Update Yolov5-Deepsort/demo.py

Yolov5-Deepsort/demo.py

@@ -61,6 +61,52 @@ def main():
     videoWriter.release()
     cv2.destroyAllWindows()
 
+
+def app_main(video):
+    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
+
 if __name__ == '__main__':
     
     main()