Yolov5_Deepsort/demo.py

from Yolov5_Deepsort.AIDetector_pytorch import Detector
import imutils
import rich
import cv2
import time

def main():

    name = 'demo'

    det = Detector()
    cap = cv2.VideoCapture('mot.mp4')
    fps = int(cap.get(5))
    print('fps:', fps)
    t = int(1000/fps)
    frame_count = 0
    total_time = 0

    videoWriter = None

    while True:

        start_time = time.time()

        # try:
        _, im = cap.read()
        if im is None:
            break
        #rich.print("im:",im)
        result = det.feedCap(im)
        #rich.print(result)
        result = result['frame']
        result = imutils.resize(result, height=500)
        if videoWriter is None:
            fourcc = cv2.VideoWriter_fourcc(
                'm', 'p', '4', 'v')  # opencv3.0
            videoWriter = cv2.VideoWriter(
                'result.mp4', fourcc, fps, (result.shape[1], result.shape[0]))

        videoWriter.write(result)
        cv2.imshow(name, result)
        cv2.waitKey(t)

        end_time = time.time()
        total_time += (end_time - start_time)
        frame_count +=1

        if frame_count > 0:
            processing_fps = frame_count / total_time
            rich.print('Processing fps:', processing_fps)


        if cv2.getWindowProperty(name, cv2.WND_PROP_AUTOSIZE) < 1:
            # 点x退出
            break
        # except Exception as e:
        #     print(e)
        #     break

    cap.release()
    videoWriter.release()
    cv2.destroyAllWindows()

def app_main(video):
    try:
        name = 'demo'
        det = Detector()
        cap = cv2.VideoCapture(video)
        fps = int(cap.get(cv2.CAP_PROP_FPS))
        print('fps:', fps)
        t = int(1000 / fps)
        frame_count = 0
        total_time = 0

        videoWriter = None
        output_filename = 'result.mp4'

        while True:
            start_time = time.time()

            ret, im = cap.read()
            if not ret:
                break

            result = det.feedCap(im)
            result = result['frame']
            result = imutils.resize(result, height=500)

            if videoWriter is None:
                fourcc = cv2.VideoWriter_fourcc(*'mp4v')
                videoWriter = cv2.VideoWriter(
                    output_filename, fourcc, fps, (result.shape[1], result.shape[0]))

            videoWriter.write(result)

            end_time = time.time()
            total_time += (end_time - start_time)
            frame_count += 1

            if frame_count > 0:
                processing_fps = frame_count / total_time
                print('Processing fps:', processing_fps)

        cap.release()
        videoWriter.release()

        return output_filename
    except Exception as e:
        print(f"Error in app_main: {e}")
        return None

if __name__ == '__main__':
    
    main()