Update app.py

app.py

@@ -4,121 +4,96 @@ from PIL import Image
 import os 
 import yolov9
 
-
-# Load your custom-trained model
-model_name='./best.pt'
-#model = YOLO(model_name)
-# Images
-
 def yolov9_inference(img_path, 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
-    
-    # Load the model
-    #model_path = download_models(model_id)
     model = yolov9.load('./best.pt')
-    
-    # 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 app():
-    with gr.Blocks():
+    with gr.Blocks() as demo:
+        gr.HTML(
+            """
+            <div style="text-align: center; max-width: 800px; margin: 0 auto;">
+                <h1 style="color: #2c3e50; font-size: 2.5em; margin-bottom: 0.5em;">YOLOv9: Manhole Detector</h1>
+                <p style="color: #34495e; font-size: 1.2em; margin-bottom: 1em;">Detect manholes in images using YOLOv9</p>
+            </div>
+            """
+        )
         with gr.Row():
-            with gr.Column():
-                img_path = gr.Image(type="filepath", label="Image")
-                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")
-
+            with gr.Column(scale=1, min_width=300):
+                img_path = gr.Image(type="filepath", label="Upload Image")
+                with gr.Box():
+                    gr.Markdown("### Model Parameters")
+                    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("Detect Manholes", variant="primary")
+            
+            with gr.Column(scale=1, min_width=300):
+                output_numpy = gr.Image(type="numpy", label="Detection Result")
+        
         yolov9_infer.click(
             fn=yolov9_inference,
-            inputs=[
-                img_path,
-                image_size,
-                conf_threshold,
-                iou_threshold,
-            ],
-            outputs=[output_numpy])
+            inputs=[img_path, image_size, conf_threshold, iou_threshold],
+            outputs=[output_numpy]
+        )
+        
         gr.Examples(
             examples=[
-                [
-                    "./openmanhole.jpg",  #openmanhole.jpg
-
-                    640,
-                    0.4,
-                    0.5,
-                ],
-                [
-                    "./images.jpeg", #images.jpeg
-
-                    640,
-                    0.4,
-                    0.5,
-                ],
+                ["./openmanhole.jpg", 640, 0.4, 0.5],
+                ["./images.jpeg", 640, 0.4, 0.5],
             ],
             fn=yolov9_inference,
-            inputs=[
-                img_path,
-                image_size,
-                conf_threshold,
-                iou_threshold,
-            ],
+            inputs=[img_path, image_size, conf_threshold, iou_threshold],
             outputs=[output_numpy],
             cache_examples=True,
         )
-        
+    
+    return demo
+
+css = """
+body {
+    font-family: 'Arial', sans-serif;
+    background-color: #f0f3f6;
+}
+.gradio-container {
+    max-width: 900px !important;
+    margin-left: auto !important;
+    margin-right: auto !important;
+    background-color: white;
+    border-radius: 10px;
+    box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
+    padding: 20px;
+}
+.gr-button {
+    background-color: #3498db !important;
+    border: none !important;
+}
+.gr-button:hover {
+    background-color: #2980b9 !important;
+}
+.gr-box {
+    border-radius: 8px;
+    border: 1px solid #e0e0e0;
+    padding: 15px;
+    margin-top: 10px;
+    background-color: #f9f9f9;
+}
+.gr-form {
+    background-color: white;
+    padding: 15px;
+    border-radius: 8px;
+    box-shadow: 0 2px 4px rgba(0, 0, 0, 0.05);
+}
+"""
+
+demo = gr.Blocks(css=css)
 
+with demo:
+    app()
 
-gradio_app = gr.Blocks()
-with gradio_app:
-    gr.HTML(
-        """
-    <h1 style='text-align: center'>
-    YOLOv9: Manhole Detector
-    </h1>
-    """)
-    with gr.Row():
-        with gr.Column():
-            app()
-gradio_app.launch(debug=True, share=True)
+if __name__ == "__main__":
+    demo.launch(debug=True, share=True)