Update app.py

app.py

@@ -2,125 +2,121 @@ import gradio as gr
 import spaces
 from huggingface_hub import hf_hub_download
 
+class YOLOv9App:
+    def __init__(self):
+        self.gradio_app = gr.Blocks()
 
-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):
-    """
-    Load a YOLOv9 model, configure it, perform inference on an image, and optionally adjust 
-    the input size and apply test time augmentation.
-    
-    :param model_path: Path to the YOLOv9 model file.
-    :param conf_threshold: Confidence threshold for NMS.
-    :param iou_threshold: IoU threshold for NMS.
-    :param img_path: Path to the image file.
-    :param size: Optional, input size for inference.
-    :return: A tuple containing the detections (boxes, scores, categories) and the results object for further actions like displaying.
-    """
-    # Import YOLOv9
-    import yolov9
-    
-    # Load the model
-    model_path = download_models(model_id)
-    model = yolov9.load(model_path, device="cpu")
-    
-    # Set model parameters
-    model.conf = conf_threshold
-    model.iou = iou_threshold
-    
-    # Perform inference
-    results = model(img_path, size=image_size)
-
-    # Optionally, show detection bounding boxes on image
-    output = results.render()
-    
-    return output[0]
+    def download_models(self, model_id):
+        hf_hub_download("merve/yolov9", filename=f"{model_id}", local_dir="./")
+        return f"./{model_id}"
 
+    @spaces.GPU
+    def yolov9_inference(self, img_path, model_id, image_size, conf_threshold, iou_threshold):
+        """
+        Load a YOLOv9 model, configure it, perform inference on an image, and optionally adjust 
+        the input size and apply test time augmentation.
+        
+        :param model_path: Path to the YOLOv9 model file.
+        :param conf_threshold: Confidence threshold for NMS.
+        :param iou_threshold: IoU threshold for NMS.
+        :param img_path: Path to the image file.
+        :param size: Optional, input size for inference.
+        :return: A tuple containing the detections (boxes, scores, categories) and the results object for further actions like displaying.
+        """
+        # Import YOLOv9
+        import yolov9
+        
+        # Load the model
+        model_path = self.download_models(model_id)
+        model = yolov9.load(model_path, device="cpu")
+        
+        # Set model parameters
+        model.conf = conf_threshold
+        model.iou = iou_threshold
+        
+        # Perform inference
+        results = model(img_path, 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="Submit")
+        # Optionally, show detection bounding boxes on image
+        output = results.render()
+        
+        return output[0]
 
-            with gr.Column():
-                output_numpy = gr.Image(type="numpy",label="Output")
+    def app(self):
+        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="Submit")
 
-        yolov9_infer.click(
-            fn=yolov9_inference,
-            inputs=[
-                img_path,
-                model_path,
-                image_size,
-                conf_threshold,
-                iou_threshold,
-            ],
-            outputs=[output_numpy],
-        )   
+                with gr.Column():
+                    output_numpy = gr.Image(type="numpy", label="Output")
 
-# Definindo a aplicação Gradio
-gradio_app = gr.Blocks()
+            yolov9_infer.click(
+                fn=self.yolov9_inference,
+                inputs=[
+                    img_path,
+                    model_path,
+                    image_size,
+                    conf_threshold,
+                    iou_threshold,
+                ],
+                outputs=[output_numpy],
+            )
 
-with gradio_app:
-    gr.HTML(
-        """
-        <h1 style='text-align: center'>
-        YOLOv9 Base Model
-        </h1>
-        """
-    )
-    gr.HTML(
-        """
-        <h3 style='text-align: center'>
-        Aplicação para ajudar nos resgates do RS
-        </h3>
-        """
-    )
-    with gr.Row():
-        with gr.Column():
-            # Carregar uma imagem de entrada e exibir a imagem processada
-            image_input = gr.Image(type="pil", label="Carregue uma imagem")
-            image_output = gr.Image(label="Resultado da Detecção")
-            
-            # Botão para acionar o processamento da imagem
-            process_button = gr.Button("Processar Imagem")
-            
-            # Conectar o botão à função de processamento
-            process_button.click(fn=process_image, inputs=image_input, outputs=image_output)
+    def launch(self):
+        with self.gradio_app:
+            gr.HTML(
+                """
+                <h1 style='text-align: center'>
+                YOLOv9 Base Model
+                </h1>
+                """
+            )
+            gr.HTML(
+                """
+                <h3 style='text-align: center'>
+                Aplicação para ajudar nos resgates do RS
+                </h3>
+                """
+            )
+            with gr.Row():
+                with gr.Column():
+                    self.app()
+        self.gradio_app.launch(debug=True)
 
-gradio_app.launch(debug=True)
+# Criação e execução da aplicação
+if __name__ == "__main__":
+    yolov9_app = YOLOv9App()
+    yolov9_app.launch()