Update app.py

app.py

@@ -1,12 +1,98 @@
+import re
+import gradio as gr
+import torch
+from PIL import Image
+from transformers import (
+    pipeline,
+    AutoProcessor,
+    AutoModelForVision2Seq,
+    AutoTokenizer,
+    AutoModelForSeq2SeqLM,
+)
+
+# Auto-detect CPU/GPU
+DEVICE = 0 if torch.cuda.is_available() else -1
+
+# 1) BLIP captioner
+processor = AutoProcessor.from_pretrained("Salesforce/blip-image-captioning-large")
+blip_model = AutoModelForVision2Seq.from_pretrained("Salesforce/blip-image-captioning-large")
+caption_pipe = pipeline(
+    task="image-to-text",
+    model=blip_model,
+    tokenizer=processor.tokenizer,
+    image_processor=processor.image_processor,
+    device=DEVICE,
+)
+
+# 2) FLAN-T5 for text-to-text
+FLAN_MODEL = "google/flan-t5-large"
+flan_tokenizer = AutoTokenizer.from_pretrained(FLAN_MODEL)
+flan_model = AutoModelForSeq2SeqLM.from_pretrained(FLAN_MODEL)
+
+category_pipe = pipeline(
+    "text2text-generation",
+    model=flan_model,
+    tokenizer=flan_tokenizer,
+    device=DEVICE,
+    max_new_tokens=32,
+    do_sample=True,
+    temperature=1.0,
+)
+
+analysis_pipe = pipeline(
+    "text2text-generation",
+    model=flan_model,
+    tokenizer=flan_tokenizer,
+    device=DEVICE,
+    max_new_tokens=256,
+    do_sample=True,
+    temperature=1.0,
+)
+
+suggestion_pipe = pipeline(
+    "text2text-generation",
+    model=flan_model,
+    tokenizer=flan_tokenizer,
+    device=DEVICE,
+    max_new_tokens=256,
+    do_sample=True,
+    temperature=1.0,
+)
+
+expansion_pipe = pipeline(
+    "text2text-generation",
+    model=flan_model,
+    tokenizer=flan_tokenizer,
+    device=DEVICE,
+    max_new_tokens=128,
+    do_sample=False,
+)
+
+# Example gallery helper
+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",
+    ]
+
+# Main processing function
 def process(image: Image):
     if image is None:
-        return "", "", "", get_recommendations()
+        return "", "", "", "", get_recommendations()
 
     # 1) BLIP caption
     caption_res = caption_pipe(image, max_new_tokens=64)
     raw_caption = caption_res[0]["generated_text"].strip()
 
-    # 1a) Expand caption if too short
+    # 1a) Expand if too short
     if len(raw_caption.split()) < 3:
         exp = expansion_pipe(f"Expand into a detailed description: {raw_caption}")
         desc = exp[0]["generated_text"].strip()
@@ -29,39 +115,54 @@ def process(image: Image):
     sentences = re.split(r'(?<=[.!?])\s+', ana_raw)
     analysis = " ".join(sentences[:5])
 
-    # 4) Five bullet-point suggestions, filter for unique & not empty
+    # 4) Five bullet-point suggestions
     sug_prompt = (
         f"Description: {desc}\n\n"
-        "Suggest five unique, practical improvements for this ad. Each must address a different aspect (such as message, visuals, CTA, targeting, layout, or design). Each suggestion must be only one sentence and start with '- '. Do NOT repeat suggestions."
+        "Provide five distinct improvement suggestions, each starting with '- '."
     )
     sug_raw = suggestion_pipe(sug_prompt)[0]["generated_text"].strip()
-    bullets = []
-    seen = set()
-    for line in sug_raw.splitlines():
-        if line.startswith("-"):
-            suggestion = line.strip()
-            # Remove duplicates and ignore empty lines
-            if suggestion and suggestion not in seen:
-                bullets.append(suggestion)
-                seen.add(suggestion)
-        elif line.strip():
-            suggestion = "- " + line.strip()
-            if suggestion and suggestion not in seen:
-                bullets.append(suggestion)
-                seen.add(suggestion)
-        if len(bullets) == 5:
-            break
-    # Add defaults if needed
-    defaults = [
-        "- Make the main headline more eye-catching.",
-        "- Add a clear and visible call-to-action button.",
-        "- Use contrasting colors for better readability.",
-        "- Highlight the unique selling point of the product.",
-        "- Simplify the design to reduce clutter."
-    ]
-    for default in defaults:
-        if len(bullets) < 5 and default not in seen:
-            bullets.append(default)
+    bullets = [l for l in sug_raw.splitlines() if l.startswith("-")]
+    if len(bullets) < 5:
+        for line in sug_raw.splitlines():
+            if len(bullets) >= 5:
+                break
+            if line and not line.startswith("-"):
+                bullets.append("- " + line)
+    while len(bullets) < 5:
+        bullets.append("- Improve visual appeal and clarity.")
     suggestions = "\n".join(bullets[:5])
 
-    return cat_out, analysis, suggestions, get_recommendations()
+    return raw_caption, cat_out, analysis, suggestions, get_recommendations()
+
+# Gradio UI
+def main():
+    with gr.Blocks(title="Smart Ad Analyzer") as demo:
+        gr.Markdown("## 📢 Smart Ad Analyzer")
+        gr.Markdown(
+            "Upload an ad image to get:\n"
+            "- 🔍 **BLIP Caption**\n"
+            "- 📂 **Ad Category**\n"
+            "- 📊 **Five-sentence Analysis**\n"
+            "- 🚀 **Five Improvement Suggestions**\n"
+            "- 📸 **Example Ads**"
+        )
+        with gr.Row():
+            inp = gr.Image(type='pil', label='Upload Ad Image')
+            with gr.Column():
+                cap_out = gr.Textbox(label='🔍 BLIP Caption', 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', variant='primary')
+        gallery = gr.Gallery(label='Example Ads')
+        btn.click(
+            fn=process,
+            inputs=[inp],
+            outputs=[cap_out, cat_out, ana_out, sug_out, gallery],
+        )
+        gr.Markdown('Made by Simon Thalmay')
+    return demo
+
+if __name__ == "__main__":
+    demo = main()
+    demo.launch()