| import os |
| import sys |
| import json |
| import math |
| import random |
| import argparse |
| from pathlib import Path |
| import numpy as np |
| from PIL import Image |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from torch.utils.data import Dataset, DataLoader |
| DATA_ROOT = Path('dataset') |
| TRAIN_IMG = DATA_ROOT / 'train' / 'images' |
| TRAIN_SC = DATA_ROOT / 'train' / 'scribbles' |
| TRAIN_GT = DATA_ROOT / 'train' / 'ground_truth' |
| TEST_IMG = DATA_ROOT / 'test1' / 'images' |
| TEST_SC = DATA_ROOT / 'test1' / 'scribbles' |
| TEST_PRED = DATA_ROOT / 'test1' / 'predictions' |
| TRAIN_H = int(os.environ.get('TRAIN_H', '384')) |
| TRAIN_W = int(os.environ.get('TRAIN_W', '512')) |
| ORIG_H, ORIG_W = (375, 500) |
| CKPT_DIR = Path(os.environ.get('CKPT_DIR', 'runs_global_unet')) |
| CKPT_DIR.mkdir(exist_ok=True) |
|
|
| def list_train_pairs(): |
| pairs = [] |
| for img_path in sorted(TRAIN_IMG.iterdir()): |
| if img_path.name.startswith('.'): |
| continue |
| stem = img_path.stem |
| sc_path = TRAIN_SC / f'{stem}.png' |
| gt_path = TRAIN_GT / f'{stem}.png' |
| if sc_path.exists() and gt_path.exists(): |
| pairs.append((stem, img_path, sc_path, gt_path)) |
| return pairs |
|
|
| def list_test_pairs(): |
| pairs = [] |
| for img_path in sorted(TEST_IMG.iterdir()): |
| if img_path.name.startswith('.'): |
| continue |
| stem = img_path.stem |
| sc_path = TEST_SC / f'{stem}.png' |
| if sc_path.exists(): |
| pairs.append((stem, img_path, sc_path)) |
| return pairs |
|
|
| def list_pseudo_pairs(pseudo_label_method='v3v4'): |
| pairs = [] |
| for setname in ['test1', 'test2']: |
| img_dir = Path(f'dataset/{setname}/images') |
| sc_dir = Path(f'dataset/{setname}/scribbles') |
| gt_dir = Path(f'dataset/{setname}/predictions_{pseudo_label_method}') |
| if not gt_dir.exists(): |
| continue |
| for ip in sorted(img_dir.iterdir()): |
| if ip.name.startswith('.'): |
| continue |
| stem = ip.stem |
| sp = sc_dir / f'{stem}.png' |
| gp = gt_dir / f'{stem}.png' |
| if sp.exists() and gp.exists(): |
| pairs.append((stem, ip, sp, gp)) |
| return pairs |
|
|
| def load_palette(): |
| any_gt = next(TRAIN_GT.glob('*.png')) |
| return Image.open(any_gt).getpalette() |
|
|
| def encode_scribble(sc): |
| bg_ch = (sc == 0).astype(np.float32) |
| fg_ch = (sc == 1).astype(np.float32) |
| return np.stack([bg_ch, fg_ch], axis=0) |
|
|
| def random_affine(img, sc, gt, rng): |
| H, W = img.shape[:2] |
| angle = rng.uniform(-12, 12) |
| scale = rng.uniform(0.85, 1.2) |
| tx = rng.uniform(-0.05, 0.05) * W |
| ty = rng.uniform(-0.05, 0.05) * H |
| cx, cy = (W / 2, H / 2) |
| a = math.radians(angle) |
| cos_a, sin_a = (math.cos(a) * scale, math.sin(a) * scale) |
| M = np.array([[cos_a, -sin_a, (1 - cos_a) * cx + sin_a * cy + tx], [sin_a, cos_a, (1 - cos_a) * cy - sin_a * cx + ty]], dtype=np.float32) |
| import cv2 |
| img_a = cv2.warpAffine(img, M, (W, H), flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_REFLECT_101) |
| sc_a = cv2.warpAffine(sc, M, (W, H), flags=cv2.INTER_NEAREST, borderMode=cv2.BORDER_CONSTANT, borderValue=255) |
| gt_a = cv2.warpAffine(gt, M, (W, H), flags=cv2.INTER_NEAREST, borderMode=cv2.BORDER_CONSTANT, borderValue=0) |
| return (img_a, sc_a, gt_a) |
|
|
| def color_jitter(img, rng): |
| img_f = img.astype(np.float32) / 255.0 |
| img_f = img_f * rng.uniform(0.8, 1.2) |
| mean = img_f.mean(axis=(0, 1), keepdims=True) |
| img_f = (img_f - mean) * rng.uniform(0.8, 1.2) + mean |
| if rng.random() < 0.7: |
| gray = img_f.mean(axis=2, keepdims=True) |
| img_f = img_f * rng.uniform(0.7, 1.3) + gray * (1 - rng.uniform(0.7, 1.3)) |
| img_f = np.clip(img_f, 0, 1) |
| return (img_f * 255).astype(np.uint8) |
|
|
| class ScribbleSegDataset(Dataset): |
|
|
| def __init__(self, pairs, train=True, image_size=(TRAIN_H, TRAIN_W), cutmix_p=0.0): |
| self.pairs = pairs |
| self.train = train |
| self.H, self.W = image_size |
| self.cutmix_p = cutmix_p |
| self.mean = np.array([0.485, 0.456, 0.406], dtype=np.float32) |
| self.std = np.array([0.229, 0.224, 0.225], dtype=np.float32) |
|
|
| def __len__(self): |
| return len(self.pairs) |
|
|
| def _load_one(self, idx): |
| import cv2 |
| stem, img_p, sc_p, gt_p = self.pairs[idx] |
| img = np.array(Image.open(img_p).convert('RGB')) |
| sc = np.array(Image.open(sc_p).convert('L')) |
| gt = np.array(Image.open(gt_p)) |
| if img.shape[:2] != (self.H, self.W): |
| img = cv2.resize(img, (self.W, self.H), interpolation=cv2.INTER_LINEAR) |
| sc = cv2.resize(sc, (self.W, self.H), interpolation=cv2.INTER_NEAREST) |
| gt = cv2.resize(gt, (self.W, self.H), interpolation=cv2.INTER_NEAREST) |
| return (stem, img, sc, gt) |
|
|
| def __getitem__(self, idx): |
| stem, img, sc, gt = self._load_one(idx) |
| rng = random.Random() |
| if self.train: |
| if rng.random() < 0.5: |
| img = img[:, ::-1, :].copy() |
| sc = sc[:, ::-1].copy() |
| gt = gt[:, ::-1].copy() |
| img, sc, gt = random_affine(img, sc, gt, rng) |
| img = color_jitter(img, rng) |
| if rng.random() < 0.3: |
| drop_mask = (sc != 255) & (np.random.rand(*sc.shape) < 0.3) |
| sc = sc.copy() |
| sc[drop_mask] = 255 |
| if self.cutmix_p > 0 and rng.random() < self.cutmix_p: |
| j = rng.randint(0, len(self.pairs) - 1) |
| _, img2, sc2, gt2 = self._load_one(j) |
| rh = rng.randint(int(0.3 * self.H), int(0.6 * self.H)) |
| rw = rng.randint(int(0.3 * self.W), int(0.6 * self.W)) |
| ry = rng.randint(0, self.H - rh) |
| rx = rng.randint(0, self.W - rw) |
| img = img.copy() |
| sc = sc.copy() |
| gt = gt.copy() |
| img[ry:ry + rh, rx:rx + rw] = img2[ry:ry + rh, rx:rx + rw] |
| sc[ry:ry + rh, rx:rx + rw] = sc2[ry:ry + rh, rx:rx + rw] |
| gt[ry:ry + rh, rx:rx + rw] = gt2[ry:ry + rh, rx:rx + rw] |
| img_f = img.astype(np.float32) / 255.0 |
| img_f = (img_f - self.mean) / self.std |
| img_t = torch.from_numpy(img_f.transpose(2, 0, 1)) |
| sc_enc = encode_scribble(sc) |
| sc_t = torch.from_numpy(sc_enc) |
| x = torch.cat([img_t, sc_t], dim=0) |
| gt_bin = (gt > 0).astype(np.float32) |
| y = torch.from_numpy(gt_bin) |
| return (x, y, stem) |
|
|
| class ConvBlock(nn.Module): |
|
|
| def __init__(self, in_ch, out_ch): |
| super().__init__() |
| self.block = nn.Sequential(nn.Conv2d(in_ch, out_ch, 3, padding=1, bias=False), nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True), nn.Conv2d(out_ch, out_ch, 3, padding=1, bias=False), nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True)) |
|
|
| def forward(self, x): |
| return self.block(x) |
|
|
| class UNet(nn.Module): |
|
|
| def __init__(self, in_ch=5, base=48, out_ch=1): |
| super().__init__() |
| c1, c2, c3, c4, c5 = (base, base * 2, base * 4, base * 8, base * 16) |
| self.enc1 = ConvBlock(in_ch, c1) |
| self.enc2 = ConvBlock(c1, c2) |
| self.enc3 = ConvBlock(c2, c3) |
| self.enc4 = ConvBlock(c3, c4) |
| self.bottleneck = ConvBlock(c4, c5) |
| self.pool = nn.MaxPool2d(2) |
| self.up4 = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False) |
| self.dec4 = ConvBlock(c5 + c4, c4) |
| self.up3 = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False) |
| self.dec3 = ConvBlock(c4 + c3, c3) |
| self.up2 = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False) |
| self.dec2 = ConvBlock(c3 + c2, c2) |
| self.up1 = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False) |
| self.dec1 = ConvBlock(c2 + c1, c1) |
| self.head = nn.Conv2d(c1, out_ch, 1) |
| self._init_weights() |
|
|
| def _init_weights(self): |
| for m in self.modules(): |
| if isinstance(m, nn.Conv2d): |
| nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') |
| if m.bias is not None: |
| nn.init.zeros_(m.bias) |
| elif isinstance(m, nn.BatchNorm2d): |
| nn.init.ones_(m.weight) |
| nn.init.zeros_(m.bias) |
|
|
| def forward(self, x): |
| e1 = self.enc1(x) |
| e2 = self.enc2(self.pool(e1)) |
| e3 = self.enc3(self.pool(e2)) |
| e4 = self.enc4(self.pool(e3)) |
| b = self.bottleneck(self.pool(e4)) |
| d4 = self.dec4(torch.cat([self.up4(b), e4], 1)) |
| d3 = self.dec3(torch.cat([self.up3(d4), e3], 1)) |
| d2 = self.dec2(torch.cat([self.up2(d3), e2], 1)) |
| d1 = self.dec1(torch.cat([self.up1(d2), e1], 1)) |
| return self.head(d1) |
|
|
| def soft_dice_loss(logits, target, eps=1e-06): |
| p = torch.sigmoid(logits).squeeze(1) |
| inter = (p * target).sum(dim=(1, 2)) |
| denom = p.sum(dim=(1, 2)) + target.sum(dim=(1, 2)) |
| dice = (2 * inter + eps) / (denom + eps) |
| return 1 - dice.mean() |
|
|
| def combined_loss(logits, target): |
| bce = F.binary_cross_entropy_with_logits(logits.squeeze(1), target) |
| dice = soft_dice_loss(logits, target) |
| return 0.5 * bce + 0.5 * dice |
|
|
| def compute_iou(pred_bin, gt_bin, cls): |
| p = pred_bin == cls |
| g = gt_bin == cls |
| inter = np.logical_and(p, g).sum() |
| union = np.logical_or(p, g).sum() |
| return inter / union if union > 0 else 0.0 |
|
|
| def evaluate_predictions(preds, gts): |
| bg, fg = ([], []) |
| for p, g in zip(preds, gts): |
| bg.append(compute_iou(p, g, 0)) |
| fg.append(compute_iou(p, g, 1)) |
| bg = np.mean(bg) |
| fg = np.mean(fg) |
| return (bg, fg, (bg + fg) / 2) |
|
|
| def train_one_fold(train_pairs, val_pairs, epochs, batch_size, lr, fold_id, device, base=48, cutmix_p=0.0): |
| train_ds = ScribbleSegDataset(train_pairs, train=True, cutmix_p=cutmix_p) |
| val_ds = ScribbleSegDataset(val_pairs, train=False) |
| train_dl = DataLoader(train_ds, batch_size=batch_size, shuffle=True, num_workers=4, pin_memory=True, drop_last=True) |
| val_dl = DataLoader(val_ds, batch_size=batch_size, shuffle=False, num_workers=2, pin_memory=True) |
| model = UNet(in_ch=5, base=base, out_ch=1).to(device) |
| n_params = sum((p.numel() for p in model.parameters())) |
| print(f'[fold {fold_id}] U-Net params: {n_params / 1000000.0:.2f}M (base={base}), train={len(train_ds)}, val={len(val_ds)}') |
| opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.0001) |
| sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=epochs, eta_min=lr / 30) |
| scaler = torch.amp.GradScaler('cuda') |
| best_miou = -1.0 |
| best_state = None |
| log = [] |
| patience = 25 |
| bad_epochs = 0 |
| for epoch in range(epochs): |
| model.train() |
| train_loss = 0.0 |
| n = 0 |
| for x, y, _ in train_dl: |
| x, y = (x.to(device, non_blocking=True), y.to(device, non_blocking=True)) |
| opt.zero_grad(set_to_none=True) |
| with torch.amp.autocast('cuda', dtype=torch.float16): |
| logits = model(x) |
| loss = combined_loss(logits, y) |
| scaler.scale(loss).backward() |
| scaler.step(opt) |
| scaler.update() |
| train_loss += loss.item() * x.size(0) |
| n += x.size(0) |
| train_loss /= n |
| sched.step() |
| model.eval() |
| all_p, all_g = ([], []) |
| with torch.no_grad(): |
| for x, y, _ in val_dl: |
| x = x.to(device, non_blocking=True) |
| with torch.amp.autocast('cuda', dtype=torch.float16): |
| logits = model(x) |
| p = (torch.sigmoid(logits).squeeze(1).float().cpu().numpy() > 0.5).astype(np.uint8) |
| g = y.numpy().astype(np.uint8) |
| for i in range(p.shape[0]): |
| all_p.append(p[i]) |
| all_g.append(g[i]) |
| bg, fg, miou = evaluate_predictions(all_p, all_g) |
| log.append({'epoch': epoch, 'loss': train_loss, 'val_bg': bg, 'val_fg': fg, 'val_miou': miou}) |
| print(f'[fold {fold_id} ep {epoch:03d}] loss={train_loss:.4f} val: bg={bg:.4f} fg={fg:.4f} mIoU={miou:.4f} lr={sched.get_last_lr()[0]:.2e}') |
| if miou > best_miou: |
| best_miou = miou |
| best_state = {k: v.detach().cpu().clone() for k, v in model.state_dict().items()} |
| bad_epochs = 0 |
| else: |
| bad_epochs += 1 |
| if bad_epochs >= patience: |
| print(f'[fold {fold_id}] early stopping at epoch {epoch} (best mIoU={best_miou:.4f})') |
| break |
| fold_dir = CKPT_DIR / f'fold_{fold_id}' |
| fold_dir.mkdir(exist_ok=True) |
| torch.save(best_state, fold_dir / 'best.pth') |
| with open(fold_dir / 'log.json', 'w') as f: |
| json.dump(log, f, indent=2) |
| return best_miou |
|
|
| def cmd_train(args): |
| set_seed(args.seed) |
| device = torch.device(f'cuda:{args.gpu}' if torch.cuda.is_available() else 'cpu') |
| print(f'Device: {device}') |
| pairs = list_train_pairs() |
| print(f'Training pairs: {len(pairs)}') |
| pseudo_pairs = [] |
| if getattr(args, 'pseudo_method', None): |
| pseudo_pairs = list_pseudo_pairs(args.pseudo_method) |
| print(f'Pseudo-labeled pairs ({args.pseudo_method}): {len(pseudo_pairs)}') |
| rng = np.random.RandomState(args.seed) |
| indices = np.arange(len(pairs)) |
| rng.shuffle(indices) |
| if args.folds == 1: |
| n_val = max(1, len(pairs) // 5) |
| splits = [(indices[n_val:], indices[:n_val])] |
| else: |
| fold_arr = np.array_split(indices, args.folds) |
| splits = [] |
| for k in range(args.folds): |
| val_idx = fold_arr[k] |
| train_idx = np.concatenate([fold_arr[i] for i in range(args.folds) if i != k]) |
| splits.append((train_idx, val_idx)) |
| fold_mious = [] |
| for k, (train_idx, val_idx) in enumerate(splits): |
| train_pairs = [pairs[i] for i in train_idx] |
| if pseudo_pairs: |
| train_pairs = train_pairs + pseudo_pairs |
| val_pairs = [pairs[i] for i in val_idx] |
| print(f'\n=== Fold {k + 1}/{len(splits)}: train={len(train_pairs)} ({len(train_idx)} real + {len(pseudo_pairs)} pseudo), val={len(val_pairs)} ===') |
| miou = train_one_fold(train_pairs, val_pairs, epochs=args.epochs, batch_size=args.batch_size, lr=args.lr, fold_id=k, device=device, base=args.base, cutmix_p=args.cutmix_p) |
| fold_mious.append(miou) |
| print('\n=== Cross-validation summary ===') |
| for k, m in enumerate(fold_mious): |
| print(f' fold {k}: {m:.4f}') |
| print(f' mean: {np.mean(fold_mious):.4f} (+/- {np.std(fold_mious):.4f})') |
|
|
| def tta_predict(model, x, device, scales=(1.0,)): |
| model.eval() |
| H, W = (x.shape[-2], x.shape[-1]) |
| probs = [] |
| with torch.no_grad(), torch.amp.autocast('cuda', dtype=torch.float16): |
| for s in scales: |
| if s == 1.0: |
| xs = x |
| else: |
| new_h = int(round(H * s / 32) * 32) |
| new_w = int(round(W * s / 32) * 32) |
| rgb = F.interpolate(x[:, :3], size=(new_h, new_w), mode='bilinear', align_corners=False) |
| sc = F.interpolate(x[:, 3:], size=(new_h, new_w), mode='nearest') |
| xs = torch.cat([rgb, sc], dim=1) |
| p1 = torch.sigmoid(model(xs)) |
| p2 = torch.sigmoid(model(torch.flip(xs, dims=[3]))) |
| p2 = torch.flip(p2, dims=[3]) |
| p = (p1 + p2) / 2 |
| if p.shape[-2:] != (H, W): |
| p = F.interpolate(p, size=(H, W), mode='bilinear', align_corners=False) |
| probs.append(p) |
| return (sum(probs) / len(probs)).squeeze().float().cpu().numpy() |
|
|
| def cmd_predict(args): |
| import cv2 |
| device = torch.device(f'cuda:{args.gpu}' if torch.cuda.is_available() else 'cpu') |
| fold_dirs = sorted(CKPT_DIR.glob('fold_*')) |
| fold_dirs = [f for f in fold_dirs if (f / 'best.pth').exists()] |
| if not fold_dirs: |
| print('No trained models found.') |
| sys.exit(1) |
| print(f'Ensembling {len(fold_dirs)} folds.') |
| models = [] |
| for fd in fold_dirs: |
| m = UNet(in_ch=5, base=args.base, out_ch=1).to(device) |
| m.load_state_dict(torch.load(fd / 'best.pth', map_location=device)) |
| m.eval() |
| models.append(m) |
| palette = load_palette() |
| test_pairs = list_test_pairs() |
| TEST_PRED.mkdir(parents=True, exist_ok=True) |
| mean = np.array([0.485, 0.456, 0.406], dtype=np.float32) |
| std = np.array([0.229, 0.224, 0.225], dtype=np.float32) |
| for stem, img_p, sc_p in test_pairs: |
| img = np.array(Image.open(img_p).convert('RGB')) |
| sc = np.array(Image.open(sc_p).convert('L')) |
| img_r = cv2.resize(img, (TRAIN_W, TRAIN_H), interpolation=cv2.INTER_LINEAR) |
| sc_r = cv2.resize(sc, (TRAIN_W, TRAIN_H), interpolation=cv2.INTER_NEAREST) |
| img_f = (img_r.astype(np.float32) / 255.0 - mean) / std |
| img_t = torch.from_numpy(img_f.transpose(2, 0, 1)) |
| sc_t = torch.from_numpy(encode_scribble(sc_r)) |
| x = torch.cat([img_t, sc_t], dim=0).unsqueeze(0).to(device) |
| prob_sum = None |
| for m in models: |
| p = tta_predict(m, x, device, scales=(0.7, 1.0, 1.3)) |
| prob_sum = p if prob_sum is None else prob_sum + p |
| prob = prob_sum / len(models) |
| prob_full = cv2.resize(prob, (ORIG_W, ORIG_H), interpolation=cv2.INTER_LINEAR) |
| pred = (prob_full > 0.5).astype(np.uint8) |
| pred_snap = pred.copy() |
| pred_snap[sc == 0] = 0 |
| pred_snap[sc == 1] = 1 |
| out_img = Image.fromarray(pred_snap.astype(np.uint8), mode='P') |
| out_img.putpalette(palette) |
| out_img.save(TEST_PRED / f'{stem}.png') |
| print(f'Wrote {len(test_pairs)} predictions to {TEST_PRED}') |
|
|
| def cmd_eval_train(args): |
| import cv2 |
| device = torch.device(f'cuda:{args.gpu}' if torch.cuda.is_available() else 'cpu') |
| pairs = list_train_pairs() |
| rng = np.random.RandomState(args.seed) |
| indices = np.arange(len(pairs)) |
| rng.shuffle(indices) |
| folds = np.array_split(indices, args.folds) |
| mean = np.array([0.485, 0.456, 0.406], dtype=np.float32) |
| std = np.array([0.229, 0.224, 0.225], dtype=np.float32) |
| train_pred_dir = Path('dataset/train/predictions') |
| if args.save: |
| train_pred_dir.mkdir(exist_ok=True) |
| palette = load_palette() |
| all_p, all_g = ([], []) |
| for k in range(args.folds): |
| ckpt = CKPT_DIR / f'fold_{k}' / 'best.pth' |
| if not ckpt.exists(): |
| print(f'skip fold {k} - no checkpoint') |
| continue |
| model = UNet(in_ch=5, base=args.base, out_ch=1).to(device) |
| model.load_state_dict(torch.load(ckpt, map_location=device)) |
| model.eval() |
| val_idx = folds[k] |
| for i in val_idx: |
| stem, img_p, sc_p, gt_p = pairs[i] |
| img = np.array(Image.open(img_p).convert('RGB')) |
| sc = np.array(Image.open(sc_p).convert('L')) |
| gt = (np.array(Image.open(gt_p)) > 0).astype(np.uint8) |
| img_r = cv2.resize(img, (TRAIN_W, TRAIN_H), interpolation=cv2.INTER_LINEAR) |
| sc_r = cv2.resize(sc, (TRAIN_W, TRAIN_H), interpolation=cv2.INTER_NEAREST) |
| img_f = (img_r.astype(np.float32) / 255.0 - mean) / std |
| x = torch.cat([torch.from_numpy(img_f.transpose(2, 0, 1)), torch.from_numpy(encode_scribble(sc_r))], 0).unsqueeze(0).to(device) |
| prob = tta_predict(model, x, device, scales=(0.7, 1.0, 1.3)) |
| prob_full = cv2.resize(prob, (ORIG_W, ORIG_H), interpolation=cv2.INTER_LINEAR) |
| pred = (prob_full > 0.5).astype(np.uint8) |
| pred[sc == 0] = 0 |
| pred[sc == 1] = 1 |
| all_p.append(pred) |
| all_g.append(gt) |
| if args.save: |
| out_img = Image.fromarray(pred.astype(np.uint8), mode='P') |
| out_img.putpalette(palette) |
| out_img.save(train_pred_dir / f'{stem}.png') |
| if args.folds == 1: |
| break |
| bg, fg, miou = evaluate_predictions(all_p, all_g) |
| print(f'Held-out CV: bg={bg:.4f} fg={fg:.4f} mIoU={miou:.4f} (n={len(all_p)} images)') |
| if args.save: |
| print(f'Saved {len(all_p)} train predictions to {train_pred_dir}') |
|
|
| def set_seed(seed): |
| random.seed(seed) |
| np.random.seed(seed) |
| torch.manual_seed(seed) |
| torch.cuda.manual_seed_all(seed) |
|
|
| def main(): |
| p = argparse.ArgumentParser() |
| sub = p.add_subparsers(dest='cmd') |
| pt = sub.add_parser('train') |
| pt.add_argument('--epochs', type=int, default=120) |
| pt.add_argument('--batch-size', type=int, default=8) |
| pt.add_argument('--lr', type=float, default=0.001) |
| pt.add_argument('--folds', type=int, default=1) |
| pt.add_argument('--seed', type=int, default=42) |
| pt.add_argument('--gpu', type=int, default=0) |
| pt.add_argument('--base', type=int, default=48, help='U-Net base channel count') |
| pt.add_argument('--ckpt-suffix', type=str, default='', help='Suffix for runs_global_unet dir') |
| pt.add_argument('--cutmix-p', type=float, default=0.0, help='Probability of CutMix per sample') |
| pt.add_argument('--pseudo-method', type=str, default='', help="If set (e.g. 'v3v4'), use that method's predictions on test1+test2 as additional pseudo-labeled training data.") |
| pp = sub.add_parser('predict') |
| pp.add_argument('--gpu', type=int, default=0) |
| pp.add_argument('--base', type=int, default=48) |
| pe = sub.add_parser('eval') |
| pe.add_argument('--folds', type=int, default=1) |
| pe.add_argument('--seed', type=int, default=42) |
| pe.add_argument('--gpu', type=int, default=0) |
| pe.add_argument('--base', type=int, default=48) |
| pe.add_argument('--save', action='store_true', help='Save out-of-fold predictions to dataset/train/predictions/') |
| args = p.parse_args() |
| if args.cmd == 'train': |
| cmd_train(args) |
| elif args.cmd == 'predict': |
| cmd_predict(args) |
| elif args.cmd == 'eval': |
| cmd_eval_train(args) |
| else: |
| p.print_help() |
| if __name__ == '__main__': |
| main() |
