Update app.py

app.py

@@ -1,93 +1,142 @@
-# app.py
-
+import logging
 import re
 import gradio as gr
 from PIL import Image
 from transformers import (
-    AutoProcessor,
-    AutoModelForVision2Seq,
+    BlipProcessor,
+    BlipForConditionalGeneration,
     pipeline,
 )
 
-# 1 – BLIP-large for image captioning
-processor = AutoProcessor.from_pretrained("Salesforce/blip-image-captioning-large")
-model     = AutoModelForVision2Seq.from_pretrained("Salesforce/blip-image-captioning-large")
-
-def generate_caption(image: Image) -> str:
-    inputs = processor(images=image, return_tensors="pt")
-    outputs = model.generate(**inputs)
-    return processor.tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-# 2 – Flan-T5 pipelines
-def make_pipe(model_name, max_tokens):
-    return pipeline(
-        "text2text-generation",
-        model=model_name,
-        tokenizer=model_name,
-        max_new_tokens=max_tokens,
-        do_sample=True,
-        temperature=1.0,
-    )
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+
+# 1) BLIP captioner (large model for richer captions)
+caption_processor = BlipProcessor.from_pretrained(
+    "Salesforce/blip-image-captioning-large",
+    use_fast=False
+)
+caption_model = BlipForConditionalGeneration.from_pretrained(
+    "Salesforce/blip-image-captioning-large"
+)
 
-cat_pipe = make_pipe("google/flan-t5-small", 80)
-ana_pipe = make_pipe("google/flan-t5-small", 200)
-sug_pipe = make_pipe("google/flan-t5-small", 200)
+caption_pipe = pipeline(
+    task="image-to-text",
+    model=caption_model,
+    processor=caption_processor,
+    device=-1,
+    max_length=64,
+    do_sample=False,
+)
+
+# 2) Flan-T5 pipelines for category, analysis, suggestions
+FLAN_MODEL = "google/flan-t5-large"
+category_pipe = pipeline(
+    "text2text-generation",
+    model=FLAN_MODEL,
+    tokenizer=FLAN_MODEL,
+    max_new_tokens=32,
+    do_sample=True,
+    temperature=1.0,
+)
+analysis_pipe = pipeline(
+    "text2text-generation",
+    model=FLAN_MODEL,
+    tokenizer=FLAN_MODEL,
+    max_new_tokens=256,
+    do_sample=True,
+    temperature=1.0,
+)
+suggestion_pipe = pipeline(
+    "text2text-generation",
+    model=FLAN_MODEL,
+    tokenizer=FLAN_MODEL,
+    max_new_tokens=256,
+    do_sample=True,
+    temperature=1.0,
+)
 
-# 3 – Recommendation gallery
-def get_recs():
+def get_recommendations():
     return [
         "https://i.imgur.com/InC88PP.jpeg",
         "https://i.imgur.com/7BHfv4T.png",
         "https://i.imgur.com/wp3Wzc4.jpeg",
         "https://i.imgur.com/5e2xOA4.jpeg",
         "https://i.imgur.com/txjRk98.jpeg",
+        "https://i.imgur.com/rQ4AYl0.jpeg",
+        "https://i.imgur.com/bDzwD04.jpeg",
+        "https://i.imgur.com/fLMngXI.jpeg",
+        "https://i.imgur.com/nYEJzxt.png",
+        "https://i.imgur.com/Xj92Cjv.jpeg",
     ]
 
-# 4 – Full workflow
 def process(image: Image):
-    caption = generate_caption(image)
-
-    # category
-    raw_cat = cat_pipe(f"Caption: {caption}\nLabel this ad in one phrase:")[0]["generated_text"]
-    category = raw_cat.strip().splitlines()[0]
-
-    # analysis
-    raw_ana = ana_pipe(
-        f"Caption: {caption}\nWrite exactly five sentences explaining what this ad communicates and its emotional impact."
-    )[0]["generated_text"]
-    sentences = re.split(r'(?<=[.!?])\s+', raw_ana.strip())
+    # Step 1: Generate BLIP caption
+    caption = caption_pipe(image)[0]["generated_text"].strip()
+    logging.info(f"RAW CAPTION: {caption}")
+
+    # Step 2: Category label
+    cat_prompt = (
+        f"Caption: {caption}\n"
+        "Provide a single concise category label for this ad (e.g. 'Food Ad', 'Fitness Promotion'):"
+    )
+    raw_cat = category_pipe(cat_prompt)[0]["generated_text"].strip()
+    category = raw_cat.splitlines()[0]
+    logging.info(f"RAW CATEGORY: {raw_cat}")
+
+    # Step 3: Five-sentence Analysis
+    ana_prompt = (
+        f"Caption: {caption}\n"
+        "Write exactly five sentences explaining what this ad communicates and its emotional impact."
+    )
+    raw_ana = analysis_pipe(ana_prompt)[0]["generated_text"].strip()
+    sentences = re.split(r'(?<=[.!?])\s+', raw_ana)
     analysis = " ".join(sentences[:5])
+    logging.info(f"RAW ANALYSIS: {raw_ana}")
 
-    # suggestions
-    raw_sug = sug_pipe(
-        f"Caption: {caption}\nSuggest five distinct improvements as bullets, each starting with '- '."
-    )[0]["generated_text"]
-    bullets = [l for l in raw_sug.splitlines() if l.strip().startswith("-")]
-    if len(bullets) < 5:
-        lines = [l.strip() for l in raw_sug.splitlines() if l.strip()]
-        bullets = [("- " + lines[i]) for i in range(min(5, len(lines)))]
-    suggestions = "\n".join(bullets[:5])
-
-    return category, analysis, suggestions, get_recs()
+    # Step 4: Five bullet-point Suggestions
+    sug_prompt = (
+        f"Caption: {caption}\n"
+        "Suggest five distinct improvements for this ad. "
+        "Each suggestion must start with '- ' and be one actionable sentence."
+    )
+    raw_sug = suggestion_pipe(sug_prompt)[0]["generated_text"].strip()
+    lines = [l for l in raw_sug.splitlines() if l.strip().startswith("-")]
+    if len(lines) < 5:
+        extras = [l.strip() for l in raw_sug.splitlines() if l.strip()]
+        for ex in extras:
+            if len(lines) >= 5:
+                break
+            prefix = "- " if not ex.startswith("-") else ""
+            lines.append(prefix + ex.lstrip("- ").strip())
+    suggestions = "\n".join(lines[:5])
+    logging.info(f"RAW SUGGESTIONS:\n{raw_sug}")
+
+    return caption, category, analysis, suggestions, get_recommendations()
 
-# 5 – Gradio UI
 with gr.Blocks(theme=gr.themes.Default(primary_hue="blue")) as demo:
     gr.Markdown("## 📢 Smart Ad Analyzer")
     gr.Markdown(
-        "Upload an image ad to get: a Category, five-sentence Analysis, "
-        "five bullet-point Suggestions, and Example Ads."
+        "Upload an image ad to get: a **BLIP Caption** (debug), a **Category**, a "
+        "**five-sentence Analysis**, **five bullet-point Suggestions**, and **Example Ads**."
     )
 
     with gr.Row():
-        inp = gr.Image(type="pil", label="Upload Ad Image")
+        img       = gr.Image(type="pil", label="Upload Ad Image")
         with gr.Column():
-            out_cat = gr.Textbox(label="Ad Category", interactive=False)
-            out_ana = gr.Textbox(label="Ad Analysis", lines=5, interactive=False)
-            out_sug = gr.Textbox(label="Improvement Suggestions", lines=5, interactive=False)
-            btn     = gr.Button("Analyze Ad", size="sm")
+            debug_cap = gr.Textbox(label="🔍 BLIP Caption (debug)", interactive=False)
+            cat_out   = gr.Textbox(label="Ad Category", interactive=False)
+            ana_out   = gr.Textbox(label="Ad Analysis", lines=5, interactive=False)
+            sug_out   = gr.Textbox(label="Improvement Suggestions", lines=5, interactive=False)
+            btn       = gr.Button("Analyze Ad")
 
     gallery = gr.Gallery(label="Example Ads", show_label=True)
-    btn.click(process, inputs=[inp], outputs=[out_cat, out_ana, out_sug, gallery])
+
+    btn.click(
+        fn=process,
+        inputs=[img],
+        outputs=[debug_cap, cat_out, ana_out, sug_out, gallery],
+    )
 
     gr.Markdown("Made by Simon Thalmay")