Update Yolov5_Deepsort/demo.py

Yolov5_Deepsort/demo.py

@@ -62,49 +62,52 @@ def main():
     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
+    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__':