Add BLEU

app.py

@@ -2,8 +2,15 @@ import gradio as gr
 from PIL import Image
 import torch
 from transformers import VisionEncoderDecoderModel, AutoTokenizer, ViTFeatureExtractor, AutoImageProcessor, AutoModelForImageClassification
+from nltk.translate.bleu_score import sentence_bleu, SmoothingFunction
+import nltk
 import warnings
 
+try:
+    nltk.data.find("tokenizers/punkt")
+except LookupError:
+    nltk.download("punkt")
+
 warnings.filterwarnings("ignore", category=UserWarning)
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -38,6 +45,17 @@ def classify_image(image):
             results[name] = label
     return results
 
+def preprocess_caption(text):
+    text = str(text).lower()
+    for term in ["magnetic resonance imaging", "magnetic resonance image"]:
+        text = text.replace(term, "mri")
+    for term in ["computed tomography"]:
+        text = text.replace(term, "ct")
+    text = text.replace("t1-weighted", "t1").replace("t1w1", "t1").replace("t1ce", "t1")
+    text = text.replace("t2-weighted", "t2").replace("t2w", "t2").replace("t2/flair", "flair")
+    text = text.replace("tumour", "tumor").replace("lesions", "lesion").replace("-", " ")
+    return text.split()
+
 def generate_captions(image, keywords):
     pixel_values = feature_extractor(images=image, return_tensors="pt").pixel_values.to(device)
 
@@ -61,7 +79,7 @@ def generate_captions(image, keywords):
 
     return caption1, caption2
 
-def run_pipeline(image):
+def run_pipeline(image, actual_caption):
     classification = classify_image(image)
     keywords = list(classification.values())
     caption1, caption2 = generate_captions(image, keywords)
@@ -73,30 +91,44 @@ def run_pipeline(image):
         + (f"Tumor Type: {classification.get('tumor_type')}" if "tumor_type" in classification else "")
     )
 
-    return classification_text, caption1, caption2
+    if actual_caption.strip():
+        ref = [preprocess_caption(actual_caption)]
+        hyp1 = preprocess_caption(caption1)
+        hyp2 = preprocess_caption(caption2)
+        smooth = SmoothingFunction().method1
+        bleu1 = f"{sentence_bleu(ref, hyp1, smoothing_function=smooth):.2f}"
+        bleu2 = f"{sentence_bleu(ref, hyp2, smoothing_function=smooth):.2f}"
+    else:
+        bleu1 = "-"
+        bleu2 = "-"
+
+    return classification_text, caption1, caption2, bleu1, bleu2
 
 with gr.Blocks(theme=gr.themes.Soft(primary_hue="pink"), css=custom_css) as demo:
     gr.Markdown(
         """
         <link href="https://fonts.googleapis.com/css2?family=Poppins&display=swap" rel="stylesheet">
         <h1 style='text-align: center;'>🧠 Brain Hierarchical Classification + Captioning</h1>
-        <p style='text-align: center;'>Upload an MRI/CT brain image. The system will classify the image (plane, modality, abnormality, tumor type) and generate two captions.</p>
+        <p style='text-align: center;'>Upload an MRI/CT brain image. The system will classify the image and generate captions. Optionally, provide ground truth to see BLEU scores.</p>
         """,
         elem_id="title"
     )
     with gr.Row():
         with gr.Column():
             image_input = gr.Image(type="pil", label="🖼️ Upload Brain MRI/CT")
+            actual_caption = gr.Textbox(label="💬 Ground Truth Caption (optional)")
             btn = gr.Button("🚀 Submit")
         with gr.Column():
             cls_box = gr.Textbox(label="📋 Classification Result", lines=4)
             cap1_box = gr.Textbox(label="📝 Caption without Keyword Integration", lines=4)
             cap2_box = gr.Textbox(label="🧠 Caption with Keyword Integration", lines=4)
+            bleu1_box = gr.Textbox(label="📊 BLEU (No Keyword)", lines=1)
+            bleu2_box = gr.Textbox(label="📈 BLEU (With Keyword)", lines=1)
 
     btn.click(
         fn=run_pipeline,
-        inputs=[image_input],
-        outputs=[cls_box, cap1_box, cap2_box]
+        inputs=[image_input, actual_caption],
+        outputs=[cls_box, cap1_box, cap2_box, bleu1_box, bleu2_box]
     )
 
 demo.launch()

setup.sh

@@ -0,0 +1,2 @@
+#!/bin/bash
+python -m nltk.downloader punkt