Capstone Project — Programuoki

Code

capstone.py — the baseline

What you just watched is capstone.py. Here you get it as a starting point: complete, working, and with defaults we managed to justify.

Your job

Create capstone.py and test_capstone.py (on the right). Both are complete.

All the important parts are old acquaintances: make_loaders (lesson 9), build_resnet (lesson 10), fit (lesson 7), predict_all, confusion_matrix, print_report and predict_image (lesson 11). Only these things are new:

--augment / --no-augment — off by default. Why — in the next step;
--patience — how many epochs to wait when the validation loss stops improving (0 = never stop);
--out and --summary — where to save the model and the summary. Different runs must get different names, otherwise one overwrites the other;
the test/ folder is opened at the very start, before training: with a wrong --test-dir the script stops within seconds, not after minutes of training;
at the end the demo shows up to five random test/ photos, in any format ImageFolder recognises.

The test

test_capstone.py does not measure whether the model is good (its data is random noise). It checks the plumbing:

$ python -m unittest test_capstone -v
test_augmentation_can_be_switched_on_and_off ... ok
test_augmentation_is_off_by_default ... ok
test_a_missing_test_folder_fails_before_any_training ... ok
test_a_test_folder_with_other_classes_is_refused ... ok
test_runs_end_to_end_and_writes_its_outputs ... ok
test_the_demo_prints_one_line_per_picture_up_to_five ... ok

Ran 6 tests in 4.208s

OK

It creates a tiny folder (2 classes, 10 training images and 4 test/ images each, the latter .jpg), runs the whole script with --epochs 1 --size 64, and checks that:

model.pt was saved with the classes and size, and summary.json with five keys;
a missing test/ folder stops the script before training;
a test/ with different classes is refused with the message test folder must have the same class folders;
augmentation is off by default.

A trap this test catches. A demo that only looks for *.png crashes on a folder of .jpg photos (phone photos usually are) with ValueError: Sample larger than population or is negative — after training and evaluation have finished. That is why the test's test/ photos are .jpg: we checked that version, and the test catches it.

Run it on your data

python capstone.py --data data --test-dir test --out noaug.pt --summary noaug.json

On our stand-in it took 151.5 seconds (another time — 221 s). Yours depends on the number of images and your computer; on a CPU it is a few minutes. Compare with the previous step and do not expect the same numbers: other classes, other images, another split.

When it does not pass

AssertionError: test folder must have the same class folders — the folder names in test/ do not match data/. Check the spelling and whether an extra folder appeared.
FileNotFoundError for test/ — a wrong --test-dir. You now see it at once, not after training.
ValueError: Sample larger than population or is negative — if you changed the demo's file search back to *.png and your photos are .jpg.

capstone.py CREATE

Create this file in ./

"""capstone.py — the whole project in one run: data -> augment -> transfer -> evaluate -> demo."""

import argparse
import json
import random
import time
from pathlib import Path

import torch
from PIL import Image
from torchvision.datasets.folder import IMG_EXTENSIONS

from data import make_eval_loader, make_loaders
from evaluate import confusion_matrix, predict_all, print_report
from predict import predict_image
from train import evaluate, fit
from transfer import build_resnet, count_trainable


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser()
    parser.add_argument("--data", default="data", help="your training folder")
    parser.add_argument("--test-dir", default="test", help="NEW images, touched once")
    parser.add_argument("--epochs", type=int, default=8)
    parser.add_argument("--lr", type=float, default=0.001)
    parser.add_argument("--size", type=int, default=224)
    parser.add_argument("--augment", action=argparse.BooleanOptionalAction, default=False)
    parser.add_argument("--patience", type=int, default=2, help="0 = never stop early")
    parser.add_argument("--out", default="model.pt")
    parser.add_argument("--summary", default="summary.json")
    return parser


def main() -> None:
    args = build_parser().parse_args()

    start = time.time()
    torch.manual_seed(42)
    random.seed(42)

    # 1. data (augmentation is opt-in: --augment; step 4 shows why it is off)
    train_loader, val_loader, classes = make_loaders(
        args.data, batch_size=32, augment=args.augment, size=args.size, rgb=True
    )
    print(f"classes: {classes}")
    print(f"train images: {len(train_loader.dataset)}, val images: {len(val_loader.dataset)}")

    # open the test folder NOW, so a wrong --test-dir fails before minutes of training
    test_loader, test_classes = make_eval_loader(args.test_dir, size=args.size)
    assert test_classes == classes, "test folder must have the same class folders"

    # 2. transfer learning: pretrained ResNet-18, last block + new head trainable
    model = build_resnet(len(classes), finetune_last_block=True)
    trainable, total = count_trainable(model)
    print(f"trainable parameters: {trainable:,} of {total:,}")
    history = fit(
        model, train_loader, args.epochs, args.lr, val_loader, patience=args.patience or None
    )

    # 3. honest validation metrics (the model picked its best epoch on these)
    _, val_acc = evaluate(model, val_loader)
    preds, targets = predict_all(model, val_loader)
    print(f"\nvalidation accuracy: {val_acc:.1%}")
    print_report(confusion_matrix(targets, preds, len(classes)), classes)

    # 4. the held-out test folder: images no decision was ever based on
    preds, targets = predict_all(model, test_loader)
    test_acc = (preds == targets).float().mean().item()
    print(f"\nTEST accuracy: {test_acc:.1%} on {len(targets)} new images")
    print_report(confusion_matrix(targets, preds, len(classes)), classes)

    # 5. save the model and a machine-readable summary
    torch.save(
        {"state_dict": model.state_dict(), "classes": classes, "size": args.size},
        args.out,
    )
    summary = {
        "classes": classes,
        "epochs_run": len(history),
        "val_accuracy": round(val_acc, 4),
        "test_accuracy": round(test_acc, 4),
        "seconds": round(time.time() - start, 1),
    }
    Path(args.summary).write_text(json.dumps(summary, indent=2))
    print(f"\nsaved {args.out} and {args.summary}")

    # 6. demo: up to five random test images through the same path a user's photo takes
    model.eval()
    paths = sorted(
        p for p in Path(args.test_dir).rglob("*") if str(p).lower().endswith(IMG_EXTENSIONS)
    )
    for path in random.sample(paths, min(5, len(paths))):
        ranked = predict_image(model, classes, args.size, Image.open(path))
        name, prob = ranked[0]
        print(f"{path.parent.name}/{path.name}: predicted {name} ({prob:.0%})")


if __name__ == "__main__":
    main()

test_capstone.py CREATE

Create this file in ./

"""test_capstone.py — a smoke test of the whole pipeline on a tiny made-up dataset.

It does NOT measure how good your model is (the data is random noise). It checks
the plumbing: the script runs end to end, writes what it promises to write, and
refuses a test folder whose classes do not match.
Run: python -m unittest test_capstone -v
"""

import json
import os
import tempfile
import unittest
from contextlib import redirect_stdout
from io import StringIO
from unittest import mock

import torch
from PIL import Image

import capstone


def make_folder(root: str, per_class: int, ext: str, classes=("bright", "dark")) -> None:
    """Tiny 28x28 grey images: 'bright' ones around 200, 'dark' ones around 50."""
    gen = torch.Generator().manual_seed(1)
    for cls in classes:
        os.makedirs(os.path.join(root, cls))
        base = 200 if cls == classes[0] else 50
        for k in range(per_class):
            noise = torch.randint(-20, 20, (28, 28), generator=gen)
            pixels = (base + noise).clamp(0, 255).to(torch.uint8).numpy()
            Image.fromarray(pixels).save(os.path.join(root, cls, f"{k:03d}{ext}"))


class TestParser(unittest.TestCase):
    def test_augmentation_is_off_by_default(self):
        self.assertFalse(capstone.build_parser().parse_args([]).augment)

    def test_augmentation_can_be_switched_on_and_off(self):
        parser = capstone.build_parser()
        self.assertTrue(parser.parse_args(["--augment"]).augment)
        self.assertFalse(parser.parse_args(["--no-augment"]).augment)


class TestPipeline(unittest.TestCase):
    def setUp(self):
        self.tmp = tempfile.TemporaryDirectory()
        self.data = os.path.join(self.tmp.name, "data")
        self.test = os.path.join(self.tmp.name, "test")
        make_folder(self.data, per_class=10, ext=".png")
        make_folder(self.test, per_class=4, ext=".jpg")  # jpg on purpose
        self.out = os.path.join(self.tmp.name, "model.pt")
        self.summary = os.path.join(self.tmp.name, "summary.json")

    def tearDown(self):
        self.tmp.cleanup()

    def run_capstone(self, test_dir: str) -> str:
        argv = [
            "capstone.py", "--data", self.data, "--test-dir", test_dir,
            "--epochs", "1", "--size", "64", "--patience", "0",
            "--out", self.out, "--summary", self.summary,
        ]
        buf = StringIO()
        with mock.patch("sys.argv", argv), redirect_stdout(buf):
            capstone.main()
        return buf.getvalue()

    def test_runs_end_to_end_and_writes_its_outputs(self):
        output = self.run_capstone(self.test)
        self.assertIn("TEST accuracy", output)
        with open(self.summary) as f:
            summary = json.load(f)
        self.assertEqual(summary["classes"], ["bright", "dark"])
        self.assertEqual(
            set(summary), {"classes", "epochs_run", "val_accuracy", "test_accuracy", "seconds"}
        )
        self.assertTrue(0.0 <= summary["test_accuracy"] <= 1.0)
        ckpt = torch.load(self.out, map_location="cpu")
        self.assertEqual(ckpt["classes"], ["bright", "dark"])
        self.assertEqual(ckpt["size"], 64)

    def test_the_demo_prints_one_line_per_picture_up_to_five(self):
        output = self.run_capstone(self.test)
        demo = [l for l in output.splitlines() if ": predicted " in l]
        self.assertEqual(len(demo), 5)  # 8 test images, the demo shows five

    def test_a_missing_test_folder_fails_before_any_training(self):
        with mock.patch("capstone.fit", side_effect=AssertionError("it trained first")):
            with self.assertRaises(FileNotFoundError):
                self.run_capstone(os.path.join(self.tmp.name, "no_such_folder"))

    def test_a_test_folder_with_other_classes_is_refused(self):
        other = os.path.join(self.tmp.name, "other")
        make_folder(other, per_class=2, ext=".png", classes=("cat", "dog"))
        with self.assertRaises(AssertionError):
            self.run_capstone(other)


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