update output

app.py

@@ -15,20 +15,21 @@ import spaces
 hf_token = os.getenv("HF_TOKEN")
 login(token=hf_token, add_to_git_credential=True)
 
-def detect_brain_tumor(image, debug: bool = False) -> str:
+import numpy as np
+from PIL import Image
+
+def detect_brain_tumor(image, seg_input, debug: bool = False):
     """
-    Detects a brain tumor in the given image and saves the image with bounding boxes.
+    Detects a brain tumor in the given image and returns the annotated image.
 
     Parameters:
-        image: The input image (as provided by Gradio).
+        image: The input image (as numpy array provided by Gradio).
+        seg_input: The segmentation input (not used in this function, but required for Gradio).
         debug (bool): Flag to enable logging for debugging purposes.
 
     Returns:
-        str: Path to the saved output image.
+        tuple: (numpy array of image, list of (label, score, bbox) tuples)
     """
-    # Generate a unique output filename
-    output_path = f"./output/tumor_detection_{int(time.time())}.jpg"
-    
     if debug:
         print(f"Image received")
 
@@ -43,12 +44,20 @@ def detect_brain_tumor(image, debug: bool = False) -> str:
     if debug:
         print("Bounding boxes overlaid on the image")
 
-    # Step 4: Save the resulting image
-    save_image(image_with_bboxes, output_path)
-    if debug:
-        print(f"Image saved to {output_path}")
-    
-    return output_path
+    # Prepare annotations for AnnotatedImage output
+    annotations = []
+    for detection in detections:
+        label = detection['label']
+        score = detection['score']
+        bbox = detection['bbox']
+        x1, y1, x2, y2 = bbox
+        annotations.append((f"{label} {score:.2f}", (x1, y1, x2-x1, y2-y1)))
+
+    # Convert image to numpy array if it's not already
+    if isinstance(image_with_bboxes, Image.Image):
+        image_with_bboxes = np.array(image_with_bboxes)
+
+    return (image_with_bboxes, annotations)
 
 INTRO_TEXT="# 🔬🧠 CellVision AI -- Intelligent Cell Imaging Analysis 🤖🧫"
 IMAGE_PROMPT="Are these cells healthy or cancerous?"
@@ -59,7 +68,7 @@ with gr.Blocks(css="style.css") as demo:
         with gr.Row():
             with gr.Column():
                 image = gr.Image(type="numpy")
-                seg_input = gr.Text(label="Entities to Segment/Detect")
+                seg_input = gr.Text(label="Entities to Segment/Detect", value="detect brain tumor")
         
             with gr.Column():
                 annotated_image = gr.AnnotatedImage(label="Output")