Data & Augmentation — Programuoki

Code

data.py — fill in two holes

Now what you saw in step 2 is written down as data.py — and you fill in two missing pieces. These are exactly the two places where data mistakes are the quietest.

What is in the file

class_counts — how many pictures in each class. Ready.
build_transform — how a PIL picture becomes a tensor: gray → [1, 28, 28] → normalised. Ready.
split_indices — a stratified split: every class gives val_frac of its pictures to validation. The loop is written, one line is missing.
mean_std — pixel mean and std over the indices you name. Ready.
make_loaders — ties it together and returns (train_loader, val_loader, classes). One line is missing.

Your job

Create data.py and test_data.py (see the right panel). data.py runs without crashing but is wrong — two places are marked HOLE 1 and HOLE 2:

HOLE 1 — how many pictures of a class go to validation? val_frac of them, rounded, but never fewer than one. Right now it is 0.
HOLE 2 — the normalisation mean and std. mean_std(plain, ...) takes a list of picture indices — which ones? Right now it is 0.0, 1.0 (changes nothing).

Nothing else needs to change.

test_data.py is the verdict. It builds a tiny folder of pictures in a temporary directory itself (three classes of 30 pictures, each with its own average brightness), so it needs neither your photographs nor FashionMNIST.

Run it before filling anything in

$ python -m unittest test_data
FF..FF
======================================================================
FAIL: test_a_tiny_class_still_gets_a_validation_image (test_data.DataTest.test_a_tiny_class_still_gets_a_validation_image)
...
AssertionError: Lists differ: [] != [0, 1]
======================================================================
FAIL: test_batches_have_the_shape_the_cnn_expects (test_data.DataTest.test_batches_have_the_shape_the_cnn_expects)
...
AssertionError: 90 != 72
======================================================================
FAIL: test_split_is_stratified_and_disjoint (test_data.DataTest.test_split_is_stratified_and_disjoint)
...
AssertionError: 0 != 18 : 20 % of 90 images
======================================================================
FAIL: test_training_pixels_are_normalised_with_training_statistics (test_data.DataTest.test_training_pixels_are_normalised_with_training_statistics)
...
AssertionError: 0.4701877236366272 != 0.0 within 4 places (0.4701877236366272 difference)
----------------------------------------------------------------------
Ran 6 tests in 0.6s

FAILED (failures=4)

(Shortened: the Traceback bodies are cut. The first run can take longer — that is the torchvision import.)

Read the messages from the top; they tell a story: AssertionError: 90 != 72 — the training set holds all 90 pictures because nothing was given to validation; 0 != 18 — validation is empty; 0.47 != 0.0 — the training pixel mean is 0.47 because normalisation does nothing.

The two . — class_counts and the repeatability test — already pass because they do not depend on these holes.

When it does not pass

Only test_a_tiny_class_still_gets_a_validation_image fails, the rest pass — you wrote round(len(idx) * val_frac) or int(...). For a class of three pictures 10 % rounds to 0 and the class is left without validation. You need max(1, ...).
test_training_pixels_are_normalised… fails by a tiny amount (e.g. 0.0007 != 0.0) — HOLE 2 uses more than the training pictures. With all pictures the difference is 0.0007; with the validation ones, 0.0036. That does not look like much — but the test is strict on purpose: validation pictures must not influence a single training number.
AssertionError: 72 != 18 — you flipped it: val_frac taken for training rather than validation.

data.py CREATE

Create this file in ./

"""data.py — your own images, read from folders: data/<class>/*.png."""

from collections import Counter

import torch
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms


def class_counts(root: str) -> dict[str, int]:
    """How many images each class folder holds."""
    ds = datasets.ImageFolder(root)
    counts = Counter(ds.targets)
    return {name: counts[i] for i, name in enumerate(ds.classes)}


def build_transform(
    mean: float | None = None, std: float | None = None
) -> transforms.Compose:
    steps: list = [transforms.Grayscale(num_output_channels=1)]
    steps.append(transforms.ToTensor())  # PIL image -> [1, H, W] floats in 0..1
    if mean is not None:
        steps.append(transforms.Normalize([mean], [std]))
    return transforms.Compose(steps)


def split_indices(
    targets: list[int], val_frac: float, seed: int
) -> tuple[list[int], list[int]]:
    """A stratified split: every class contributes val_frac of its images."""
    gen = torch.Generator().manual_seed(seed)
    train_idx: list[int] = []
    val_idx: list[int] = []
    for cls in sorted(set(targets)):
        idx = [i for i, t in enumerate(targets) if t == cls]
        order = torch.randperm(len(idx), generator=gen).tolist()
        # HOLE 1: how many images of this class go to validation?
        # val_frac of them, rounded - but never fewer than one.
        n_val = 0
        val_idx += [idx[j] for j in order[:n_val]]
        train_idx += [idx[j] for j in order[n_val:]]
    return train_idx, val_idx


def mean_std(ds: datasets.ImageFolder, idx: list[int]) -> tuple[float, float]:
    """Pixel mean and std over the TRAINING images only."""
    pixels = torch.cat([ds[i][0].flatten() for i in idx])
    return pixels.mean().item(), pixels.std().item()


def make_loaders(
    root: str,
    batch_size: int = 32,
    val_frac: float = 0.2,
    seed: int = 42,
) -> tuple[DataLoader, DataLoader, list[str]]:
    """Return (train_loader, val_loader, class_names)."""
    plain = datasets.ImageFolder(root, transform=build_transform())
    train_idx, val_idx = split_indices(plain.targets, val_frac, seed)
    # HOLE 2: the mean and std used to normalise. Measure them with
    # mean_std() - on WHICH images?
    mean, std = 0.0, 1.0

    ds = datasets.ImageFolder(root, transform=build_transform(mean, std))
    train_loader = DataLoader(
        Subset(ds, train_idx), batch_size=batch_size, shuffle=True
    )
    val_loader = DataLoader(Subset(ds, val_idx), batch_size=batch_size)
    return train_loader, val_loader, plain.classes

test_data.py CREATE

Create this file in ./

"""test_data.py — the verdict. Run: python -m unittest -v

The tests build their own tiny image folders in a temporary directory, so they
need neither your photographs nor FashionMNIST.
"""

import random
import tempfile
import unittest
from pathlib import Path

import torch
from PIL import Image

from data import class_counts, make_loaders, split_indices


def make_folder(root: Path, counts: dict[str, int], seed: int = 0) -> str:
    """Write counts[name] noisy 28x28 grayscale PNGs into root/<name>/.

    Each class has its own average brightness, so the classes really differ.
    """
    rng = random.Random(seed)
    for k, (name, n) in enumerate(counts.items()):
        (root / name).mkdir(parents=True)
        base = 60 + 60 * k
        for i in range(n):
            img = Image.new("L", (28, 28))
            img.putdata([min(255, max(0, base + rng.randint(-40, 40))) for _ in range(784)])
            img.save(root / name / f"{i:04d}.png")
    return str(root)


class DataTest(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.tmp = tempfile.TemporaryDirectory()
        base = Path(cls.tmp.name)
        cls.root = make_folder(base / "even", {"alpha": 30, "beta": 30, "gamma": 30})

    @classmethod
    def tearDownClass(cls):
        cls.tmp.cleanup()

    def test_class_counts(self):
        self.assertEqual(class_counts(self.root), {"alpha": 30, "beta": 30, "gamma": 30})

    def test_split_is_stratified_and_disjoint(self):
        targets = [0] * 30 + [1] * 30 + [2] * 30
        train_idx, val_idx = split_indices(targets, 0.2, seed=42)
        self.assertEqual(len(val_idx), 18, "20 % of 90 images")
        self.assertEqual(set(train_idx) & set(val_idx), set(), "an image is in both splits")
        self.assertEqual(sorted(train_idx + val_idx), list(range(90)))
        for cls in (0, 1, 2):
            in_val = sum(1 for i in val_idx if targets[i] == cls)
            self.assertEqual(in_val, 6, f"class {cls}: validation images")

    def test_split_is_reproducible(self):
        targets = [0] * 30 + [1] * 30
        self.assertEqual(split_indices(targets, 0.2, 42), split_indices(targets, 0.2, 42))
        self.assertNotEqual(split_indices(targets, 0.2, 42), split_indices(targets, 0.2, 7))

    def test_a_tiny_class_still_gets_a_validation_image(self):
        targets = [0, 0, 0, 1, 1, 1]  # 10 % of 3 images rounds to 0
        _, val_idx = split_indices(targets, 0.1, seed=0)
        self.assertEqual(sorted(targets[i] for i in val_idx), [0, 1])

    def test_batches_have_the_shape_the_cnn_expects(self):
        train_loader, val_loader, classes = make_loaders(self.root, batch_size=8)
        xb, yb = next(iter(train_loader))
        self.assertEqual(tuple(xb.shape), (8, 1, 28, 28))
        self.assertEqual(yb.dtype, torch.int64)
        self.assertEqual(classes, ["alpha", "beta", "gamma"])
        self.assertEqual(len(train_loader.dataset), 72)
        self.assertEqual(len(val_loader.dataset), 18)

    def test_training_pixels_are_normalised_with_training_statistics(self):
        train_loader, _, _ = make_loaders(self.root, batch_size=8)
        pixels = torch.cat([xb.flatten() for xb, _ in train_loader])
        self.assertAlmostEqual(pixels.mean().item(), 0.0, places=4)
        self.assertAlmostEqual(pixels.std().item(), 1.0, places=3)


if __name__ == "__main__":
    unittest.main()