Coverage for cafl4ds/data/sources.py: 56%
63 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-07-17 14:39 +0000
1"""Data sources: produce ``(images, labels)`` tensors for the streams to order.
3A :class:`DataSource` is the *raw material* a stream orders into eras — it decouples the
4stream/ordering logic from where the pixels come from. Two sources exist in Phase 0:
6* :class:`STL10Source` — the real STL-10 labeled split (Coates et al. 2011), resized tiny for
7 CPU. Labels are carried only so the stream can build class-blocked ordering and held-out
8 eval sets; they never reach the SSL update.
9* :class:`SyntheticSource` — class-structured Gaussian blobs, network-free, for fast unit
10 tests and the fastest smoke runs.
12Both return images as ``float32`` ``[N, C, H, W]`` in ``[0, 1]`` and integer labels ``[N]``.
13Adding BDD100K/ZOD later means adding a new source, not touching the stream.
14"""
16from __future__ import annotations
18from abc import ABC, abstractmethod
19from pathlib import Path
21import torch
22import torch.nn.functional as F # noqa: N812 - conventional alias
23from loguru import logger
24from torchvision.datasets import STL10
27class DataSource(ABC):
28 """Produces ``(images, labels)`` tensors for a stream to order into eras."""
30 @abstractmethod
31 def load(self) -> tuple[torch.Tensor, torch.Tensor]:
32 """Load the full dataset into memory.
34 Returns:
35 A tuple ``(images, labels)`` with ``images`` of shape ``[N, C, H, W]``
36 (``float32`` in ``[0, 1]``) and integer ``labels`` of shape ``[N]``.
37 """
39 @property
40 @abstractmethod
41 def num_classes(self) -> int:
42 """Number of distinct classes in the source."""
45class STL10Source(DataSource):
46 """The real STL-10 labeled split, resized to a tiny CPU-friendly size."""
48 def __init__(
49 self,
50 root: str,
51 split: str = "train",
52 img_size: int = 32,
53 max_per_class: int | None = None,
54 ) -> None:
55 """Configure the STL-10 source.
57 Args:
58 root: Directory holding the downloaded ``stl10_binary`` (torchvision layout).
59 split: Which labeled split to load (``"train"`` or ``"test"``).
60 img_size: Side length to bilinearly resize the 96px images to.
61 max_per_class: If set, keep at most this many images per class (tiny runs).
62 """
63 self.root = root
64 self.split = split
65 self.img_size = img_size
66 self.max_per_class = max_per_class
68 @property
69 def num_classes(self) -> int:
70 """STL-10 has 10 classes."""
71 return 10
73 def load(self) -> tuple[torch.Tensor, torch.Tensor]:
74 """Load, resize, and (optionally) per-class-subsample the STL-10 split.
76 Returns:
77 ``(images, labels)`` with images ``[N, 3, img_size, img_size]`` in ``[0, 1]``.
79 Raises:
80 FileNotFoundError: If the STL-10 binaries are not present under ``root``.
81 """
82 if not (Path(self.root) / "stl10_binary").is_dir():
83 raise FileNotFoundError(
84 f"STL-10 binaries not found under {self.root}. Download once with "
85 "torchvision.datasets.STL10(root=..., split=..., download=True)."
86 )
87 ds = STL10(root=self.root, split=self.split, download=False)
88 images = torch.from_numpy(ds.data).float() / 255.0 # [N, 3, 96, 96]
89 labels = torch.from_numpy(ds.labels).long()
90 if self.max_per_class is not None:
91 images, labels = _subsample_per_class(images, labels, self.max_per_class)
92 images = F.interpolate(images, size=self.img_size, mode="bilinear", align_corners=False, antialias=True)
93 logger.info(f"STL10Source: loaded {images.shape[0]} images ({self.split}) at {self.img_size}px")
94 return images, labels
97class SyntheticSource(DataSource):
98 """Class-structured Gaussian images — network-free, for tests and fast smoke runs."""
100 def __init__(
101 self,
102 num_classes: int = 4,
103 per_class: int = 64,
104 img_size: int = 16,
105 channels: int = 3,
106 noise: float = 0.3,
107 seed: int = 0,
108 ) -> None:
109 """Configure the synthetic source.
111 Args:
112 num_classes: Number of classes to generate.
113 per_class: Images per class.
114 img_size: Image side length.
115 channels: Number of channels.
116 noise: Standard deviation of the per-pixel Gaussian noise around each class mean.
117 seed: RNG seed for reproducibility.
118 """
119 self._num_classes = num_classes
120 self.per_class = per_class
121 self.img_size = img_size
122 self.channels = channels
123 self.noise = noise
124 self.seed = seed
126 @property
127 def num_classes(self) -> int:
128 """Return the configured class count."""
129 return self._num_classes
131 def load(self) -> tuple[torch.Tensor, torch.Tensor]:
132 """Generate class-structured images (a distinct random mean pattern per class).
134 Returns:
135 ``(images, labels)`` with images ``[N, C, img_size, img_size]`` in ``[0, 1]``,
136 where a class's images cluster around its own random pattern (so probes are
137 learnable and the effective rank is meaningful).
138 """
139 g = torch.Generator().manual_seed(self.seed)
140 shape = (self.channels, self.img_size, self.img_size)
141 images, labels = [], []
142 for c in range(self._num_classes):
143 mean = torch.rand(shape, generator=g)
144 block = mean.unsqueeze(0) + self.noise * torch.randn(self.per_class, *shape, generator=g)
145 images.append(block.clamp_(0.0, 1.0))
146 labels.append(torch.full((self.per_class,), c, dtype=torch.long))
147 return torch.cat(images), torch.cat(labels)
150def _subsample_per_class(
151 images: torch.Tensor, labels: torch.Tensor, max_per_class: int
152) -> tuple[torch.Tensor, torch.Tensor]:
153 """Keep at most ``max_per_class`` images of each class (first-occurring, order-preserving).
155 Args:
156 images: All images ``[N, C, H, W]``.
157 labels: All labels ``[N]``.
158 max_per_class: Cap on images retained per class.
160 Returns:
161 The subsampled ``(images, labels)``.
162 """
163 keep: list[int] = []
164 counts: dict[int, int] = {}
165 for i, y in enumerate(labels.tolist()):
166 if counts.get(y, 0) < max_per_class:
167 keep.append(i)
168 counts[y] = counts.get(y, 0) + 1
169 idx = torch.tensor(keep, dtype=torch.long)
170 return images[idx], labels[idx]