opencv video

mydemo.py

import cv2
import sys
import argparse
import time

from Processor import Processor
from Visualizer import Visualizer

 
def gstreamer_pipeline(
    capture_width=1024,
    capture_height=1024,
    display_width=608,
    display_height=608,
    framerate=2,
    flip_method=0,#
):
    return (
        "nvarguscamerasrc ! "
        "video/x-raw(memory:NVMM), "
        "width=(int)%d, height=(int)%d, "
        "format=(string)NV12, framerate=(fraction)%d/1 ! "
        "nvvidconv flip-method=%d ! "
        "video/x-raw, width=(int)%d, height=(int)%d, format=(string)BGRx ! "
        "videoconvert ! "
        "video/x-raw, format=(string)BGR ! appsink"
        % (
            capture_width,
            capture_height,
            framerate,
            flip_method,
            display_width,
            display_height,
        )
    )
 

def cli():
    desc = 'Run TensorRT yolov5 visualizer'
    parser = argparse.ArgumentParser(description=desc)
    parser.add_argument('-model',  help='trt engine file located in ./models', required=False)
    parser.add_argument('-image', help='image file path', required=False)
    args = parser.parse_args()
    model = args.model or 'yolov3.trt'
    img = args.image or 'dogs.jpeg'
    return { 'model': model, 'image': img }

def main():
    # parse arguments
    args = cli()

    # setup processor and visualizer
    processor = Processor(model=args['model'])
    visualizer = Visualizer()

    # fetch input
    # print('image arg', args['image'])
    # img = cv2.imread('inputs/{}'.format(args['image']))

    cap = cv2.VideoCapture(gstreamer_pipeline(flip_method=2), cv2.CAP_GSTREAMER)
 
    while cap.isOpened():
        flag, img = cap.read()  
        # inference
        t0 =time.time()
        output = processor.detect(img) 
        # img = cv2.resize(img, (640, 640))

        # object visualization
        # object_grids = processor.extract_object_grids(output)
        # visualizer.draw_object_grid(img, object_grids, 0.1)

        # class visualization
        # class_grids = processor.extract_class_grids(output)
        # visualizer.draw_class_grid(img, class_grids, 0.01)

        # bounding box visualization
        t1 =time.time()
        boxes = processor.extract_boxes(output)
        t2 =time.time()
        boxedimg = visualizer.draw_boxes(img, boxes)
        t3 =time.time()

        print("detect time ",t1-t0)
        print("extract boxes time ",t2-t1) 
        print("draw boxes time ",t3-t2)
        print("total time ",t3-t0)

        cv2.imshow("CSI Camera", boxedimg)
        # cv2.imshow("CSI Camera", img)
        kk = cv2.waitKey(1)
        # do other things
 
        if kk == ord('q'):  # 按下 q 键，退出
             break
 
    cap.release()
    cv2.destroyAllWindows()
    

    # final results
    # boxes, confs, classes = processor.post_process(output)
    # visualizer.draw_results(img, boxes, confs, classes)

if __name__ == '__main__':
    main()