minor update to yolo.py code

glMa7 · Jan 14, 2019 · 026930a · 026930a
1 parent 3eb897a
commit 026930a
Showing 1 changed file with 5 additions and 4 deletions.
diff --git a/yolo/yolo.py b/yolo/yolo.py
@@ -172,12 +172,14 @@ def detect_img(yolo):
         try:
             image = Image.open(img)
         except:
-            print('*** Open Error! Try again!')
-            continue
+            if img == 'q' or img == 'Q':
+                break
+            else:
+                print('*** Open Error! Try again!')
+                continue
         else:
             res_image, _ = yolo.detect_image(image)
             res_image.show()
-
     yolo.close_session()
 
 
@@ -247,6 +249,5 @@ def detect_video(model, video_path=None, output=None):
     vid.release()
     out.release()
     cv2.destroyAllWindows()
-
     # close the session
     model.close_session()