Update app.py

app.py

@@ -75,7 +75,7 @@ async def upload_image(file: UploadFile = File(...)):
 
 @app.post("/upload-video/")
 async def upload_video(file: UploadFile = File(...)):
-    """Upload a video, process it frame by frame, and return detection results."""
+    """Upload a video, process it frame by frame, and return the processed video."""
     if model is None:
         return {"error": "Model not loaded. Please upload '12x.pt' to run detection."}
 
@@ -88,30 +88,42 @@ async def upload_video(file: UploadFile = File(...)):
         os.remove(temp_video_path)
         return {"error": "Could not open video file"}
 
-    frame_results = []
-    frame_interval = 5  # Process every 5th frame for efficiency
+    # Get video properties
+    fps = int(cap.get(cv2.CAP_PROP_FPS))
+    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+
+    # Output video
+    output_video_path = temp_video_path.replace(".mp4", "_processed.mp4")
+    out = cv2.VideoWriter(output_video_path, fourcc, fps, (width, height))
 
+    frame_interval = 5  # Process every 5th frame for efficiency
     frame_index = 0
+
     while True:
         ret, frame = cap.read()
         if not ret:
             break
 
         if frame_index % frame_interval == 0:
-            predictions, object_count = process_frame(frame)
-            frame_results.append({
-                "frame_index": frame_index,
-                "object_count": object_count,
-                "detections": predictions
-            })
+            predictions, _ = process_frame(frame)
+
+            # Draw bounding boxes on the frame
+            for pred in predictions:
+                x1, y1, x2, y2 = map(int, pred["bbox"])
+                label = f"{pred['class']} ({pred['confidence']:.2f})"
+                cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
+                cv2.putText(frame, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
 
+        out.write(frame)
         frame_index += 1
 
     cap.release()
-    os.remove(temp_video_path)  # Clean up temporary file
-
-    return {"video_results": frame_results}
+    out.release()
+    os.remove(temp_video_path)  # Clean up temp file
 
+    return FileResponse(output_video_path, media_type="video/mp4", filename="processed_video.mp4")
 
 @app.get("/")
 def home():