code/train_redseg.py

"""Phase 2: train a small U-Net red-mark segmenter (pure torch, no torchvision).

Binary red vs not. BCE+Dice loss, heavy photometric augmentation so the model keys on
relative red-ness + spatial context (not absolute colour). Val IoU on the gold synthetic
set. Saves best weights to redseg_best.pt.
"""
import os, glob, random, time, numpy as np
from PIL import Image
import torch, torch.nn as nn, torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader

DATA = os.environ.get("DATA", "/home/ubuntu/Projects/AutoEval/red_seg/data")
OUT = os.environ.get("OUTDIR", "/home/ubuntu/Projects/AutoEval/red_seg")
EPOCHS = int(os.environ.get("EPOCHS", "40"))
CROP = 256; BS = int(os.environ.get("BS", "16"))
dev = "cuda" if torch.cuda.is_available() else "cpu"
random.seed(0); np.random.seed(0); torch.manual_seed(0)


def rgb_to_hsv_shift(a, dh):
    import colorsys
    # vectorized hue shift via numpy (approx): operate in HSV
    r, g, b = a[..., 0] / 255., a[..., 1] / 255., a[..., 2] / 255.
    mx = a.max(2) / 255.; mn = a.min(2) / 255.; df = mx - mn + 1e-9
    h = np.zeros_like(mx)
    m = mx == r; h[m] = ((g - b) / df)[m] % 6
    m = mx == g; h[m] = ((b - r) / df)[m] + 2
    m = mx == b; h[m] = ((r - g) / df)[m] + 4
    h = (h / 6.0 + dh) % 1.0
    s = df / (mx + 1e-9); v = mx
    i = (h * 6).astype(int) % 6; f = h * 6 - i
    p = v * (1 - s); q = v * (1 - f * s); t = v * (1 - (1 - f) * s)
    out = np.zeros_like(a, float)
    for idx, (R, G, B) in enumerate([(v, t, p), (q, v, p), (p, v, t), (p, q, v), (t, p, v), (v, p, q)]):
        mm = i == idx
        out[..., 0][mm] = R[mm]; out[..., 1][mm] = G[mm]; out[..., 2][mm] = B[mm]
    return np.clip(out * 255, 0, 255)


class RedSet(Dataset):
    def __init__(self, split, aug):
        self.ii = sorted(glob.glob(f"{DATA}/{split}/img/*.png"))
        self.aug = aug
    def __len__(self): return len(self.ii)
    def __getitem__(self, k):
        ip = self.ii[k]; mp = ip.replace("/img/", "/mask/")
        a = np.asarray(Image.open(ip).convert("RGB")).astype(np.float32)
        m = (np.asarray(Image.open(mp).convert("L")) > 127).astype(np.float32)
        H, W = m.shape
        if self.aug:
            y = random.randint(0, H - CROP); x = random.randint(0, W - CROP)
            a = a[y:y + CROP, x:x + CROP]; m = m[y:y + CROP, x:x + CROP]
            if random.random() < 0.5: a = a[:, ::-1].copy(); m = m[:, ::-1].copy()
            if random.random() < 0.5: a = a[::-1].copy(); m = m[::-1].copy()
            kk = random.randint(0, 3); a = np.rot90(a, kk).copy(); m = np.rot90(m, kk).copy()
            a *= random.uniform(0.7, 1.3)                       # brightness
            a = (a - 128) * random.uniform(0.7, 1.3) + 128      # contrast
            if random.random() < 0.6: a = rgb_to_hsv_shift(np.clip(a,0,255), random.uniform(-0.06, 0.06))
            a += np.random.randn(*a.shape).astype(np.float32) * random.uniform(0, 8)  # noise
            a = np.clip(a, 0, 255)
        else:
            a = a[:CROP * 2, :CROP * 2]; m = m[:CROP * 2, :CROP * 2]  # full-ish for val
        a = torch.from_numpy(a.transpose(2, 0, 1) / 255.0).float()
        return a, torch.from_numpy(m[None]).float()


def conv(i, o): return nn.Sequential(nn.Conv2d(i, o, 3, padding=1), nn.BatchNorm2d(o), nn.ReLU(True),
                                     nn.Conv2d(o, o, 3, padding=1), nn.BatchNorm2d(o), nn.ReLU(True))
class UNet(nn.Module):
    def __init__(self, c=32):
        super().__init__()
        self.d1 = conv(3, c); self.d2 = conv(c, c*2); self.d3 = conv(c*2, c*4); self.d4 = conv(c*4, c*8)
        self.p = nn.MaxPool2d(2)
        self.u3 = nn.ConvTranspose2d(c*8, c*4, 2, 2); self.c3 = conv(c*8, c*4)
        self.u2 = nn.ConvTranspose2d(c*4, c*2, 2, 2); self.c2 = conv(c*4, c*2)
        self.u1 = nn.ConvTranspose2d(c*2, c, 2, 2); self.c1 = conv(c*2, c)
        self.o = nn.Conv2d(c, 1, 1)
    def forward(self, x):
        e1 = self.d1(x); e2 = self.d2(self.p(e1)); e3 = self.d3(self.p(e2)); e4 = self.d4(self.p(e3))
        d = self.c3(torch.cat([self.u3(e4), e3], 1))
        d = self.c2(torch.cat([self.u2(d), e2], 1))
        d = self.c1(torch.cat([self.u1(d), e1], 1))
        return self.o(d)


def dice_loss(p, t):
    p = torch.sigmoid(p); n = (2 * (p * t).sum((2, 3)) + 1) / ((p + t).sum((2, 3)) + 1)
    return 1 - n.mean()

@torch.no_grad()
def evaluate(net, dl):
    net.eval(); I = U = 0.0
    for a, m in dl:
        a, m = a.to(dev), m.to(dev)
        p = (torch.sigmoid(net(a)) > 0.5).float()
        I += (p * m).sum().item(); U += ((p + m) >= 1).float().sum().item()
    return I / (U + 1e-9)


def main():
    tr = DataLoader(RedSet("train", True), BS, shuffle=True, num_workers=8, drop_last=True)
    vs = DataLoader(RedSet("valsyn", False), 4, num_workers=4)
    print(f"device={dev} train={len(tr.dataset)} valsyn={len(vs.dataset)}", flush=True)
    net = UNet().to(dev); opt = torch.optim.Adam(net.parameters(), 1e-3)
    sch = torch.optim.lr_scheduler.CosineAnnealingLR(opt, EPOCHS)
    best = 0.0
    for ep in range(EPOCHS):
        net.train(); t0 = time.time(); tot = 0
        for a, m in tr:
            a, m = a.to(dev), m.to(dev)
            p = net(a); loss = F.binary_cross_entropy_with_logits(p, m) + dice_loss(p, m)
            opt.zero_grad(); loss.backward(); opt.step(); tot += loss.item()
        sch.step(); iou = evaluate(net, vs)
        if iou > best:
            best = iou; torch.save(net.state_dict(), f"{OUT}/redseg_best.pt")
        print(f"ep{ep+1}/{EPOCHS} loss={tot/len(tr):.4f} valsyn_IoU={iou:.3f} best={best:.3f} {time.time()-t0:.0f}s", flush=True)
    print(f"=== train done best_valsyn_IoU={best:.3f} -> {OUT}/redseg_best.pt ===", flush=True)


if __name__ == "__main__":
    main()