Update app.py

app.py

@@ -4,12 +4,15 @@ from detect_deepsort import run_deepsort
 from detect_strongsort import run_strongsort
 from detect import run
 import os
+from PIL import Image
+import numpy as np
 import threading
-import torch
+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
@@ -21,23 +24,25 @@ def yolov9_inference(model_id, img_path=None, vid_path=None, tracking_algorithm
     input_path = None
     output_path = None
     if img_path is not None:
-        #_, img_extension = os.path.splitext(img_path)
-        #if img_extension.lower() in img_extensions:
+        # Convert the numpy array to an image
+        img = Image.fromarray(img_path)
+        img_path = 'output.png'
+        # Save the image
+        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_extension = os.path.splitext(vid_path)
         #if vid_extension.lower() in vid_extensions:
         input_path = vid_path
         print(input_path)
         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:
@@ -51,78 +56,46 @@ def yolov9_inference(model_id, img_path=None, vid_path=None, tracking_algorithm
 
 
 
-def stop_processing():
-    global should_continue
-    should_continue = False 
-    return "Stop..."
-
-def app():
-    with gr.Blocks():
-        with gr.Row():
-            with gr.Column():
-                gr.HTML("<h2>Input Parameters</h2>")
-                img_path = gr.File(label="Image")
-                vid_path = gr.File(label="Video")
-                model_id = gr.Dropdown(
-                    label="Model",
-                    choices=[
-                        "yolov9_e_trained-converted.pt",
-                        "our-converted.pt",
-                        "last_best_model.pt"
-                    ],
-                    value="our-converted.pt"
-                    
-                )
-                tracking_algorithm = gr.Dropdown(
-                    label= "Tracking Algorithm",
-                    choices=[
-                        "None",
-                        "deep_sort",
-                        "strong_sort"
-                    ],
-                    value="None"
-                )
-                yolov9_infer = gr.Button(value="Inference")
-                stop_button = gr.Button(value="Stop")
-            with gr.Column():
-                gr.HTML("<h2>Output</h2>")
-                output_image = gr.Image(type="numpy",label="Output Image")
-                output_video = gr.Video(label="Output Video")
-                output_path = gr.Textbox(label="Output path")
+def app(model_id, img_path, vid_path, tracking_algorithm):
+    return yolov9_inference(model_id, img_path, vid_path, tracking_algorithm)
 
-        yolov9_infer.click(
-            fn=yolov9_inference,
-            inputs=[
-                model_id,
-                img_path,
-                vid_path,
-                tracking_algorithm
+iface = gr.Interface(
+    fn=app, 
+    inputs=[
+        gr.Dropdown(
+            label="Model",
+            choices=[
+                "last_best_model.pt",
+                "best_model-converted.pt",
+                "yolov9_e_trained.pt",
+                "best_model-converted-reparams.pt"
+            ],
+            value="last_best_model.pt"
+        ),
+        gr.Image(label="Image"),
+        gr.Video(label="Video"),
+        gr.Dropdown(
+            label= "Tracking Algorithm",
+            choices=[
+                "None",
+                "deep_sort",
+                "strong_sort"
             ],
-            outputs=[output_image, output_video, output_path],
+            value="None"
         )
-        stop_button.click(stop_processing)
-
-
-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)
-
+    ], 
+    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)