training/local_dataset.py

"""Read the locally downloaded hoho22k_2026 webdataset tar files directly.

Bypasses the datasets library entirely -- no version conflicts, no network.
Each tar file contains per-scene entries:
  {order_id}.{field}_{image_id}.npy   -- per-image numpy arrays (K, R, t, pose_only_in_colmap)
  {order_id}.{ade|depth|gestalt}_{image_id}.png  -- per-image PIL images
  {order_id}.colmap.zip               -- COLMAP reconstruction bytes
  {order_id}.wf_vertices.npy          -- (N, 3) float32
  {order_id}.wf_edges.npy             -- (M, 2) int64
  {order_id}.wf_classifications.npy   -- (M,) int64

Returns sample dicts compatible with predict_wireframe / convert_entry_to_human_readable.
"""
import glob
import io
import tarfile

import numpy as np
from PIL import Image


VECTOR_FIELDS = {'K', 'R', 't', 'pose_only_in_colmap'}
IMAGE_FIELDS   = {'ade', 'depth', 'gestalt'}
GLOBAL_FIELDS  = {'colmap', 'wf_vertices', 'wf_edges', 'wf_classifications'}


def _assemble_scene(order_id, raw):
    """raw: {field_key_with_ext: bytes}"""
    # Collect image_ids from any per-image field
    image_ids = set()
    for key in raw:
        for prefix in [f'{f}_' for f in VECTOR_FIELDS | IMAGE_FIELDS]:
            if key.startswith(prefix):
                img_id = key[len(prefix):]
                img_id = img_id.rsplit('.', 1)[0]  # strip extension
                image_ids.add(img_id)
    image_ids = sorted(image_ids)

    sample = {'order_id': order_id, 'image_ids': image_ids}
    for f in VECTOR_FIELDS | IMAGE_FIELDS:
        sample[f] = []

    for img_id in image_ids:
        for field in VECTOR_FIELDS:
            key = f'{field}_{img_id}.npy'
            if key in raw:
                sample[field].append(np.load(io.BytesIO(raw[key])))
        for field in IMAGE_FIELDS:
            key = f'{field}_{img_id}.png'
            if key in raw:
                sample[field].append(Image.open(io.BytesIO(raw[key])).copy())

    # Global fields
    if 'colmap.zip' in raw:
        sample['colmap'] = raw['colmap.zip']
    for field in ('wf_vertices', 'wf_edges', 'wf_classifications'):
        key = f'{field}.npy'
        if key in raw:
            sample[field] = np.load(io.BytesIO(raw[key])).tolist()

    return sample


def iter_tar(tar_path):
    """Yield assembled scene dicts from a single tar file."""
    scenes = {}
    with tarfile.open(tar_path, 'r') as tf:
        for member in tf.getmembers():
            if not member.isfile():
                continue
            name = member.name
            dot = name.index('.')
            order_id  = name[:dot]
            field_key = name[dot + 1:]
            f = tf.extractfile(member)
            if f is None:
                continue
            if order_id not in scenes:
                scenes[order_id] = {}
            scenes[order_id][field_key] = f.read()

    for order_id, raw in scenes.items():
        yield _assemble_scene(order_id, raw)


def iter_split(dataset_dir, split='validation'):
    """Yield all scenes from a dataset split, sorted by tar file name."""
    tars = sorted(glob.glob(f'{dataset_dir}/data/{split}/*.tar'))
    if not tars:
        raise FileNotFoundError(f'No tar files found in {dataset_dir}/data/{split}/')
    for tar_path in tars:
        yield from iter_tar(tar_path)