Update app.py

app.py

@@ -1,64 +1,72 @@
 import gradio as gr
 from huggingface_hub import hf_hub_download
-import torch
 
 def download_models(model_id):
-    """
-    Downloads a model file from Hugging Face Hub to a specified local directory.
-
-    Parameters:
-    - model_id (str): Identifier of the model to download.
+    model_file_path = hf_hub_download("merve/yolov9", filename=model_id)
+    return model_file_path
 
-    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):
+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:
-    - 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
+    # Import YOLOv9
+    import yolov9
+    
+    # Load the model
     model_path = download_models(model_id)
-    model = YOLOv9(model_path, conf_threshold=conf_threshold, iou_threshold=iou_threshold, img_size=image_size)
-    model.eval()  # Set the model to evaluation mode
+    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)
 
-    # Load image
-    img = Image.open(img_path).convert("RGB")
-    img = np.array(img)
+    # Optionally, show detection bounding boxes on image
+    output = results.render()
+    
+    return output[0]
 
-    # Perform inference
-    results = model.predict(img, size=image_size)
+def app():
+    with gr.Blocks() as blocks:
+        with gr.Row():
+            with gr.Column():
+                img_path = gr.Image(type="filepath", label="Image")
+                model_id = 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("Inference")
 
-    # Extract results and visualize
-    output_image = model.visualize(results, img)
-    return output_image
+            with gr.Column():
+                output_image = gr.Image(type="numpy", label="Output")
 
-# Example Gradio interface setup (simplified for demonstration purposes)
-def gradio_interface(img_path):
-    return yolov9_inference(img_path)
+        yolov9_infer.click(
+            fn=yolov9_inference,
+            inputs=[img_path, model_id, image_size, conf_threshold, iou_threshold],
+            outputs=[output_image]
+        )
 
-iface = gr.Interface(fn=gradio_interface, inputs="image", outputs="image", title="YOLOv9 Object Detection")
-iface.launch()
+    return blocks
 
+gradio_app = app()
 
+# Display a title using HTML, centered.
+gradio_app[''].add(
+    gr.HTML("""
+        <h1 style='text-align: center; margin-bottom: 20px;'>
+        YOLOv9 from PipYoloV9 on my data
+        </h1>
+    """)
+)
 
+# Launch the Gradio app, enabling debug mode for detailed error logs and server information.
+gradio_app.launch(debug=True)