front end updated

app.py

@@ -180,29 +180,49 @@ model = initialize_model()
 @torch.inference_mode()
 @torch.autocast(device_type="cuda", dtype=torch.bfloat16)
 def process_image(image_path, text_prompts, modality):
-    image = read_rgb(image_path)
-    text_prompts = [prompt.strip() for prompt in text_prompts.split(',')]
-    
-    # Run inference
-    pred_masks = interactive_infer_image(model, Image.fromarray(image), text_prompts)
-    
-    # Prepare outputs
-    results = []
-    dice_scores = []
-    p_values = []
+    try:
+        # Input validation
+        if not image_path:
+            raise ValueError("Please upload an image")
+        if not text_prompts or text_prompts.strip() == "":
+            raise ValueError("Please enter prompts for analysis")
+        if not modality:
+            raise ValueError("Please select a modality")
+
+        image = read_rgb(image_path)
+        text_prompts = [prompt.strip() for prompt in text_prompts.split(',')]
+        
+        # Run inference
+        pred_masks = interactive_infer_image(model, Image.fromarray(image), text_prompts)
+        
+        # Prepare outputs
+        results = []
+        p_values = []
+
+        for i, prompt in enumerate(text_prompts):
+            # Calculate p-value for the selected modality
+            print("PROMPT: ", prompt, flush=True)   
+            p_value = check_mask_stats(image, pred_masks[i] * 255, modality, prompt)
+            p_values.append(f"P-value for '{prompt}' ({modality}): {p_value:.4f}")
 
-    for i, prompt in enumerate(text_prompts):
-        # Calculate p-value for the selected modality
-        print("PROMPT: ", prompt, flush=True)   
-        p_value = check_mask_stats(image, pred_masks[i] * 255, modality, prompt)
-        p_values.append(f"P-value for '{prompt}' ({modality}): {p_value:.4f}")
+            # Overlay predictions on the image
+            overlay_image = image.copy()
+            overlay_image[pred_masks[i] > 0.5] = [255, 0, 0]  # Highlight predictions in red
+            results.append(overlay_image)
 
-        # Overlay predictions on the image
-        overlay_image = image.copy()
-        overlay_image[pred_masks[i] > 0.5] = [255, 0, 0]  # Highlight predictions in red
-        results.append(overlay_image)
+            return results, "\n".join(p_values)
 
-    return results, p_values
+    except ValueError as ve:
+        # Handle validation errors
+        return None, f"⚠️ Input Error: {str(ve)}"
+    except torch.cuda.OutOfMemoryError:
+        # Handle CUDA out of memory errors
+        return None, "⚠️ Error: GPU memory exceeded. Please try with a smaller image."
+    except Exception as e:
+        # Handle all other errors
+        error_msg = f"��️ An error occurred: {str(e)}"
+        print(f"Error details: {str(e)}", flush=True)  # For logging
+        return None, error_msg
 
 # Define Gradio interface
 with gr.Blocks() as demo:
@@ -210,7 +230,11 @@ with gr.Blocks() as demo:
     with gr.Row():
         with gr.Column():
             image_input = gr.Image(type="filepath", label="Input Image")
-            prompts_input = gr.Textbox(lines=2, placeholder="Enter prompts separated by commas...", label="Prompts")
+            prompts_input = gr.Textbox(
+                lines=2, 
+                placeholder="Enter prompts separated by commas...", 
+                label="Prompts"
+            )
             modality_dropdown = gr.Dropdown(
                 choices=list(BIOMEDPARSE_MODES.keys()),
                 value=list(BIOMEDPARSE_MODES.keys())[0],
@@ -219,28 +243,18 @@ with gr.Blocks() as demo:
             submit_btn = gr.Button("Submit")
         with gr.Column():
             output_gallery = gr.Gallery(label="Findings")
-            pvalue_output = gr.Textbox(label="Confidence (P-values)", interactive=False)
+            pvalue_output = gr.Textbox(
+                label="Results", 
+                interactive=False,
+                show_label=True
+            )
     
+    # Add error handling for the submit button
     submit_btn.click(
-        process_image,
+        fn=process_image,
         inputs=[image_input, prompts_input, modality_dropdown],
-        outputs=[output_gallery, pvalue_output]
+        outputs=[output_gallery, pvalue_output],
+        api_name="process"
     )
-    # with gr.Row():
-    #         gr.Examples(
-    #             fn=process_image,
-    #             examples=IMAGE_PROCESSING_EXAMPLES,
-    #             inputs=[
-    #                 image_processing_mode_dropdown_component,
-    #                 image_processing_image_input_component,
-    #                 image_processing_text_input_component
-    #             ],
-    #             outputs=[
-    #                 image_processing_image_output_component,
-    #                 image_processing_text_output_component
-    #             ],
-    #             run_on_click=True
-    #         )
-
-# Launch the app
+
 demo.launch()