tools/view_sampler.py

"""Pluggable view selection strategies for TACO dataset.

Selects which of the 12 allocentric cameras are used as context (input)
vs target (supervision) views for novel view synthesis training.
"""

from abc import ABC, abstractmethod

import numpy as np


class ViewSampler(ABC):
    """Base class for view sampling strategies.

    Given a list of available camera IDs and desired counts, returns
    disjoint sets of context and target camera IDs.
    """

    @abstractmethod
    def sample(
        self,
        camera_ids: list[str],
        n_context: int,
        n_target: int,
        rng: np.random.Generator | None = None,
    ) -> tuple[list[str], list[str]]:
        """Select context and target camera IDs.

        Args:
            camera_ids: Available camera IDs for this sequence.
            n_context: Number of context (input) views.
            n_target: Number of target (supervision) views.
            rng: Optional numpy random generator for reproducibility.

        Returns:
            (context_ids, target_ids) — disjoint lists of camera ID strings.

        Raises:
            ValueError: If n_context + n_target > len(camera_ids).
        """
        ...


class RandomViewSampler(ViewSampler):
    """Random split of cameras into context/target (no overlap)."""

    def sample(self, camera_ids, n_context, n_target, rng=None):
        if n_context + n_target > len(camera_ids):
            raise ValueError(
                f"Requested {n_context}+{n_target}={n_context+n_target} views "
                f"but only {len(camera_ids)} cameras available"
            )
        if rng is None:
            rng = np.random.default_rng()
        chosen = rng.choice(len(camera_ids), size=n_context + n_target, replace=False)
        ids = [camera_ids[i] for i in chosen]
        return ids[:n_context], ids[n_context:]


class FixedViewSampler(ViewSampler):
    """Predefined context/target camera assignments.

    Args:
        context_cameras: List of camera IDs to always use as context.
        target_cameras: List of camera IDs to always use as target.
            If None, all remaining cameras become targets.
    """

    def __init__(self, context_cameras: list[str], target_cameras: list[str] | None = None):
        self.context_cameras = context_cameras
        self.target_cameras = target_cameras

    def sample(self, camera_ids, n_context, n_target, rng=None):
        available = set(camera_ids)
        context = [c for c in self.context_cameras if c in available][:n_context]
        if self.target_cameras is not None:
            target = [c for c in self.target_cameras if c in available and c not in context][:n_target]
        else:
            remaining = [c for c in camera_ids if c not in context]
            target = remaining[:n_target]
        return context, target


class InterpolationViewSampler(ViewSampler):
    """Target views are spatially between context views (for NVS interpolation).

    Uses camera ID ordering as a proxy for spatial arrangement (cameras are
    arranged in a ring around the scene). Context views are evenly spaced,
    and target views are chosen from the gaps between them.

    Args:
        camera_ring_order: Optional explicit ordering of camera IDs around the ring.
            If None, uses sorted camera_ids as the ring order.
    """

    def __init__(self, camera_ring_order: list[str] | None = None):
        self.camera_ring_order = camera_ring_order

    def sample(self, camera_ids, n_context, n_target, rng=None):
        if n_context + n_target > len(camera_ids):
            raise ValueError(
                f"Requested {n_context}+{n_target}={n_context+n_target} views "
                f"but only {len(camera_ids)} cameras available"
            )

        ring = self.camera_ring_order or sorted(camera_ids)
        # Filter to available cameras while preserving ring order
        ring = [c for c in ring if c in camera_ids]

        if rng is None:
            rng = np.random.default_rng()

        n = len(ring)
        # Pick evenly-spaced context indices along the ring
        offset = rng.integers(0, n)
        step = n / n_context
        context_indices = [int((offset + i * step) % n) for i in range(n_context)]

        # Target views: those between consecutive context views
        context_set = set(context_indices)
        target_pool = [i for i in range(n) if i not in context_set]

        # Score each target candidate by how "between" context views it is
        # (prefer midpoints of gaps)
        sorted_ctx = sorted(context_indices)
        gap_mids = []
        for k in range(len(sorted_ctx)):
            a = sorted_ctx[k]
            b = sorted_ctx[(k + 1) % len(sorted_ctx)]
            if b <= a:
                b += n
            mid = ((a + b) / 2) % n
            gap_mids.append(mid)

        def min_dist_to_gap_mid(idx):
            return min(abs(idx - m) for m in gap_mids)

        target_pool.sort(key=min_dist_to_gap_mid)
        target_indices = target_pool[:n_target]

        context = [ring[i] for i in context_indices]
        target = [ring[i] for i in target_indices]
        return context, target


class ExtrapolationViewSampler(ViewSampler):
    """Target views are outside the context view hull (harder NVS task).

    Context views are clustered in a contiguous arc on the camera ring,
    and target views are chosen from outside that arc.
    """

    def sample(self, camera_ids, n_context, n_target, rng=None):
        if n_context + n_target > len(camera_ids):
            raise ValueError(
                f"Requested {n_context}+{n_target}={n_context+n_target} views "
                f"but only {len(camera_ids)} cameras available"
            )

        ring = sorted(camera_ids)
        n = len(ring)

        if rng is None:
            rng = np.random.default_rng()

        # Pick a contiguous arc of n_context cameras
        start = rng.integers(0, n)
        context_indices = [(start + i) % n for i in range(n_context)]
        context_set = set(context_indices)

        # Target views: outside the context arc
        target_pool = [i for i in range(n) if i not in context_set]
        if len(target_pool) < n_target:
            raise ValueError(
                f"Only {len(target_pool)} cameras outside context arc, "
                f"but need {n_target} targets"
            )
        chosen = rng.choice(len(target_pool), size=n_target, replace=False)
        target_indices = [target_pool[i] for i in chosen]

        context = [ring[i] for i in context_indices]
        target = [ring[i] for i in target_indices]
        return context, target