Update app.py

app.py

@@ -362,103 +362,99 @@ def run_inference(
     if image is None:
         return None, "Please upload an image first."
 
-    try:
-        # Convert numpy array to PIL Image if needed
-        if isinstance(image, np.ndarray):
-            image = Image.fromarray(image)
-
-        image_width, image_height = image.size
-
-        # Parse visual prompts if needed
-        visual_prompt_boxes = []
-        if task_selection == "Visual Prompting":
-            # Check if we have predefined visual prompt boxes from examples
-            if hasattr(image, "_example_visual_prompts"):
-                visual_prompt_boxes = image._example_visual_prompts
-            elif visual_prompt_data is not None and "points" in visual_prompt_data:
-                visual_prompt_boxes = parse_visual_prompt(visual_prompt_data["points"])
-
-        # Determine task type and categories based on task selection
-        if task_selection == "OCR":
-            # For OCR, use the selected output format to determine task type
-            task_type = OCR_OUTPUT_FORMATS[ocr_output_format]["task_type"]
-            task_key = task_type.value
-            # Use granularity level to determine categories
-            categories_list = [OCR_GRANULARITY_LEVELS[ocr_granularity]["categories"]]
-        elif task_selection == "Visual Prompting":
-            # For visual prompting, we don't need explicit categories
-            task_key = "visual_prompting"
-            categories_list = ["object"]
 
-            # Check if visual prompt boxes are provided
-            if not visual_prompt_boxes:
-                return (
-                    None,
-                    "Please draw bounding boxes on the image to provide visual examples for Visual Prompting task.",
-                )
-        elif task_selection == "Keypoint":
-            task_key = "keypoint"
-            categories_list = [keypoint_type] if keypoint_type else ["person"]
-        else:
-            # For other tasks, get task type from demo config
-            demo_config = DEMO_TASK_CONFIGS[task_selection]
-            task_type = demo_config["task_type"]
-            task_key = task_type.value
-
-            # Split categories by comma and clean up
-            categories_list = [
-                cat.strip() for cat in categories.split(",") if cat.strip()
-            ]
-            if not categories_list:
-                categories_list = ["object"]
-
-        # Run inference
-        if task_selection == "Visual Prompting":
-            results = rex_model.inference(
-                images=image,
-                task=task_key,
-                categories=categories_list,
-                visual_prompt_boxes=visual_prompt_boxes,
-            )
-        elif task_selection == "Keypoint":
-            results = rex_model.inference(
-                images=image,
-                task=task_key,
-                categories=categories_list,
-                keypoint_type=keypoint_type if keypoint_type else "person",
-            )
-        else:
-            results = rex_model.inference(
-                images=image, task=task_key, categories=categories_list
-            )
+    # Convert numpy array to PIL Image if needed
+    if isinstance(image, np.ndarray):
+        image = Image.fromarray(image)
 
-        result = results[0]
-
-        # Check if inference was successful
-        if not result.get("success", False):
-            error_msg = result.get("error", "Unknown error occurred during inference")
-            return None, f"Inference failed: {error_msg}"
-
-        # Get predictions and raw output
-        predictions = result["extracted_predictions"]
-        raw_output = result["raw_output"]
-
-        # Create visualization
-        try:
-            vis_image = RexOmniVisualize(
-                image=image,
-                predictions=predictions,
-                font_size=font_size,
-                draw_width=draw_width,
-                show_labels=show_labels,
+    image_width, image_height = image.size
+
+    # Parse visual prompts if needed
+    visual_prompt_boxes = []
+    if task_selection == "Visual Prompting":
+        # Check if we have predefined visual prompt boxes from examples
+        if hasattr(image, "_example_visual_prompts"):
+            visual_prompt_boxes = image._example_visual_prompts
+        elif visual_prompt_data is not None and "points" in visual_prompt_data:
+            visual_prompt_boxes = parse_visual_prompt(visual_prompt_data["points"])
+
+    # Determine task type and categories based on task selection
+    if task_selection == "OCR":
+        # For OCR, use the selected output format to determine task type
+        task_type = OCR_OUTPUT_FORMATS[ocr_output_format]["task_type"]
+        task_key = task_type.value
+        # Use granularity level to determine categories
+        categories_list = [OCR_GRANULARITY_LEVELS[ocr_granularity]["categories"]]
+    elif task_selection == "Visual Prompting":
+        # For visual prompting, we don't need explicit categories
+        task_key = "visual_prompting"
+        categories_list = ["object"]
+
+        # Check if visual prompt boxes are provided
+        if not visual_prompt_boxes:
+            return (
+                None,
+                "Please draw bounding boxes on the image to provide visual examples for Visual Prompting task.",
             )
-            return vis_image, raw_output
-        except Exception as e:
-            return image, f"Visualization failed: {str(e)}\n\nRaw output:\n{raw_output}"
+    elif task_selection == "Keypoint":
+        task_key = "keypoint"
+        categories_list = [keypoint_type] if keypoint_type else ["person"]
+    else:
+        # For other tasks, get task type from demo config
+        demo_config = DEMO_TASK_CONFIGS[task_selection]
+        task_type = demo_config["task_type"]
+        task_key = task_type.value
 
-    except Exception as e:
-        return None, f"Error during inference: {str(e)}"
+        # Split categories by comma and clean up
+        categories_list = [
+            cat.strip() for cat in categories.split(",") if cat.strip()
+        ]
+        if not categories_list:
+            categories_list = ["object"]
+
+    # Run inference
+    if task_selection == "Visual Prompting":
+        results = rex_model.inference(
+            images=image,
+            task=task_key,
+            categories=categories_list,
+            visual_prompt_boxes=visual_prompt_boxes,
+        )
+    elif task_selection == "Keypoint":
+        results = rex_model.inference(
+            images=image,
+            task=task_key,
+            categories=categories_list,
+            keypoint_type=keypoint_type if keypoint_type else "person",
+        )
+    else:
+        results = rex_model.inference(
+            images=image, task=task_key, categories=categories_list
+        )
+
+    result = results[0]
+
+    # Check if inference was successful
+    if not result.get("success", False):
+        error_msg = result.get("error", "Unknown error occurred during inference")
+        return None, f"Inference failed: {error_msg}"
+
+    # Get predictions and raw output
+    predictions = result["extracted_predictions"]
+    raw_output = result["raw_output"]
 
+    # Create visualization
+    try:
+        vis_image = RexOmniVisualize(
+            image=image,
+            predictions=predictions,
+            font_size=font_size,
+            draw_width=draw_width,
+            show_labels=show_labels,
+        )
+        return vis_image, raw_output
+    except Exception as e:
+        return image, f"Visualization failed: {str(e)}\n\nRaw output:\n{raw_output}"
 
 def update_interface(task_selection):
     """Update interface based on task selection"""