Koala_Smart_City/draw_box.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# @Created : 2021/04/21
# @Author  : Koala
# @FileName: draw_box.py

import copy
import cv2
import glob
import numpy as np
import argparse

drawing = False
ix, iy = -1, -1
rects = []
frame = np.zeros([0])


def draw_circle(event, x, y, flags, param):
    global ix, iy, drawing, frame, rects
    img = frame.copy()
    if event == cv2.EVENT_LBUTTONDOWN:
        drawing = True
        ix, iy = x, y 
    elif event == cv2.EVENT_LBUTTONUP:
        if drawing:
            img = frame.copy()
            drawing = False  
            ix,x = (x,ix) if x <ix else (ix,x)
            iy,y = (y,iy) if y <iy else (iy,y)
            if (x-ix)*(y-iy) > 10:
                rects.append([ix, iy, x, y])
                print([ix, iy, x, y])
    elif event == cv2.EVENT_RBUTTONDOWN:
        img = frame.copy()
        if rects:
            del rects[-1]
        drawing = False
    elif event == cv2.EVENT_MOUSEMOVE and drawing:
        cv2.rectangle(img, (ix, iy), (x, y), (0, 0, 255), 1)
    for rect in rects:
        cv2.rectangle(img, (rect[0], rect[1]), (rect[2], rect[3]), (0, 255, 0), 1)
    cv2.imshow('img', img)

def draw_in_video(video_path, filename=''):
    global ix, iy, drawing, frame, rects
    print('读取视频文件:', video_path)
    rects.clear()
    cv2.namedWindow('img')
    cv2.setMouseCallback('img', draw_circle)
    print(video_path)
    cap = cv2.VideoCapture(video_path)
    W = cap.get(cv2.CAP_PROP_FRAME_WIDTH)
    H = cap.get(cv2.CAP_PROP_FRAME_HEIGHT)
    print(f'视频大小: WxH={W}x{H}')
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        time_ = cap.get(cv2.CAP_PROP_POS_MSEC)

        img = frame.copy()
        for rect in rects:
            cv2.rectangle(img, (rect[0], rect[1]), (rect[2], rect[3]), (0, 255, 0), 1)
        cv2.putText(img, 'time:' + str(int(time_)), (100, 100), 2, 2, (255, 255, 0))
        cv2.imshow('img', img)
        k = cv2.waitKey(40)

        if k == 27 or k == ord('q') or k == ord('Q'):  # 退出视频
            print('esc------------')
            break
        elif k == 32:  # 如果按下空格键
            print('block------------')
            while True:
                cv2.imshow('img', img)
                k = cv2.waitKey(0)  # 等待调用鼠标回调函数框出目标
                if k == 32:  # 框完目标，再次按下空格键，摄像头捕捉的画面重新播放
                    break
    if filename:
        with open(filename, 'w') as f:
            for (x1, y1, x2, y2) in rects:
                f.write(f"{x1} {y1} {x2} {y2}\n")
        print(f'已保存警戒框在文件"{filename}"')


if __name__ == '__main__':
    parser = argparse.ArgumentParser('画警戒框')
    parser.add_argument('--input', default='2.mp4', help='视频的路径')
    parser.add_argument('-o', '--output', default='sources/stall_alert_boxes.txt', help='保存框的路径')
    cfg = parser.parse_args()
    draw_in_video(cfg.input, cfg.output)