CNNs for chest X-ray and retinal images

import numpy as np
import torch
from torch import nn

def make_images(n, seed, marker_with_lesion=0.0):
    """Synthetic 28x28 'scans': noisy background; half contain a faint round lesion at a random place.
    marker_with_lesion: chance that lesion images also carry a bright corner marker (a shortcut)."""
    rng = np.random.default_rng(seed)
    yy, xx = np.mgrid[:28, :28]
    images = rng.normal(0, 1, (n, 28, 28))
    labels = rng.integers(0, 2, n)
    for i in np.where(labels == 1)[0]:
        cy, cx, r = rng.integers(6, 22), rng.integers(6, 22), rng.uniform(2.5, 4.0)
        images[i] += 1.0 * ((yy - cy) ** 2 + (xx - cx) ** 2 < r ** 2)
        if rng.random() < marker_with_lesion:
            images[i, :3, :3] += 4.0                     # e.g. a portable-scanner label
    return torch.tensor(images, dtype=torch.float32).unsqueeze(1), torch.tensor(labels)

class SmallCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(1, 8, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),      # 28 -> 14
            nn.Conv2d(8, 16, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),     # 14 -> 7
        )
        self.head = nn.Sequential(nn.Flatten(), nn.Linear(16 * 7 * 7, 2))
    def forward(self, x):
        return self.head(self.features(x))

def train(model, X, y, epochs=8, seed=0):
    torch.manual_seed(seed)
    opt = torch.optim.Adam(model.parameters(), lr=1e-3)
    for _ in range(epochs):
        for idx in torch.randperm(len(X)).split(64):
            opt.zero_grad()
            nn.functional.cross_entropy(model(X[idx]), y[idx]).backward()
            opt.step()
    return model

def accuracy(model, X, y):
    with torch.no_grad():
        return (model(X).argmax(1) == y).float().mean().item()

X_train, y_train = make_images(4000, seed=1)
X_test, y_test = make_images(1000, seed=2)
torch.manual_seed(0)
cnn = train(SmallCNN(), X_train, y_train)
print(f"CNN test accuracy: {accuracy(cnn, X_test, y_test):.3f}")
print(f"Parameters: {sum(p.numel() for p in cnn.parameters()):,}")

# CNN test accuracy: 0.946
# Parameters: 2,818

from sklearn.linear_model import LogisticRegression

flat = LogisticRegression(max_iter=2000).fit(X_train.reshape(len(X_train), -1).numpy(), y_train.numpy())
print(f"Logistic regression on raw pixels: "
      f"{flat.score(X_test.reshape(len(X_test), -1).numpy(), y_test.numpy()):.3f}")

# Logistic regression on raw pixels: 0.738

X_dev, y_dev = make_images(4000, seed=3, marker_with_lesion=0.9)
torch.manual_seed(0)
shortcut_cnn = train(SmallCNN(), X_dev, y_dev)
X_same, y_same = make_images(1000, seed=4, marker_with_lesion=0.9)
X_new, y_new = make_images(1000, seed=5, marker_with_lesion=0.0)
print(f"Same hospital (marker present): {accuracy(shortcut_cnn, X_same, y_same):.3f}")
print(f"New hospital (no marker):       {accuracy(shortcut_cnn, X_new, y_new):.3f}")

# Same hospital (marker present): 0.984
# New hospital (no marker):       0.813