Update app.py

app.py

@@ -1,16 +1,25 @@
-import torch
 import cv2
 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')
 
-# Load the YOLOv8 model
-model = YOLO('yolov8n-seg.pt').to(device)
+# 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)
 
-# 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():
@@ -30,9 +39,7 @@ def process_video(input_video_path):
 
         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]
+        results = model(frame)[0]
         for result in results.boxes.data.tolist():
             x1, y1, x2, y2, score, class_id = result
             if score > threshold:
@@ -46,17 +53,25 @@ def process_video(input_video_path):
     cap.release()
     out.release()
 
-# Define the input and output interfaces for Gradio
+# 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",
+# )
+
 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)",
+    title="Animal detector using YOLOv8 NANO for Videos",
 )
 
-# Launch the interface
+#  = gr.Interface([ interface_video], title="Animal Detector using YOLOv8 NANO")
 demo.launch()