update app.py

app.py

@@ -16,7 +16,7 @@ from sklearn.preprocessing import normalize
 st.set_page_config(
     page_title="Image Captioning Refinement Fusion System",
     layout="wide",
-    initial_sidebar_state="expanded"
+    initial_sidebar_bar="expanded"
 )
 
 JINA_KEY = os.environ.get("JINA_KEY", "")
@@ -87,20 +87,6 @@ def image_to_data_uri(image: Image.Image) -> str:
     b64 = base64.b64encode(raw).decode()
     return f"data:image/jpeg;base64,{b64}"
 
-# ============================================================================
-# STEP 1 — FLORENCE-2-LARGE: 5 DISTINCT CAPTION APPROACHES
-#
-# Cap 1: <CAPTION> greedy
-#        → single concise sentence, primary subject only
-# Cap 2: <CAPTION> sampling temp=1.0
-#        → alt-text accessibility style, concise but different phrasing
-# Cap 3: <DETAILED_CAPTION> temp=0.7
-#        → paragraph describing foreground, background, colors
-# Cap 4: <DETAILED_CAPTION> temp=1.1
-#        → focuses on mood, atmosphere, implied action
-# Cap 5: <MORE_DETAILED_CAPTION> temp=0.8
-#        → exhaustive breakdown of every visible element
-# ============================================================================
 def generate_captions_florence(image: Image.Image, florence_proc, florence_mod) -> list:
 
     captions   = []
@@ -246,25 +232,25 @@ def majority_voting(captions, itm, jina, cosine) -> tuple:
 def fuse_captions(cap1: str, cap2: str, qwen_tok, qwen_mod) -> str:
 
     system_prompt = (
-        "You write image captions. "
+        "You write detailed image captions. "
         "You will receive two captions of the same image. "
-        "Your job is to combine them into one detailed caption. "
-        "Include ALL specific details you find: "
-        "the clothing colors and style of each person, "
+        "Combine them into one complete, detailed caption. "
+        "Include ALL visible details: "
+        "clothing colors and style of each person, "
         "what each person looks like and what they are doing, "
-        "the objects and surroundings visible around them, "
-        "and the setting or background of the scene. "
-        "Write 5 to 6 sentences. Use simple, clear, everyday words. "
-        "Do NOT summarize or shorten — keep every specific detail. "
-        "Only include what is clearly visible. "
+        "objects and surroundings visible in the scene, "
+        "and the background or setting. "
+        "Write 3 to 4 complete sentences. "
+        "Always finish the last sentence properly — never leave it incomplete. "
+        "Use simple, clear, everyday words. "
         "Return ONLY the caption, nothing else."
     )
 
     user_prompt = (
         f"Caption A: {cap1}\n"
         f"Caption B: {cap2}\n\n"
-        "Write a detailed caption that includes all the clothing, "
-        "people, objects and background in details:"
+        "Write a detailed caption covering all clothing, "
+        "people, objects and background:"
     )
 
     try:
@@ -282,8 +268,11 @@ def fuse_captions(cap1: str, cap2: str, qwen_tok, qwen_mod) -> str:
         with torch.no_grad():
             generated_ids = qwen_mod.generate(
                 **model_inputs,
-                max_new_tokens=120,
-                do_sample=False
+                max_new_tokens     = 220,
+                do_sample          = False,
+                repetition_penalty = 1.1,
+                eos_token_id       = qwen_tok.eos_token_id,
+                pad_token_id       = qwen_tok.eos_token_id
             )
 
         output_ids = generated_ids[0][len(model_inputs.input_ids[0]):]
@@ -293,6 +282,16 @@ def fuse_captions(cap1: str, cap2: str, qwen_tok, qwen_mod) -> str:
             if fused.lower().startswith(prefix.lower()):
                 fused = fused[len(prefix):].strip()
 
+        # Safety — trim to last complete sentence if still cut off
+        if fused and not fused.endswith((".", "!", "?")):
+            last_stop = max(
+                fused.rfind("."),
+                fused.rfind("!"),
+                fused.rfind("?")
+            )
+            if last_stop > len(fused) // 2:
+                fused = fused[:last_stop + 1].strip()
+
         return fused if fused else cap1
 
     except Exception as e:
@@ -340,10 +339,6 @@ def compute_caption_quality(image, final_caption, blip_proc, blip_itm) -> tuple:
     avg_score = round((itm_score + cosine_score) / 2, 4)
     return avg_score, round(itm_score, 4), round(cosine_score, 4)
 
-# ============================================================================
-# GAUGE — updated to match reference style
-# Bright saturated zone colors, sharp black needle, clean arc, no dark shades
-# ============================================================================
 def render_gauge(score, itm, cosine, placeholder):
 
     if score >= 0.75:
@@ -421,9 +416,7 @@ def render_gauge(score, itm, cosine, placeholder):
                 unsafe_allow_html=True
             )
 
-# ============================================================================
-# SIDEBAR — pipeline steps + live accuracy section (session_state)
-# ============================================================================
+# ── SIDEBAR ────────────────────────────────────────────────────────
 with st.sidebar:
     st.title("Image Captioning Refinement Fusion")
     st.markdown("---")
@@ -450,8 +443,6 @@ Caption fusion
     st.markdown("---")
     st.markdown("**Local:** Florence-2, BLIP ITM, Qwen2.5")
     st.markdown("**API:** Jina")
-
-    # ── accuracy section ──────────────
     st.markdown("---")
     st.markdown("### Caption Quality Metrics")
     st.markdown("""
@@ -464,9 +455,7 @@ Measures embedding distance
 between image and caption.
     """)
 
-# ============================================================================
-# MAIN UI
-# ============================================================================
+# ── MAIN UI ────────────────────────────────────────────────────────
 st.title("Image Captioning Refinement Fusion System")
 st.markdown("Upload an image to generate a refined, grounded caption.")
 st.markdown("---")
@@ -499,7 +488,9 @@ if uploaded_file is not None:
             status   = st.empty()
 
             status.info("Step 1/6: Generating captions with Florence-2-Large...")
-            captions = generate_captions_florence(input_image, florence_proc, florence_mod)
+            captions = generate_captions_florence(
+                input_image, florence_proc, florence_mod
+            )
             progress.progress(16)
 
             with st.expander("5 Generated Captions", expanded=True):
@@ -507,7 +498,9 @@ if uploaded_file is not None:
                     st.write(f"**{i+1}.** {cap}")
 
             status.info("Step 2/6: Computing BLIP ITM scores...")
-            itm_scores = compute_itm_scores(input_image, captions, blip_proc, blip_itm)
+            itm_scores = compute_itm_scores(
+                input_image, captions, blip_proc, blip_itm
+            )
             progress.progress(32)
 
             status.info("Step 3/6: Computing Jina Reranker scores...")
@@ -515,7 +508,9 @@ if uploaded_file is not None:
             progress.progress(50)
 
             status.info("Step 4/6: Computing Cosine Similarity scores...")
-            cosine_scores = compute_cosine_scores(input_image, captions, blip_proc, blip_itm)
+            cosine_scores = compute_cosine_scores(
+                input_image, captions, blip_proc, blip_itm
+            )
             progress.progress(66)
 
             scores_df = pd.DataFrame({
@@ -560,7 +555,6 @@ if uploaded_file is not None:
                 input_image, final, blip_proc, blip_itm
             )
 
-            # Store in session_state so sidebar updates on rerender
             st.session_state.avg_score = avg_score
             st.session_state.itm_q     = itm_q
             st.session_state.cosine_q  = cosine_q