Update app.py

app.py

@@ -3,40 +3,21 @@
 import re
 import gradio as gr
 from PIL import Image
-from transformers import (
-    Blip2Processor,
-    Blip2ForConditionalGeneration,
-    pipeline
+from transformers import pipeline
+
+# Single pipeline: BLIP-2 + Flan-T5-XL for image-to-text
+pipe = pipeline(
+    "image-to-text",
+    model="Salesforce/blip2-flan-t5-xl",
+    tokenizer="Salesforce/blip2-flan-t5-xl",
+    do_sample=True,
+    temperature=1.0,
+    top_k=50,
+    top_p=0.95,
+    max_new_tokens=512
 )
 
-# 1) BLIP-2 for richer image captions
-processor = Blip2Processor.from_pretrained("Salesforce/blip2-opt-2.7b")
-model     = Blip2ForConditionalGeneration.from_pretrained("Salesforce/blip2-opt-2.7b")
-
-def generate_caption(image: Image) -> str:
-    inputs = processor(images=image, return_tensors="pt")
-    outputs = model.generate(**inputs)
-    return processor.decode(outputs[0], skip_special_tokens=True)
-
-# 2) Helper to build Flan-T5-small text pipelines (temperature=1.0)
-def make_pipeline(model_name: str, max_tokens: int):
-    return pipeline(
-        "text2text-generation",
-        model=model_name,
-        tokenizer=model_name,
-        max_new_tokens=max_tokens,
-        do_sample=True,
-        temperature=1.0,
-        top_k=50,
-        top_p=0.95
-    )
-
-# 3) Pipelines: category, analysis, suggestions
-category_generator   = make_pipeline("google/flan-t5-small", 100)
-analysis_generator   = make_pipeline("google/flan-t5-small", 500)
-suggestion_generator = make_pipeline("google/flan-t5-small", 500)
-
-# Hardcoded example ads for gallery
+# Hard-coded example-ad URLs
 def get_recommendations():
     return [
         "https://i.imgur.com/InC88PP.jpeg",
@@ -51,69 +32,60 @@ def get_recommendations():
         "https://i.imgur.com/Xj92Cjv.jpeg",
     ]
 
-# Step B: Flan interprets caption into concise category
-def generate_category(caption: str) -> str:
-    prompt = f"Caption: {caption}\nProvide a concise category label for this ad."
-    raw = category_generator(prompt)[0]["generated_text"].strip()
-    return raw.splitlines()[0]
-
-# Step C: Flan produces exactly five-sentence analysis
-def generate_analysis(caption: str) -> str:
+def process(image: Image):
+    # A single prompt that asks for exactly what you need
     prompt = (
-        f"Caption: {caption}\n"
-        "In exactly five sentences, explain what this ad communicates and its emotional impact."
+        "You are a smart ad analyst. Given the following ad image, output:\n"
+        "Category: <one concise label>\n"
+        "Analysis: <exactly five sentences explaining what it communicates and its emotional impact>\n"
+        "Suggestions:\n"
+        "- <bullet 1>\n"
+        "- <bullet 2>\n"
+        "- <bullet 3>\n"
+        "- <bullet 4>\n"
+        "- <bullet 5>\n"
     )
-    raw = analysis_generator(prompt)[0]["generated_text"].strip()
-    sentences = re.split(r'(?<=[.!?])\s+', raw)
-    return " ".join(sentences[:5])
+    raw = pipe(image, prompt=prompt)[0]["generated_text"]
 
-# Step D: Flan suggests five actionable bullet-point improvements
-def generate_suggestions(caption: str) -> str:
-    prompt = (
-        f"Caption: {caption}\n"
-        "Suggest five distinct improvements as bullet points. Each line must start with '- '."
-    )
-    raw = suggestion_generator(prompt)[0]["generated_text"].strip()
-    lines = [l for l in raw.splitlines() if l.strip().startswith('- ')]
-    if len(lines) < 5:
-        all_lines = [l.strip() for l in raw.splitlines() if l.strip()]
-        lines = [
-            ('- ' + all_lines[i]) if not all_lines[i].startswith('- ') else all_lines[i]
-            for i in range(min(5, len(all_lines)))
-        ]
-    return "\n".join(lines[:5])
+    # Parse out the three sections
+    cat_match = re.search(r"Category:(.*)Analysis:", raw, re.S)
+    ana_match = re.search(r"Analysis:(.*)Suggestions:", raw, re.S)
+    sug_match = re.search(r"Suggestions:(.*)", raw, re.S)
 
-# Orchestrator: process image through all steps
-def process(image: Image):
-    caption     = generate_caption(image)
-    category    = generate_category(caption)
-    analysis    = generate_analysis(caption)
-    suggestions = generate_suggestions(caption)
-    recs        = get_recommendations()
-    return category, analysis, suggestions, recs
+    category    = cat_match.group(1).strip() if cat_match else ""
+    analysis    = ana_match.group(1).strip() if ana_match else ""
+    suggestions = sug_match.group(1).strip() if sug_match else ""
+
+    # Ensure suggestions each start with '-'
+    bullets = [line.strip() for line in suggestions.splitlines() if line.strip()]
+    if len(bullets) < 5:
+        bullets = bullets + ["- (no bullet returned)"] * (5 - len(bullets))
+    suggestions = "\n".join(bullets[:5])
+
+    return category, analysis, suggestions, get_recommendations()
 
-# Gradio UI layout
+# Build 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 see: an Ad Category, a five-sentence Analysis, "
-        "five bullet-point Suggestions, and Example Ads."
+        "Upload an image ad to see: an **Ad Category**, a **five-sentence Analysis**, "
+        "**five bullet-point Suggestions**, and **Example Ads**."
     )
 
     with gr.Row():
         image_input = gr.Image(type="pil", label="Upload Ad Image")
         with gr.Column():
-            category_out   = gr.Textbox(label="Ad Category", interactive=False)
-            analysis_out   = gr.Textbox(label="Ad Analysis", lines=5, interactive=False)
-            suggestion_out = gr.Textbox(label="Improvement Suggestions", lines=5, interactive=False)
-            btn            = gr.Button("Analyze Ad", size="sm", variant="primary")
+            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", size="sm", variant="primary")
 
-    recommendation_gallery = gr.Gallery(label="Recommended Example Ads", show_label=True)
+    gallery = gr.Gallery(label="Recommended Example Ads", show_label=True)
 
     btn.click(
         fn=process,
         inputs=[image_input],
-        outputs=[category_out, analysis_out, suggestion_out, recommendation_gallery]
+        outputs=[cat_out, ana_out, sug_out, gallery]
     )
 
     gr.Markdown("Made by Simon Thalmay")