Update app.py

app.py

@@ -12,7 +12,7 @@ caption_pipe = pipeline(
     do_sample=False,
 )
 
-# 2) Flan-T5 for text‐to‐text
+# 2) Flan-T5 for text‑to‑text
 FLAN = "google/flan-t5-large"
 category_pipe = pipeline(
     "text2text-generation",
@@ -38,7 +38,16 @@ suggestion_pipe = pipeline(
     do_sample=True,
     temperature=1.0,
 )
+# 3) Caption expander when BLIP is too short
+expansion_pipe = pipeline(
+    "text2text-generation",
+    model=FLAN,
+    tokenizer=FLAN,
+    max_new_tokens=128,
+    do_sample=False,
+)
 
+# Example gallery helper returns 10 example ad URLs
 def get_recommendations():
     return [
         "https://i.imgur.com/InC88PP.jpeg",
@@ -46,38 +55,49 @@ def get_recommendations():
         "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",
     ]
 
+
 def process(image: Image):
     # 1) BLIP caption
     caption = caption_pipe(image)[0]["generated_text"].strip()
 
-    # 2) Single‐label category
+    # 1a) Expand if too short
+    if len(caption.split()) < 3:
+        exp_prompt = f"Expand this into a detailed description: {caption}"
+        desc = expansion_pipe(exp_prompt)[0]["generated_text"].strip()
+    else:
+        desc = caption
+
+    # 2) Single‑label category
     cat_prompt = (
-        f"Caption: {caption}\n\n"
-        "Give me one concise category label for this ad (e.g. 'Fitness', 'Food'):"
+        f"Description: {desc}\n\n"
+        "Provide a concise category label for this ad (e.g. 'Fitness', 'Food'):"
     )
     category = category_pipe(cat_prompt)[0]["generated_text"].strip().splitlines()[0]
 
-    # 3) Five‐sentence analysis
+    # 3) Five‑sentence analysis
     ana_prompt = (
-        f"Caption: {caption}\n\n"
+        f"Description: {desc}\n\n"
         "Write exactly five sentences explaining what this ad communicates and its emotional impact."
     )
     raw_ana = analysis_pipe(ana_prompt)[0]["generated_text"].strip()
-    # ensure exactly five sentences
     sents = re.split(r'(?<=[.!?])\s+', raw_ana)
     analysis = " ".join(sents[:5])
 
-    # 4) Five bullet‐point suggestions
+    # 4) Five bullet‑point suggestions
     sug_prompt = (
-        f"Caption: {caption}\n\n"
+        f"Description: {desc}\n\n"
         "Suggest five distinct improvements for this ad. "
         "Each suggestion must start with '- ' and be one sentence."
     )
     raw_sug = suggestion_pipe(sug_prompt)[0]["generated_text"].strip()
     bullets = [l for l in raw_sug.splitlines() if l.strip().startswith("-")]
-    # pad/truncate to 5
     if len(bullets) < 5:
         extras = [l.strip() for l in raw_sug.splitlines() if l.strip()]
         for ex in extras:
@@ -93,7 +113,7 @@ with gr.Blocks() as demo:
     gr.Markdown(
         "Upload an image ad to get:\n"
         "- **BLIP Caption** (debug)\n"
-        "- **Category**\n"
+        "- **Ad Category**\n"
         "- **Five-sentence Analysis**\n"
         "- **Five bullet-point Suggestions**\n"
         "- **Example Ads**"
@@ -108,7 +128,7 @@ with gr.Blocks() as demo:
             sug_out = gr.Textbox(label="Improvement Suggestions", lines=5, interactive=False)
             btn     = gr.Button("Analyze Ad")
 
-    gallery = gr.Gallery(label="Example Ads").style(grid=[5], height="auto")
+    gallery = gr.Gallery(label="Example Ads")
 
     btn.click(
         fn=process,
@@ -119,4 +139,4 @@ with gr.Blocks() as demo:
     gr.Markdown("Made by Simon Thalmay")
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()