Diabetic foot-ulcer image classification

import numpy as np
import torch
from torch import nn

def wound_photos(n, seed, dark_share):
    """Synthetic 28x28 'foot photos'. Skin tone sets the overall brightness; an ulcer is a darker,
    textured patch. dark_share: fraction of photos with darker skin tones."""
    rng = np.random.default_rng(seed)
    yy, xx = np.mgrid[:28, :28]
    dark = rng.random(n) < dark_share
    tone = np.where(dark, rng.uniform(-1.6, -0.8, n), rng.uniform(0.8, 1.6, n))
    images = tone[:, None, None] + rng.normal(0, 0.6, (n, 28, 28))
    labels = rng.integers(0, 2, n)
    for i in np.where(labels == 1)[0]:
        cy, cx, r = rng.integers(7, 21), rng.integers(7, 21), rng.uniform(3, 5)
        patch = (yy - cy) ** 2 + (xx - cx) ** 2 < r ** 2
        images[i][patch] -= 1.2 + rng.normal(0, 0.4, patch.sum())
    to_t = lambda a: torch.tensor(a, dtype=torch.float32)
    return to_t(images).unsqueeze(1), torch.tensor(labels), dark

class SmallCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Conv2d(1, 8, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
            nn.Conv2d(8, 16, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
            nn.Flatten(), nn.Linear(16 * 7 * 7, 2))
    def forward(self, x):
        return self.net(x)

def train(X, y, epochs=8):
    torch.manual_seed(0)
    model = SmallCNN()
    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_by_tone(model, X, y, dark):
    with torch.no_grad():
        correct = (model(X).argmax(1) == y).numpy()
    return correct[~dark].mean(), correct[dark].mean()

X_tr, y_tr, dark_tr = wound_photos(4000, seed=1, dark_share=0.05)       # 95% lighter skin in training
X_te, y_te, dark_te = wound_photos(2000, seed=2, dark_share=0.5)        # balanced test set
model = train(X_tr, y_tr)
light_acc, dark_acc = accuracy_by_tone(model, X_te, y_te, dark_te)
print(f"Trained on 5% darker skin:  lighter skin {light_acc:.1%}, darker skin {dark_acc:.1%}")

# Trained on 5% darker skin:  lighter skin 99.7%, darker skin 71.6%

X_bal, y_bal, _ = wound_photos(4000, seed=3, dark_share=0.5)            # diverse training set
light_acc, dark_acc = accuracy_by_tone(train(X_bal, y_bal), X_te, y_te, dark_te)
print(f"Trained on 50% darker skin: lighter skin {light_acc:.1%}, darker skin {dark_acc:.1%}")

def normalise(X):                                                         # per-image standardisation
    return (X - X.mean(dim=(2, 3), keepdim=True)) / X.std(dim=(2, 3), keepdim=True)
light_acc, dark_acc = accuracy_by_tone(train(normalise(X_tr), y_tr), normalise(X_te), y_te, dark_te)
print(f"5% darker skin, normalised: lighter skin {light_acc:.1%}, darker skin {dark_acc:.1%}")

# Trained on 50% darker skin: lighter skin 99.2%, darker skin 97.6%
# 5% darker skin, normalised: lighter skin 99.9%, darker skin 100.0%

# Wound area (cm2) measured weekly from segmented photos for three synthetic patients
areas = {"patient A": [4.0, 3.1, 2.4, 1.7, 1.1], "patient B": [3.5, 3.4, 3.3, 3.0, 2.9],
         "patient C": [6.2, 5.6, 4.0, 3.5, 3.0]}
for name, a in areas.items():
    reduction = 1 - a[4] / a[0]                        # percent area reduction after 4 weeks
    flag = "on track" if reduction >= 0.5 else "REVIEW: under 50% reduction at 4 weeks"
    print(f"{name}: {reduction:.0%} smaller after 4 weeks -> {flag}")

# patient A: 72% smaller after 4 weeks -> on track
# patient B: 17% smaller after 4 weeks -> REVIEW: under 50% reduction at 4 weeks
# patient C: 52% smaller after 4 weeks -> on track