Update app.py

app.py

@@ -13,17 +13,46 @@ model.eval()
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model.to(device)
 
+# def generate_caption(image):
+#     # Preprocess
+#     inputs = processor(image, return_tensors="pt").to(device)
+
+#     # Generate caption
+#     with torch.no_grad():
+#         output = model.generate(**inputs, max_new_tokens=50)
+
+#     # Decode
+#     caption = processor.decode(output[0], skip_special_tokens=True)
+#     return caption
+
 def generate_caption(image):
-    # Preprocess
-    inputs = processor(image, return_tensors="pt").to(device)
+    prompt = "Radiology report:"
+    inputs = processor(images=image, text=prompt, return_tensors="pt").to(device)
+
+    output = model.generate(
+        **inputs,
+        max_length=256,
+        num_beams=5,
+        repetition_penalty=1.2,
+        length_penalty=1.0,
+        early_stopping=True
+    )
+
+    caption = processor.batch_decode(output, skip_special_tokens=True)[0]
+
+    # Optional: format as a full report
+    report = f"""
+    **Radiology Report**
+
+    **Findings**: {caption}
+
+    **Impression**: {caption}
+
+    **Recommendation**: Clinical correlation advised. Consider MRI if necessary.
+    """
+    return report.strip()
 
-    # Generate caption
-    with torch.no_grad():
-        output = model.generate(**inputs, max_new_tokens=50)
 
-    # Decode
-    caption = processor.decode(output[0], skip_special_tokens=True)
-    return caption
 
 # Gradio UI
 demo = gr.Interface(