Fix BadZipFile skip + _missing_uids harden

code/meshai_train/dataset.py

@@ -0,0 +1,221 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import numpy as np
+import torch
+from PIL import Image
+from torch.utils.data import Dataset
+
+HF_PREPROCESSED_DEFAULT = "HayrettinIscan/MeshAI-Preprocessed-4K"
+
+
+def _uid_hf_prefix(uid: str) -> str:
+    uid = str(uid)
+    return f"preprocessed/{uid[:2]}/{uid}"
+
+
+class PreprocessedMeshDataset(Dataset):
+    """Loads geometry_latent.npz + render PNGs (local cache or HF)."""
+
+    def __init__(
+        self,
+        *,
+        token: str | None,
+        data_root: Path | None = None,
+        hf_repo: str = HF_PREPROCESSED_DEFAULT,
+        limit: int | None = None,
+        render_size: int = 128,
+        max_views: int = 4,
+    ) -> None:
+        self.token = token
+        self.data_root = data_root
+        self.hf_repo = hf_repo
+        self.render_size = render_size
+        self.max_views = max_views
+        self._missing_uids: set[str] = set()
+        self.objects = self._load_index(limit)
+
+    def _load_index(self, limit: int | None) -> list[dict]:
+        rows: list[dict] = []
+        roots = [p for p in (self.data_root, Path("data/preprocessed")) if p and p.exists()]
+        for root in roots:
+            for latent in root.rglob("geometry_latent.npz"):
+                uid = latent.parent.name
+                if len(uid) >= 8:
+                    rows.append({"uid": uid, "object_id": uid, "status": "ready"})
+            if rows:
+                break
+        if not rows:
+            local_summary = Path("data/preprocessed/preprocessed_objects.json")
+            if local_summary.exists():
+                payload = json.loads(local_summary.read_text(encoding="utf-8"))
+                rows = [r for r in payload.get("objects", []) if r.get("status", "ready") == "ready"]
+        if not rows:
+            from huggingface_hub import hf_hub_download
+
+            path = hf_hub_download(
+                repo_id=self.hf_repo,
+                filename="preprocessed/preprocessed_objects.json",
+                repo_type="dataset",
+                token=self.token,
+            )
+            payload = json.loads(Path(path).read_text(encoding="utf-8"))
+            rows = [r for r in payload.get("objects", []) if r.get("status", "ready") == "ready"]
+        if limit is not None:
+            rows = rows[:limit]
+        if not rows:
+            raise RuntimeError("Preprocessed object listesi bos.")
+        return rows
+
+    def __len__(self) -> int:
+        return len(self.objects)
+
+    def _resolve_file(self, uid: str, name: str) -> Path:
+        candidates = []
+        if self.data_root:
+            candidates.append(self.data_root / uid[:2] / uid / name)
+            candidates.append(self.data_root / uid / name)
+        candidates.append(Path("data/preprocessed") / uid[:2] / uid / name)
+        candidates.append(Path("data/preprocessed") / uid / name)
+        for local in candidates:
+            if local.exists():
+                return local
+        if not self.token:
+            raise FileNotFoundError(f"Yerel dosya yok ve HF_TOKEN yok: {uid}/{name}")
+        from huggingface_hub import hf_hub_download
+
+        rel = f"{_uid_hf_prefix(uid)}/{name}"
+        try:
+            return Path(
+                hf_hub_download(
+                    repo_id=self.hf_repo,
+                    filename=rel,
+                    repo_type="dataset",
+                    token=self.token,
+                )
+            )
+        except Exception as exc:
+            # 404 / EntryNotFound / RemoteEntryNotFound — manifestte olup HF'de olmayan UID
+            msg = str(exc).lower()
+            if "404" in msg or "entrynotfound" in msg.replace(" ", "") or "not found" in msg:
+                raise FileNotFoundError(f"HF eksik: {rel}") from exc
+            raise
+
+    def _load_npz(self, uid: str):
+        """Incomplete HF cache can yield BadZipFile; delete + force redownload once."""
+        import zipfile
+
+        path = self._resolve_file(uid, "geometry_latent.npz")
+        try:
+            return np.load(path)
+        except (zipfile.BadZipFile, OSError, ValueError) as exc:
+            try:
+                path.unlink(missing_ok=True)
+            except OSError:
+                pass
+            if not self.token:
+                raise FileNotFoundError(f"Bozuk latent (token yok, yeniden indirilemez): {uid}") from exc
+            from huggingface_hub import hf_hub_download
+
+            rel = f"{_uid_hf_prefix(uid)}/geometry_latent.npz"
+            path = Path(
+                hf_hub_download(
+                    repo_id=self.hf_repo,
+                    filename=rel,
+                    repo_type="dataset",
+                    token=self.token,
+                    force_download=True,
+                )
+            )
+            try:
+                return np.load(path)
+            except (zipfile.BadZipFile, OSError, ValueError) as exc2:
+                raise FileNotFoundError(f"Bozuk/eksik latent: {rel}") from exc2
+
+    def _getitem_one(self, idx: int) -> dict:
+        row = self.objects[idx]
+        uid = str(row.get("uid") or row.get("object_id"))
+        if uid in self._missing_uids:
+            raise FileNotFoundError(f"cached missing uid: {uid}")
+        try:
+            z = self._load_npz(uid)
+        except FileNotFoundError:
+            self._missing_uids.add(uid)
+            raise
+        hist = z["vertex_hist"].astype(np.float32)
+        stats = z["normal_stats"].astype(np.float32)
+        voxel = z["voxel_occ_32"].astype(np.float32).reshape(-1) / 255.0
+
+        views: list[np.ndarray] = []
+        for i in range(self.max_views):
+            for suffix in (f"view_{i:02d}_tex.png", f"view_{i:02d}.png"):
+                try:
+                    img_path = self._resolve_file(uid, f"renders/{suffix}")
+                    img = Image.open(img_path).convert("RGB").resize(
+                        (self.render_size, self.render_size), Image.BILINEAR
+                    )
+                    arr = np.asarray(img, dtype=np.float32) / 255.0
+                    views.append(arr.transpose(2, 0, 1))
+                    break
+                except Exception:
+                    continue
+        if not views:
+            views.append(np.zeros((3, self.render_size, self.render_size), dtype=np.float32))
+        # Collate batch_size>1 icin view sayisini sabitle (eksik render pad).
+        while len(views) < self.max_views:
+            views.append(views[-1].copy())
+        views = views[: self.max_views]
+
+        return {
+            "uid": uid,
+            "idx": idx,
+            "geom_in": np.concatenate([hist, stats]),
+            "voxel_tgt": voxel,
+            "views": np.stack(views, axis=0),
+            "quality_score": float(row.get("quality_score", 0.0)),
+        }
+
+    def __getitem__(self, idx: int) -> dict:
+        # Manifest bazen HF'de olmayan / bozuk UID icerir; sonraki ornege kay.
+        import zipfile
+
+        if not hasattr(self, "_missing_uids"):
+            self._missing_uids = set()
+        n = len(self.objects)
+        last_err: Exception | None = None
+        for offset in range(n):
+            cur = (idx + offset) % n
+            try:
+                return self._getitem_one(cur)
+            except (FileNotFoundError, zipfile.BadZipFile, OSError, ValueError, KeyError) as exc:
+                last_err = exc
+                try:
+                    uid = str(self.objects[cur].get("uid") or self.objects[cur].get("object_id") or "")
+                    if uid:
+                        self._missing_uids.add(uid)
+                except Exception:
+                    pass
+                continue
+        raise RuntimeError(f"Erisilebilir preprocessed ornek yok (son: {last_err})")
+
+
+def collate_preprocessed(batch: list[dict]) -> dict:
+    # Savunmaci pad: eski checkpoint/yamali dataset karisik shape getirebilir.
+    max_v = max(int(b["views"].shape[0]) for b in batch)
+    views_list = []
+    for b in batch:
+        v = b["views"]
+        if v.shape[0] < max_v:
+            pad = np.repeat(v[-1:], max_v - v.shape[0], axis=0)
+            v = np.concatenate([v, pad], axis=0)
+        views_list.append(v[:max_v])
+    return {
+        "uid": [b["uid"] for b in batch],
+        "idx": torch.tensor([b["idx"] for b in batch], dtype=torch.long),
+        "geom_in": torch.tensor(np.stack([b["geom_in"] for b in batch]), dtype=torch.float32),
+        "voxel_tgt": torch.tensor(np.stack([b["voxel_tgt"] for b in batch]), dtype=torch.float32),
+        "views": torch.tensor(np.stack(views_list), dtype=torch.float32),
+        "quality_score": torch.tensor([b["quality_score"] for b in batch], dtype=torch.float32),
+    }