Update app.py

app.py

@@ -1,11 +1,12 @@
 # app.py
 import gradio as gr
-from PIL import Image, ImageOps
+from PIL import Image
 import numpy as np
 import pandas as pd
 import io
 import os
 import cv2
+import glob
 from pathlib import Path
 from typing import List, Dict, Any, Tuple
 
@@ -34,21 +35,6 @@ def _to_np(img: Image.Image) -> np.ndarray:
     return np.array(_ensure_rgb(img))
 
 
-def _thumbnail(img: Image.Image, max_side: int = 320) -> Image.Image:
-    """
-    Create a *small* thumbnail for gallery display so you don't need to scroll.
-    Keeps aspect ratio; pads to square so grids look neat.
-    """
-    img = _ensure_rgb(img.copy())
-    img.thumbnail((max_side, max_side))
-    # Pad to square with white (looks nicer for product photos)
-    w, h = img.size
-    side = max(w, h)
-    bg = Image.new("RGB", (side, side), (255, 255, 255))
-    bg.paste(img, ((side - w) // 2, (side - h) // 2))
-    return bg
-
-
 def _hsv_hist_features(img: Image.Image) -> np.ndarray:
     """Return simple features for matching and scoring.
     - Hue histogram (18 bins)
@@ -83,9 +69,7 @@ def _complementary_hue_score(q_hist: np.ndarray, w_hist: np.ndarray) -> float:
     q_h = q_hist[:hb]
     w_h = w_hist[:hb]
     q_shift = np.roll(q_h, hb // 2)
-    # cosine similarity on shifted hues
-    denom = (np.linalg.norm(q_shift) * np.linalg.norm(w_h) + 1e-8)
-    hue_sim = float(np.dot(q_shift, w_h) / denom)
+    hue_sim = float(np.dot(q_shift, w_h) / (np.linalg.norm(q_shift) * np.linalg.norm(w_h) + 1e-8))
 
     # Encourage pairing items with different edge density (texture contrast)
     q_ed = q_hist[-1]
@@ -270,9 +254,7 @@ def _get_embedder() -> _Embedder:
 # ----------------------------
 # State schema:
 # {
-#   "wardrobe": [ {"id": int, "name": str, "image": PIL.Image, "thumb": PIL.Image,
-#                  "features": np.ndarray, "embedding": np.ndarray, "category": str,
-#                  "rating": int|None} ],
+#   "wardrobe": [ {"id": int, "name": str, "image": PIL.Image, "features": np.ndarray, "embedding": np.ndarray, "rating": int|None} ],
 #   "selected_idx": int|None
 # }
 
@@ -280,8 +262,7 @@ def _blank_state() -> Dict[str, Any]:
     return {"wardrobe": [], "selected_idx": None}
 
 
-# ----------------------------
-# Wardrobe management
+
 # ----------------------------
 
 def add_wardrobe(files: List[Any], state: Dict[str, Any]):
@@ -289,13 +270,22 @@ def add_wardrobe(files: List[Any], state: Dict[str, Any]):
         state = _blank_state()
     next_id = 0 if not state["wardrobe"] else max(w["id"] for w in state["wardrobe"]) + 1
     if files is None:
-        return state, _render_gallery(state), _ratings_df(state)
+        return state, _render_gallery(state), _ratings_df(state), ""
 
     scorer = _get_scorer()
     embedder = _get_embedder()
 
+    warnings = []
+    added = 0
+    allowed_exts = {".png", ".jpg", ".jpeg"}
+
     for f in files:
         try:
+            fname = os.path.basename(getattr(f, 'name', f))
+            ext = Path(fname).suffix.lower()
+            if ext not in allowed_exts:
+                warnings.append(fname)
+                continue
             img = Image.open(f.name if hasattr(f, "name") else f)
             img = _ensure_rgb(img)
             feats = _hsv_hist_features(img)
@@ -306,50 +296,84 @@ def add_wardrobe(files: List[Any], state: Dict[str, Any]):
                 category = _ALIAS_TO_CANON.get(alias, "shirt")
             except Exception:
                 category = "shirt"
-            name = os.path.basename(getattr(f, 'name', f))
+            name = fname
             rating = scorer.predict_1to100(img)
             state["wardrobe"].append({
                 "id": next_id,
                 "name": name,
                 "image": img,
-                "thumb": _thumbnail(img),
                 "features": feats,
                 "embedding": emb,
                 "category": category,
                 "rating": int(rating),
             })
+            added += 1
             next_id += 1
         except Exception:
             continue
 
     gallery = _render_gallery(state)
-    return state, gallery, _ratings_df(state)
+    status_lines = []
+    if added:
+        status_lines.append(f"✅ Added {added} item(s) to your wardrobe.")
+    if warnings:
+        status_lines.append(
+            f"⚠️ Skipped {len(warnings)} file(s) (not PNG/JPG): " + ", ".join(warnings[:5]) + ("..." if len(warnings) > 5 else "")
+        )
+        status_lines.append("Please upload .png, .jpg, or .jpeg files.")
+    status_msg = "\n\n".join(status_lines)
+    return state, gallery, _ratings_df(state), status_msg
 
 
 def add_wardrobe_from_dir(example_dir: str, state: Dict[str, Any]):
-    """Load all images in a folder into the wardrobe and auto-rate/classify them."""
+    """Load all images in a folder into the wardrobe and auto-rate/classify them.
+    Used by gr.Examples. Accepts relative paths in the Space repo.
+    """
     if not example_dir:
-        return state, _render_gallery(state), _ratings_df(state)
+        return state, _render_gallery(state), _ratings_df(state), ""
     p = Path(example_dir)
-    patterns = ["*.jpg", "*.jpeg", "*.png", "*.webp", "*.bmp", "*.tif", "*.tiff"]
+    patterns = ["*.jpg", "*.jpeg", "*.png"]  # keep examples aligned with allowed types
     files = []
     for pat in patterns:
         files.extend([str(x) for x in p.glob(pat)])
-    return add_wardrobe(files, state)
+    st, gal, df = add_wardrobe(files, state)[:3]
+    # add_wardrobe returns 4 now; reuse message-less return for examples
+    return st, gal, df, "Loaded example wardrobe."
 
 
 def clear_wardrobe(state: Dict[str, Any]):
     state = _blank_state()
-    return state, [], _ratings_df(state)
+    return state, [], _ratings_df(state), "Wardrobe cleared."
 
 
-def _render_gallery(state: Dict[str, Any]) -> List[Tuple[Image.Image, str]]:
-    """Return list of (thumbnail, caption)."""
-    out = []
-    for w in state.get("wardrobe", []):
-        caption = f"#{w['id']} · {w['name']} · {w.get('category','?')} · {w.get('rating','-')}/100"
-        out.append((w["thumb"], caption))
-    return out
+def _render_gallery(state: Dict[str, Any]) -> List[Image.Image]:
+    return [w["image"] for w in state.get("wardrobe", [])]
+
+
+def on_select_item(item_label: str, state: Dict[str, Any]):
+    if not item_label:
+        return state, None, gr.update(value=50)
+    # label format: "#<id> · <name>"
+    try:
+        item_id = int(item_label.split(" ")[0][1:])
+    except Exception:
+        return state, None, gr.update(value=50)
+
+    idx = next((i for i, w in enumerate(state["wardrobe"]) if w["id"] == item_id), None)
+    state["selected_idx"] = idx
+    if idx is None:
+        return state, None, gr.update(value=50)
+    w = state["wardrobe"][idx]
+    current_rating = w["rating"] if w["rating"] is not None else 50
+    return state, w["image"], gr.update(value=int(current_rating))
+
+
+def save_rating(rating: int, state: Dict[str, Any]):
+    idx = state.get("selected_idx", None)
+    if idx is None:
+        return state, _ratings_df(state)
+    state["wardrobe"][idx]["rating"] = int(rating)
+    return state, _ratings_df(state)
 
 
 def _ratings_df(state: Dict[str, Any]) -> pd.DataFrame:
@@ -373,6 +397,25 @@ def export_ratings(state: Dict[str, Any]):
     df.to_csv(buf, index=False)
     buf.seek(0)
     return buf
+    buf = io.BytesIO()
+    df.to_csv(buf, index=False)
+    buf.seek(0)
+    return buf
+
+
+def import_ratings(file_obj, state: Dict[str, Any]):
+    # Deprecated in auto-rating flow; keep no-op for compatibility
+    return state, _ratings_df(state)
+    try:
+        df = pd.read_csv(file_obj.name if hasattr(file_obj, "name") else file_obj)
+        names_to_rating = {str(row["name"]): int(row["rating"]) if not pd.isna(row["rating"]) else None
+                               for _, row in df.iterrows()}
+        for w in state.get("wardrobe", []):
+            if w["name"] in names_to_rating:
+                w["rating"] = names_to_rating[w["name"]]
+    except Exception:
+        pass
+    return state, _ratings_df(state)
 
 
 # ----------------------------
@@ -381,9 +424,7 @@ def export_ratings(state: Dict[str, Any]):
 
 def rate_and_recommend(query_img: Image.Image, top_k: int, matching_mode: str, state: Dict[str, Any]):
     if query_img is None:
-        return 0, "No image provided.", [], pd.DataFrame(columns=[
-            "rank", "name", "category", "model_rating", "match_score"
-        ])
+        return 0, "Please upload a PNG/JPG image to get a rating and matches.", []
 
     query_img = _ensure_rgb(query_img)
 
@@ -432,7 +473,7 @@ def rate_and_recommend(query_img: Image.Image, top_k: int, matching_mode: str, s
         else:  # Complementary color + style
             final = 0.5 * ((cos + 1.0) / 2.0) + 0.5 * comp
 
-        # Quality prior (prefer items you generally like)
+        # Quality prior
         if w.get("rating") is not None:
             qual = 0.5 + 0.5 * (w["rating"] / 100.0)
             final *= qual
@@ -450,177 +491,147 @@ def rate_and_recommend(query_img: Image.Image, top_k: int, matching_mode: str, s
             final = 0.5 * ((cos + 1.0) / 2.0) + 0.5 * comp
             candidates.append((final, w))
 
-    # Rank and prepare outputs
     candidates.sort(key=lambda x: x[0], reverse=True)
-    chosen = candidates[: max(0, top_k)]
-
-    if not chosen:
-        txt = f"**Predicted wear score:** {pred}/100  \n" \
-              f"**Detected category:** {qcat}  \n" \
-              f"**Results:** No compatible matches found in your wardrobe."
-        return pred, txt, [], pd.DataFrame(columns=[
-            "rank", "name", "category", "model_rating", "match_score"
-        ])
+    top = candidates[: max(0, top_k)]
+    recs = []
+    for score, w in top:
+        caption = f"{w['name']} · {w.get('category','?')} · match {int(round(100*score))}%"
+        recs.append((w["image"], caption))
 
-    # Human-friendly summary
-    top_names = ", ".join([f"{w['name']} ({w.get('category')})" for _, w in chosen])
+    if len(recs) == 0:
+        txt = f"**Predicted wear score:** {pred}/100\n\n_No compatible matches found in your wardrobe._"
+        return pred, txt, []
+
+    top_names = ", ".join([f"{w['name']} ({w.get('category')})" for _, w in top])
     txt = (
         f"**Predicted wear score:** {pred}/100  \n"
-        f"**Detected category:** `{qcat}`  \n"
-        f"**Matching mode:** `{matching_mode}`  \n"
-        f"**Suggested pairings:** {top_names}  \n"
-        f"_Scoring blends style similarity (CLIP) with color/texture complement and your model ratings._"
+        f"**Detected category:** {qcat}  \n"
+        f"**Top suggestions:** {top_names}  \n"
+        f"_Matching mode:_ {matching_mode.lower()} with category filtering and quality prior."
     )
+    return pred, txt, recs
 
-    # Gallery (thumbnails) + table
-    rec_gallery = [(w["thumb"], f"{w['name']} · {w.get('category')} · {w.get('rating','-')}/100")
-                   for _, w in chosen]
+    # (Unreached legacy block kept to minimize the overall diff)
+    try:
+        scorer = _get_scorer()
+        pred = scorer.predict_1to100(query_img)
+    except Exception:
+        pred = 50
 
-    rec_table = pd.DataFrame([{
-        "rank": i + 1,
-        "name": w["name"],
-        "category": w.get("category"),
-        "model_rating": w.get("rating"),
-        "match_score": round(float(s), 3)
-    } for i, (s, w) in enumerate(chosen)])
+    qfeat = _hsv_hist_features(query_img)
+    candidates = []
+    for w in state.get("wardrobe", []):
+        comp = _complementary_hue_score(qfeat, w["features"])
+        user_w = 1.0
+        if w["rating"] is not None:
+            user_w = 0.5 + 0.5 * (w["rating"] / 100.0)
+        final = comp * user_w
+        candidates.append((final, w))
 
-    return pred, txt, rec_gallery, rec_table
+    candidates.sort(key=lambda x: x[0], reverse=True)
+    recs = [w["image"] for _, w in candidates[: max(0, top_k)]]
+
+    if len(recs) == 0:
+        txt = f"Predicted rating: {pred}/100. No matches found in your wardrobe."
+        return pred, txt, []
+
+    top_names = ", ".join([w["name"] for _, w in candidates[: max(0, top_k)]])
+    txt = (
+        f"Predicted rating: {pred}/100. Suggested pairings from your wardrobe: {top_names}."
+        f"Logic: complementary hues + texture contrast + your cached ratings."
+    )
+    return pred, txt, recs
 
 
 # ----------------------------
 # Gradio UI
 # ----------------------------
-
-APP_DESCRIPTION = """
-# Wardrobe Rater + Recommender
-
-**What this app does (1 minute):**
-- **Upload** photos of items in your wardrobe (tops, pants, jackets, shoes).
-- The app **auto-rates** each item with your model and tags its category.
-- When you upload a **new item** (e.g., a shopping photo), it:
-  1) Predicts how likely **you** are to wear it (1–100),
-  2) Suggests the **best matches** from your wardrobe using style similarity and color/texture complement.
-
-**How to use it:**
-1. Go to **“1) Wardrobe Manager”** and upload several wardrobe images (front-on, good lighting).
-2. Then open **“2) Rate + Recommend New Item”**, upload your candidate item, pick the number of matches, and click **Rate + Recommend**.
-3. Review the **score**, the **explanation**, the **top matches** (thumbnails), and the **table**.
-"""
-
-WARDROBE_TIPS = """
-**What to upload:** clear product-style photos (JPG/PNG).  
-Avoid group photos or cluttered backgrounds when possible.  
-You can click thumbnails to preview full size.
-"""
-
-QUERY_TIPS = """
-**Upload a single item** you’re considering (e.g., a screenshot or photo from a listing).  
-Good lighting and centered framing helps the detector and embeddings.
-"""
-
-with gr.Blocks(title="Wardrobe Rater + Recommender",
-               css="""
-               .gradio-container {max-width: 1100px}
-               .small-note {font-size: 0.9em; color: #4b5563}
-               """) as demo:
-    gr.Markdown(APP_DESCRIPTION)
+with gr.Blocks(title="Wardrobe Rater + Recommender", css="""
+.gradio-container {max-width: 1200px}
+""") as demo:
+    gr.Markdown(
+        "# Wardrobe Rater + Recommender\n"
+        "**What this app does:** Scores how likely you are to wear a new item (1–100) and suggests compatible pieces from your wardrobe.  \n"
+        "**How to use it:** (1) Upload a few wardrobe images first. (PNG/JPG only.) (2) Go to *Rate + Recommend* and upload a new item to get a score and matches."
+    )
 
     app_state = gr.State(_blank_state())
 
     with gr.Tab("1) Wardrobe Manager"):
-        gr.Markdown("### Upload your wardrobe")
         with gr.Row():
             wardrobe_uploader = gr.File(
-                label="Upload wardrobe images",
-                file_types=["image"],
-                file_count="multiple"
+                label="Upload wardrobe images (PNG/JPG)",
+                file_types=[".png", ".jpg", ".jpeg"],  # enforce png/jpg
+                file_count="multiple",
+                info="Tip: crop to the item; solid backgrounds work well."
             )
-        gr.Markdown(WARDROBE_TIPS, elem_classes=["small-note"])
         with gr.Row():
-            add_btn = gr.Button("➕ Add to wardrobe (auto-rate + auto-category)", variant="primary")
-            clear_btn = gr.Button("🗑️ Clear wardrobe", variant="secondary")
-
+            add_btn = gr.Button("Add to wardrobe (auto-rate + auto-category)")
+            clear_btn = gr.Button("Clear wardrobe")
+        status_md = gr.Markdown("")  # status / reminders (e.g., wrong file types)
         gallery = gr.Gallery(
-            label="Current wardrobe (click to preview)",
+            label="Current wardrobe",
             columns=6,
             height=220,
-            allow_preview=True,
-            show_label=True
+            object_fit="contain",
+            allow_preview=False
         )
+        ratings_table = gr.Dataframe(headers=["id", "name", "category", "model_rating"], interactive=False)
 
-        gr.Markdown("### Item summary")
-        ratings_table = gr.Dataframe(
-            headers=["id", "name", "category", "model_rating"],
-            interactive=False,
-            wrap=True
-        )
-
-        # Optional Example loader (kept hidden input textbox)
         gr.Markdown("### Or load an example wardrobe")
         example_dir = gr.Textbox(label="Example folder path", value="examples/wardrobe_basic", visible=False)
         gr.Examples(
             examples=[["examples/wardrobe_basic", None]],
             inputs=[example_dir, app_state],
-            outputs=[app_state, gallery, ratings_table],
+            outputs=[app_state, gallery, ratings_table, status_md],
             fn=add_wardrobe_from_dir,
             cache_examples=False,
             run_on_click=True,
         )
 
     with gr.Tab("2) Rate + Recommend New Item"):
-        gr.Markdown("### Upload a candidate item")
         with gr.Row():
             query_img = gr.Image(
-                label="Upload or paste image (single item)",
-                sources=["upload", "webcam", "clipboard"],
-                type="pil"
+                label="Upload or take photo (PNG/JPG)",
+                sources=["upload", "webcam","clipboard"],
+                type="pil",
+                image_mode="RGB",
+                info="Upload a clear photo of the item to score and match."
             )
-            controls_col = gr.Column()
-        with controls_col:
             topk = gr.Slider(1, 6, value=3, step=1, label="# of matches to return")
             matching_mode = gr.Radio(
                 ["Complementary color+style", "Similar style"],
                 value="Complementary color+style",
                 label="Matching mode"
             )
-            go_btn = gr.Button("⭐ Rate + Recommend", variant="primary")
-        gr.Markdown(QUERY_TIPS, elem_classes=["small-note"])
-
+        go_btn = gr.Button("Rate + Recommend")
         with gr.Row():
             pred_score = gr.Number(label="Predicted wear score (1–100)")
-        rec_text = gr.Markdown()  # human-readable summary
-
-        gr.Markdown("### Top matches (click to preview)")
+        rec_text = gr.Markdown()
         rec_gallery = gr.Gallery(
+            label="Matches in your wardrobe",
             columns=6,
             height=220,
-            allow_preview=True,
-            show_label=False
-        )
-
-        gr.Markdown("### Details")
-        rec_table = gr.Dataframe(
-            headers=["rank", "name", "category", "model_rating", "match_score"],
-            interactive=False,
-            wrap=True
+            object_fit="contain",
+            allow_preview=False
         )
 
     # --- Wiring ---
     add_btn.click(
         add_wardrobe,
         inputs=[wardrobe_uploader, app_state],
-        outputs=[app_state, gallery, ratings_table]
+        outputs=[app_state, gallery, ratings_table, status_md]
     )
     clear_btn.click(
         clear_wardrobe,
         inputs=[app_state],
-        outputs=[app_state, gallery, ratings_table]
+        outputs=[app_state, gallery, ratings_table, status_md]
     )
 
     go_btn.click(
         rate_and_recommend,
         inputs=[query_img, topk, matching_mode, app_state],
-        outputs=[pred_score, rec_text, rec_gallery, rec_table]
+        outputs=[pred_score, rec_text, rec_gallery]
     )
 
 # Lightweight tests. Run only when RUN_TESTS=1
@@ -643,14 +654,9 @@ if __name__ == "__main__":
         # Test recommend path with small wardrobe
         st = _blank_state()
         img_b = solid(32)
-        st["wardrobe"].append({
-            "id":0, "name":"test.png", "image":img_b, "thumb":_thumbnail(img_b),
-            "features":_hsv_hist_features(img_b), "embedding":_get_embedder().embed(img_b),
-            "rating":50, "category":"pants"
-        })
-        pred, txt, recs, tbl = rate_and_recommend(solid(200), 1, "Similar style", st)
+        st["wardrobe"].append({"id":0, "name":"test.png", "image":img_b, "features":_hsv_hist_features(img_b), "embedding":_get_embedder().embed(img_b), "rating":50})
+        pred, txt, recs = rate_and_recommend(solid(200), 1, "Similar style", st)
         assert isinstance(pred, int) and isinstance(txt, str) and isinstance(recs, list)
-        assert isinstance(tbl, pd.DataFrame)
         print("Tests passed.")
     else:
         demo.launch()