Update app.py

app.py

@@ -1,18 +1,18 @@
-import spaces
+#import spaces
 import gradio as gr
 from detect_deepsort import run_deepsort
 from detect_strongsort import run_strongsort
 from detect import run
 import os
-import torch
 from PIL import Image
+import cv2
 import numpy as np
 import threading
-import cv2
 
 should_continue = True
 
-@spaces.GPU(duration=120)
+
+#@spaces.GPU(duration=120)
 def yolov9_inference(model_id, img_path=None, vid_path=None, tracking_algorithm = None):
     global should_continue
     img_extensions = ['.jpg', '.jpeg', '.png', '.gif']  # Add more image extensions if needed
@@ -31,7 +31,7 @@ def yolov9_inference(model_id, img_path=None, vid_path=None, tracking_algorithm
         img.save(img_path)
         input_path = img_path
         print(input_path)
-        output_path = run(weights=model_id, imgsz=(image_size,image_size), conf_thres=conf_threshold, iou_thres=iou_threshold, source=input_path, device='0', hide_conf= True)
+        output_path = run(weights=model_id, imgsz=(image_size,image_size), conf_thres=conf_threshold, iou_thres=iou_threshold, source=input_path, device='cpu', hide_conf= True)
     elif vid_path is not None:
         vid_name = 'output.mp4'
 
@@ -68,64 +68,91 @@ def yolov9_inference(model_id, img_path=None, vid_path=None, tracking_algorithm
         out.release()
         input_path = vid_name
         if tracking_algorithm == 'deep_sort':
-            output_path = run_deepsort(weights=model_id, imgsz=(image_size,image_size), conf_thres=conf_threshold, iou_thres=iou_threshold, source=input_path, device='0', draw_trails=True)
+            output_path = run_deepsort(weights=model_id, imgsz=(image_size,image_size), conf_thres=conf_threshold, iou_thres=iou_threshold, source=input_path, device='cpu', draw_trails=True)
         elif tracking_algorithm == 'strong_sort':
-            device_strongsort = torch.device('cuda:0')
-            output_path = run_strongsort(yolo_weights=model_id, imgsz=(image_size,image_size), conf_thres=conf_threshold, iou_thres=iou_threshold, source=input_path, device=device_strongsort, strong_sort_weights = "osnet_x0_25_msmt17.pt", hide_conf= True)
+            output_path = run_strongsort(yolo_weights=model_id, imgsz=(image_size,image_size), conf_thres=conf_threshold, iou_thres=iou_threshold, source=input_path, device='cpu', strong_sort_weights = "osnet_x0_25_msmt17.pt", hide_conf= True)
         else: 
-            output_path =  run(weights=model_id, imgsz=(image_size,image_size), conf_thres=conf_threshold, iou_thres=iou_threshold, source=input_path, device='0', hide_conf= True)
+            output_path =  run(weights=model_id, imgsz=(image_size,image_size), conf_thres=conf_threshold, iou_thres=iou_threshold, source=input_path, device='cpu', hide_conf= True)
         # Assuming output_path is the path to the output file
     _, output_extension = os.path.splitext(output_path)
     if output_extension.lower() in img_extensions:
-        output_image = output_path  # Load the image file here
-        output_video = None
+        output = output_path  # Load the image file here
     elif output_extension.lower() in vid_extensions:
-        output_image = None
-        output_video = output_path  # Load the video file here
-
-    return output_image, output_video, output_path
+        output = output_path  # Load the video file here
+
+    return output, output_path
+
+
+def app():
+    with gr.Blocks(title="YOLOv9: Real-time Object Detection", css=".gradio-container {background:lightyellow;}"):
+        with gr.Row():
+            with gr.Column():
+                gr.HTML("<h2>Input Parameters</h2>")
+                img_path = gr.Image(label="Image", height = 370, width = 600)
+                vid_path = gr.Video(label="Video", height = 370, width = 600)
+                model_id = gr.Dropdown(
+                    label="Model",
+                    choices=[
+                        "our-converted.pt",
+                        "yolov9_e_trained-converted.pt"
+                    ],
+                    value="our-converted.pt"
+                    
+                )
+                tracking_algorithm = gr.Dropdown(
+                    label= "Tracking Algorithm",
+                    choices=[
+                        "None",
+                        "deep_sort",
+                        "strong_sort"
+                    ],
+                    value="None"
+                )
+                gr.Examples(['camera1_A_133.png'], inputs=img_path,label = "Image Example")
+                gr.Examples(['test.mp4'], inputs=vid_path, label = "Video Example")
+                yolov9_infer = gr.Button(value="Inference")
+            with gr.Column():
+                gr.HTML("<h2>Output</h2>")
+                if img_path is not None:
+                    output_image = gr.Image(type="numpy",label="Output")
+                    output = output_image
+                else:
+                    output_video = gr.Video(label="Output")
+                    output = output_video
+                output_path = gr.Textbox(label="Output path")
+
+        yolov9_infer.click(
+            fn=yolov9_inference,
+            inputs=[
+                model_id,
+                img_path,
+                vid_path,
+                tracking_algorithm
+            ],
+            outputs=[output, output_path],
+        )
 
 
+gradio_app = gr.Blocks()
+with gradio_app:
+    gr.HTML(
+        """
+    <h1 style='text-align: center'>
+    YOLOv9: Real-time Object Detection
+    </h1>
+    """)
+    css = """
+    body {
+        background-color: #f0f0f0;
+    }
+    h1 {
+        color: #4CAF50;
+    }
+    """
+    with gr.Row():
+        with gr.Column():
+            app()
+
+gradio_app.launch(debug=True)
 
-def app(model_id, img_path, vid_path, tracking_algorithm):
-    return yolov9_inference(model_id, img_path, vid_path, tracking_algorithm)
 
-iface = gr.Interface(
-    fn=app, 
-    inputs=[
-        gr.Dropdown(
-            label="Model",
-            choices=[
-                "our-converted.pt",
-                "yolov9_e_trained-converted.pt",
-                "last_best_model.pt"
-            ],
-            value="our-converted.pt"
-        ),
-        gr.Image(label="Image"),
-        gr.Video(label="Video"),
-        gr.Dropdown(
-            label= "Tracking Algorithm",
-            choices=[
-                "None",
-                "deep_sort",
-                "strong_sort"
-            ],
-            value="None"
-        )
-    ], 
-    outputs=[
-        gr.Image(type="numpy",label="Output Image"),
-        gr.Video(label="Output Video"),
-        gr.Textbox(label="Output path")
-    ],
-    examples=[
-        ["last_best_model.pt", "camera1_A_133.png", None, "deep_sort"],
-        ["last_best_model.pt", None, "test.mp4", "strong_sort"]
-    ],
-    title='YOLOv9: Real-time Object Detection',
-    description='This is a real-time object detection system using YOLOv9.',
-    theme='huggingface'
-)
-
-iface.launch(debug=True)