Update app.py

app.py

@@ -1,24 +1,14 @@
+import torch
 import cv2
 import numpy as np
 import gradio as gr
 from ultralytics import YOLO
 
-model = YOLO('yolov8n-seg.pt')
+# Check if CUDA (GPU support) is available
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
-# 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)
+# Load the YOLOv8 model
+model = YOLO(r'E:\CV-project\MAIN-DATA\GRADIO\Yolo-v8-Video\yolov8x-seg.pt').to(device)
 
 def process_video(input_video_path):
     cap = cv2.VideoCapture(input_video_path)
@@ -37,9 +27,15 @@ def process_video(input_video_path):
         if not ret:
             break
 
+        # Resize frame to match the expected input shape of the model
+        resized_frame = cv2.resize(frame, (640, 640))
+
+        # Convert resized frame to torch tensor and move it to GPU
+        frame_tensor = torch.from_numpy(resized_frame).permute(2, 0, 1).unsqueeze(0).float().to(device) / 255.0
+
         threshold = 0.1
         frame_copy = frame.copy()
-        results = model(frame)[0]
+        results = model(frame_tensor)[0]
         for result in results.boxes.data.tolist():
             x1, y1, x2, y2, score, class_id = result
             if score > threshold:
@@ -52,26 +48,21 @@ def process_video(input_video_path):
 
     cap.release()
     out.release()
+    cv2.destroyAllWindows()
 
-# inputs_image = [gr.Image(label="Input Image")]
-# outputs_image = [gr.Image( label="Output Image")]
-
-# interface_image = gr.Interface(
-#     fn=show_preds_image,
-#     inputs=inputs_image,
-#     outputs=outputs_image,
-#     title="Animal detector using YOLOv8 NANO for Images",
-# )
+    return "output_video.mp4"
 
+# 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",
+    title="Animal detector using YOLOv8 NANO for Videos (GPU)",
 )
 
-#  = gr.Interface([ interface_video], title="Animal Detector using YOLOv8 NANO")
+# Launch the interface
 demo.launch()