Update app.py

app.py

@@ -12,6 +12,7 @@ import logging
 import matplotlib.pyplot as plt
 import gradio as gr
 from PIL import ImageDraw
+from enum import Enum
 
 # Load environment variables
 # load_dotenv()
@@ -20,6 +21,11 @@ from PIL import ImageDraw
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
+class EnumModel(str, Enum):
+    lion90 = "lion90"
+    dulcet72 = "dulcet72"
+    default = "default"
+
 class YOLODetector:
     def __init__(self, api_url: str, stride: int = 32):
         self.api_url = api_url
@@ -47,8 +53,8 @@ class YOLODetector:
             logger.warning(f'WARNING: Image size {imgsz} must be multiple of max stride {s}, updating to {new_size}')
         return new_size
 
-    def detect(self, image_path: str, conf_thres: float = 0.25, iou_thres: float = 0.45) -> Dict:
-        params = {'conf_thres': conf_thres, 'iou_thres': iou_thres}
+    def detect(self, image_path: str, conf_thres: float = 0.25, iou_thres: float = 0.45, model: EnumModel = EnumModel.dulcet72) -> Dict:
+        params = {'conf_thres': conf_thres, 'iou_thres': iou_thres, 'model': model.value}
         
         try:
             with open(image_path, 'rb') as image_file:
@@ -94,7 +100,7 @@ def plot_detections(image_path: str, detections: List[Dict]):
     plt.tight_layout()
     plt.show()
 
-def process_image(image, conf_thres, iou_thres):
+def process_image(image, conf_thres, iou_thres, model):
     api_url = os.getenv('YOLO_API_URL', 'https://akc8p0uym6ktml-8000.proxy.runpod.net/detect')
     detector = YOLODetector(api_url)
     
@@ -103,7 +109,7 @@ def process_image(image, conf_thres, iou_thres):
         temp_file_path = temp_file.name
     
     processed_image_path = detector.preprocess_image(temp_file_path)
-    results = detector.detect(processed_image_path, conf_thres=conf_thres, iou_thres=iou_thres)
+    results = detector.detect(processed_image_path, conf_thres=conf_thres, iou_thres=iou_thres, model=model)
     
     img = Image.open(processed_image_path)
     img_draw = img.copy()
@@ -123,24 +129,27 @@ def process_image(image, conf_thres, iou_thres):
 def main():
     with gr.Blocks() as demo:
         gr.Markdown("# YOLO Object Detection")
-        gr.Markdown("Upload an image and adjust the confidence and IOU thresholds to detect objects.")
+        gr.Markdown("Upload an image, select a model, and adjust the confidence and IOU thresholds to detect objects.")
         
         with gr.Row():
             with gr.Column():
                 image_input = gr.Image(type="pil", label="Input Image")
+                model_dropdown = gr.Dropdown(choices=[model.value for model in EnumModel], value=EnumModel.dulcet72.value, label="Model")
                 conf_slider = gr.Slider(minimum=0.01, maximum=1.0, value=0.25, label="Confidence Threshold")
                 iou_slider = gr.Slider(minimum=0.01, maximum=1.0, value=0.45, label="IOU Threshold")
             with gr.Column():
                 image_output = gr.Image(type="pil", label="Detection Result")
-        def update_output(image, conf_thres, iou_thres):
+        
+        def update_output(image, model, conf_thres, iou_thres):
             if image is None:
                 return None
-            return process_image(image, conf_thres, iou_thres)
+            return process_image(image, conf_thres, iou_thres, EnumModel(model))
         
-        inputs = [image_input, conf_slider, iou_slider]
+        inputs = [image_input, model_dropdown, conf_slider, iou_slider]
         outputs = image_output
 
         image_input.change(fn=update_output, inputs=inputs, outputs=outputs)
+        model_dropdown.change(fn=update_output, inputs=inputs, outputs=outputs)
         conf_slider.change(fn=update_output, inputs=inputs, outputs=outputs)
         iou_slider.change(fn=update_output, inputs=inputs, outputs=outputs)