Update app.py

app.py

@@ -1,49 +1,53 @@
-import torch
 import cv2
-import numpy as np
-
-# Load your model (assuming it is a PyTorch model)
-model = torch.load('./data/model.pt', weights_only=False)
+import torch
+import gradio as gr
+from ultralytics import YOLO
 
-model.eval()
+# Load YOLOv8 model
+model = YOLO('./data/model.pt')  # Path to your model
 
-# Open video file (input video)
-input_video = cv2.VideoCapture('input_video.mp4')
+# Define the function that processes the uploaded video
+def process_video(video):
+    # Read the uploaded video file
+    input_video = cv2.VideoCapture(video.name)  # 'video' here is the uploaded video file
 
-# Get the frame width, height, and frames per second (fps) from the input video
-frame_width = int(input_video.get(cv2.CAP_PROP_FRAME_WIDTH))
-frame_height = int(input_video.get(cv2.CAP_PROP_FRAME_HEIGHT))
-fps = input_video.get(cv2.CAP_PROP_FPS)
+    # Get frame width, height, and fps from input video
+    frame_width = int(input_video.get(cv2.CAP_PROP_FRAME_WIDTH))
+    frame_height = int(input_video.get(cv2.CAP_PROP_FRAME_HEIGHT))
+    fps = input_video.get(cv2.CAP_PROP_FPS)
 
-# Define the output video writer
-fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # You can change this to any codec
-output_video = cv2.VideoWriter('output_video.mp4', fourcc, fps, (frame_width, frame_height))
+    # Define output video writer
+    output_video_path = "/mnt/data/output_video.mp4"  # Path to save the output video
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+    output_video = cv2.VideoWriter(output_video_path, fourcc, fps, (frame_width, frame_height))
 
-while True:
-    # Read a frame from the input video
-    ret, frame = input_video.read()
-    if not ret:
-        break  # End of video
+    while True:
+        # Read a frame from the video
+        ret, frame = input_video.read()
+        if not ret:
+            break  # End of video
 
-    # Preprocess the frame if necessary (depends on your model)
-    # For example, convert to tensor and normalize if required
-    frame_tensor = torch.tensor(frame).float().unsqueeze(0)  # Add batch dimension
+        # Perform inference on the frame
+        results = model(frame)
 
-    # Pass the frame through the model
-    with torch.no_grad():
-        output = model(frame_tensor)  # Adjust based on your model's requirements
+        # Extract annotated image from results
+        annotated_frame = results.render()[0]  # This will give the frame with bounding boxes
 
-    # Postprocess the output if necessary (depends on your model's output format)
-    output_frame = output.squeeze(0).cpu().numpy()  # Remove batch dimension and convert to NumPy
+        # Write the annotated frame to the output video
+        output_video.write(annotated_frame)
 
-    # Convert the model output to a valid image format (if necessary)
-    output_frame = np.uint8(output_frame)
+    # Release resources
+    input_video.release()
+    output_video.release()
 
-    # Write the frame to the output video
-    output_video.write(output_frame)
+    return output_video_path
 
-# Release resources
-input_video.release()
-output_video.release()
+# Create a Gradio interface for video upload
+iface = gr.Interface(fn=process_video, 
+                     inputs=gr.inputs.Video(label="Upload Video"), 
+                     outputs="file", 
+                     title="YOLOv8 Object Detection on Video",
+                     description="Upload a video for object detection using YOLOv8")
 
-cv2.destroyAllWindows()
+# Launch the interface
+iface.launch()