Update app.py

app.py

@@ -1,81 +1,62 @@
-import gradio as gr
+import torch
 import cv2
-import os
- 
+import numpy as np
+import gradio as gr
 from ultralytics import YOLO
 
+# Check if CUDA (GPU support) is available
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
-model = YOLO('yolov8n-seg.pt')
-path  = [['3891186464_00d76e10a2_z.jpg'], ['images (1).jpeg']]
-video_path = [['sheep.mp4']]
+# Load the YOLOv8 model
+model = YOLO('yolov8n-seg.pt').to(device)
 
+# Define the function to process the video on GPU
+def process_video(input_video_path):
+    cap = cv2.VideoCapture(input_video_path)
+    if not cap.isOpened():
+        print("Error: Couldn't open the video file.")
+        return
 
-def show_preds_image(image_path):
-    image = cv2.imread(image_path)
-    image_copy=image.copy()
-    threshold = 0.1
-    results = model(image)[0]
-    for result in results.boxes.data.tolist():
-                x1, y1, x2, y2, score, class_id = result
-
-                if score > threshold:
-                    cv2.rectangle(image_copy, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 4)
-                    cv2.putText(image_copy, results.names[int(class_id)].upper(), (int(x1), int(y1 - 10)),
-                                cv2.FONT_HERSHEY_SIMPLEX, 1.3, (0, 255, 0), 3, cv2.LINE_AA)
-                    cv2.putText(image_copy, str(score), (int(x1), int(y2 + 10)),
-                                cv2.FONT_HERSHEY_SIMPLEX, 1.3, (0, 0, 255), 3, cv2.LINE_AA)
-    return cv2.cvtColor(image_copy, cv2.COLOR_BGR2RGB)
- 
-inputs_image = [
-    gr.components.Image(type="filepath", label="Input Image"),
-]
-outputs_image = [
-    gr.components.Image(type="numpy", label="Output Image"),
-]
-interface_image = gr.Interface(
-    fn=show_preds_image,
-    inputs=inputs_image,
-    outputs=outputs_image,
-    title="Animal detector using YOLOV8 NANO",
-    examples=path,
-    cache_examples=False,
-)
+    frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+    fps = cap.get(cv2.CAP_PROP_FPS)
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+    out = cv2.VideoWriter("output_video.mp4", fourcc, fps, (frame_width, frame_height))
 
-def show_preds_video(video_path):
-    cap = cv2.VideoCapture(video_path)
-    while(cap.isOpened()):
+    while cap.isOpened():
         ret, frame = cap.read()
-        if ret:
-            threshold = 0.1
-            frame_copy = frame.copy()
-            results = model(frame)[0]
-            for result in results.boxes.data.tolist():
-                x1, y1, x2, y2, score, class_id = result
-
-                if score > threshold:
-                    cv2.rectangle(frame_copy, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 4)
-                    cv2.putText(frame_copy, results.names[int(class_id)].upper(), (int(x1), int(y1 - 10)),
-                                cv2.FONT_HERSHEY_SIMPLEX, 1.3, (0, 255, 0), 3, cv2.LINE_AA)
-                    cv2.putText(frame_copy, str(score), (int(x1), int(y2 + 10)),
-                                cv2.FONT_HERSHEY_SIMPLEX, 1.3, (0, 0, 255), 3, cv2.LINE_AA)
-            
-            yield cv2.cvtColor(frame_copy, cv2.COLOR_BGR2RGB)
-inputs = [
-    gr.Video(label='Input Video')]
-outputs = gr.Video(label='Output Video')
-
-
-
-interface_video = gr.Interface(
-    fn=show_preds_video,
-    inputs=inputs,
-    outputs=outputs,
-    title="Cattle detector using YOLOV8 NANO",
-    examples=video_path,
-    cache_examples=False,
+        if not ret:
+            break
+
+        threshold = 0.1
+        frame_copy = frame.copy()
+        # Convert frame to torch tensor and move it to GPU
+        frame_tensor = torch.from_numpy(frame_copy).permute(2, 0, 1).unsqueeze(0).float().to(device) / 255.0
+        results = model(frame_tensor)[0]
+        for result in results.boxes.data.tolist():
+            x1, y1, x2, y2, score, class_id = result
+            if score > threshold:
+                cv2.rectangle(frame_copy, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 4)
+                cv2.putText(frame_copy, results.names[int(class_id)].upper(), (int(x1), int(y1 - 10)),
+                            cv2.FONT_HERSHEY_SIMPLEX, 1.3, (0, 255, 0), 3, cv2.LINE_AA)
+                cv2.putText(frame_copy, str(score), (int(x1), int(y2 + 10)),
+                            cv2.FONT_HERSHEY_SIMPLEX, 1.3, (0, 0, 255), 3, cv2.LINE_AA)
+        out.write(frame_copy)
+
+    cap.release()
+    out.release()
+
+# Define the input and output interfaces for Gradio
+inputs_video = gr.Video(label="Input Video")
+outputs_video = gr.Video(label="Output Video")
+
+# Create the Gradio interface
+demo = gr.Interface(
+    fn=process_video,
+    inputs=inputs_video,
+    outputs=outputs_video,
+    title="Animal detector using YOLOv8 NANO for Videos (GPU)",
 )
 
-gr.TabbedInterface(
-    [interface_image, interface_video],
-    tab_names=['Image inference', 'Video inference']
-).launch()
+# Launch the interface
+demo.launch()