update app.py

app.py

@@ -4,8 +4,10 @@ import gradio as gr
 from typing import *
 from pillow_heif import register_heif_opener
 register_heif_opener()
+from PIL import Image
+import numpy as np
 import vision_agent as va
-from vision_agent.tools import register_tool
+from vision_agent.tools import owl_v2, overlay_bounding_boxes, save_image
 
 from vision_agent.tools import load_image, owl_v2, overlay_bounding_boxes, save_image
 
@@ -21,7 +23,7 @@ def detect_brain_tumor(image, debug: bool = False) -> str:
     Detects a brain tumor in the given image and saves the image with bounding boxes.
 
     Parameters:
-        image: The input image (as a PIL Image or numpy array).
+        image: The input image (can be PIL Image, numpy array, or file path).
         debug (bool): Flag to enable logging for debugging purposes.
 
     Returns:
@@ -30,14 +32,28 @@ def detect_brain_tumor(image, debug: bool = False) -> str:
     # Generate a unique output filename
     output_path = f"./output/tumor_detection_{int(time.time())}.jpg"
     
-    # Step 1: Load the image (not needed if image is already a PIL Image or numpy array)
-    # image = load_image(image_path)
+    # Ensure image is in the correct format
+    if isinstance(image, str):
+        # If image is a file path
+        image = Image.open(image)
+    elif isinstance(image, np.ndarray):
+        # If image is already a numpy array
+        image = Image.fromarray(image)
+    elif not isinstance(image, Image.Image):
+        raise ValueError("Unsupported image type. Please provide a PIL Image, numpy array, or file path.")
+
+    # Convert to RGB if it's not
+    image = image.convert('RGB')
+    
+    # Convert PIL Image to numpy array for owl_v2
+    image_array = np.array(image)
+
     if debug:
-        print(f"Image loaded")
+        print(f"Image loaded and converted to numpy array of shape {image_array.shape}")
 
     # Step 2: Detect brain tumor using owl_v2
     prompt = "detect brain tumor"
-    detections = owl_v2(prompt, image)
+    detections = owl_v2(prompt, image_array)
     if debug:
         print(f"Detections: {detections}")
 
@@ -53,7 +69,6 @@ def detect_brain_tumor(image, debug: bool = False) -> str:
     
     return output_path
 
-
 # Example usage (uncomment to run):
 # detect_brain_tumor("/content/drive/MyDrive/kaggle/datasets/brain-tumor-image-dataset-semantic-segmentation_old/train_categories/1385_jpg.rf.3c67cb92e2922dba0e6dba86f69df40b.jpg", "/content/drive/MyDrive/kaggle/datasets/brain-tumor-image-dataset-semantic-segmentation_old/output/1385_jpg.rf.3c67cb92e2922dba0e6dba86f69df40b.jpg", debug=True)