Update app.py

app.py

@@ -1,156 +1,64 @@
 import gradio as gr
-import spaces
 from huggingface_hub import hf_hub_download
-
+import torch
 
 def download_models(model_id):
-    hf_hub_download("merve/yolov9", filename=f"{model_id}", local_dir=f"./")
-    return f"./{model_id}"
-
-@spaces.GPU
-def yolov9_inference(img_path, model_id, image_size, conf_threshold, iou_threshold):
-	"""
-	Performs object detection using a YOLOv9 model. This function loads a specified YOLOv9 model,
-	configures it based on the provided parameters, and carries out inference on a given image. 
-	Additionally, it allows for optional modification of the input size and the application of 
-	test time augmentation to potentially improve detection accuracy.
-
-	Parameters:
-	    model_path (str): The file path to the pre-trained YOLOv9 model.
-	    conf_threshold (float): The confidence threshold used during Non-Maximum Suppression (NMS)
-		to filter detections. Detections with confidence levels below this threshold are discarded.
-	    iou_threshold (float): The Intersection over Union (IoU) threshold applied in NMS. Detections 
-		with IoU values above this threshold are considered overlapping and are hence merged.
-	    img_path (str): The file path to the image on which inference is to be performed.
-	    size (int, optional): The input size for inference. Altering this may affect the accuracy and 
-		speed of the detection process.
-
-	Returns:
-	    tuple: A tuple containing the detection results. This includes the bounding boxes (boxes),
-	    confidence scores (scores), and detected categories (categories), alongside a results object
-	    for further processing, such as visualization.
-	"""
-	
-    # Import YOLOv9
-    import yolov9
-    
-    # Load the model
+    """
+    Downloads a model file from Hugging Face Hub to a specified local directory.
+
+    Parameters:
+    - model_id (str): Identifier of the model to download.
+
+    Returns:
+    - str: Path to the downloaded model file.
+    """
+    model_path = hf_hub_download(repo_id="merve/yolov9", filename=model_id)
+    return model_path
+
+def yolov9_inference(img_path, model_id="model_weights.pth", image_size=640, conf_threshold=0.25, iou_threshold=0.45):
+    """
+    Performs object detection using a YOLOv9 model. This function loads a specified YOLOv9 model,
+    configures it based on the provided parameters, and carries out inference on a given image. 
+    Additionally, it allows for optional modification of the input size and the application of 
+    test time augmentation to potentially improve detection accuracy.
+
+    Parameters:
+    - img_path (str): The file path to the image on which inference is to be performed.
+    - model_id (str): Identifier of the model to use.
+    - image_size (int): The input size for inference.
+    - conf_threshold (float): The confidence threshold used during Non-Maximum Suppression.
+    - iou_threshold (float): The Intersection over Union threshold applied in NMS.
+
+    Returns:
+    - Image: An image with detection bounding boxes drawn on it.
+    """
+    # Import YOLOv9 and torch only when the function is called to save on initial script load time
+    from yolov9 import YOLOv9
+    from PIL import Image
+    import numpy as np
+
+    # Download and load the model
     model_path = download_models(model_id)
-    model = yolov9.load(model_path, device="cuda:0")
-    
-    # Set model parameters
-    model.conf = conf_threshold
-    model.iou = iou_threshold
-    
-    # Perform inference
-    results = model(img_path, size=image_size)
+    model = YOLOv9(model_path, conf_threshold=conf_threshold, iou_threshold=iou_threshold, img_size=image_size)
+    model.eval()  # Set the model to evaluation mode
 
-    # Optionally, show detection bounding boxes on image
-    output = results.render()
-    
-    return output[0]
+    # Load image
+    img = Image.open(img_path).convert("RGB")
+    img = np.array(img)
 
+    # Perform inference
+    results = model.predict(img, size=image_size)
 
-def app():
-    with gr.Blocks():
-        with gr.Row():
-            with gr.Column():
-                img_path = gr.Image(type="filepath", label="Image")
-                model_path = gr.Dropdown(
-                    label="Model",
-                    choices=[
-                        "gelan-c.pt",
-                        "gelan-e.pt",
-                        "yolov9-c.pt",
-                        "yolov9-e.pt",
-                    ],
-                    value="gelan-e.pt",
-                )
-                image_size = gr.Slider(
-                    label="Image Size",
-                    minimum=320,
-                    maximum=1280,
-                    step=32,
-                    value=640,
-                )
-                conf_threshold = gr.Slider(
-                    label="Confidence Threshold",
-                    minimum=0.1,
-                    maximum=1.0,
-                    step=0.1,
-                    value=0.4,
-                )
-                iou_threshold = gr.Slider(
-                    label="IoU Threshold",
-                    minimum=0.1,
-                    maximum=1.0,
-                    step=0.1,
-                    value=0.5,
-                )
-                yolov9_infer = gr.Button(value="Inference")
-
-            with gr.Column():
-                output_numpy = gr.Image(type="numpy",label="Output")
-
-        yolov9_infer.click(
-            fn=yolov9_inference,
-            inputs=[
-                img_path,
-                model_path,
-                image_size,
-                conf_threshold,
-                iou_threshold,
-            ],
-            outputs=[output_numpy],
-        )
-        
-        gr.Examples(
-            examples=[
-                [
-                    "image_data/IMG_3352.JPG",
-                    "gelan-e.pt",
-                    640,
-                    0.4,
-                    0.5,
-                ],
-                [
-                    "image_data/IMG_3353.JPG",
-                    "yolov9-c.pt",
-                    640,
-                    0.4,
-                    0.5,
-                ],
-            ],
-            fn=yolov9_inference,
-            inputs=[
-                img_path,
-                model_path,
-                image_size,
-                conf_threshold,
-                iou_threshold,
-            ],
-            outputs=[output_numpy],
-            cache_examples=True,
-        )
-
-# Initialize a Gradio Blocks application.
-gradio_app = gr.Blocks()
+    # Extract results and visualize
+    output_image = model.visualize(results, img)
+    return output_image
 
-with gradio_app:
-    # Display a title using HTML, centered.
-    gr.HTML("""
-        <h1 style='text-align: center; margin-bottom: 20px;'>
-        YOLOv9 from PipYoloV9 on my data
-        </h1>
-    """)
+# Example Gradio interface setup (simplified for demonstration purposes)
+def gradio_interface(img_path):
+    return yolov9_inference(img_path)
 
-    # Assuming 'app()' represents your main app functionality.
-    # If 'app()' involves Gradio interface components, directly place those components here.
-    # Example: gr.Image() for image input, gr.Button() for a button, etc.
-    # Replace 'app()' with your actual Gradio components or function call.
-    app()
+iface = gr.Interface(fn=gradio_interface, inputs="image", outputs="image", title="YOLOv9 Object Detection")
+iface.launch()
 
-# Launch the Gradio app, enabling debug mode for detailed error logs and server information.
-gradio_app.launch(debug=True)