dataset.py

import os
import struct
import torch
import random
from torch import Tensor
from pathlib import Path
from typing import List, Optional, Sequence, Union, Any, Callable, Dict
from torchvision.datasets.folder import default_loader
from pytorch_lightning import LightningDataModule
from torch.utils.data import DataLoader, Dataset, IterableDataset
from torchvision import transforms
from torchvision.datasets import CelebA
import zipfile
import numpy as np
import numpy.lib.format as npy_format


def _npz_member_data_offset(fh, zip_path: str, member: str) -> int:
    """Return the absolute byte offset of ``member``'s data inside ``zip_path``.

    Parses the ZIP local file header directly (rather than trusting the
    central-directory entry) so the offset is correct regardless of any
    local-vs-central header differences.
    """
    with zipfile.ZipFile(zip_path) as z:
        zi = z.getinfo(member)
    fh.seek(zi.header_offset)
    sig = fh.read(4)
    if sig != b'PK\x03\x04':
        raise ValueError(f"Bad local header signature in {zip_path}:{member}")
    fh.read(2)                       # version needed
    fh.read(2)                       # general purpose bit flag
    fh.read(2)                       # compression method
    fh.read(8)                       # last mod time/date
    fh.read(8)                       # crc-32, compressed size, uncompressed size
    name_len = struct.unpack('<H', fh.read(2))[0]
    extra_len = struct.unpack('<H', fh.read(2))[0]
    fh.read(name_len)
    fh.read(extra_len)
    return fh.tell()


def _read_npy_header_at(fh, offset: int):
    """Parse a ``.npy`` header located at ``offset`` in ``fh``.

    Returns ``(array_data_offset, dtype, shape, fortran_order)``.
    """
    fh.seek(offset)
    major, _ = npy_format.read_magic(fh)
    if major == 1:
        shape, fortran_order, dtype = npy_format.read_array_header_1_0(fh)
    elif major == 2:
        shape, fortran_order, dtype = npy_format.read_array_header_2_0(fh)
    else:
        raise ValueError(f"Unsupported npy format version {major}")
    return fh.tell(), dtype, shape, fortran_order


def _npz_member_memmap(zip_path: str,
                       member: str,
                       cache_dir: Optional[Path] = None) -> Optional[np.memmap]:
    """Memory-map array ``member`` from a (possibly STORED) ``.npz`` file.

    For STORED members the array is mapped directly inside the ``.npz`` (no
    extra disk, no copy). For compressed members the member is extracted to
    ``cache_dir`` (validated by size) and the extracted ``.npy`` is mapped.
    Returns ``None`` if ``member`` is not present.
    """
    try:
        with zipfile.ZipFile(zip_path) as z:
            info = z.getinfo(member)
    except KeyError:
        return None

    with open(zip_path, 'rb') as fh:
        data_offset = _npz_member_data_offset(fh, zip_path, member)
        arr_offset, dtype, shape, fortran_order = _read_npy_header_at(fh, data_offset)

    if info.compress_type == zipfile.ZIP_STORED:
        return np.memmap(zip_path, dtype=dtype, mode='r',
                         offset=arr_offset, shape=shape,
                         order='F' if fortran_order else 'C')

    if cache_dir is None:
        raise ValueError(
            f"Member '{member}' in {zip_path} is compressed and cannot be "
            f"memory-mapped in place; provide lazy_cache_dir to enable the "
            f"disk-cache fallback."
        )
    cache_dir = Path(cache_dir)
    cache_dir.mkdir(parents=True, exist_ok=True)
    extracted = cache_dir / f"{Path(zip_path).stem}__{member.replace('/', '_')}.npy"
    need_extract = (not extracted.exists()
                    or extracted.stat().st_size != info.file_size)
    if need_extract:
        with zipfile.ZipFile(zip_path) as z, z.open(member) as src, \
                open(extracted, 'wb') as dst:
            dst.write(src.read())
    return np.memmap(extracted, dtype=dtype, mode='r', shape=shape,
                     order='F' if fortran_order else 'C')


class LazyConcat:
    """Read-only virtual concatenation of arrays along axis 0.

    Behaves like ``np.concatenate(arrays, axis=0)`` for the access patterns
    used by this module (integer indexing, contiguous slices, ``len``,
    ``shape``, ``dtype``) but never materializes the full concatenation:
    integer and contiguous-slice indexing return views into the underlying
    arrays (typically ``np.memmap``). Fancy/non-contiguous indexing
    materializes the result.
    """

    def __init__(self, arrays: Sequence[np.ndarray]):
        self.arrays = list(arrays)
        if len(self.arrays) == 0:
            self._offsets = np.array([0], dtype=np.int64)
            self._shape = (0,)
            self.dtype = np.dtype(np.float32)
            return
        lengths = np.array([int(a.shape[0]) for a in self.arrays], dtype=np.int64)
        self._offsets = np.concatenate([[0], np.cumsum(lengths)])
        self._shape = (int(self._offsets[-1]),) + tuple(self.arrays[0].shape[1:])
        self.dtype = self.arrays[0].dtype

    @property
    def shape(self) -> tuple:
        return self._shape

    def __len__(self) -> int:
        return self._shape[0]

    def __array__(self, dtype=None) -> np.ndarray:
        out = np.concatenate([np.asarray(a) for a in self.arrays], axis=0)
        return out if dtype is None else out.astype(dtype)

    def _locate(self, idx: int):
        i = np.searchsorted(self._offsets, idx, side='right') - 1
        return i, int(idx - self._offsets[i])

    def __getitem__(self, idx):
        n = self._shape[0]
        if isinstance(idx, (int, np.integer)):
            if idx < 0:
                idx += n
            if idx < 0 or idx >= n:
                raise IndexError(f"index {idx} out of range for length {n}")
            i, local = self._locate(idx)
            return self.arrays[i][local]

        if isinstance(idx, slice):
            start, stop, step = idx.indices(n)
            if step == 1 and stop > start:
                i0, local0 = self._locate(start)
                i1, local1 = self._locate(stop - 1)
                if i0 == i1:
                    return self.arrays[i0][local0:local1 + 1]
                pieces = [self.arrays[i0][local0:]]
                for i in range(i0 + 1, i1):
                    pieces.append(self.arrays[i][:])
                pieces.append(self.arrays[i1][:local1 + 1])
                if len(pieces) == 1:
                    return pieces[0]
                return LazyConcat(pieces)
            return np.asarray(self)[start:stop:step]

        idx_arr = np.asarray(idx)
        if idx_arr.dtype == bool:
            idx_arr = np.nonzero(idx_arr)
        return np.asarray(self)[idx_arr]


# Add your custom dataset class here
class AnimalAI(Dataset):
    def __init__(self,
                 data,
                 transform: Callable,
                 indices: Optional[np.ndarray] = None,
                 **kwargs):
        self.transforms = transform
        self.im = data['s']
        self.rew = data['r']
        self.rst = data.get('reset', np.zeros(len(self.im), dtype=np.float32))
        self.indices = indices

    def __len__(self):
        return len(self.indices) if self.indices is not None else len(self.im)

    def __getitem__(self, idx):
        i = self.indices[idx] if self.indices is not None else idx
        img = self.im[i]
        rew = self.rew[i]
        rst = self.rst[i]

        if self.transforms is not None:
            img = self.transforms(np.array(img, dtype=np.float32))

        return img, rew.astype(np.float32), rst.astype(np.float32)


class MyCelebA(CelebA):
    """
    A work-around to address issues with pytorch's celebA dataset class.
    
    Download and Extract
    URL : https://drive.google.com/file/d/1m8-EBPgi5MRubrm6iQjafK2QMHDBMSfJ/view?usp=sharing
    """
    
    def _check_integrity(self) -> bool:
        return True
    

class OxfordPets(Dataset):
    """
    URL = https://www.robots.ox.ac.uk/~vgg/data/pets/
    """
    def __init__(self, 
                 data_path: str, 
                 split: str,
                 transform: Callable,
                **kwargs):
        self.data_dir = Path(data_path) / "OxfordPets"        
        self.transforms = transform
        imgs = sorted([f for f in self.data_dir.iterdir() if f.suffix == '.jpg'])
        
        self.imgs = imgs[:int(len(imgs) * 0.75)] if split == "train" else imgs[int(len(imgs) * 0.75):]
    
    def __len__(self):
        return len(self.imgs)
    
    def __getitem__(self, idx):
        img = default_loader(self.imgs[idx])
        
        if self.transforms is not None:
            img = self.transforms(img)
        
        return img, 0.0 # dummy datat to prevent breaking 


class AnimalAIEpisodeParallel(IterableDataset):
    def __init__(self,
                 data,
                 transform: Callable,
                 batch_size: int,
                 **kwargs):
        super().__init__()
        self.data = data
        self.transforms = transform
        self.batch_size = batch_size

    def _split_into_episodes(self):
        resets = self.data.get('reset', np.zeros(len(self.data['s']), dtype=np.float32))
        indices = np.where(resets > 0)[0]
        if len(indices) == 0 or indices[0] != 0:
            indices = np.concatenate([[0], indices])
        episodes = []
        for i, start in enumerate(indices):
            end = indices[i + 1] if i + 1 < len(indices) else len(self.data['s'])
            if end - start > 0:
                episodes.append({
                    's': self.data['s'][start:end],
                    'r': self.data['r'][start:end],
                })
        return episodes

    def __iter__(self):
        episodes = self._split_into_episodes()
        random.shuffle(episodes)

        effective_batch_size = min(self.batch_size, len(episodes))
        if effective_batch_size == 0:
            return

        track_episodes = [[] for _ in range(effective_batch_size)]
        for i, ep in enumerate(episodes):
            track_episodes[i % effective_batch_size].append(ep)

        track_ep_idx = [0] * effective_batch_size
        track_frame_idx = [0] * effective_batch_size

        while True:
            batch_imgs = []
            batch_rews = []
            batch_resets = []
            valid = True

            for i in range(effective_batch_size):
                while True:
                    if track_ep_idx[i] >= len(track_episodes[i]):
                        valid = False
                        break
                    ep = track_episodes[i][track_ep_idx[i]]
                    pos = track_frame_idx[i]
                    if pos < len(ep['s']):
                        break
                    track_ep_idx[i] += 1
                    track_frame_idx[i] = 0

                if not valid:
                    break

                ep = track_episodes[i][track_ep_idx[i]]
                pos = track_frame_idx[i]
                img = ep['s'][pos]
                rew = ep['r'][pos]
                rst = 1.0 if pos == 0 else 0.0

                if self.transforms is not None:
                    img = self.transforms(img.astype(np.float32))
                else:
                    img = torch.from_numpy(img.astype(np.float32))

                batch_imgs.append(img)
                batch_rews.append(rew)
                batch_resets.append(rst)
                track_frame_idx[i] += 1

            if not valid or len(batch_imgs) < effective_batch_size:
                break

            yield (torch.stack(batch_imgs),
                   torch.tensor(batch_rews, dtype=torch.float32),
                   torch.tensor(batch_resets, dtype=torch.float32))


class VAEDataset(LightningDataModule):
    """
    PyTorch Lightning data module 

    Args:
        data_dir: root directory of your dataset.
        train_batch_size: the batch size to use during training.
        val_batch_size: the batch size to use during validation.
        patch_size: the size of the crop to take from the original images.
        num_workers: the number of parallel workers to create to load data
            items (see PyTorch's Dataloader documentation for more details).
        pin_memory: whether prepared items should be loaded into pinned memory
            or not. This can improve performance on GPUs.
        lazy: if True, the large 's' arrays are memory-mapped from disk
            instead of loaded into RAM. They are mapped directly inside the
            existing (uncompressed) ``.npz`` files with no extra disk usage;
            compressed members are extracted to ``lazy_cache_dir`` on first
            use. Set to False to load everything into RAM (old behavior).
        lazy_cache_dir: directory for the disk-cache fallback used when a
            ``.npz`` member is compressed and cannot be mapped in place.
            Defaults to ``<data_path>/.lazycache``. Only populated on demand.
    """

    def __init__(
        self,
        data_path: str,
        train_batch_size: int = 8,
        val_batch_size: int = 8,
        patch_size: Union[int, Sequence[int]] = (256, 256),
        num_workers: int = 0,
        pin_memory: bool = False,
        shuffle: bool = True,
        trace_decay: float = 0.0,
        test_data_path: Optional[str] = None,
        lazy: bool = True,
        lazy_cache_dir: Optional[str] = None,
        **kwargs,
    ):
        super().__init__()

        self.data_dir = data_path
        self.train_batch_size = train_batch_size
        self.val_batch_size = val_batch_size
        self.patch_size = patch_size
        self.num_workers = num_workers
        self.pin_memory = pin_memory
        self.shuffle = shuffle
        self.trace_decay = trace_decay
        self.test_data_path = test_data_path
        self.lazy = lazy
        self.lazy_cache_dir = lazy_cache_dir

    def _resolve_cache_dir(self, data_path: str) -> Optional[Path]:
        if not self.lazy:
            return None
        if self.lazy_cache_dir is not None:
            return Path(self.lazy_cache_dir)
        p = Path(data_path)
        base = p if p.is_dir() else p.parent
        return base / ".lazycache"

    def _load_data_from_path(self, data_path: str) -> dict:
        data_dir = Path(data_path)
        if data_dir.is_dir():
            files = sorted(f for f in data_dir.iterdir() if f.is_file())
            if self.lazy:
                cache_dir = self._resolve_cache_dir(data_path)
                s_arrays, r_arrays, reset_arrays = [], [], []
                has_reset = False
                for f in files:
                    s = _npz_member_memmap(str(f), 's.npy', cache_dir)
                    if s is None:
                        continue
                    s_arrays.append(s)
                    with np.load(f, allow_pickle=True) as d:
                        if 'r' in d:
                            r_arrays.append(np.asarray(d['r']))
                        if 'reset' in d:
                            has_reset = True
                            reset_arrays.append(np.asarray(d['reset']))
                out: Dict[str, Any] = {
                    's': LazyConcat(s_arrays) if len(s_arrays) > 1 else s_arrays[0],
                    'r': np.concatenate(r_arrays, axis=0) if r_arrays else np.zeros(0, dtype=np.float32),
                }
                if has_reset:
                    out['reset'] = np.concatenate(reset_arrays, axis=0)
                return out
            all_data = [np.load(f, allow_pickle=True) for f in files]
            data_keys = ['s', 'r'] + (['reset'] if any('reset' in d for d in all_data) else [])
            return {k: np.concatenate([d[k] for d in all_data if k in d], axis=0) for k in data_keys}
        else:
            if self.lazy:
                cache_dir = self._resolve_cache_dir(str(data_dir.parent))
                with np.load(data_dir, allow_pickle=True) as d:
                    keys = list(d.files)
                out = {}
                for k in keys:
                    m = _npz_member_memmap(str(data_dir), f'{k}.npy', cache_dir)
                    out[k] = m if m is not None else np.asarray(np.load(data_dir, allow_pickle=True)[k])
                return out
            return dict(np.load(data_dir, allow_pickle=True))

    def setup(self, stage: Optional[str] = None) -> None:
        train_transforms = transforms.Compose([lambda x: torch.from_numpy(x)])
        
        val_transforms = transforms.Compose([lambda x: torch.from_numpy(x)])

        data = self._load_data_from_path(self.data_dir)

        data_keys = ['s', 'r']
        if 'reset' in data:
            data_keys.append('reset')

        if self.test_data_path is not None:
            test_data = self._load_data_from_path(self.test_data_path)
            if self.trace_decay > 0:
                train_data = data
                self.train_dataset = AnimalAIEpisodeParallel(
                    train_data,
                    transform=train_transforms,
                    batch_size=self.train_batch_size,
                )
            else:
                train_data = data
                self.train_dataset = AnimalAI(
                    train_data,
                    transform=train_transforms,
                )
            self.val_dataset = AnimalAI(
                test_data,
                transform=val_transforms,
            )
            return

        if self.trace_decay > 0:
            all_resets = data.get('reset', np.zeros(len(data['s']), dtype=np.float32))
            ep_starts = np.where(all_resets > 0)[0]
            if len(ep_starts) == 0 or ep_starts[0] != 0:
                ep_starts = np.concatenate([[0], ep_starts])

            n = len(data['s'])
            n_train = int(n * 0.75)
            train_ep_idx = np.searchsorted(ep_starts, n_train, side='right')
            train_start_idx = 0
            train_end_idx = ep_starts[train_ep_idx] if train_ep_idx < len(ep_starts) else n
            val_start_idx = train_end_idx
            val_end_idx = n

            train_data = {name: data[name][train_start_idx:train_end_idx] for name in data_keys}
            test_data = {name: data[name][val_start_idx:val_end_idx] for name in data_keys}
            val_indices = None

            self.train_dataset = AnimalAIEpisodeParallel(
                train_data,
                transform=train_transforms,
                batch_size=self.train_batch_size,
            )
        else:
            split_indexes = np.arange(len(data['s']))
            np.random.shuffle(split_indexes)
            cut = int(len(split_indexes) * 0.75)
            val_indices = split_indexes[cut:]

            self.train_dataset = AnimalAI(
                data,
                transform=train_transforms,
                indices=split_indexes[:cut],
            )

        self.val_dataset = AnimalAI(
            test_data,
            transform=val_transforms,
            indices=val_indices,
        )

    def train_dataloader(self) -> DataLoader:
        if self.trace_decay > 0:
            return DataLoader(
                self.train_dataset,
                batch_size=1,
                num_workers=0,
                pin_memory=self.pin_memory,
                collate_fn=lambda x: x[0],
            )
        return DataLoader(
            self.train_dataset,
            batch_size=self.train_batch_size,
            num_workers=self.num_workers,
            shuffle=self.shuffle,
            pin_memory=self.pin_memory,
        )

    def val_dataloader(self) -> Union[DataLoader, List[DataLoader]]:
        return DataLoader(
            self.val_dataset,
            batch_size=self.val_batch_size,
            num_workers=self.num_workers,
            shuffle=False,
            pin_memory=self.pin_memory,
        )
    
    def test_dataloader(self) -> Union[DataLoader, List[DataLoader]]:
        return DataLoader(
            self.val_dataset,
            batch_size=144,
            num_workers=self.num_workers,
            shuffle=True,
            pin_memory=self.pin_memory,
        )