Update app.py

app.py

@@ -23,13 +23,13 @@ for i, url in enumerate(file_urls):
 
 colors = {
     0: (255, 0, 0),    # Red for class 0
-    1: (0, 255, 0),    # Green for class 1
+    1: (0, 128, 0),    # Green (dark) for class 1
     2: (0, 0, 255),    # Blue for class 2
     3: (255, 255, 0),  # Yellow for class 3
     4: (255, 0, 255),  # Magenta for class 4
     5: (0, 255, 255),  # Cyan for class 5
     6: (128, 0, 0),    # Maroon for class 6
-    7: (0, 128, 0),    # Green (dark) for class 7
+    7: (0, 225, 0),    # Green for class 7
 }
 
 model = YOLO('modelbest.pt')
@@ -95,41 +95,61 @@ interface_image = gr.Interface(
 )
 
 def show_preds_video(video_path):
+    # Open the input video
     cap = cv2.VideoCapture(video_path)
+    
+    # Get video properties
+    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+    fps = int(cap.get(cv2.CAP_PROP_FPS))
+    
+    # Define the codec and create a VideoWriter object
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # 'mp4v' for .mp4 format
+    out = cv2.VideoWriter('output_video.mp4', fourcc, fps, (width, height))
+    
     while cap.isOpened():
         ret, frame = cap.read()
-        if ret:
-            frame_copy = frame.copy()
-            outputs = model.predict(source=frame)
-            results = outputs[0].cpu().numpy()
-
-            for i, det in enumerate(results.boxes.xyxy):
-                class_id = int(results.boxes.cls[i])
-                label = model.names[class_id]
-
-                # Get the bounding box coordinates
-                x1, y1, x2, y2 = int(det[0]), int(det[1]), int(det[2]), int(det[3])
-                
-                # Draw the bounding box with the specified color
-                color = colors.get(class_id, (0, 0, 255))
-                cv2.rectangle(frame_copy, (x1, y1), (x2, y2), color, 2, cv2.LINE_AA)
-                
-                # Calculate text size and position
-                label_size, _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.75, 2)
-                text_x = x1 + (x2 - x1) // 2 - label_size[0] // 2
-                text_y = y1 + (y2 - y1) // 2 + label_size[1] // 2
-
-                # Draw the label text
-                cv2.putText(frame_copy, label, (text_x, text_y), cv2.FONT_HERSHEY_SIMPLEX, 0.75, color, 2, cv2.LINE_AA)
-
-            yield cv2.cvtColor(frame_copy, cv2.COLOR_BGR2RGB)
+        if not ret:
+            break
+        
+        frame_copy = frame.copy()
+        outputs = model.predict(source=frame)
+        results = outputs[0].cpu().numpy()
 
+        for i, det in enumerate(results.boxes.xyxy):
+            class_id = int(results.boxes.cls[i])
+            label = model.names[class_id]
+
+            # Get the bounding box coordinates
+            x1, y1, x2, y2 = int(det[0]), int(det[1]), int(det[2]), int(det[3])
+            
+            # Draw the bounding box with the specified color
+            color = colors.get(class_id, (0, 0, 255))
+            cv2.rectangle(frame_copy, (x1, y1), (x2, y2), color, 2, cv2.LINE_AA)
             
+            # Calculate text size and position
+            label_size, _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.75, 2)
+            text_x = x1 + (x2 - x1) // 2 - label_size[0] // 2
+            text_y = y1 + (y2 - y1) // 2 + label_size[1] // 2
+
+            # Draw the label text
+            cv2.putText(frame_copy, label, (text_x, text_y), cv2.FONT_HERSHEY_SIMPLEX, 0.75, color, 2, cv2.LINE_AA)
+
+        # Write the frame to the output video
+        out.write(frame_copy)
+    
+    # Release everything
+    cap.release()
+    out.release()
+
+    return 'output_video.mp4'
+
+# Updated Gradio interface
 inputs_video = [
     gr.Video(format="mp4", label="Input Video"),
 ]
 outputs_video = [
-    gr.Image(type="numpy", label="Output Image"),
+    gr.Video(label="Output Video"),
 ]
 interface_video = gr.Interface(
     fn=show_preds_video,
@@ -139,12 +159,12 @@ interface_video = gr.Interface(
     examples=video_path,
     cache_examples=False,
 )
-
 gr.TabbedInterface(
     [interface_image, interface_video],
     tab_names=['Image inference', 'Video inference']
 ).queue().launch()
 
+
 # import gradio as gr
 # import cv2
 # import requests