Update app.py

app.py

@@ -71,34 +71,32 @@ def detect_video(video):
     # Create a video capture object
     cap = cv2.VideoCapture(video)
 
+    nb_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    frame_h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+    frame_w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    fps = cap.get(cv2.CAP_PROP_FPS)
+    
     # Process frames in a loop
     while cap.isOpened():
         ret, frame = cap.read()
         if not ret:
             break
-
-        # Expand dimensions since model expects images to have shape: [1, None, None, 3]
-        image_np_expanded = np.expand_dims(frame, axis=0)
-
-        # Run inference
-        output_dict = detection_model(image_np_expanded)
-
-        # Extract detections
-        boxes = output_dict['detection_boxes'][0].numpy()
-        scores = output_dict['detection_scores'][0].numpy()
-        classes = output_dict['detection_classes'][0].numpy().astype(np.int64)
-
-        # Draw bounding boxes and labels
+    for i in tqdm(range(nb_frames)):
+        ret, image_np = video_reader.read()
+        input_tensor = tf.convert_to_tensor(np.expand_dims(image_np, 0), dtype=tf.uint8)
+        results = detection_model(input_tensor)
         image_np_with_detections = viz_utils.visualize_boxes_and_labels_on_image_array(
-            frame,
-            boxes,
-            classes,
-            scores,
-            category_index,
-            use_normalized_coordinates=True,
-            max_boxes_to_draw=20,
-            min_score_thresh=.5,
-            agnostic_mode=False)
+                  image_np,
+                  results['detection_boxes'][0].numpy(),
+                  (results['detection_classes'][0].numpy()+ label_id_offset).astype(int),
+                  results['detection_scores'][0].numpy(),
+                  category_index,
+                  use_normalized_coordinates=True,
+                  max_boxes_to_draw=200,
+                  min_score_thresh=.50,
+                  agnostic_mode=False,
+                  line_thickness=2)
+
 
         # Yield the processed frame
         yield image_np_with_detections