Update app.py

app.py

@@ -75,11 +75,15 @@ def load_sample_phrase():
     return [save_temp_image(sample["frontal"]), sample["phrase"]]
 
 def generate_report(frontal_path, lateral_path, indication, technique, comparison,
-                   prior_frontal_path, prior_lateral_path, prior_report, grounding):
+                    prior_frontal_path, prior_lateral_path, prior_report, grounding):
     """Generate radiology report with authentication check"""
     if not MODEL_STATE["authenticated"]:
         return "⚠️ Please authenticate with your Hugging Face token first!"
     
+    # Check that mandatory image paths are provided.
+    if not frontal_path or not lateral_path:
+        return "❌ Please upload both the frontal and lateral images for the current study."
+    
     try:
         current_frontal = Image.open(frontal_path)
         current_lateral = Image.open(lateral_path)
@@ -117,6 +121,10 @@ def ground_phrase(frontal_path, phrase):
     if not MODEL_STATE["authenticated"]:
         return "⚠️ Please authenticate with your Hugging Face token first!"
     
+
+    if not frontal_path:
+        return "❌ Please upload the frontal image for phrase grounding."
+    
     try:
         frontal = Image.open(frontal_path)
         processed = MODEL_STATE["processor"].format_and_preprocess_phrase_grounding_input(
@@ -125,6 +133,9 @@ def ground_phrase(frontal_path, phrase):
             return_tensors="pt"
         ).to("cpu")
         
+
+        processed.pop("image_sizes", None)
+        
         outputs = MODEL_STATE["model"].generate(
             **processed,
             max_new_tokens=150,
@@ -187,12 +198,12 @@ with gr.Blocks(title="MAIRA-2 Medical Assistant") as demo:
             sample_btn.click(
                 load_sample_findings,
                 outputs=[frontal, lateral, indication, technique, comparison,
-                       prior_frontal, prior_lateral, prior_report, grounding]
+                         prior_frontal, prior_lateral, prior_report, grounding]
             )
             generate_btn.click(
                 generate_report,
                 inputs=[frontal, lateral, indication, technique, comparison,
-                      prior_frontal, prior_lateral, prior_report, grounding],
+                        prior_frontal, prior_lateral, prior_report, grounding],
                 outputs=report_output
             )