from __future__ import annotations import json import zipfile 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: 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): 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: {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)) 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: 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: 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), }