more fix

data/polyvore.py

@@ -118,9 +118,7 @@ class PolyvoreOutfitTripletDataset(Dataset):
             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)
+        img_path = PolyvoreTripletDataset._find_image_path(self, item_id)
         return Image.open(img_path).convert("RGB")
 
     def __len__(self) -> int:

scripts/prepare_polyvore.py

@@ -3,7 +3,79 @@ import json
 import random
 import argparse
 from pathlib import Path
-from typing import Dict, Any, List, Set
+from typing import Dict, Any, List, Set, Union
+
+
+def _normalize_outfits(obj: Union[List[Any], Dict[str, Any]]) -> List[Dict[str, Any]]:
+    """Normalize various Polyvore JSON formats into a list of {"items": [id,...]} dicts.
+
+    Accepts:
+    - List of objects where each object may be:
+      - {"items": [id,...]} already
+      - {"items": [{"item_id": id}...]} (extract item_id or id)
+      - {"set_id": ..., "items": [...]}
+      - List of ids directly
+    - Dict mapping outfit_id -> list of item ids or an object with items.
+    """
+    result: List[Dict[str, Any]] = []
+    if isinstance(obj, dict):
+        # values could be list of ids or dicts with items
+        values = list(obj.values())
+        for v in values:
+            if isinstance(v, list):
+                # list of ids or list of dicts
+                if len(v) > 0 and isinstance(v[0], dict):
+                    items = []
+                    for it in v:
+                        if isinstance(it, dict):
+                            iid = it.get("item_id") or it.get("id") or it.get("itemId")
+                            if iid is not None:
+                                items.append(str(iid))
+                    if items:
+                        result.append({"items": items})
+                else:
+                    result.append({"items": [str(x) for x in v]})
+            elif isinstance(v, dict):
+                if "items" in v:
+                    itm = v["items"]
+                    if isinstance(itm, list):
+                        if itm and isinstance(itm[0], dict):
+                            items = []
+                            for it in itm:
+                                iid = it.get("item_id") or it.get("id") or it.get("itemId")
+                                if iid is not None:
+                                    items.append(str(iid))
+                            if items:
+                                result.append({"items": items})
+                        else:
+                            result.append({"items": [str(x) for x in itm]})
+        return result
+    if isinstance(obj, list):
+        for e in obj:
+            if isinstance(e, dict):
+                if "items" in e:
+                    itm = e["items"]
+                    if isinstance(itm, list):
+                        if itm and isinstance(itm[0], dict):
+                            items = []
+                            for it in itm:
+                                iid = it.get("item_id") or it.get("id") or it.get("itemId")
+                                if iid is not None:
+                                    items.append(str(iid))
+                            if items:
+                                result.append({"items": items})
+                        else:
+                            result.append({"items": [str(x) for x in itm]})
+                else:
+                    # some variants use different key names but include list of item ids
+                    for k in ("good", "outfit", "products"):
+                        if k in e and isinstance(e[k], list):
+                            result.append({"items": [str(x) for x in e[k]]})
+                            break
+            elif isinstance(e, list):
+                result.append({"items": [str(x) for x in e]})
+        return result
+    return result
 
 
 def load_outfits_json(root: str, split: str) -> List[Dict[str, Any]]:
@@ -16,30 +88,33 @@ def load_outfits_json(root: str, split: str) -> List[Dict[str, Any]]:
     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")
+                raw = json.load(f)
+            data = _normalize_outfits(raw)
+            if data:
+                return data
+    raise FileNotFoundError(f"Could not find usable {split} split 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
+    # Prefer official splits if present
     merged: List[Dict[str, Any]] = []
     for sp in ["train", "valid", "test"]:
         try:
             merged.extend(load_outfits_json(root, sp))
         except FileNotFoundError:
             continue
-    return merged
+    if merged:
+        return merged
+    # Fallback: common aggregated files
+    for name in ("outfits.json", "all.json", "data.json"):
+        p = os.path.join(root, name)
+        if os.path.exists(p):
+            with open(p, "r") as f:
+                raw = json.load(f)
+            data = _normalize_outfits(raw)
+            if data:
+                return data
+    return []
 
 
 def collect_all_items(outfits: List[Dict[str, Any]]) -> List[str]:
@@ -141,7 +216,20 @@ def main() -> None:
     out_dir = args.out or os.path.join(args.root, "splits")
     Path(out_dir).mkdir(parents=True, exist_ok=True)
 
-    if args.random_split:
+    # Prefer official splits; if missing, optionally create random split
+    splits = {}
+    found_any_official = False
+    for split in ["train", "valid", "test"]:
+        try:
+            data = load_outfits_json(args.root, split)
+            splits[split] = data
+            if data:
+                found_any_official = True
+        except FileNotFoundError as e:
+            print(f"Skipping {split}: {e}")
+            splits[split] = []
+
+    if not found_any_official and 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.")
@@ -155,14 +243,6 @@ def main() -> None:
             "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:

train_resnet.py

@@ -30,6 +30,35 @@ def main() -> None:
     if device == "cuda":
         torch.backends.cudnn.benchmark = True
 
+    # Ensure splits exist; if missing, prepare from official splits
+    splits_dir = os.path.join(args.data_root, "splits")
+    triplet_path = os.path.join(splits_dir, "train.json")
+    if not os.path.exists(triplet_path):
+        os.makedirs(splits_dir, exist_ok=True)
+        try:
+            from scripts.prepare_polyvore import main as prepare_main
+            import sys
+            argv_bak = sys.argv
+            try:
+                # First try using official splits (no random)
+                sys.argv = ["prepare_polyvore.py", "--root", args.data_root]
+                prepare_main()
+            finally:
+                sys.argv = argv_bak
+        except Exception:
+            # As a fallback, try random split on any available aggregate file
+            try:
+                from scripts.prepare_polyvore import main as prepare_main
+                import sys
+                argv_bak = sys.argv
+                try:
+                    sys.argv = ["prepare_polyvore.py", "--root", args.data_root, "--random_split"]
+                    prepare_main()
+                finally:
+                    sys.argv = argv_bak
+            except Exception:
+                pass
+
     dataset = PolyvoreTripletDataset(args.data_root, split="train")
 
     loader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True, num_workers=4, pin_memory=(device=="cuda"))

train_vit_triplet.py

@@ -45,6 +45,33 @@ def main() -> None:
     if device == "cuda":
         torch.backends.cudnn.benchmark = True
 
+    # Ensure outfit triplets exist
+    splits_dir = os.path.join(args.data_root, "splits")
+    trip_path = os.path.join(splits_dir, "outfit_triplets_train.json")
+    if not os.path.exists(trip_path):
+        os.makedirs(splits_dir, exist_ok=True)
+        try:
+            from scripts.prepare_polyvore import main as prepare_main
+            import sys
+            argv_bak = sys.argv
+            try:
+                sys.argv = ["prepare_polyvore.py", "--root", args.data_root]
+                prepare_main()
+            finally:
+                sys.argv = argv_bak
+        except Exception:
+            try:
+                from scripts.prepare_polyvore import main as prepare_main
+                import sys
+                argv_bak = sys.argv
+                try:
+                    sys.argv = ["prepare_polyvore.py", "--root", args.data_root, "--random_split"]
+                    prepare_main()
+                finally:
+                    sys.argv = argv_bak
+            except Exception:
+                pass
+
     dataset = PolyvoreOutfitTripletDataset(args.data_root, split="train")
 
     def collate(batch):

utils/data_fetch.py

@@ -36,42 +36,43 @@ def ensure_dataset_ready() -> Optional[str]:
     root = os.path.abspath(os.path.join(os.getcwd(), "data", "Polyvore"))
     Path(root).mkdir(parents=True, exist_ok=True)
 
-    # If already present, ensure images are unzipped and return
+    # If images are already present, don't return early; still ensure metadata JSONs exist
     _unzip_images_if_needed(root)
-    if os.path.isdir(os.path.join(root, "images")):
-        return root
 
     # Download the HF dataset snapshot into root
     try:
         # Only fetch what's needed to run and prepare splits
         allow = [
             "images.zip",
-            "images/*.jpg",
-            "images/*.jpeg",
-            "images/*.png",
             "train.json",
             "valid.json",
             "test.json",
-            "fill_in_blank_*.json",
-            "compatibility_*.txt",
             "polyvore_item_metadata.json",
             "polyvore_outfit_titles.json",
             "categories.csv",
         ]
+        # Explicit ignores to prevent huge downloads (>10GB)
         ignore = [
             "**/*hglmm*",
-            "disjoint/*",
-            "nondisjoint/*",
-            "*/large/*",
+            "disjoint/**",
+            "nondisjoint/**",
+            "*/large/**",
+            "**/*.tar",
+            "**/*.tar.gz",
+            "**/*.7z",
         ]
-        snapshot_download(
-            "Stylique/Polyvore",
-            repo_type="dataset",
-            local_dir=root,
-            local_dir_use_symlinks=False,
-            allow_patterns=allow,
-            ignore_patterns=ignore,
-        )
+        need_meta = not all(os.path.exists(os.path.join(root, f)) for f in [
+            "train.json", "valid.json", "test.json", "categories.csv"
+        ])
+        if need_meta or not os.path.isdir(os.path.join(root, "images")):
+            snapshot_download(
+                "Stylique/Polyvore",
+                repo_type="dataset",
+                local_dir=root,
+                local_dir_use_symlinks=False,
+                allow_patterns=allow,
+                ignore_patterns=ignore,
+            )
     except Exception as e:  # pragma: no cover
         print(f"Failed to download Stylique/Polyvore dataset: {e}")
         return None