Bug saving webcam output fixed.

ksanjeevan · Aug 22, 2018 · fa96f91 · fa96f91
1 parent be9da6c
commit fa96f91
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 15 deletions.
diff --git a/confs/config_coco.json b/confs/config_coco.json
@@ -4,7 +4,7 @@
         "grid_size":        13,
         "true_box_buffer":  10,
         "iou_threshold":    0.5,  
-        "nms_threshold":    0.4
+        "nms_threshold":    0.3
     },
     "config_path" : {
         "labels":           "models/coco/labels_coco.txt",

diff --git a/yolov2.py b/yolov2.py
@@ -20,7 +20,7 @@
 from net.netgen import YoloDataGenerator
 
 
-CAM_WIDTH = 1038
+CAM_WIDTH = 1024
 CAM_HEIGHT = 576
 
 class YoloV2(object):
@@ -92,20 +92,19 @@ def inference(self, path):
 
 
     def _video_params(self, name):
-
-        cap = cv2.VideoCapture(name)    
-        if name == 0:        
-            cap.set(cv2.CAP_PROP_FRAME_HEIGHT, CAM_HEIGHT)
-            cap.set(cv2.CAP_PROP_FRAME_WIDTH, CAM_WIDTH)
-            size = (CAM_WIDTH, CAM_HEIGHT)
-        else:
-            video_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
-            video_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-            size = (video_width, video_height)
 
-        video_len = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+        cap = cv2.VideoCapture(name)    
 
+        video_len = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
         fps = round(cap.get(cv2.CAP_PROP_FPS))
+
+        if name == 0:        
+            cap.set(cv2.CAP_PROP_FRAME_WIDTH, CAM_WIDTH)
+            cap.set(cv2.CAP_PROP_FRAME_HEIGHT, CAM_HEIGHT)
+
+        video_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+        video_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+        size = (video_width, video_height)
 
         return cap, size, video_len, fps
 
@@ -122,7 +121,6 @@ def video_inference(self, filename):
 
             boxes, scores, _, labels = self.inf_model.predict(frame)
             frame_pred = draw_boxes(frame, (boxes, scores, labels))
-
             writer.write(frame_pred)
 
             if cv2.waitKey(1) & 0xFF == ord('q'):
@@ -141,6 +139,7 @@ def cam_inference(self, fname):
         cap, size, _, fps = self._video_params(0)
 
         fourcc = cv2.VideoWriter_fourcc('m','p','4','v')
+
         if fname: writer = cv2.VideoWriter(fname, fourcc, fps, size)
 
         while(cap.isOpened()):
@@ -158,7 +157,8 @@ def cam_inference(self, fname):
                     break
 
         cap.release()
-        if fname: 
+
+        if fname:
             writer.release()
             if YoloParams.STORE_GIF: generate_gif(fname)
         cv2.destroyAllWindows()