intial starter

Dockerfile

@@ -0,0 +1,31 @@
+FROM python:3.11-slim
+
+ENV PYTHONDONTWRITEBYTECODE=1 \
+    PYTHONUNBUFFERED=1 \
+    PIP_NO_CACHE_DIR=1 \
+    HF_HUB_ENABLE_HF_TRANSFER=1
+
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    build-essential \
+    git \
+    curl \
+    ca-certificates \
+    libgomp1 \
+    && rm -rf /var/lib/apt/lists/*
+
+WORKDIR /app
+
+COPY recommendation/requirements.txt /app/requirements.txt
+RUN pip install --upgrade pip && pip install -r /app/requirements.txt
+
+COPY recommendation /app/
+
+EXPOSE 8000
+EXPOSE 7860
+
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "8000"]
+
+
+
+
+

README.md

@@ -1,4 +1,3 @@
----
 title: Recomendation
 emoji: 🏆
 colorFrom: purple
@@ -6,7 +5,4 @@ colorTo: green
 sdk: gradio
 sdk_version: 5.44.1
 app_file: app.py
-pinned: false
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+pinned: false

app.py

@@ -0,0 +1,263 @@
+import os
+import base64
+import io
+from typing import List, Optional, Any, Dict
+
+import gradio as gr
+import numpy as np
+import requests
+import torch
+from fastapi import FastAPI, Header, HTTPException
+from pydantic import BaseModel
+from PIL import Image
+from starlette.staticfiles import StaticFiles
+import threading
+
+from inference import InferenceService
+from utils.data_fetch import ensure_dataset_ready
+
+
+AI_API_KEY = os.getenv("AI_API_KEY")
+
+
+def require_api_key(x_api_key: Optional[str]):
+    if AI_API_KEY and x_api_key != AI_API_KEY:
+        raise HTTPException(status_code=401, detail="Invalid API key")
+
+
+class EmbedRequest(BaseModel):
+    image_urls: Optional[List[str]] = None
+    images_base64: Optional[List[str]] = None
+
+
+class Item(BaseModel):
+    id: str
+    embedding: Optional[List[float]] = None
+    category: Optional[str] = None
+    image_url: Optional[str] = None
+
+
+class ComposeRequest(BaseModel):
+    items: List[Item]
+    context: Optional[Dict[str, Any]] = None
+
+
+app = FastAPI(title="Dressify Recommendation Service")
+service = InferenceService()
+
+# Non-blocking bootstrap: fetch data, prepare splits, and train if needed in background
+BOOT_STATUS = "idle"
+
+
+def _background_bootstrap():
+    global BOOT_STATUS
+    try:
+        BOOT_STATUS = "preparing-dataset"
+        ds_root = ensure_dataset_ready()
+        if not ds_root:
+            BOOT_STATUS = "dataset-not-prepared"
+            return
+
+        # Prepare 70/10/10 splits if missing
+        splits_dir = os.path.join(ds_root, "splits")
+        need_prepare = not (
+            os.path.isfile(os.path.join(splits_dir, "train.json")) or
+            os.path.isfile(os.path.join(splits_dir, "outfit_triplets_train.json"))
+        )
+        if need_prepare:
+            BOOT_STATUS = "creating-splits"
+            os.makedirs(splits_dir, exist_ok=True)
+            from scripts.prepare_polyvore import main as prepare_main
+            os.environ.setdefault("PYTHONWARNINGS", "ignore")
+            import sys
+            argv_bak = sys.argv
+            try:
+                sys.argv = ["prepare_polyvore.py", "--root", ds_root, "--random_split"]
+                prepare_main()
+            finally:
+                sys.argv = argv_bak
+
+        # Train if checkpoints are absent
+        export_dir = os.getenv("EXPORT_DIR", "models/exports")
+        os.makedirs(export_dir, exist_ok=True)
+        resnet_ckpt = os.path.join(export_dir, "resnet_item_embedder_best.pth")
+        vit_ckpt = os.path.join(export_dir, "vit_outfit_model_best.pth")
+        import subprocess
+        if not os.path.exists(resnet_ckpt):
+            BOOT_STATUS = "training-resnet"
+            subprocess.run([
+                "python", "train_resnet.py", "--data_root", ds_root, "--epochs", "3",
+                "--out", os.path.join(export_dir, "resnet_item_embedder.pth")
+            ], check=False)
+        if not os.path.exists(vit_ckpt):
+            BOOT_STATUS = "training-vit"
+            subprocess.run([
+                "python", "train_vit_triplet.py", "--data_root", ds_root, "--epochs", "3",
+                "--export", os.path.join(export_dir, "vit_outfit_model.pth")
+            ], check=False)
+        service.reload_models()
+        BOOT_STATUS = "ready"
+    except Exception as e:
+        BOOT_STATUS = f"error: {e}"
+
+
+threading.Thread(target=_background_bootstrap, daemon=True).start()
+
+
+@app.get("/health")
+def health() -> dict:
+    return {"status": "ok", "device": service.device, "resnet": service.resnet_version, "vit": service.vit_version}
+
+
+@app.post("/embed")
+def embed(req: EmbedRequest, x_api_key: Optional[str] = Header(None)) -> dict:
+    require_api_key(x_api_key)
+    images: List[Image.Image] = []
+    if req.image_urls:
+        for url in req.image_urls:
+            resp = requests.get(url, timeout=20)
+            resp.raise_for_status()
+            images.append(Image.open(io.BytesIO(resp.content)).convert("RGB"))
+    if req.images_base64:
+        for b64 in req.images_base64:
+            images.append(Image.open(io.BytesIO(base64.b64decode(b64))).convert("RGB"))
+    if not images:
+        raise HTTPException(status_code=400, detail="No images provided")
+    embs = service.embed_images(images)
+    return {"embeddings": [e.tolist() for e in embs], "model_version": service.resnet_version}
+
+
+@app.post("/compose")
+def compose(req: ComposeRequest, x_api_key: Optional[str] = Header(None)) -> dict:
+    require_api_key(x_api_key)
+    items = [
+        {
+            "id": it.id,
+            "embedding": np.array(it.embedding, dtype=np.float32) if it.embedding is not None else None,
+            "category": it.category,
+            "image_url": it.image_url,
+        }
+        for it in req.items
+    ]
+    outfits = service.compose_outfits(items, context=req.context or {})
+    return {"outfits": outfits, "version": service.vit_version}
+
+
+@app.get("/artifacts")
+def artifacts() -> dict:
+    # list exported model artifacts for download
+    export_dir = os.getenv("EXPORT_DIR", "models/exports")
+    files = []
+    if os.path.isdir(export_dir):
+        for fn in os.listdir(export_dir):
+            if fn.endswith((".pth", ".pt", ".onnx", ".ts", ".json")):
+                files.append({
+                    "name": fn,
+                    "path": f"{export_dir}/{fn}",
+                    "url": f"/files/{fn}",
+                })
+    return {"artifacts": files}
+
+
+# --------- Gradio UI ---------
+
+def _load_images_from_files(files: List[str]) -> List[Image.Image]:
+    images: List[Image.Image] = []
+    for fp in files:
+        try:
+            with Image.open(fp) as im:
+                images.append(im.convert("RGB"))
+        except Exception:
+            continue
+    return images
+
+
+def gradio_embed(files: List[str]):
+    if not files:
+        return "[]"
+    images = _load_images_from_files(files)
+    if not images:
+        return "[]"
+    embs = service.embed_images(images)
+    return str([e.tolist() for e in embs])
+
+
+def gradio_compose(files: List[str], occasion: str, weather: str, num_outfits: int):
+    if not files:
+        return []
+    images = _load_images_from_files(files)
+    if not images:
+        return []
+    embs = service.embed_images(images)
+    items = [
+        {"id": f"item_{i}", "embedding": embs[i], "category": None, "image_url": None}
+        for i in range(len(embs))
+    ]
+    results = service.compose_outfits(items, context={"occasion": occasion, "weather": weather, "num_outfits": int(num_outfits)})
+    # Render as a simple markdown summary
+    lines = []
+    for r in results:
+        lines.append(f"score={r['score']:.3f}, items={r['item_ids']}")
+    return "\n".join(lines)
+
+
+with gr.Blocks() as demo:
+    gr.Markdown("# Dressify Recommendations – HF Test UI")
+    with gr.Tab("Embed"):
+        inp = gr.Files(label="Upload Items (multiple images)")
+        out = gr.Textbox(label="Embeddings (JSON)")
+        btn = gr.Button("Compute Embeddings")
+        btn.click(fn=gradio_embed, inputs=inp, outputs=out)
+    with gr.Tab("Compose"):
+        inp2 = gr.Files(label="Upload Wardrobe (multiple images)")
+        occasion = gr.Dropdown(choices=["casual", "business", "formal", "sport"], value="casual", label="Occasion")
+        weather = gr.Dropdown(choices=["any", "hot", "mild", "cold", "rain"], value="any", label="Weather")
+        num_outfits = gr.Slider(minimum=1, maximum=10, step=1, value=3, label="Num outfits")
+        out2 = gr.Textbox(label="Recommendations")
+        btn2 = gr.Button("Generate")
+        btn2.click(fn=gradio_compose, inputs=[inp2, occasion, weather, num_outfits], outputs=out2)
+    with gr.Tab("Downloads"):
+        gr.Markdown("Download trained artifacts from models/exports")
+        file_list = gr.JSON(label="Artifacts JSON")
+        def list_artifacts_for_ui():
+            export_dir = os.getenv("EXPORT_DIR", "models/exports")
+            files = []
+            if os.path.isdir(export_dir):
+                for fn in os.listdir(export_dir):
+                    if fn.endswith((".pth", ".pt", ".onnx", ".ts", ".json")):
+                        files.append({
+                            "name": fn,
+                            "path": f"{export_dir}/{fn}",
+                            "url": f"/files/{fn}",
+                        })
+            return {"artifacts": files}
+        refresh = gr.Button("Refresh")
+        refresh.click(fn=lambda: list_artifacts_for_ui(), inputs=[], outputs=file_list)
+    with gr.Tab("Status"):
+        gr.Markdown("Startup & training status")
+        status = gr.Textbox(label="Status", value=lambda: BOOT_STATUS)
+        refresh_status = gr.Button("Refresh Status")
+        refresh_status.click(fn=lambda: BOOT_STATUS, inputs=[], outputs=status)
+
+
+try:
+    # Mount Gradio onto FastAPI root path
+    app = gr.mount_gradio_app(app, demo, path="/")
+except Exception:
+    # In case mounting fails in certain runners, we still want FastAPI to be available
+    pass
+
+# Mount static files for direct artifact download
+export_dir = os.getenv("EXPORT_DIR", "models/exports")
+os.makedirs(export_dir, exist_ok=True)
+try:
+    app.mount("/files", StaticFiles(directory=export_dir), name="files")
+except Exception:
+    pass
+
+
+if __name__ == "__main__":
+    # Local testing
+    demo.launch()
+
+

configs/default.yaml

@@ -0,0 +1,33 @@
+# Default training/eval configuration
+
+seed: 42
+image_size: 224
+
+resnet:
+  backbone: resnet50
+  embedding_dim: 512
+  batch_size: 64
+  epochs: 20
+  lr: 0.001
+  weight_decay: 0.0001
+  use_imagenet_pretrained: true
+
+vit:
+  embedding_dim: 512
+  num_layers: 6
+  num_heads: 8
+  ff_multiplier: 4
+  dropout: 0.1
+  batch_size: 64
+  epochs: 30
+  lr: 0.0005
+  weight_decay: 0.05
+
+paths:
+  polyvore_root: /kaggle/input/polyvore-outfits
+  resnet_ckpt: models/exports/resnet_item_embedder.pth
+  vit_ckpt: models/exports/vit_outfit_model.pth
+  export_dir: models/exports
+
+
+

data/__init__.py

@@ -0,0 +1,5 @@
+
+
+
+
+

data/polyvore.py

@@ -0,0 +1,127 @@
+import os
+import json
+from typing import List, Dict, Any, Tuple
+
+import torch
+from torch.utils.data import Dataset
+from PIL import Image
+
+from utils.transforms import build_train_transforms
+
+
+class PolyvoreTripletDataset(Dataset):
+    """
+    Creates (anchor, positive, negative) image triplets for training the ResNet embedder.
+    Assumes a JSON list or multiple files that describe compatible pairs/sets and item image paths.
+
+    Expected structure (customize as needed):
+      root/
+        images/<item_id>.jpg
+        splits/train.json  # [{"anchor": id, "positive": id, "negative": id}, ...]
+    """
+
+    def __init__(self, root: str, split: str = "train", image_size: int = 224) -> None:
+        super().__init__()
+        self.root = root
+        self.split = split
+        self.transforms = build_train_transforms(image_size=image_size)
+        triplet_path = os.path.join(root, "splits", f"{split}.json")
+        if not os.path.exists(triplet_path):
+            raise FileNotFoundError(f"Triplet file not found: {triplet_path}")
+        with open(triplet_path, "r") as f:
+            self.samples: List[Dict[str, Any]] = json.load(f)
+
+    def _load_image(self, item_id: str) -> Image.Image:
+        # Customize if images are arranged differently
+        img_path = os.path.join(self.root, "images", f"{item_id}.jpg")
+        if not os.path.exists(img_path):
+            raise FileNotFoundError(img_path)
+        return Image.open(img_path).convert("RGB")
+
+    def __len__(self) -> int:
+        return len(self.samples)
+
+    def __getitem__(self, idx: int):
+        s = self.samples[idx]
+        a = self._load_image(str(s["anchor"]))
+        p = self._load_image(str(s["positive"]))
+        n = self._load_image(str(s["negative"]))
+        return self.transforms(a), self.transforms(p), self.transforms(n)
+
+
+class PolyvoreOutfitDataset(Dataset):
+    """
+    Produces (tokens, label) where tokens is a sequence of item embeddings or images preprocessed downstream.
+    For simplicity here we return a list of image tensors to be embedded externally or pre-embedded offline.
+
+    Expected structure:
+      root/
+        images/<item_id>.jpg
+        splits/outfits_train.json  # [{"items": [id1,id2,...], "label": 1}, ...]
+    """
+
+    def __init__(self, root: str, split: str = "train", image_size: int = 224) -> None:
+        super().__init__()
+        self.root = root
+        self.split = split
+        self.transforms = build_train_transforms(image_size=image_size)
+        outfit_path = os.path.join(root, "splits", f"outfits_{split}.json")
+        if not os.path.exists(outfit_path):
+            raise FileNotFoundError(f"Outfit file not found: {outfit_path}")
+        with open(outfit_path, "r") as f:
+            self.samples: List[Dict[str, Any]] = json.load(f)
+        # enforce outfit slot constraints: require at least upper, bottom, shoes, accessory if metadata available
+        # If metadata isn't available, we will rely on count >= 3 and let model learn; here, keep as-is.
+
+    def _load_image(self, item_id: str) -> Image.Image:
+        img_path = os.path.join(self.root, "images", f"{item_id}.jpg")
+        if not os.path.exists(img_path):
+            raise FileNotFoundError(img_path)
+        return Image.open(img_path).convert("RGB")
+
+    def __len__(self) -> int:
+        return len(self.samples)
+
+    def __getitem__(self, idx: int):
+        s = self.samples[idx]
+        imgs = [self.transforms(self._load_image(str(i))) for i in s["items"]]
+        label = torch.tensor(int(s.get("label", 1)), dtype=torch.long)
+        # Returns list of tensors; training loop can embed then pack to (N,D)
+        return imgs, label
+
+
+class PolyvoreOutfitTripletDataset(Dataset):
+    """
+    Outfit-level triplets for ViT triplet training: (good1, good2, bad).
+    Expects file `outfit_triplets_<split>.json` with entries:
+      {"good_a": [id...], "good_b": [id...], "bad": [id...]}
+    """
+
+    def __init__(self, root: str, split: str = "train", image_size: int = 224) -> None:
+        super().__init__()
+        self.root = root
+        self.split = split
+        self.transforms = build_train_transforms(image_size=image_size)
+        trip_path = os.path.join(root, "splits", f"outfit_triplets_{split}.json")
+        if not os.path.exists(trip_path):
+            raise FileNotFoundError(f"Outfit triplet file not found: {trip_path}")
+        with open(trip_path, "r") as f:
+            self.samples: List[Dict[str, Any]] = json.load(f)
+
+    def _load_image(self, item_id: str) -> Image.Image:
+        img_path = os.path.join(self.root, "images", f"{item_id}.jpg")
+        if not os.path.exists(img_path):
+            raise FileNotFoundError(img_path)
+        return Image.open(img_path).convert("RGB")
+
+    def __len__(self) -> int:
+        return len(self.samples)
+
+    def __getitem__(self, idx: int):
+        s = self.samples[idx]
+        ga = [self.transforms(self._load_image(str(i))) for i in s["good_a"]]
+        gb = [self.transforms(self._load_image(str(i))) for i in s["good_b"]]
+        bd = [self.transforms(self._load_image(str(i))) for i in s["bad"]]
+        return ga, gb, bd
+
+

inference.py

@@ -0,0 +1,156 @@
+import os
+from typing import List, Dict, Any
+
+import numpy as np
+import torch
+import torch.nn as nn
+from PIL import Image
+
+from utils.transforms import build_inference_transform
+from models.resnet_embedder import ResNetItemEmbedder
+from models.vit_outfit import OutfitCompatibilityModel
+
+
+def _get_device() -> str:
+    if torch.cuda.is_available():
+        return "cuda"
+    if torch.backends.mps.is_available():
+        return "mps"
+    return "cpu"
+
+
+class InferenceService:
+    def __init__(self) -> None:
+        self.device = _get_device()
+        self.transform = build_inference_transform()
+        self.embed_dim = int(os.getenv("EMBED_DIM", "512"))
+        self.resnet_version = "resnet_v1"
+        self.vit_version = "vit_v1"
+
+        self.resnet = self._load_resnet().to(self.device).eval()
+        self.vit = self._load_vit().to(self.device).eval()
+
+        for m in [self.resnet, self.vit]:
+            for p in m.parameters():
+                p.requires_grad_(False)
+
+    def _load_resnet(self) -> nn.Module:
+        strategy = os.getenv("MODEL_LOAD_STRATEGY", "state_dict")
+        ckpt_path = os.getenv("RESNET_CHECKPOINT", "models/exports/resnet_item_embedder.pth")
+        model = ResNetItemEmbedder(embedding_dim=self.embed_dim)
+        if strategy == "random":
+            return model
+        # prefer best if present
+        best_path = os.path.join(os.path.dirname(ckpt_path), "resnet_item_embedder_best.pth")
+        if os.path.exists(best_path):
+            ckpt_to_use = best_path
+        else:
+            ckpt_to_use = ckpt_path
+        if os.path.exists(ckpt_to_use):
+            state = torch.load(ckpt_to_use, map_location="cpu")
+            # accept either full state_dict or {"state_dict": ...}
+            state_dict = state.get("state_dict", state) if isinstance(state, dict) else state
+            missing, unexpected = model.load_state_dict(state_dict, strict=False)
+            if len(unexpected) == 0:
+                return model
+        return model
+
+    def _load_vit(self) -> nn.Module:
+        strategy = os.getenv("MODEL_LOAD_STRATEGY", "state_dict")
+        ckpt_path = os.getenv("VIT_CHECKPOINT", "models/exports/vit_outfit_model.pth")
+        model = OutfitCompatibilityModel(embedding_dim=self.embed_dim)
+        if strategy == "random":
+            return model
+        best_path = os.path.join(os.path.dirname(ckpt_path), "vit_outfit_model_best.pth")
+        ckpt_to_use = best_path if os.path.exists(best_path) else ckpt_path
+        if os.path.exists(ckpt_to_use):
+            state = torch.load(ckpt_to_use, map_location="cpu")
+            state_dict = state.get("state_dict", state) if isinstance(state, dict) else state
+            model.load_state_dict(state_dict, strict=False)
+            return model
+        return model
+
+    def reload_models(self) -> None:
+        """Reload weights from current checkpoint locations (used after background training)."""
+        self.resnet = self._load_resnet().to(self.device).eval()
+        self.vit = self._load_vit().to(self.device).eval()
+        for m in [self.resnet, self.vit]:
+            for p in m.parameters():
+                p.requires_grad_(False)
+
+    @torch.inference_mode()
+    def embed_images(self, images: List[Image.Image]) -> List[np.ndarray]:
+        if len(images) == 0:
+            return []
+        batch = torch.stack([self.transform(img) for img in images]).to(self.device)
+        emb = self.resnet(batch)
+        emb = nn.functional.normalize(emb, dim=-1)
+        return [e.detach().cpu().numpy().astype(np.float32) for e in emb]
+
+    @torch.inference_mode()
+    def compose_outfits(self, items: List[Dict[str, Any]], context: Dict[str, Any]) -> List[Dict[str, Any]]:
+        # Ensure embeddings
+        proc_items: List[Dict[str, Any]] = []
+        for it in items:
+            e = it.get("embedding")
+            if e is None and it.get("image") is not None:
+                # Not used in Gradio path, but kept for completeness
+                emb = self.embed_images([it["image"]])[0]
+            elif e is None:
+                # If missing embedding and image, skip
+                continue
+            else:
+                emb = np.asarray(e, dtype=np.float32)
+            proc_items.append({"id": it.get("id"), "embedding": emb, "category": it.get("category")})
+
+        if len(proc_items) < 2:
+            return []
+
+        # Candidate generation: enforce minimum slots (upper, bottom, shoes, accessory) if categories provided
+        rng = np.random.default_rng(42)
+        num_outfits = int(context.get("num_outfits", 3))
+        min_size, max_size = 3, 5
+        candidates: List[List[int]] = []
+        ids = list(range(len(proc_items)))
+        # slot-aware sampling if categories exist
+        def has_cat(i: int, cat_prefix: str) -> bool:
+            c = (proc_items[i].get("category") or "").lower()
+            return cat_prefix in c
+
+        uppers = [i for i in ids if any(k in (proc_items[i].get("category") or "").lower() for k in ["top", "shirt", "tshirt", "blouse", "jacket", "hoodie"]) ]
+        bottoms = [i for i in ids if any(k in (proc_items[i].get("category") or "").lower() for k in ["pant", "trouser", "jean", "skirt", "short"]) ]
+        shoes = [i for i in ids if "shoe" in (proc_items[i].get("category") or "").lower()]
+        accs = [i for i in ids if any(k in (proc_items[i].get("category") or "").lower() for k in ["watch", "belt", "ring", "bracelet", "accessor"]) ]
+
+        for _ in range(num_outfits * 10):
+            if uppers and bottoms and shoes and accs:
+                subset = [rng.choice(uppers).item(), rng.choice(bottoms).item(), rng.choice(shoes).item(), rng.choice(accs).item()]
+                # optional: add one more random
+                remain = list(set(ids) - set(subset))
+                if remain and rng.random() < 0.5:
+                    subset.append(rng.choice(remain).item())
+            else:
+                k = rng.integers(min_size, max_size + 1)
+                subset = rng.choice(ids, size=int(k), replace=False).tolist()
+            candidates.append(subset)
+
+        # Score using ViT
+        def score_subset(idx_subset: List[int]) -> float:
+            embs = torch.tensor(np.stack([proc_items[i]["embedding"] for i in idx_subset]), dtype=torch.float32, device=self.device)
+            embs = embs.unsqueeze(0)  # (1, N, D)
+            s = self.vit.score_compatibility(embs).item()
+            return float(s)
+
+        scored = [(subset, score_subset(subset)) for subset in candidates]
+        scored.sort(key=lambda x: x[1], reverse=True)
+        topk = scored[:num_outfits]
+        results = [
+            {
+                "item_ids": [proc_items[i]["id"] for i in subset],
+                "score": float(score),
+            }
+            for subset, score in topk
+        ]
+        return results
+
+

models/__init__.py

@@ -0,0 +1,5 @@
+
+
+
+
+

models/resnet_embedder.py

@@ -0,0 +1,35 @@
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torchvision.models as tvm
+
+
+class ResNetItemEmbedder(nn.Module):
+    def __init__(self, embedding_dim: int = 512, backbone: str = "resnet50", pretrained: bool = True) -> None:
+        super().__init__()
+        if backbone == "resnet50":
+            model = tvm.resnet50(weights=tvm.ResNet50_Weights.DEFAULT if pretrained else None)
+            feat_dim = 2048
+        elif backbone == "resnet101":
+            model = tvm.resnet101(weights=tvm.ResNet101_Weights.DEFAULT if pretrained else None)
+            feat_dim = 2048
+        else:
+            raise ValueError(f"Unsupported backbone: {backbone}")
+
+        # Remove classifier, keep global average pooling output
+        modules = list(model.children())[:-1]  # drop fc
+        self.backbone = nn.Sequential(*modules)
+        self.proj = nn.Linear(feat_dim, embedding_dim)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # x: (B, 3, H, W)
+        feats = self.backbone(x)  # (B, C, 1, 1)
+        feats = feats.flatten(1)  # (B, C)
+        emb = self.proj(feats)    # (B, D)
+        return emb
+
+
+
+
+

models/vit_outfit.py

@@ -0,0 +1,40 @@
+from typing import Optional
+
+import torch
+import torch.nn as nn
+
+
+class OutfitCompatibilityModel(nn.Module):
+    def __init__(self, embedding_dim: int = 512, num_layers: int = 4, num_heads: int = 8, ff_multiplier: int = 4, dropout: float = 0.1) -> None:
+        super().__init__()
+        encoder_layer = nn.TransformerEncoderLayer(
+            d_model=embedding_dim,
+            nhead=num_heads,
+            dim_feedforward=ff_multiplier * embedding_dim,
+            dropout=dropout,
+            batch_first=True,
+            activation="gelu",
+            norm_first=True,
+        )
+        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
+        self.compatibility_head = nn.Sequential(
+            nn.LayerNorm(embedding_dim),
+            nn.Linear(embedding_dim, embedding_dim // 2),
+            nn.GELU(),
+            nn.Linear(embedding_dim // 2, 1),
+        )
+
+    def forward(self, tokens: torch.Tensor) -> torch.Tensor:
+        # tokens: (B, N, D) sequence of item embeddings
+        h = self.encoder(tokens)  # (B, N, D)
+        pooled = h.mean(dim=1)    # (B, D)
+        score = self.compatibility_head(pooled)  # (B, 1)
+        return score.squeeze(-1)  # (B,)
+
+    def score_compatibility(self, tokens: torch.Tensor) -> torch.Tensor:
+        return self.forward(tokens)
+
+
+
+
+

requirements.txt

@@ -0,0 +1,17 @@
+fastapi==0.111.0
+uvicorn[standard]==0.30.1
+gradio==5.44.1
+pydantic==2.8.2
+python-multipart==0.0.9
+Pillow==10.4.0
+numpy==2.0.1
+torch==2.3.1
+torchvision==0.18.1
+timm==1.0.8
+requests==2.32.3
+scikit-learn==1.5.1
+faiss-cpu==1.8.0.post1
+pyyaml==6.0.2
+open-clip-torch==2.26.1
+huggingface-hub==0.24.5
+kagglehub[pandas-datasets]

scripts/prepare_polyvore.py

@@ -0,0 +1,187 @@
+import os
+import json
+import random
+import argparse
+from pathlib import Path
+from typing import Dict, Any, List, Set
+
+
+def load_outfits_json(root: str, split: str) -> List[Dict[str, Any]]:
+    candidates = [
+        os.path.join(root, f"{split}.json"),
+        os.path.join(root, f"{split}_no_dup.json"),
+        os.path.join(root, "splits", f"{split}.json"),
+        os.path.join(root, "splits", f"{split}_no_dup.json"),
+    ]
+    for p in candidates:
+        if os.path.exists(p):
+            with open(p, "r") as f:
+                data = json.load(f)
+            # Expect list where each item has key "items" listing item ids
+            return data
+    raise FileNotFoundError(f"Could not find {split}.json in {root} or {root}/splits")
+
+
+def try_load_any_outfits(root: str) -> List[Dict[str, Any]]:
+    candidates = [
+        os.path.join(root, "outfits.json"),
+        os.path.join(root, "all.json"),
+        os.path.join(root, "data.json"),
+    ]
+    for p in candidates:
+        if os.path.exists(p):
+            with open(p, "r") as f:
+                return json.load(f)
+    # As a last resort, merge available splits
+    merged: List[Dict[str, Any]] = []
+    for sp in ["train", "valid", "test"]:
+        try:
+            merged.extend(load_outfits_json(root, sp))
+        except FileNotFoundError:
+            continue
+    return merged
+
+
+def collect_all_items(outfits: List[Dict[str, Any]]) -> List[str]:
+    s: Set[str] = set()
+    for o in outfits:
+        for it in o.get("items", []):
+            s.add(str(it))
+    return sorted(s)
+
+
+def build_triplets(outfits: List[Dict[str, Any]], all_items: List[str], max_triplets: int = 200000) -> List[Dict[str, str]]:
+    rng = random.Random(42)
+    all_items_set = set(all_items)
+    triplets: List[Dict[str, str]] = []
+    for o in outfits:
+        items = [str(i) for i in o.get("items", [])]
+        if len(items) < 2:
+            continue
+        local_set = set(items)
+        for i in range(len(items) - 1):
+            a = items[i]
+            p = items[i + 1]
+            # pick a negative not in this outfit
+            negatives = list(all_items_set - local_set)
+            if not negatives:
+                continue
+            n = rng.choice(negatives)
+            triplets.append({"anchor": a, "positive": p, "negative": n})
+            if len(triplets) >= max_triplets:
+                return triplets
+    return triplets
+
+
+def build_outfit_pairs(outfits: List[Dict[str, Any]], num_negatives_per_pos: int = 1) -> List[Dict[str, Any]]:
+    rng = random.Random(123)
+    all_items = collect_all_items(outfits)
+    all_set = set(all_items)
+    pairs: List[Dict[str, Any]] = []
+    # Positive samples
+    for o in outfits:
+        items = [str(i) for i in o.get("items", [])]
+        if len(items) < 2:
+            continue
+        pairs.append({"items": items, "label": 1})
+        # Negative by corrupting one item
+        for _ in range(num_negatives_per_pos):
+            if not items:
+                continue
+            idx = rng.randrange(len(items))
+            neg_pool = list(all_set - set(items))
+            if not neg_pool:
+                continue
+            neg_item = rng.choice(neg_pool)
+            neg_items = items.copy()
+            neg_items[idx] = neg_item
+            pairs.append({"items": neg_items, "label": 0})
+    return pairs
+
+
+def build_outfit_triplets(outfits: List[Dict[str, Any]], num_triplets: int = 200000) -> List[Dict[str, Any]]:
+    rng = random.Random(999)
+    # Collect only valid positive outfits (len >= 3 or ideally slot-complete)
+    pos = [o for o in outfits if len(o.get("items", [])) >= 3]
+    all_items = collect_all_items(outfits)
+    all_set = set(all_items)
+    triplets: List[Dict[str, Any]] = []
+    for _ in range(num_triplets):
+        if len(pos) < 2:
+            break
+        ga = rng.choice(pos)
+        gb = rng.choice(pos)
+        # Ensure ga != gb
+        if ga is gb:
+            continue
+        # Create bad by corrupting one item in ga
+        items_ga = [str(i) for i in ga.get("items", [])]
+        if not items_ga:
+            continue
+        corrupt_idx = rng.randrange(len(items_ga))
+        neg_pool = list(all_set - set(items_ga))
+        if not neg_pool:
+            continue
+        neg_item = rng.choice(neg_pool)
+        bad = items_ga.copy()
+        bad[corrupt_idx] = neg_item
+        triplets.append({"good_a": items_ga, "good_b": [str(i) for i in gb.get("items", [])], "bad": bad})
+    return triplets
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--root", type=str, required=True, help="Polyvore dataset root")
+    ap.add_argument("--out", type=str, default=None, help="Output directory for splits (default: <root>/splits)")
+    ap.add_argument("--max_triplets", type=int, default=200000)
+    ap.add_argument("--neg_per_pos", type=int, default=1)
+    ap.add_argument("--random_split", action="store_true", help="Create 70/10/10 random split if official splits are missing")
+    args = ap.parse_args()
+
+    out_dir = args.out or os.path.join(args.root, "splits")
+    Path(out_dir).mkdir(parents=True, exist_ok=True)
+
+    if args.random_split:
+        all_outfits = try_load_any_outfits(args.root)
+        if not all_outfits:
+            raise FileNotFoundError("No outfits found to split. Provide official splits or an outfits.json file.")
+        rng = random.Random(2024)
+        rng.shuffle(all_outfits)
+        n = len(all_outfits)
+        n_train = int(0.7 * n)
+        n_valid = int(0.1 * n)
+        splits = {
+            "train": all_outfits[:n_train],
+            "valid": all_outfits[n_train:n_train + n_valid],
+            "test": all_outfits[n_train + n_valid:],
+        }
+    else:
+        splits = {}
+        for split in ["train", "valid", "test"]:
+            try:
+                splits[split] = load_outfits_json(args.root, split)
+            except FileNotFoundError as e:
+                print(f"Skipping {split}: {e}")
+                splits[split] = []
+
+    for split, outfits in splits.items():
+        if not outfits:
+            continue
+        all_items = collect_all_items(outfits)
+        triplets = build_triplets(outfits, all_items, max_triplets=args.max_triplets)
+        pairs = build_outfit_pairs(outfits, num_negatives_per_pos=args.neg_per_pos)
+
+        with open(os.path.join(out_dir, f"{split}.json"), "w") as f:
+            json.dump(triplets, f)
+        with open(os.path.join(out_dir, f"outfits_{split}.json"), "w") as f:
+            json.dump(pairs, f)
+        triplets_o = build_outfit_triplets(outfits)
+        with open(os.path.join(out_dir, f"outfit_triplets_{split}.json"), "w") as f:
+            json.dump(triplets_o, f)
+        print(f"Wrote {split}: {len(triplets)} item-triplets, {len(pairs)} outfit-pairs, {len(triplets_o)} outfit-triplets -> {out_dir}")
+
+
+if __name__ == "__main__":
+    main()
+
+

train_resnet.py

@@ -0,0 +1,80 @@
+import os
+import argparse
+from typing import Tuple
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import DataLoader
+
+from data.polyvore import PolyvoreTripletDataset
+from models.resnet_embedder import ResNetItemEmbedder
+from utils.export import ensure_export_dir
+import json
+
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser()
+    p.add_argument("--data_root", type=str, default=os.getenv("POLYVORE_ROOT", "/kaggle/input/polyvore-outfits"))
+    p.add_argument("--epochs", type=int, default=20)
+    p.add_argument("--batch_size", type=int, default=64)
+    p.add_argument("--lr", type=float, default=1e-3)
+    p.add_argument("--embedding_dim", type=int, default=512)
+    p.add_argument("--out", type=str, default="models/exports/resnet_item_embedder.pth")
+    return p.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    device = "cuda" if torch.cuda.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu")
+
+    dataset = PolyvoreTripletDataset(args.data_root, split="train")
+
+    loader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True, num_workers=4, pin_memory=True)
+    model = ResNetItemEmbedder(embedding_dim=args.embedding_dim).to(device)
+    optimizer = optim.AdamW(model.parameters(), lr=args.lr, weight_decay=1e-4)
+    criterion = nn.TripletMarginLoss(margin=0.2, p=2)
+
+    export_dir = ensure_export_dir(os.path.dirname(args.out) or "models/exports")
+    best_loss = float("inf")
+    history = []
+    for epoch in range(args.epochs):
+        model.train()
+        running = 0.0
+        steps = 0
+        for batch in loader:
+            # Expect batch as (anchor, positive, negative)
+            anchor, positive, negative = batch
+            anchor = anchor.to(device)
+            positive = positive.to(device)
+            negative = negative.to(device)
+            emb_a = model(anchor)
+            emb_p = model(positive)
+            emb_n = model(negative)
+            loss = criterion(emb_a, emb_p, emb_n)
+            optimizer.zero_grad(set_to_none=True)
+            loss.backward()
+            optimizer.step()
+            running += loss.item()
+            steps += 1
+        avg_loss = running / max(1, steps)
+        out_path = args.out
+        if not out_path.startswith("models/"):
+            out_path = os.path.join(export_dir, os.path.basename(args.out))
+        torch.save({"state_dict": model.state_dict()}, out_path)
+        print(f"Epoch {epoch+1}/{args.epochs} avg_triplet_loss={avg_loss:.4f} saved -> {out_path}")
+        history.append({"epoch": epoch + 1, "avg_triplet_loss": avg_loss})
+        if avg_loss < best_loss:
+            best_loss = avg_loss
+            torch.save({"state_dict": model.state_dict()}, os.path.join(export_dir, "resnet_item_embedder_best.pth"))
+
+    # write metrics
+    metrics_path = os.path.join(export_dir, "resnet_metrics.json")
+    with open(metrics_path, "w") as f:
+        json.dump({"best_triplet_loss": best_loss, "history": history}, f)
+
+
+if __name__ == "__main__":
+    main()
+
+

train_vit_triplet.py

@@ -0,0 +1,145 @@
+import os
+import argparse
+from typing import List
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import DataLoader
+
+from data.polyvore import PolyvoreOutfitTripletDataset
+from models.vit_outfit import OutfitCompatibilityModel
+from models.resnet_embedder import ResNetItemEmbedder
+from utils.export import ensure_export_dir
+import json
+
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser()
+    p.add_argument("--data_root", type=str, default=os.getenv("POLYVORE_ROOT", "/kaggle/input/polyvore-outfits"))
+    p.add_argument("--epochs", type=int, default=30)
+    p.add_argument("--batch_size", type=int, default=32)
+    p.add_argument("--lr", type=float, default=5e-4)
+    p.add_argument("--embedding_dim", type=int, default=512)
+    p.add_argument("--triplet_margin", type=float, default=0.3)
+    p.add_argument("--export", type=str, default="models/exports/vit_outfit_model.pth")
+    p.add_argument("--eval_every", type=int, default=1)
+    return p.parse_args()
+
+
+def embed_outfit(imgs: List[torch.Tensor], embedder: ResNetItemEmbedder, device: str, max_len: int = 4) -> torch.Tensor:
+    if len(imgs) == 0:
+        return torch.zeros((max_len, embedder.proj.out_features), device=device)
+    k = min(len(imgs), max_len)
+    x = torch.stack(imgs[:k], dim=0).to(device)
+    with torch.no_grad():
+        e = embedder(x)  # (k, D)
+    tokens = torch.zeros((max_len, e.shape[-1]), device=device)
+    tokens[:k] = e
+    return tokens
+
+
+def main() -> None:
+    args = parse_args()
+    device = "cuda" if torch.cuda.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu")
+
+    dataset = PolyvoreOutfitTripletDataset(args.data_root, split="train")
+
+    def collate(batch):
+        return batch  # variable length handled inside training loop
+
+    loader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True, num_workers=4, pin_memory=True, collate_fn=collate)
+
+    model = OutfitCompatibilityModel(embedding_dim=args.embedding_dim).to(device)
+    embedder = ResNetItemEmbedder(embedding_dim=args.embedding_dim).to(device).eval()
+    for p in embedder.parameters():
+        p.requires_grad_(False)
+
+    optimizer = optim.AdamW(model.parameters(), lr=args.lr, weight_decay=5e-2)
+    triplet = nn.TripletMarginWithDistanceLoss(distance_function=lambda x, y: 1 - nn.functional.cosine_similarity(x, y), margin=args.triplet_margin)
+
+    export_dir = ensure_export_dir(os.path.dirname(args.export) or "models/exports")
+    best_loss = float("inf")
+    hist = []
+    for epoch in range(args.epochs):
+        model.train()
+        for batch in loader:
+            # batch: List[(ga_imgs, gb_imgs, bd_imgs)]
+            anchor_tokens = []
+            positive_tokens = []
+            negative_tokens = []
+            for ga, gb, bd in batch:
+                ta = embed_outfit(ga, embedder, device)
+                tb = embed_outfit(gb, embedder, device)
+                tn = embed_outfit(bd, embedder, device)
+                anchor_tokens.append(ta.unsqueeze(0))
+                positive_tokens.append(tb.unsqueeze(0))
+                negative_tokens.append(tn.unsqueeze(0))
+            A = torch.cat(anchor_tokens, dim=0)  # (B, N, D)
+            P = torch.cat(positive_tokens, dim=0)
+            N = torch.cat(negative_tokens, dim=0)
+
+            # get outfit-level embeddings via ViT encoder pooled output
+            ea = model.encoder(A).mean(dim=1)
+            ep = model.encoder(P).mean(dim=1)
+            en = model.encoder(N).mean(dim=1)
+            loss = triplet(ea, ep, en)
+            optimizer.zero_grad(set_to_none=True)
+            loss.backward()
+            optimizer.step()
+
+        # Simple validation using a subset of training data as a proxy if no val split here
+        # For true 70/10/10, prepare_polyvore.py will create outfit_triplets_valid.json
+        val_path = os.path.join(args.data_root, "splits", "outfit_triplets_valid.json")
+        val_loss = None
+        if os.path.exists(val_path) and (epoch + 1) % args.eval_every == 0:
+            val_ds = PolyvoreOutfitTripletDataset(args.data_root, split="valid")
+            val_loader = DataLoader(val_ds, batch_size=args.batch_size, shuffle=False, num_workers=2, collate_fn=lambda x: x)
+            model.eval()
+            losses = []
+            with torch.no_grad():
+                for vbatch in val_loader:
+                    anchor_tokens = []
+                    positive_tokens = []
+                    negative_tokens = []
+                    for ga, gb, bd in vbatch:
+                        ta = embed_outfit(ga, embedder, device)
+                        tb = embed_outfit(gb, embedder, device)
+                        tn = embed_outfit(bd, embedder, device)
+                        anchor_tokens.append(ta.unsqueeze(0))
+                        positive_tokens.append(tb.unsqueeze(0))
+                        negative_tokens.append(tn.unsqueeze(0))
+                    A = torch.cat(anchor_tokens, dim=0)
+                    P = torch.cat(positive_tokens, dim=0)
+                    N = torch.cat(negative_tokens, dim=0)
+                    ea = model.encoder(A).mean(dim=1)
+                    ep = model.encoder(P).mean(dim=1)
+                    en = model.encoder(N).mean(dim=1)
+                    l = triplet(ea, ep, en).item()
+                    losses.append(l)
+            val_loss = sum(losses) / max(1, len(losses))
+
+        out_path = args.export
+        if not out_path.startswith("models/"):
+            out_path = os.path.join(export_dir, os.path.basename(args.export))
+        torch.save({"state_dict": model.state_dict()}, out_path)
+        if val_loss is not None:
+            print(f"Epoch {epoch+1}/{args.epochs} triplet_loss={loss.item():.4f} val_triplet_loss={val_loss:.4f} saved -> {out_path}")
+            hist.append({"epoch": epoch + 1, "triplet_loss": float(loss.item()), "val_triplet_loss": float(val_loss)})
+            if val_loss < best_loss:
+                best_loss = val_loss
+                torch.save({"state_dict": model.state_dict()}, os.path.join(export_dir, "vit_outfit_model_best.pth"))
+        else:
+            print(f"Epoch {epoch+1}/{args.epochs} triplet_loss={loss.item():.4f} saved -> {out_path}")
+            hist.append({"epoch": epoch + 1, "triplet_loss": float(loss.item())})
+
+    metrics_path = os.path.join(export_dir, "vit_metrics.json")
+    payload = {"best_val_triplet_loss": best_loss if best_loss != float("inf") else None, "history": hist}
+    with open(metrics_path, "w") as f:
+        json.dump(payload, f)
+
+
+if __name__ == "__main__":
+    main()
+
+

utils/__init__.py

@@ -0,0 +1,5 @@
+
+
+
+
+

utils/data_fetch.py

@@ -0,0 +1,148 @@
+import os
+import shutil
+import zipfile
+from pathlib import Path
+from typing import Optional, List
+
+import requests
+
+try:
+    from huggingface_hub import snapshot_download  # type: ignore
+except Exception:  # pragma: no cover
+    snapshot_download = None
+
+try:
+    import kagglehub  # type: ignore
+    from kagglehub import KaggleDatasetAdapter  # type: ignore
+except Exception:  # pragma: no cover
+    kagglehub = None
+    KaggleDatasetAdapter = None
+
+
+def _download_zip(url: str, dest_dir: str) -> str:
+    os.makedirs(dest_dir, exist_ok=True)
+    local_zip = os.path.join(dest_dir, "dataset.zip")
+    with requests.get(url, stream=True, timeout=60) as r:
+        r.raise_for_status()
+        with open(local_zip, "wb") as f:
+            for chunk in r.iter_content(chunk_size=1024 * 1024):
+                if chunk:
+                    f.write(chunk)
+    with zipfile.ZipFile(local_zip, "r") as zf:
+        zf.extractall(dest_dir)
+    os.remove(local_zip)
+    return dest_dir
+
+
+def _unzip_inner_archives(root: str) -> None:
+    """Find and extract any zip files inside root (e.g., images.zip)."""
+    for dirpath, _dirnames, filenames in os.walk(root):
+        for fn in filenames:
+            if fn.lower().endswith(".zip"):
+                zpath = os.path.join(dirpath, fn)
+                try:
+                    with zipfile.ZipFile(zpath, "r") as zf:
+                        zf.extractall(dirpath)
+                    # keep original zip to avoid repeated work? remove to save disk
+                    try:
+                        os.remove(zpath)
+                    except Exception:
+                        pass
+                except Exception as e:  # pragma: no cover
+                    print(f"Failed to unzip inner archive {zpath}: {e}")
+
+
+def _ensure_images_dir(root: str) -> None:
+    """Ensure a stable images/ path exists under root. Create a symlink if needed."""
+    images_root = os.path.join(root, "images")
+    if os.path.isdir(images_root):
+        return
+    # Try to find a folder with many jpg/png files
+    candidate_dirs: List[str] = []
+    for dirpath, dirnames, filenames in os.walk(root):
+        if dirpath == root:
+            # skip root level files, look deeper
+            continue
+        img_files = [f for f in filenames if f.lower().endswith((".jpg", ".jpeg", ".png"))]
+        if len(img_files) > 1000:  # heuristic: big image folder
+            candidate_dirs.append(dirpath)
+    # Prefer the shallowest candidate
+    candidate_dirs.sort(key=lambda p: len(Path(p).parts))
+    if candidate_dirs:
+        src = candidate_dirs[0]
+        try:
+            os.symlink(src, images_root)
+            print(f"Created images symlink: {images_root} -> {src}")
+        except Exception:
+            # fallback: create folder and leave it empty (training will fail fast if missing)
+            os.makedirs(images_root, exist_ok=True)
+    else:
+        os.makedirs(images_root, exist_ok=True)
+
+
+def ensure_dataset_ready() -> Optional[str]:
+    """
+    Ensure Polyvore dataset is present locally.
+    Priority:
+    1) If POLYVORE_ROOT exists and has splits, return it
+    2) Try Hugging Face dataset repo (defaults to Stylique/Polyvore if not set)
+    3) If DATA_ZIP_URL is set, download and unzip
+    4) Try KaggleHub (best-effort)
+    Returns resolved root path or None if nothing done.
+    """
+    root = os.getenv("POLYVORE_ROOT", "./data/Polyvore")
+    auto_fetch = os.getenv("AUTO_FETCH_DATA", "true").lower() == "true"
+    Path(root).mkdir(parents=True, exist_ok=True)
+
+    # Already prepared?
+    if os.path.isdir(os.path.join(root, "splits")):
+        _unzip_inner_archives(root)
+        _ensure_images_dir(root)
+        return root
+    if not auto_fetch:
+        return None
+
+    # Try HF dataset repo
+    repo = os.getenv("HF_DATASET_REPO", "Stylique/Polyvore")
+    if repo and snapshot_download is not None:
+        try:
+            snapshot_download(repo, repo_type="dataset", local_dir=root)
+            _unzip_inner_archives(root)
+            _ensure_images_dir(root)
+            # If splits not provided, they'll be prepared by the caller
+            return root
+        except Exception as e:  # pragma: no cover
+            print(f"HF dataset download failed: {e}")
+
+    # Try ZIP URL
+    zip_url = os.getenv("DATA_ZIP_URL")
+    if zip_url:
+        try:
+            _download_zip(zip_url, root)
+            _unzip_inner_archives(root)
+            _ensure_images_dir(root)
+        except Exception as e:  # pragma: no cover
+            print(f"ZIP download failed: {e}")
+            return None
+
+    # Try KaggleHub (no Kaggle keys required for public datasets)
+    if kagglehub is not None and KaggleDatasetAdapter is not None:
+        try:
+            # Attempt to load core file to trigger dataset download locally
+            # User can override POLYVORE_FILE_PATH to select a specific CSV/JSON
+            file_path = os.getenv("POLYVORE_FILE_PATH", "")
+            kagglehub.load_dataset(
+                KaggleDatasetAdapter.PANDAS,
+                "dnepozitek/polyvore-outfits",
+                file_path,
+            )
+            # KaggleHub stores under ~/.cache/kagglehub/datasets/<slug>/...; copy to root if needed
+            # For simplicity, assume user will run prepare script using POLYVORE_ROOT pointing to extracted images
+            _unzip_inner_archives(root)
+            _ensure_images_dir(root)
+        except Exception as e:  # pragma: no cover
+            print(f"KaggleHub download failed: {e}")
+
+    return root
+
+

utils/export.py

@@ -0,0 +1,37 @@
+import os
+from typing import Optional
+
+import torch
+
+
+def ensure_export_dir(path: str) -> str:
+    os.makedirs(path, exist_ok=True)
+    return path
+
+
+def export_torchscript(model: torch.nn.Module, example_inputs: torch.Tensor, out_path: str) -> str:
+    model.eval()
+    traced = torch.jit.trace(model, example_inputs)
+    torch.jit.save(traced, out_path)
+    return out_path
+
+
+def export_onnx(model: torch.nn.Module, example_inputs: torch.Tensor, out_path: str, opset: int = 17) -> str:
+    model.eval()
+    torch.onnx.export(
+        model,
+        example_inputs,
+        out_path,
+        export_params=True,
+        opset_version=opset,
+        do_constant_folding=True,
+        input_names=["input"],
+        output_names=["output"],
+        dynamic_axes={"input": {0: "batch"}, "output": {0: "batch"}},
+    )
+    return out_path
+
+
+
+
+

utils/transforms.py

@@ -0,0 +1,31 @@
+from typing import Tuple
+
+import torchvision.transforms as T
+
+
+def build_inference_transform(image_size: int = 224) -> T.Compose:
+    return T.Compose(
+        [
+            T.Resize((image_size, image_size), interpolation=T.InterpolationMode.BICUBIC),
+            T.ToTensor(),
+            T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+        ]
+    )
+
+
+def build_train_transforms(image_size: int = 224) -> T.Compose:
+    return T.Compose(
+        [
+            T.Resize(int(image_size * 1.1), interpolation=T.InterpolationMode.BICUBIC),
+            T.RandomResizedCrop(image_size, scale=(0.7, 1.0), ratio=(0.75, 1.33)),
+            T.RandomHorizontalFlip(),
+            T.ColorJitter(0.2, 0.2, 0.2, 0.1),
+            T.ToTensor(),
+            T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+        ]
+    )
+
+
+
+
+