Update yolo3_predict_video.py

gittigxuy · Aug 3, 2018 · a937d91 · a937d91
1 parent 583ce99
commit a937d91
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/yolo3_predict_video.py b/yolo3_predict_video.py
@@ -45,6 +45,8 @@ def detect_video(yolo, video_path, output_path=""):
     prev_time = timer()
     while True:
         return_value, frame = vid.read()
+        if not return_value:
+            break
         image = Image.fromarray(frame)
         image = yolo.detect_image(image)
         result = np.asarray(image)
@@ -73,4 +75,4 @@ def detect_video(yolo, video_path, output_path=""):
         output_path = sys.argv[2]
         detect_video(YOLO(), video_path, output_path)#在这里调用YOLO函数，从而使用了已经训练好的
     else:
-        detect_video(YOLO(), video_path)
+        detect_video(YOLO(), video_path)