Add actual caption field and calculate bleu score

app.py

@@ -2,7 +2,12 @@ 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
+from nltk.translate.bleu_score import sentence_bleu
+import warnings
+import nltk
+nltk.download('punkt')
+
+warnings.filterwarnings("ignore", category=UserWarning)
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
@@ -36,6 +41,24 @@ def classify_image(image):
             results[name] = label
     return results
 
+# Preprocessing caption
+def preprocess_caption(text):
+    text = str(text).lower()
+    text = text.replace("magnetic resonance imaging", "mri")
+    text = text.replace("magnetic resonance image", "mri")
+    text = text.replace("computed tomography", "ct")
+    text = text.replace("t1-weighted", "t1")
+    text = text.replace("t1w1", "t1")
+    text = text.replace("t1w", "t1")
+    text = text.replace("t1ce", "t1")
+    text = text.replace("t2-weighted", "t2")
+    text = text.replace("t2w", "t2")
+    text = text.replace("t2/flair", "flair")
+    text = text.replace("tumour", "tumor")
+    text = text.replace("lesions", "lesion")
+    text = text.replace("-", " ")
+    return text.split()
+
 # Caption generation
 def generate_captions(image, keywords):
     pixel_values = feature_extractor(images=image, return_tensors="pt").pixel_values.to(device)
@@ -68,7 +91,7 @@ def run_pipeline(image, actual_caption):
     keywords = list(classification.values())
     caption1, caption2 = generate_captions(image, keywords)
 
-    # Format classification result as string
+    # Format classification result
     classification_str = (
         f"🧭 Plane: {classification.get('plane')}\n"
         f"🖼️ Modality: {classification.get('modality')}\n"
@@ -79,15 +102,20 @@ def run_pipeline(image, actual_caption):
 
     # BLEU Score calculation
     if actual_caption.strip():
-        ref = [actual_caption.lower().split()]
-        hyp = caption2.lower().split()
-        score = sentence_bleu(ref, hyp, smoothing_function=SmoothingFunction().method1)
+        ref = [preprocess_caption(actual_caption)]
+        hyp = preprocess_caption(caption2)
+        score = sentence_bleu(ref, hyp)
         bleu = f"📊 BLEU Score: {score:.2f}"
     else:
         bleu = "📊 BLEU Score: -"
 
-    # Output
-    result_text = f"{classification_str}\n\n✏️ Caption without Keywords:\n{caption1}\n\n✨ Caption with Keywords:\n{caption2}\n\n{bleu}"
+    # Final output
+    result_text = (
+        f"{classification_str}\n"
+        f"✏️ Caption without Keywords:\n{caption1}\n\n"
+        f"✨ Caption with Keywords:\n{caption2}\n\n"
+        f"{bleu}"
+    )
 
     return result_text
 
@@ -96,7 +124,7 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="pink")) as demo:
     gr.Markdown(
         """
         <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) and generate two captions, along with a BLEU score if ground truth is given.</p>
+        <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. Optionally, provide a ground truth caption to get BLEU score.</p>
         """
     )
     with gr.Row():
@@ -105,7 +133,7 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="pink")) as demo:
             actual_caption = gr.Textbox(label="🧠 Ground Truth Caption (optional)")
             btn = gr.Button("🚀 Submit")
         with gr.Column():
-            output_box = gr.Textbox(label="📝 Result", lines=20)
+            output_box = gr.Textbox(label="📋 Result", lines=20)
 
     btn.click(fn=run_pipeline, inputs=[image_input, actual_caption], outputs=output_box)
 

requirements.txt

@@ -1,3 +1,5 @@
 torch
 transformers
+gradio
 Pillow
+nltk

setup.py

@@ -0,0 +1,2 @@
+import nltk
+nltk.download('punkt')