Update app.py

app.py

@@ -1,33 +1,63 @@
-import torch
-from torchvision import transforms
-from PIL import Image
 import gradio as gr
+import torch
+from ultralytics import YOLO
+import cv2
+import tempfile
 
-# Load the YOLO model
-model = torch.hub.load('ultralytics/yolov8', 'custom', path='best.pt')  # Ensure 'best.pt' is the path to your trained model
-
-# Define a function to process the image and make predictions
-def detect_objects(image):
-    # Preprocess the image
-    transform = transforms.Compose([
-        transforms.ToTensor()
-    ])
-    image = transform(image).unsqueeze(0)  # Add batch dimension
+# Load the trained YOLOv8 model
+model = YOLO('best.pt')
 
-    # Perform inference
+def predict(image):
     results = model(image)
+    # You might want to process results to return bounding boxes, class labels, etc.
+    annotated_image = results[0].plot()  # plot the results on the image
+    return annotated_image
+
+def predict_video(video):
+    # Read the video file
+    cap = cv2.VideoCapture(video)
+    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))
+    
+    # Create a temporary file to save the output video
+    out_file = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4')
+    out_path = out_file.name
+    
+    # Define the codec and create VideoWriter object
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+    out = cv2.VideoWriter(out_path, fourcc, fps, (width, height))
+
+    while cap.isOpened():
+        ret, frame = cap.read()
+        if not ret:
+            break
+        results = model(frame)
+        annotated_frame = results[0].plot()  # plot the results on the frame
+        out.write(annotated_frame)
+    
+    cap.release()
+    out.release()
+    
+    return out_path
+
+# Create Gradio interface
+interface = gr.Interface(
+    fn=lambda img, vid: (predict(img), predict_video(vid)),
+    inputs=[
+        gr.inputs.Image(type="numpy", label="Input Image"),
+        gr.inputs.Video(label="Input Video")
+    ],
+    outputs=[
+        gr.outputs.Image(type="numpy", label="Output Image"),
+        gr.outputs.Video(label="Output Video")
+    ],
+    title="YOLOv8 Object Detection",
+    description="Upload an image or a video and get the object detection results using a YOLOv8 model."
+)
 
-    # Extract bounding boxes and labels
-    bbox_img = results.render()[0]  # This gets the image with bounding boxes drawn
-
-    return Image.fromarray(bbox_img)
-
-# Create the Gradio interface
-inputs = gr.inputs.Image(shape=(640, 480))
-outputs = gr.outputs.Image(type="pil")
+if __name__ == "__main__":
+    interface.launch()
 
-gr_interface = gr.Interface(fn=detect_objects, inputs=inputs, outputs=outputs, title="YOLO Object Detection", description="Upload an image to detect objects using a YOLO model.")
 
-# Run the Gradio app
-if __name__ == "__main__":
-    gr_interface.launch()
+