code/meshai_train/dataset.py

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),
    }