File size: 5,732 Bytes
7615398 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 | """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()
|