Update app.py

app.py

@@ -6,20 +6,16 @@ import re
 import torch
 from transformers import BlipForConditionalGeneration, BlipProcessor, T5ForConditionalGeneration, T5Tokenizer
 
-# Device: CPU-only
 device = torch.device("cpu")
 
-# Load caption model (small, CPU-friendly)
-# HF model: Salesforce/blip-image-captioning-base
+# Load models (CPU-friendly)
 processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
-model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")
-model.to(device)
+model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base").to(device)
 
-# Lightweight rewriter model (T5-small) to improve fluency/detail without big cost
 rewriter_tokenizer = T5Tokenizer.from_pretrained("t5-small")
 rewriter = T5ForConditionalGeneration.from_pretrained("t5-small").to(device)
 
-# Simple safety patterns (non-bypass)
+# Safety patterns (simple filter)
 SENSITIVE_PATTERNS = [
     r"\b(nude|naked|porn|sex|sexual|explicit|hardcore)\b",
     r"\b(blood|gore|mutilat|disembowel|organs)\b",
@@ -41,15 +37,12 @@ def is_caption_allowed(text: str):
     return SENSITIVE_RE.search(text) is None
 
 def generate_caption(img: Image.Image, max_len:int=30):
-    # prepare inputs
     inputs = processor(images=img, return_tensors="pt").to(device)
-    # generate with small beams to balance speed/quality
     out = model.generate(**inputs, max_length=max_len, num_beams=3, early_stopping=True)
     caption = processor.decode(out[0], skip_special_tokens=True).strip()
     return caption
 
 def rewrite_caption(caption: str, max_len:int=64):
-    # Use T5-small to expand/clean the caption. Prefix "paraphrase:" helps instruct T5.
     input_text = "paraphrase: " + caption
     tok = rewriter_tokenizer(input_text, return_tensors="pt", truncation=True).to(device)
     out = rewriter.generate(**tok, max_length=max_len, num_beams=2, early_stopping=True)
@@ -59,28 +52,42 @@ def rewrite_caption(caption: str, max_len:int=64):
 def describe_image(url: str, max_caption_len: int = 30, expand: bool = True):
     img, err = load_image_from_url(url)
     if err:
-        return None, err
+        return None, f"Error: {err}"
     caption = generate_caption(img, max_len=max_caption_len)
     if not is_caption_allowed(caption):
-        return img, "Caption omitted for safety."
+        # Provide a neutral, respectful safety message with next steps
+        safety_msg = ("A descriptive caption was not provided because the image may contain explicit or graphic content. "
+                      "If this is unexpected, try a different image or upload a cropped/edited version that removes sensitive content.")
+        return img, safety_msg
     if expand:
         try:
             caption = rewrite_caption(caption, max_len=64)
             if not is_caption_allowed(caption):
-                return img, "Caption omitted for safety."
+                safety_msg = ("A descriptive caption was not provided because the generated text may describe explicit or graphic content. "
+                              "Try a different image or disable the expansion option.")
+                return img, safety_msg
         except Exception:
             pass
+    # Make caption more descriptive by appending structural cues (objects, colors, setting)
+    # Quick heuristic: if short, expand with a simple template
+    if len(caption.split()) < 6:
+        caption = f"{caption}. The scene appears to contain: {caption.lower()}."
     return img, caption
 
+# Gradio UI: image left, caption right
 with gr.Blocks() as demo:
-    gr.Markdown("## CPU-optimized Image Captioning (BLIP-base + T5-small rewriter)")
+    gr.Markdown("## Image captioning — image on the left, descriptive caption on the right (CPU-optimized)")
     with gr.Row():
-        url_in = gr.Textbox(label="Image URL", placeholder="https://example.com/photo.jpg")
-        max_len = gr.Slider(minimum=10, maximum=60, value=30, label="Max caption token length")
-        expand_chk = gr.Checkbox(label="Expand/rewite caption (slower, more natural)", value=True)
-    img_out = gr.Image(type="pil", label="Image")
-    caption_out = gr.Textbox(label="Caption")
-    go = gr.Button("Load & Describe")
+        with gr.Column(scale=1):
+            url_in = gr.Textbox(label="Image URL", placeholder="https://example.com/photo.jpg")
+            max_len = gr.Slider(minimum=10, maximum=60, value=30, label="Max caption length")
+            expand_chk = gr.Checkbox(label="Expand/rewite caption (slower, more natural)", value=True)
+            go = gr.Button("Load & Describe")
+        with gr.Column(scale=1):
+            img_out = gr.Image(type="pil", label="Image")
+        with gr.Column(scale=1):
+            caption_out = gr.Textbox(label="Descriptive caption", lines=6)
+
     go.click(fn=describe_image, inputs=[url_in, max_len, expand_chk], outputs=[img_out, caption_out])
 
 if __name__ == "__main__":