tweak

app.py

@@ -12,7 +12,6 @@ classes = ['zero', 'one', 'two', 'three', 'four', 'five', 'six', 'seven', 'eight
 API_KEY = os.environ.get("ROBOFLOW_API_KEY")
 ACCESS_TOKEN = os.environ.get("ACCESS_TOKEN")
 
-
 def RoboFlowGetOutlineBoxesPIL(pil_img):
     client = InferenceHTTPClient(
         api_url="https://detect.roboflow.com",
@@ -27,14 +26,20 @@ def RoboFlowGetOutlineBoxesPIL(pil_img):
         use_cache=True # cache workflow definition for 15 minutes
     )
     return result
-
-def localOutlineBox(image):
-    model_name = "./vt_dataset_yolov12_v7_weights.pt"
-    if os.path.exists(model_name):
+    
+def localOutlineBox(image, selected_model="yolo_accurate"):
+    model_paths = {
+        "yolo_accurate": "./vt_dataset_yolov12_v7_weights.pt",
+        "yolo_extra_large": "./yolo_extra_large_weights.pt"
+    }
+    
+    model_path = model_paths.get(selected_model, "./vt_dataset_yolov12_v7_weights.pt")
+    model_path = "./vt_dataset_yolov12_v7_weights.pt"
+    if os.path.exists(model_path):
         print("model exists")
     else:
         print("model is not available")
-    model = YOLO(model_name)
+    model = YOLO(model_path)
     yolo_ret = model(image, verbose=False)
     useful_ret = yolo_ret[0]
     names = model.names
@@ -294,13 +299,13 @@ def online_process_image(image):
     json_str = json.dumps(kept_box_info, indent=2)
     return boxed_img, textual, json_str
 
-def process_image(image):
+def process_image(image, selected_model = "yolo_accurate"):
     print("start processing image")
     if image is None:
         # Ensure we always return 3 values to prevent errors.
         return None, "", ""
     #pil_image = image.copy() if hasattr(image, 'copy') else Image.fromarray(image)
-    all_box_info = localOutlineBox(image)
+    all_box_info = localOutlineBox(image, selected_model)
     kept_box_info = filter_overlapping_boxes(all_box_info)
     del all_box_info
     boxed_img = drawWithAllBox_info(image, kept_box_info)
@@ -323,7 +328,15 @@ with gr.Blocks() as demo:
     # --- Main Application (initially hidden) ---
     with gr.Column(visible=False) as main_app:
         img_input = gr.Image(type="pil", label="Upload Image")
-        run_btn = gr.Button("Run Detection")
+
+        with gr.Row():
+            model_list = gr.Dropdown(
+                choices=["yolo_accurate", "yolo_extra_large"],
+                value="yolo_accurate",
+                label="Select Model",
+                info="Choose the YOLO model for detection"
+            )
+            run_btn = gr.Button("Run Detection")
         img_out = gr.Image(type="pil", label="Image with Boxes")
         text_out = gr.Textbox(label="Textual Description", lines=8)
         json_state = gr.State("")
@@ -344,9 +357,9 @@ with gr.Blocks() as demo:
             f.write(json_str)
             return f.name
 
-    def _process_and_prepare_download(image):
+    def _process_and_prepare_download(image, selected_model):
         """Processes the image, creates the JSON file, and updates the UI."""
-        boxed_img, textual, json_str = process_image(image)
+        boxed_img, textual, json_str = process_image(image, selected_model)
         filepath = create_json_file(json_str)
         
         # Use the legacy gr.update() for compatibility with older Gradio versions.
@@ -370,7 +383,7 @@ with gr.Blocks() as demo:
 
     run_btn.click(
         _process_and_prepare_download,
-        inputs=img_input,
+        inputs=[img_input, model_list],
         # The output now includes the download button itself.
         outputs=[img_out, text_out, json_state, download_btn]
     )