#!/usr/bin/env python3 import os import json import math import glob import gc import random import argparse from pathlib import Path import numpy as np import pandas as pd from PIL import Image, ImageDraw # Support running as a script (python scripts/train_node_heatmap.py) and as a module (import scripts.train_node_heatmap) try: import hardneg_utils as hardneg # when executed from within scripts/ except Exception: try: from . import hardneg_utils as hardneg # when imported as a package module except Exception: import scripts.hardneg_utils as hardneg # fallback when CWD is project root import torch import torch.nn as nn import torch.nn.functional as F from torch.utils.data import Dataset, DataLoader import torchvision.transforms.functional as TF from tqdm import tqdm from torch.utils.checkpoint import checkpoint import torchvision from torchvision import models import csv # ----------------------- # Utility # ----------------------- os.environ.setdefault('PYTORCH_CUDA_ALLOC_CONF','expandable_segments:True') def set_seed(seed: int): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) def ensure_dir(p: Path): p.mkdir(parents=True, exist_ok=True) return p def _set_bn_eval(model): import torch.nn as nn for m in model.modules(): if isinstance(m, nn.BatchNorm2d): m.eval() for p in m.parameters(): p.requires_grad = False def load_image(path: Path, in_channels: int = 1): img = Image.open(path) if in_channels == 1: img = img.convert("L") else: img = img.convert("RGB") arr = np.array(img, dtype=np.uint8) return arr def resize_keep_aspect(arr: np.ndarray, max_side: int): h, w = arr.shape[:2] if max(h, w) <= max_side: return arr, 1.0, 1.0 if h >= w: scale = max_side / float(h) else: scale = max_side / float(w) new_w = max(1, int(round(w * scale))) new_h = max(1, int(round(h * scale))) img = Image.fromarray(arr) img = img.resize((new_w, new_h), resample=Image.BILINEAR) return np.array(img, dtype=np.uint8), (new_w / w), (new_h / h) def pad_to_multiple(arr: np.ndarray, multiple: int): h, w = arr.shape[:2] pad_h = (multiple - (h % multiple)) % multiple pad_w = (multiple - (w % multiple)) % multiple if pad_h == 0 and pad_w == 0: return arr, (0, 0, 0, 0) if arr.ndim == 2: out = np.zeros((h + pad_h, w + pad_w), dtype=arr.dtype) out[:h, :w] = arr else: c = arr.shape[2] out = np.zeros((h + pad_h, w + pad_w, c), dtype=arr.dtype) out[:h, :w, :] = arr return out, (0, 0, pad_w, pad_h) # left, top, right, bottom def draw_gaussians(h, w, points, sigma, type_to_channel): C = len(type_to_channel) # 3 (tip, internal, root) Y, X = np.ogrid[:h, :w] heat = np.zeros((C, h, w), dtype=np.float32) s2 = 2 * (sigma ** 2) for (x, y, t) in points: # Skip invalid points if not np.isfinite(x) or not np.isfinite(y): continue c = type_to_channel.get(t, type_to_channel.get("internal", 1)) # Guard against OOB caused by rounding x = float(np.clip(x, 0, w - 1)) y = float(np.clip(y, 0, h - 1)) g = np.exp(-((X - x) ** 2 + (Y - y) ** 2) / s2) heat[c] = np.maximum(heat[c], g) return heat def to_tensor_image(arr: np.ndarray): # HxW -> 1xHxW; HxWx3 -> 3xHxW in [0,1] t = torch.from_numpy(arr.astype(np.float32) / 255.0) if t.ndim == 2: return t.unsqueeze(0) elif t.ndim == 3: return t.permute(2, 0, 1) else: raise ValueError("Unsupported image shape for to_tensor_image") def to_tensor_heatmap(arr: np.ndarray): # CxHxW return torch.from_numpy(arr.astype(np.float32)) def to_tensor_negmask(arr: np.ndarray): # HxW float32 map return torch.from_numpy(arr.astype(np.float32)) def random_augment(img_arr: np.ndarray, p_bright: float, p_contrast: float, p_noise: float): """Photometric jitter and light noise that do not change geometry.""" img = Image.fromarray(img_arr) if random.random() < max(0.0, p_bright): b = 0.9 + 0.2 * random.random() img = TF.adjust_brightness(img, b) if random.random() < max(0.0, p_contrast): c = 0.9 + 0.2 * random.random() img = TF.adjust_contrast(img, c) if random.random() < max(0.0, p_noise): a = np.array(img, dtype=np.float32) a += np.random.normal(0, 3.0, size=a.shape).astype(np.float32) a = np.clip(a, 0, 255).astype(np.uint8) img = Image.fromarray(a) return np.array(img, dtype=np.uint8) def random_extra_downscale(img_arr: np.ndarray, min_scale: float = 0.6, max_scale: float = 1.0, prob: float = 0.0): """Optionally downscale further by a factor in [min_scale, max_scale]. Returns (arr, scale_x, scale_y).""" if max_scale <= 0 or min_scale >= max_scale: return img_arr, 1.0, 1.0 if random.random() < max(0.0, prob): scale = random.uniform(min_scale, max_scale) h, w = img_arr.shape[:2] new_w = max(1, int(round(w * scale))) new_h = max(1, int(round(h * scale))) img = Image.fromarray(img_arr) img = img.resize((new_w, new_h), resample=Image.BILINEAR) return np.array(img, dtype=np.uint8), scale, scale return img_arr, 1.0, 1.0 def _rand_gray(min_v=0, max_v=255): return int(round(random.uniform(min_v, max_v))) def overlay_random_lines(img_arr: np.ndarray, vertical=True, n_range=(1, 6), thickness=(1, 3), intensity=(0, 30)): """Overlay thin dark or light lines. If vertical=False, draws horizontal lines.""" h, w = img_arr.shape[:2] arr = img_arr.copy() n = random.randint(n_range[0], n_range[1]) if n_range[1] >= n_range[0] else 0 for _ in range(n): t = random.randint(thickness[0], thickness[1]) val = _rand_gray(*intensity) if vertical: x = random.randint(0, w - 1) x0 = max(0, x - t // 2); x1 = min(w, x0 + t) arr[:, x0:x1] = val else: y = random.randint(0, h - 1) y0 = max(0, y - t // 2); y1 = min(h, y0 + t) arr[y0:y1, :] = val return arr def overlay_random_rectangles(img_arr: np.ndarray, n_range=(0, 3), thickness=(1, 3), intensity=(0, 40)): """Draw random rectangular borders to mimic boxes/frames/legends.""" h, w = img_arr.shape[:2] if n_range[1] <= 0: return img_arr arr = img_arr.copy() for _ in range(random.randint(n_range[0], n_range[1])): x0 = random.randint(0, max(0, w - 5)) y0 = random.randint(0, max(0, h - 5)) x1 = random.randint(x0 + 3, w) y1 = random.randint(y0 + 3, h) t = random.randint(thickness[0], thickness[1]) val = _rand_gray(*intensity) # top/bottom arr[y0:y0+t, x0:x1] = val arr[y1-t:y1, x0:x1] = val # left/right arr[y0:y1, x0:x0+t] = val arr[y0:y1, x1-t:x1] = val return arr def overlay_quadrilateral_shift(img_arr: np.ndarray, delta_range=(-40, 40)): """Random quadrilateral region with intensity shift.""" h, w = img_arr.shape[:2] if random.random() < 0.5: return img_arr # pick 4 points pts = [(random.randint(0, w-1), random.randint(0, h-1)) for _ in range(4)] mask = Image.new('L', (w, h), 0) ImageDraw.Draw(mask).polygon(pts, outline=255, fill=255) m = np.array(mask, dtype=np.uint8) delta = random.randint(delta_range[0], delta_range[1]) arr = img_arr.astype(np.int16) arr = np.where(m > 0, arr + delta, arr) return np.clip(arr, 0, 255).astype(np.uint8) def overlay_node_occlusions(img_arr: np.ndarray, points, max_frac=0.3, size_px=(3, 12), shapes=("square","circle")): """Occlude regions around a random subset of nodes with small shapes.""" if len(points) == 0: return img_arr h, w = img_arr.shape[:2] k = max(1, int(round(len(points) * random.uniform(0.05, max_frac)))) idxs = random.sample(range(len(points)), k=k) arr = img_arr.copy() for i in idxs: x, y, _ = points[i] xi, yi = int(round(x)), int(round(y)) r = random.randint(size_px[0], size_px[1]) val = _rand_gray(0, 255) shape = random.choice(shapes) y0 = max(0, yi - r); y1 = min(h, yi + r + 1) x0 = max(0, xi - r); x1 = min(w, xi + r + 1) if shape == "circle": yy, xx = np.ogrid[y0:y1, x0:x1] mask = (yy - yi)**2 + (xx - xi)**2 <= r*r sub = arr[y0:y1, x0:x1] sub[mask] = val arr[y0:y1, x0:x1] = sub else: # square arr[y0:y1, x0:x1] = val return arr def apply_random_overlays(img_arr: np.ndarray, points, prob: float = 0.0): """Compose targeted overlays that do not change target geometry.""" arr = img_arr # vertical time-scale like lines if random.random() < 0.6: arr = overlay_random_lines(arr, vertical=True, n_range=(1, 8), thickness=(1, 3), intensity=(0, 40)) # occasional horizontal guide lines if random.random() < 0.3: arr = overlay_random_lines(arr, vertical=False, n_range=(1, 4), thickness=(1, 2), intensity=(0, 40)) # random border/legend-like rectangles if random.random() < 0.4: arr = overlay_random_rectangles(arr, n_range=(0, 2), thickness=(1, 3), intensity=(0, 40)) # local region color shifts if random.random() < 0.5: arr = overlay_quadrilateral_shift(arr, delta_range=(-35, 35)) # occlude near some nodes (simulate markers/pie charts) if random.random() < 0.6: arr = overlay_node_occlusions(arr, points, max_frac=0.25, size_px=(3, 10)) return arr # ----------------------- # Dataset # ----------------------- class NodeDataset(Dataset): def __init__(self, image_paths, labels_dir, max_side=1536, sigma=1.5, pad_mult=32, types=("tip", "internal", "root"), augment=False, aug_brightness_prob=0.5, aug_contrast_prob=0.5, aug_noise_prob=0.0, aug_extra_downscale_prob=0.0, aug_extra_downscale_min=0.6, aug_extra_downscale_max=1.0, overlay_aug_prob=0.0, in_channels: int = 1, internal_only: bool = False): self.image_paths = image_paths self.labels_dir = Path(labels_dir) self.max_side = max_side self.sigma = sigma self.pad_mult = pad_mult self.types = types self.type_to_channel = {t: i for i, t in enumerate(types)} self.augment = augment self.aug_brightness_prob = aug_brightness_prob self.aug_contrast_prob = aug_contrast_prob self.aug_noise_prob = aug_noise_prob self.aug_extra_downscale_prob = aug_extra_downscale_prob self.aug_extra_downscale_min = aug_extra_downscale_min self.aug_extra_downscale_max = aug_extra_downscale_max self.overlay_aug_prob = overlay_aug_prob self.internal_only = bool(internal_only) self.in_channels = in_channels # Hard-negative config defaults self.hardneg_enable = False self.hardneg_lines_count = (5,10) self.hardneg_line_thickness = (1,3) self.hardneg_line_shade = (0,100) self.hardneg_line_margin_to_node = 8 self.hardneg_penalty_radius = 3 self.hardneg_line_boost = 3.0 self.hardneg_cross_boost = 6.0 self.hardneg_rect_internal_frac = 0.1 self.hardneg_rect_side = (10,40) self.hardneg_rect_margin_to_node = 2 self.hardneg_rect_boost = 2.5 def __len__(self): return len(self.image_paths) def __getitem__(self, idx): ipath = Path(self.image_paths[idx]) base = ipath.stem # Expected CSV name csv_name = (self.labels_dir / f"{base}_node_locations.csv") img_arr = load_image(ipath, in_channels=self.in_channels) orig_h, orig_w = img_arr.shape[:2] if self.augment: img_arr = random_augment(img_arr, self.aug_brightness_prob, self.aug_contrast_prob, self.aug_noise_prob) # Resize & pad img_arr, sx, sy = resize_keep_aspect(img_arr, self.max_side) # Optional extra downscaling to simulate shrunken figures if self.augment and self.aug_extra_downscale_prob > 0.0: img_arr, sdx, sdy = random_extra_downscale(img_arr, min_scale=self.aug_extra_downscale_min, max_scale=self.aug_extra_downscale_max, prob=self.aug_extra_downscale_prob) sx *= sdx; sy *= sdy h, w = img_arr.shape[:2] img_arr, pads = pad_to_multiple(img_arr, self.pad_mult) pad_left, pad_top, pad_right, pad_bottom = pads H, W = img_arr.shape[:2] # Load points df = pd.read_csv(csv_name) # Coerce numeric columns and drop invalid rows for col in ("x","y"): if col in df.columns: df[col] = pd.to_numeric(df[col], errors='coerce') df = df.dropna(subset=["x","y"]) if set(["x","y"]).issubset(df.columns) else df # Optionally map root->internal and drop tips df["type"] = df["type"].astype(str).str.lower() if self.internal_only: df = df[df["type"] != "tip"].copy() df.loc[df["type"] == "root", "type"] = "internal" # Scale points xs = (df["x"].values * sx + pad_left).astype(np.float64) ys = (df["y"].values * sy + pad_top).astype(np.float64) ts = df["type"].tolist() points = list(zip(xs, ys, ts)) # Targeted overlays that can add distractors or occlusions if self.augment and self.overlay_aug_prob > 0.0: img_arr = apply_random_overlays(img_arr, points, prob=self.overlay_aug_prob) # Hard-negative overlays & mask neg_mask = np.zeros((H,W), dtype=np.float32) if self.augment and self.hardneg_enable: img_arr, neg_mask = hardneg.generate_hardneg_overlays( img_arr, points, line_minmax=self.hardneg_lines_count, thickness_minmax=self.hardneg_line_thickness, shade_minmax=self.hardneg_line_shade, margin_to_node=self.hardneg_line_margin_to_node, penal_radius=self.hardneg_penalty_radius, line_boost=self.hardneg_line_boost, cross_boost=self.hardneg_cross_boost, rect_internal_frac=self.hardneg_rect_internal_frac, rect_side_minmax=self.hardneg_rect_side, rect_margin_to_node=self.hardneg_rect_margin_to_node, rect_boost=self.hardneg_rect_boost) # Targets heat = draw_gaussians(H, W, points, self.sigma, self.type_to_channel) # Tensors img_t = to_tensor_image(img_arr) # 1xHxW heat_t = to_tensor_heatmap(heat) # CxHxW neg_t = to_tensor_negmask(neg_mask) # HxW meta = { "path": str(ipath), "orig_w": orig_w, "orig_h": orig_h, "scaled_w": W, "scaled_h": H, "scale_x": sx, "scale_y": sy, "pad": pads } return img_t, heat_t, neg_t, meta # ----------------------- # Model (small U-Net) # ----------------------- class DoubleConv(nn.Module): def __init__(self, in_ch, out_ch): super().__init__() self.net = nn.Sequential( nn.Conv2d(in_ch, out_ch, 3, padding=1), nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True), nn.Conv2d(out_ch, out_ch, 3, padding=1), nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True), ) def forward(self, x): return self.net(x) class UNetSmall(nn.Module): def __init__(self, in_ch=1, out_ch=3, base=32): super().__init__() self.inc = DoubleConv(in_ch, base) self.down1 = nn.Sequential(nn.MaxPool2d(2), DoubleConv(base, base*2)) self.down2 = nn.Sequential(nn.MaxPool2d(2), DoubleConv(base*2, base*4)) self.down3 = nn.Sequential(nn.MaxPool2d(2), DoubleConv(base*4, base*8)) self.down4 = nn.Sequential(nn.MaxPool2d(2), DoubleConv(base*8, base*8)) self.up1 = nn.ConvTranspose2d(base*8, base*8, 2, stride=2) self.conv1 = DoubleConv(base*16, base*4) self.up2 = nn.ConvTranspose2d(base*4, base*4, 2, stride=2) self.conv2 = DoubleConv(base*8, base*2) self.up3 = nn.ConvTranspose2d(base*2, base*2, 2, stride=2) self.conv3 = DoubleConv(base*4, base) self.up4 = nn.ConvTranspose2d(base, base, 2, stride=2) self.conv4 = DoubleConv(base*2, base) self.outc = nn.Conv2d(base, out_ch, 1) def forward(self, x): x1 = self.inc(x) # base x2 = self.down1(x1) # 2*base x3 = self.down2(x2) # 4*base x4 = self.down3(x3) # 8*base x5 = self.down4(x4) # 8*base x = self.up1(x5) x = torch.cat([x, x4], dim=1) x = self.conv1(x) x = self.up2(x) x = torch.cat([x, x3], dim=1) x = self.conv2(x) x = self.up3(x) x = torch.cat([x, x2], dim=1) x = self.conv3(x) x = self.up4(x) x = torch.cat([x, x1], dim=1) x = self.conv4(x) return self.outc(x) # ----------------------- # ResNet backbone U-Net # ----------------------- class ResNetUNet(nn.Module): def __init__(self, backbone="resnet34", in_ch=1, out_ch=3, pretrained=False, freeze_stages=0, base=64, grad_ckpt=False): super().__init__() self.backbone_name = backbone self.grad_ckpt = bool(grad_ckpt) # Instantiate backbone with or without weights (torchvision >=0.13 API) weights = None if pretrained: try: if backbone == "resnet34": from torchvision.models import ResNet34_Weights weights = ResNet34_Weights.DEFAULT elif backbone == "resnet50": from torchvision.models import ResNet50_Weights weights = ResNet50_Weights.DEFAULT except Exception: weights = "IMAGENET1K_V1" # best-effort for older APIs if backbone == "resnet34": self.enc = models.resnet34(weights=weights) ch = {"x0": 64, "l1": 64, "l2": 128, "l3": 256, "l4": 512} elif backbone == "resnet50": self.enc = models.resnet50(weights=weights) ch = {"x0": 64, "l1": 256, "l2": 512, "l3": 1024, "l4": 2048} else: raise ValueError(f"Unsupported backbone: {backbone}") # Adapt first conv for grayscale if needed if in_ch != 3: w = self.enc.conv1.weight.data self.enc.conv1 = nn.Conv2d(in_ch, self.enc.conv1.out_channels, kernel_size=7, stride=2, padding=3, bias=False) with torch.no_grad(): if w.shape[1] == 3 and in_ch == 1: self.enc.conv1.weight[:] = w.mean(dim=1, keepdim=True) elif w.shape[1] == 3 and in_ch == 2: self.enc.conv1.weight[:, :2] = w[:, :2] self.enc.conv1.weight[:, 2:] = w[:, :1] # Optionally freeze early stages stages = [ [self.enc.conv1, self.enc.bn1], # 1 [self.enc.layer1], # 2 [self.enc.layer2], # 3 [self.enc.layer3], # 4 ] # freeze_stages: 0 = none, 1 = conv1+bn1, 2 = +layer1, 3 = +layer2, 4 = +layer3 for i in range(min(freeze_stages, len(stages))): for m in stages[i]: for p in m.parameters(): p.requires_grad = False # Decoder # up from l4->l3 self.up4 = nn.ConvTranspose2d(ch["l4"], ch["l3"], 2, stride=2) self.dec4 = DoubleConv(ch["l3"] + ch["l3"], ch["l3"]) # cat(l3) # l3->l2 self.up3 = nn.ConvTranspose2d(ch["l3"], ch["l2"], 2, stride=2) self.dec3 = DoubleConv(ch["l2"] + ch["l2"], ch["l2"]) # cat(l2) # l2->l1 self.up2 = nn.ConvTranspose2d(ch["l2"], ch["l1"], 2, stride=2) self.dec2 = DoubleConv(ch["l1"] + ch["l1"], ch["l1"]) # cat(l1) # l1->x0 (1/4 -> 1/2) self.up1 = nn.ConvTranspose2d(ch["l1"], ch["x0"], 2, stride=2) self.dec1 = DoubleConv(ch["x0"] + ch["x0"], ch["x0"]) # cat(x0) # x0 -> full (1/2 -> 1) self.up0 = nn.ConvTranspose2d(ch["x0"], ch["x0"], 2, stride=2) self.dec0 = DoubleConv(ch["x0"], base) self.outc = nn.Conv2d(base, out_ch, 1) def _ckpt(self, fn, x): # Only checkpoint during training when grads are enabled and tensor requires grad if self.grad_ckpt and torch.is_grad_enabled() and isinstance(x, torch.Tensor) and x.requires_grad: return checkpoint(fn, x, use_reentrant=False) else: return fn(x) def forward(self, x): # Encoder forward with skip captures x0 = self.enc.conv1(x) x0 = self.enc.bn1(x0) x0 = self.enc.relu(x0) # 1/2, ch x0 x1 = self.enc.maxpool(x0) # 1/4 if self.grad_ckpt: l1 = self._ckpt(self.enc.layer1, x1) # 1/4 l2 = self._ckpt(self.enc.layer2, l1) # 1/8 l3 = self._ckpt(self.enc.layer3, l2) # 1/16 l4 = self._ckpt(self.enc.layer4, l3) # 1/32 else: l1 = self.enc.layer1(x1) # 1/4 l2 = self.enc.layer2(l1) # 1/8 l3 = self.enc.layer3(l2) # 1/16 l4 = self.enc.layer4(l3) # 1/32 if self.grad_ckpt: y = self._ckpt(self.up4, l4) y = torch.cat([y, l3], dim=1) y = self._ckpt(self.dec4, y) y = self._ckpt(self.up3, y) y = torch.cat([y, l2], dim=1) y = self._ckpt(self.dec3, y) y = self._ckpt(self.up2, y) y = torch.cat([y, l1], dim=1) y = self._ckpt(self.dec2, y) y = self._ckpt(self.up1, y) y = torch.cat([y, x0], dim=1) y = self._ckpt(self.dec1, y) y = self._ckpt(self.up0, y) y = self._ckpt(self.dec0, y) else: y = self.up4(l4) y = torch.cat([y, l3], dim=1) y = self.dec4(y) y = self.up3(y) y = torch.cat([y, l2], dim=1) y = self.dec3(y) y = self.up2(y) y = torch.cat([y, l1], dim=1) y = self.dec2(y) y = self.up1(y) y = torch.cat([y, x0], dim=1) y = self.dec1(y) y = self.up0(y) y = self.dec0(y) return self.outc(y) # ----------------------- # Metrics & decoding # ----------------------- def decode_peaks(logits, thresh=0.3, window=3, per_channel_topk=None, max_peaks=None, fallback_topk=0): """ logits: CxHxW tensor (raw logits). Returns list of (x,y,channel,score). Efficiently limits peaks to avoid explosion when outputs are flat. - per_channel_topk: keep at most K peaks per channel (after NMS) - max_peaks: global cap across all channels """ with torch.no_grad(): prob = torch.sigmoid(logits) C, H, W = prob.shape xs_list = [] ys_list = [] cs_list = [] ss_list = [] for c in range(C): p = prob[c:c+1, :, :] # 1xHxW # NMS via max-pool over window maxm = F.max_pool2d(p.unsqueeze(0), kernel_size=window, stride=1, padding=window//2).squeeze(0) keep = (p == maxm) & (p > thresh) # 1xHxW idx = torch.nonzero(keep[0], as_tuple=False) # Nx2 (y,x) if idx.numel() == 0 and fallback_topk and fallback_topk > 0: # Fallback: take top-K per channel regardless of threshold flat = p[0].reshape(-1) k = min(int(fallback_topk), flat.numel()) vals, inds = torch.topk(flat, k) y = (inds // W).to(torch.long) x = (inds % W).to(torch.long) idx = torch.stack([y, x], dim=1) scores = vals elif idx.numel() == 0: continue scores = p[0, idx[:, 0], idx[:, 1]] # N if per_channel_topk is not None and idx.shape[0] > per_channel_topk: vals, order = torch.topk(scores, per_channel_topk) scores = vals idx = idx[order] xs_list.append(idx[:, 1]) ys_list.append(idx[:, 0]) cs_list.append(torch.full((idx.shape[0],), c, device=idx.device, dtype=torch.long)) ss_list.append(scores) if len(xs_list) == 0: return [] xs = torch.cat(xs_list, dim=0) ys = torch.cat(ys_list, dim=0) cs = torch.cat(cs_list, dim=0) ss = torch.cat(ss_list, dim=0) if max_peaks is not None and xs.shape[0] > max_peaks: vals, order = torch.topk(ss, max_peaks) ss = vals xs = xs[order] ys = ys[order] cs = cs[order] xs = xs.detach().cpu().tolist() ys = ys.detach().cpu().tolist() cs = cs.detach().cpu().tolist() ss = ss.detach().cpu().tolist() return [(float(x), float(y), int(c), float(s)) for x, y, c, s in zip(xs, ys, cs, ss)] def greedy_match(pred_pts, gt_pts, tau): """ pred_pts: list of (x,y) gt_pts: list of (x,y) tau: distance threshold """ if len(pred_pts) == 0 and len(gt_pts) == 0: return 1.0, 1.0, 1.0, 0.0 if len(pred_pts) == 0: return 0.0, 0.0, 0.0, float('inf') if len(gt_pts) == 0: return 0.0, 0.0, 0.0, float('inf') used_gt = set() tp = 0 dsum = 0.0 for px, py in pred_pts: # find nearest unmatched gt best = None best_d2 = None for j, (gx, gy) in enumerate(gt_pts): if j in used_gt: continue dx = px - gx dy = py - gy d2 = dx*dx + dy*dy if best_d2 is None or d2 < best_d2: best_d2 = d2 best = j if best is not None and math.sqrt(best_d2) <= tau: used_gt.add(best) tp += 1 dsum += math.sqrt(best_d2) fp = len(pred_pts) - tp fn = len(gt_pts) - tp prec = tp / (tp + fp) if (tp + fp) > 0 else 0.0 rec = tp / (tp + fn) if (tp + fn) > 0 else 0.0 f1 = 2*prec*rec / (prec+rec) if (prec+rec) > 0 else 0.0 mean_err = dsum / tp if tp > 0 else float('inf') return prec, rec, f1, mean_err def greedy_match_torch(pred_pts, gt_pts, tau, device=None, max_pairs=None): """ Vectorized greedy matching using pairwise distances in torch for speed. pred_pts, gt_pts: list of (x,y) tau: distance threshold device: torch device (defaults to CPU if None) max_pairs: optional cap on maximum matches to consider """ if len(pred_pts) == 0 and len(gt_pts) == 0: return 1.0, 1.0, 1.0, 0.0 if len(pred_pts) == 0 or len(gt_pts) == 0: return 0.0, 0.0, 0.0, float('inf') dev = device if device is not None else torch.device('cpu') P = torch.tensor(pred_pts, dtype=torch.float32, device=dev) # Nx2 G = torch.tensor(gt_pts, dtype=torch.float32, device=dev) # Mx2 N, M = P.shape[0], G.shape[0] if N == 0 or M == 0: return 0.0, 0.0, 0.0, float('inf') # Pairwise Euclidean distances D = torch.cdist(P.unsqueeze(0), G.unsqueeze(0)).squeeze(0) # NxM # Mask distances beyond tau inf = torch.tensor(float('inf'), device=dev) D = torch.where(D <= tau, D, inf) tp = 0 dsum = 0.0 max_iters = min(N, M) if max_pairs is not None: max_iters = min(max_iters, int(max_pairs)) for _ in range(max_iters): v, idx = torch.min(D.view(-1), dim=0) if not torch.isfinite(v): break i = (idx // M).item() j = (idx % M).item() tp += 1 dsum += float(v.item()) # Invalidate this row and column D[i, :] = inf D[:, j] = inf fp = N - tp fn = M - tp prec = tp / (tp + fp) if (tp + fp) > 0 else 0.0 rec = tp / (tp + fn) if (tp + fn) > 0 else 0.0 f1 = 2*prec*rec / (prec+rec) if (prec+rec) > 0 else 0.0 mean_err = dsum / tp if tp > 0 else float('inf') return prec, rec, f1, mean_err # ----------------------- # Training # ----------------------- def main(): ap = argparse.ArgumentParser(description="Train heatmap-based node detector for tree figures.") ap.add_argument("--data_root", type=str, required=True, help="Root folder with images/ and labels/") ap.add_argument("--out_dir", type=str, required=True, help="Where to save model and samples") ap.add_argument("--epochs", type=int, default=40) ap.add_argument("--batch_size", type=int, default=2) ap.add_argument("--lr", type=float, default=1e-3) ap.add_argument("--max_side", type=int, default=1536, help="Resize long side to at most this many pixels") ap.add_argument("--sigma", type=float, default=1.5, help="Gaussian std in pixels (after resizing)") ap.add_argument("--pos_weight", type=float, default=25.0, help="BCE positive class weight") ap.add_argument("--val_split", type=float, default=0.15) ap.add_argument("--test_split", type=float, default=0.15) ap.add_argument("--seed", type=int, default=1337) ap.add_argument("--device", type=str, default="cuda", choices=["cuda","cpu"]) ap.add_argument("--save_every", type=int, default=10) ap.add_argument("--workers", type=int, default=0, help="DataLoader worker processes (0 in restricted envs)") ap.add_argument("--pin_memory", action="store_true", help="Enable pinned memory for DataLoader") ap.add_argument("--val_thresh", type=float, default=0.8, help="Validation peak threshold (higher avoids early dense peaks)") ap.add_argument("--val_topk", type=int, default=2000, help="Top-K peaks per channel during validation") ap.add_argument("--val_max_peaks", type=int, default=5000, help="Global peak cap per image during validation") ap.add_argument("--val_fallback_topk", type=int, default=50, help="If no peaks over threshold, take top-K per channel for metrics") ap.add_argument("--sample_thresh", type=float, default=0.6, help="Sample overlay peak threshold") ap.add_argument("--sample_topk", type=int, default=1000, help="Top-K peaks per channel for sample overlay") ap.add_argument("--nms_window_val", type=int, default=5, help="NMS window for validation decode") ap.add_argument("--nms_window_sample", type=int, default=5, help="NMS window for sample overlay decode") ap.add_argument("--val_match_device", type=str, default="cpu", choices=["cpu","cuda"], help="Device to run validation matching on (cpu avoids GPU OOM)") # Proximity inhibition (repulsion) loss ap.add_argument("--repel_lambda", type=float, default=0.0, help="Weight for repulsion loss (discourage nearby activations)") ap.add_argument("--repel_window", type=int, default=5, help="Window for repulsion loss (odd integer)") ap.add_argument("--count_loss_lambda", type=float, default=0.0, help="Weight for count penalty (L1 of mean probs vs GT heat)") ap.add_argument("--count_lambda_init", type=float, default=None, help="Start count loss weight here and ramp if provided") ap.add_argument("--count_lambda_final", type=float, default=None, help="Target final count loss weight when using schedule") ap.add_argument("--count_lambda_warmup_epochs", type=int, default=0, help="Epochs to linearly ramp count loss weight from init to final") # Logging/plots ap.add_argument("--write_csv", action="store_true", help="Write metrics CSV to out_dir/metrics.csv") ap.add_argument("--plot", action="store_true", help="Save loss/val curves as PNG in out_dir") ap.add_argument("--amp", action="store_true", help="Enable mixed-precision training to save memory") ap.add_argument("--no_aug", action="store_true", help="Disable all training-time augmentations") ap.add_argument("--aug_brightness_prob", type=float, default=0.5, help="Prob of brightness jitter") ap.add_argument("--aug_contrast_prob", type=float, default=0.5, help="Prob of contrast jitter") ap.add_argument("--aug_noise_prob", type=float, default=0.0, help="Prob of Gaussian noise jitter") ap.add_argument("--aug_extra_downscale_prob", type=float, default=0.0, help="Prob of extra downscale augmentation (0 disables)") ap.add_argument("--aug_extra_downscale_min", type=float, default=0.6, help="Min extra downscale factor") ap.add_argument("--aug_extra_downscale_max", type=float, default=1.0, help="Max extra downscale factor") ap.add_argument("--overlay_aug_prob", type=float, default=0.0, help="Probability to apply overlay augmentations (lines/rects/occlusions)") ap.add_argument("--scheduler", type=str, default="cosine", choices=["none","cosine","cosine_wr"], help="LR scheduler") ap.add_argument("--cosine_min_lr", type=float, default=1e-5, help="Min LR for cosine annealing") ap.add_argument("--cosinewr_t0", type=int, default=20, help="T0 for cosine warm restarts") ap.add_argument("--cosinewr_tmult", type=int, default=2, help="T_mult for cosine warm restarts") ap.add_argument("--early_stop_burnin", type=int, default=5, help="Epochs before monitoring early stop") ap.add_argument("--early_stop_patience", type=int, default=10, help="Epochs without improvement to stop") ap.add_argument("--early_stop_min_delta", type=float, default=0.0, help="Min val_loss improvement to reset patience") ap.add_argument("--grad_clip", type=float, default=1.0, help="Clip global grad norm; 0 disables") ap.add_argument("--nan_backoff", type=float, default=0.5, help="Multiply LR by this on NaN/inf loss (0 disables)") ap.add_argument("--logit_clip", type=float, default=20.0, help="Clamp logits to [-k,k] for loss to avoid overflow") ap.add_argument("--max_weight_multiplier", type=float, default=5.0, help="Clamp per-pixel loss weight to at most this factor") ap.add_argument("--freeze_all_bn", action="store_true", help="Keep all BatchNorm layers in eval mode during training") ap.add_argument("--normalize_weighted_loss", action="store_true", help="Normalize BCE by sum of weights for stability") ap.add_argument("--debug_nan", action="store_true", help="Print detailed tensor stats and sample path on NaN/Inf loss") ap.add_argument("--grad_ckpt", action="store_true", help="Enable gradient checkpointing to reduce memory") ap.add_argument("--channels_last", action="store_true", help="Use channels_last memory format for model and inputs") ap.add_argument("--normalize_input", action="store_true", help="Normalize inputs (ImageNet mean/std) when in_channels=3") # Loss shaping ap.add_argument("--use_separate_posneg_loss", action="store_true", help="Compute positive and negative BCE terms separately for class balance") ap.add_argument("--neg_lambda", type=float, default=1.0, help="Multiplier for negative term in separate loss") ap.add_argument("--max_neg_weight_multiplier", type=float, default=3.0, help="Clamp for negative overlay boost only") # Curriculum for negatives ap.add_argument("--neg_lambda_init", type=float, default=None, help="Start neg_lambda here and ramp to final if provided") ap.add_argument("--neg_lambda_final", type=float, default=None, help="Target final neg_lambda when using schedule") ap.add_argument("--neg_lambda_warmup_epochs", type=int, default=0, help="Epochs to linearly ramp neg_lambda from init to final") # Hard-negative overlays ap.add_argument("--hardneg_lines", action="store_true", help="Enable hard-negative overlays (lines + rectangles)") ap.add_argument("--hardneg_lines_count", type=int, nargs=2, default=[5,10], help="Min/Max lines per orientation") ap.add_argument("--hardneg_line_thickness", type=int, nargs=2, default=[1,3], help="Min/Max line thickness (px)") ap.add_argument("--hardneg_line_shade", type=int, nargs=2, default=[0,100], help="Shade range for lines (0=black..100=dark gray)") ap.add_argument("--hardneg_line_margin_to_node", type=int, default=8, help="Min distance from any node (px)") ap.add_argument("--hardneg_penalty_radius", type=int, default=3, help="Penalty radius around lines/edges (px)") ap.add_argument("--hardneg_line_boost", type=float, default=3.0, help="Negative weight boost near lines") ap.add_argument("--hardneg_cross_boost", type=float, default=6.0, help="Negative weight boost at line crossings") ap.add_argument("--hardneg_rect_internal_frac", type=float, default=0.1, help="Fraction of internal nodes to draw rectangles around") ap.add_argument("--hardneg_rect_side", type=int, nargs=2, default=[10,40], help="Min/Max rectangle side (px)") ap.add_argument("--hardneg_rect_margin_to_node", type=int, default=2, help="Min distance from rectangles to other nodes (px)") ap.add_argument("--hardneg_rect_boost", type=float, default=2.5, help="Negative weight boost on rectangle borders") ap.add_argument("--hardneg_weight_gamma", type=float, default=0.5, help="Global multiplier on hard-negative weights added to loss weighting") ap.add_argument("--hardneg_weight_gamma_init", type=float, default=None, help="Start gamma here and ramp to final if provided") ap.add_argument("--hardneg_weight_gamma_final", type=float, default=None, help="Target final gamma when using schedule") ap.add_argument("--hardneg_weight_warmup_epochs", type=int, default=0, help="Epochs to linearly ramp gamma from init to final") # Inter-class proximity penalty ap.add_argument("--prox_lambda", type=float, default=0.1, help="Weight for inter-class proximity penalty (tip/internal within 2px)") # Predict only two classes (tip, internal). Root is treated as internal during training ap.add_argument("--no_root_pred", action="store_true", help="Predict only tip and internal (map root to internal in training)") ap.add_argument("--internal_only", action="store_true", help="Predict internal nodes only; map root to internal and drop tips") # Backbone options ap.add_argument("--backbone", type=str, default="none", choices=["none","resnet34","resnet50"], help="Use pretrained backbone+UNet decoder") ap.add_argument("--pretrained", action="store_true", help="Use pretrained ImageNet weights for backbone") ap.add_argument("--freeze_stages", type=int, default=1, help="Freeze first N backbone stages (0..4)") ap.add_argument("--in_channels", type=int, default=1, help="Number of input channels (1=grayscale)") # Resume and checkpointing options ap.add_argument("--resume", type=str, default="", help="Path to checkpoint to warm-start from") ap.add_argument("--save_name", type=str, default="model", help="Base filename for checkpoints (no extension)") ap.add_argument("--save_epoch_suffix", action="store_true", help="Append _e{epoch} to checkpoint filename") args = ap.parse_args() set_seed(args.seed) data_root = Path(args.data_root) out_dir = ensure_dir(Path(args.out_dir)) samples_dir = ensure_dir(out_dir / "samples") model_path = out_dir / f"{args.save_name}.pt" img_dir = data_root / "images" lbl_dir = data_root / "labels" image_paths = sorted( [p for p in glob.glob(str(img_dir / "*")) if p.lower().endswith((".png", ".jpg", ".jpeg", ".tif", ".tiff", ".webp"))] ) # Filter to only those that have a matching CSV has_csv = [] for p in image_paths: base = Path(p).stem if (lbl_dir / f"{base}_node_locations.csv").exists(): has_csv.append(p) image_paths = has_csv # Exclude circular trees: require exactly one root with strictly smallest X among all nodes def is_non_circular(stem: str) -> bool: csv_path = lbl_dir / f"{stem}_node_locations.csv" try: df = pd.read_csv(csv_path) if "type" not in df.columns or "x" not in df.columns: return False roots = df[df["type"].astype(str).str.lower() == "root"] if len(roots) != 1: return False root_x = float(roots.iloc[0]["x"]) if not pd.isna(roots.iloc[0]["x"]) else None if root_x is None: return False others = df[df.index != roots.index[0]] if len(others) == 0: return False return root_x < float(others["x"].min()) except Exception: return False image_paths = [p for p in image_paths if is_non_circular(Path(p).stem)] if len(image_paths) == 0: raise RuntimeError("No image/CSV pairs found.") # Split train/val/test random.shuffle(image_paths) n_val = max(1, int(len(image_paths) * args.val_split)) n_test = max(0, int(len(image_paths) * args.test_split)) val_paths = image_paths[:n_val] test_paths = image_paths[n_val:n_val+n_test] if n_test > 0 else [] train_paths = image_paths[n_val+n_test:] if n_test > 0 else image_paths[n_val:] # Save splits for downstream prediction filtering splits = { "train": [Path(p).stem for p in train_paths], "val": [Path(p).stem for p in val_paths], "test": [Path(p).stem for p in test_paths], } with open(out_dir / "splits.json", "w") as f: json.dump(splits, f) # If resuming, infer class count from checkpoint and align training config if args.resume: try: tmp_state = torch.load(args.resume, map_location='cpu') sd = tmp_state.get("model", tmp_state) key = next((k for k in sd.keys() if k.endswith('outc.weight')), None) if key is not None: out_ch_loaded = sd[key].shape[0] args.no_root_pred = (out_ch_loaded == 2) except Exception: pass if args.internal_only: ds_types = ("internal",) else: ds_types = ("tip","internal") if args.no_root_pred else ("tip","internal","root") train_ds = NodeDataset(train_paths, lbl_dir, max_side=args.max_side, sigma=args.sigma, pad_mult=32, augment=(not args.no_aug), types=ds_types, in_channels=args.in_channels, internal_only=args.internal_only) # Hard-negative configuration from args if args.hardneg_lines or args.hardneg_rect_internal_frac > 0.0: train_ds.hardneg_enable = True train_ds.hardneg_lines_count = tuple(args.hardneg_lines_count) train_ds.hardneg_line_thickness = tuple(args.hardneg_line_thickness) train_ds.hardneg_line_shade = tuple(args.hardneg_line_shade) train_ds.hardneg_line_margin_to_node = int(args.hardneg_line_margin_to_node) train_ds.hardneg_penalty_radius = int(args.hardneg_penalty_radius) train_ds.hardneg_line_boost = float(args.hardneg_line_boost) train_ds.hardneg_cross_boost = float(args.hardneg_cross_boost) train_ds.hardneg_rect_internal_frac = float(args.hardneg_rect_internal_frac) train_ds.hardneg_rect_side = tuple(args.hardneg_rect_side) train_ds.hardneg_rect_margin_to_node = int(args.hardneg_rect_margin_to_node) train_ds.hardneg_rect_boost = float(args.hardneg_rect_boost) val_ds = NodeDataset(val_paths, lbl_dir, max_side=args.max_side, sigma=args.sigma, pad_mult=32, augment=False, types=ds_types, in_channels=args.in_channels, internal_only=args.internal_only) test_ds = NodeDataset(test_paths, lbl_dir, max_side=args.max_side, sigma=args.sigma, pad_mult=32, augment=False, types=ds_types, in_channels=args.in_channels, internal_only=args.internal_only) train_loader = DataLoader(train_ds, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=args.pin_memory, drop_last=True) val_loader = DataLoader(val_ds, batch_size=1, shuffle=False, num_workers=max(0, min(1, args.workers)), pin_memory=args.pin_memory) device = torch.device(args.device if (args.device == "cuda" and torch.cuda.is_available()) else "cpu") out_ch = 1 if args.internal_only else (2 if args.no_root_pred else 3) if args.backbone == "none": model = UNetSmall(in_ch=args.in_channels, out_ch=out_ch, base=32).to(device) else: model = ResNetUNet(backbone=args.backbone, in_ch=args.in_channels, out_ch=out_ch, pretrained=args.pretrained, freeze_stages=args.freeze_stages, grad_ckpt=args.grad_ckpt).to(device) # Optional channels_last memory format if args.channels_last: model = model.to(memory_format=torch.channels_last) # Warm-start from checkpoint if requested if args.resume: try: ckpt = torch.load(args.resume, map_location=device) state_dict = ckpt.get("model", ckpt) model.load_state_dict(state_dict, strict=False) print(f"[resume] Loaded weights from {args.resume}") except Exception as e: print(f"[resume] Failed to load {args.resume}: {e}") # Build per-element weight equivalent to pos_weight for channels-first tensors pos_w_vec = None if args.pos_weight and args.pos_weight > 0: pos_w = [args.pos_weight] * out_ch pos_w_vec = torch.tensor(pos_w, device=device).view(1, out_ch, 1, 1) opt = torch.optim.AdamW(model.parameters(), lr=args.lr) scaler = torch.amp.GradScaler(device="cuda", enabled=(args.amp and (device.type == "cuda"))) sched = None if args.scheduler == "cosine": try: sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=max(1, args.epochs), eta_min=float(args.cosine_min_lr)) except Exception: sched = None elif args.scheduler == "cosine_wr": try: sched = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts( opt, T_0=int(args.cosinewr_t0), T_mult=int(args.cosinewr_tmult), eta_min=float(args.cosine_min_lr)) except Exception: sched = None # Prepare logs logs = [] metrics_csv_path = out_dir / "metrics.csv" best_val = float("inf"); best_epoch = 0; epochs_no_improve = 0 if args.write_csv: with open(metrics_csv_path, "w", newline="") as f: w = csv.writer(f) w.writerow(["epoch","lr","lr_next","train_loss","val_loss","valP2","valR2","valF12","valE2","valP4","valR4","valF14","valE4","valP8","valR8","valF18","valE8"]) # Small helper for linear schedules def lin_sched(init_val, final_val, epoch_idx, warmup_epochs): if init_val is None or final_val is None or warmup_epochs is None or int(warmup_epochs) <= 0: return final_val if final_val is not None else (init_val if init_val is not None else None) t = max(0.0, min(1.0, float(epoch_idx-1) / float(max(1, int(warmup_epochs))))) return float(init_val) + t * (float(final_val) - float(init_val)) # Training loop for epoch in tqdm(range(1, args.epochs + 1), smoothing=0, desc="Epochs"): model.train() if args.freeze_all_bn: _set_bn_eval(model) # Compute scheduled weights for this epoch gamma_cur = lin_sched(getattr(args,'hardneg_weight_gamma_init',None), getattr(args,'hardneg_weight_gamma_final',None), epoch, getattr(args,'hardneg_weight_warmup_epochs',0)) if gamma_cur is None: gamma_cur = float(getattr(args, 'hardneg_weight_gamma', 0.5)) neg_lambda_cur = lin_sched(getattr(args,'neg_lambda_init',None), getattr(args,'neg_lambda_final',None), epoch, getattr(args,'neg_lambda_warmup_epochs',0)) if neg_lambda_cur is None: neg_lambda_cur = float(getattr(args,'neg_lambda',1.0)) cur_lr = float(opt.param_groups[0]["lr"]) epoch_loss = 0.0 for batch in tqdm(train_loader, smoothing=0, desc=f"Train {epoch}/{args.epochs}"): if isinstance(batch, (list,tuple)) and len(batch)==4: img_t, heat_t, neg_t, meta = batch else: img_t, heat_t, _ = batch neg_t = torch.zeros(heat_t.shape[-2:], dtype=heat_t.dtype) meta = {"path": ""} img_t = img_t.to(device, non_blocking=True) if args.channels_last: img_t = img_t.contiguous(memory_format=torch.channels_last) # Optional RGB normalization if args.normalize_input and int(args.in_channels) == 3: mean = torch.tensor([0.485, 0.456, 0.406], device=device, dtype=img_t.dtype).view(1,3,1,1) std = torch.tensor([0.229, 0.224, 0.225], device=device, dtype=img_t.dtype).view(1,3,1,1) img_t = (img_t - mean) / std heat_t = heat_t.to(device, non_blocking=True) neg_t = neg_t.to(device, non_blocking=True) if isinstance(neg_t, torch.Tensor) else torch.zeros((heat_t.shape[-2], heat_t.shape[-1]), device=device, dtype=heat_t.dtype) try: with torch.amp.autocast('cuda', enabled=scaler.is_enabled()): raw_logits = model(img_t) # Mask-out non-finite activations to block NaN gradients from propagating finite_mask = torch.isfinite(raw_logits) logits = torch.where(finite_mask, raw_logits, torch.zeros_like(raw_logits)) if args.logit_clip and float(args.logit_clip) > 0: logits = torch.clamp(logits, min=-float(args.logit_clip), max=float(args.logit_clip)) # Build weights if pos_w_vec is not None: base_pos_w = 1.0 + (pos_w_vec - 1.0) * heat_t else: base_pos_w = torch.ones_like(heat_t) B,C,H,W = heat_t.shape wneg = None if neg_t is not None: wneg = neg_t if wneg.dim()==2: wneg = wneg.unsqueeze(0).unsqueeze(0).expand(B,C,H,W) elif wneg.dim()==3: wneg = wneg.unsqueeze(1).expand(B,C,H,W) elif wneg.dim()==4 and wneg.shape[1]==1: wneg = wneg.expand(B,C,H,W) gamma = float(gamma_cur) # Compute BCE per-element in fp32 with torch.amp.autocast('cuda', enabled=False): per_elem = torch.nn.functional.binary_cross_entropy_with_logits( logits.float(), heat_t.float(), reduction='none') if bool(getattr(args, 'use_separate_posneg_loss', False)): # Positive term: emphasize GT regions; normalize by GT mass pos_mask = heat_t.float() pos_w = base_pos_w.float() pos_denom = pos_mask.sum().clamp_min(1e-6) # total GT heat mass across batch pos_loss = (per_elem * pos_w * pos_mask).sum() / pos_denom # Negative term: base=1, add hard-neg boost only on negatives, clamp only neg boost neg_mask = (1.0 - heat_t).float() if wneg is not None and gamma > 0: neg_boost = torch.clamp(gamma * wneg.float(), max=float(getattr(args,'max_neg_weight_multiplier', 3.0))) neg_w = 1.0 + neg_boost else: neg_w = torch.ones_like(heat_t).float() neg_denom = neg_mask.sum().clamp_min(1e-6) # total background mass neg_loss = (per_elem * neg_w * neg_mask).sum() / neg_denom # Use scheduled neg_lambda only; do not downscale by pos/neg mass ratio (overly weak) bce = pos_loss + float(neg_lambda_cur) * neg_loss else: # Legacy: combine weights and take mean or weighted mean weight = base_pos_w if wneg is not None and gamma > 0: weight = weight + gamma * wneg * (1.0 - heat_t) if args.max_weight_multiplier and float(args.max_weight_multiplier) > 0: weight = torch.clamp(weight, max=float(args.max_weight_multiplier)) if bool(getattr(args, 'normalize_weighted_loss', False)): w = weight.float(); denom = w.sum().clamp_min(1e-6) bce = (per_elem * w).sum() / denom else: bce = (per_elem * weight.float()).mean() repel_loss = 0.0 if args.repel_lambda > 0.0: prob = torch.sigmoid(logits) win = max(1, int(args.repel_window)) if win % 2 == 0: win += 1 maxm = torch.nn.functional.max_pool2d(prob, kernel_size=win, stride=1, padding=win//2) mask = (prob >= (maxm - 1e-6)).float() selected = prob * mask repel = (prob - selected).mean() repel_loss = args.repel_lambda * repel # Count loss schedule count_lambda_cur = lin_sched(getattr(args,'count_lambda_init',None), getattr(args,'count_lambda_final',None), epoch, getattr(args,'count_lambda_warmup_epochs',0)) if count_lambda_cur is None: count_lambda_cur = float(getattr(args,'count_loss_lambda',0.0)) count_loss = 0.0 if count_lambda_cur and float(count_lambda_cur) > 0: prob = torch.sigmoid(logits) mean_pred = prob.mean(dim=(2,3)) # BxC mean_gt = heat_t.mean(dim=(2,3)) # BxC count_loss = torch.abs(mean_pred - mean_gt).mean() * float(count_lambda_cur) prox_loss = 0.0 if args.prox_lambda and float(args.prox_lambda) > 0 and logits.shape[1] >= 2: prob = torch.sigmoid(logits) rad = 2; k = 2*rad + 1 tip_idx=0; int_idx=1 dil_tip = F.max_pool2d(heat_t[:,tip_idx:tip_idx+1], kernel_size=k, stride=1, padding=rad) dil_int = F.max_pool2d(heat_t[:,int_idx:int_idx+1], kernel_size=k, stride=1, padding=rad) prox_tip = (prob[:,tip_idx:tip_idx+1] * dil_int).mean() prox_int = (prob[:,int_idx:int_idx+1] * dil_tip).mean() prox_loss = float(args.prox_lambda) * (prox_tip + prox_int) loss = bce + repel_loss + count_loss + prox_loss # Non-finite loss guard if not torch.isfinite(loss): print("[NaN][train] non-finite loss; skipping batch") if args.debug_nan: try: # Summarize tensors def tstat(t, name): if not isinstance(t, torch.Tensor): return finite = torch.isfinite(t) nan_n = (~finite).sum().item() print(f" {name}: shape={tuple(t.shape)}, min={t[finite].min().item() if finite.any() else 'n/a'}, max={t[finite].max().item() if finite.any() else 'n/a'}, nan/inf={nan_n}") print(f" sample: {meta.get('path','')}") tstat(img_t, 'img_t') tstat(heat_t, 'heat_t') tstat(neg_t if isinstance(neg_t, torch.Tensor) else torch.tensor([]), 'neg_t') tstat(logits, 'logits') try: try: tstat(weight, 'weight') except Exception: pass try: tstat(base_pos_w, 'base_pos_w') except Exception: pass except Exception: pass try: tstat(base_pos_w, 'base_pos_w') except Exception: pass except Exception as _e: print(f" [debug_nan] failed to print stats: {_e}") if args.nan_backoff and 0.0 < float(args.nan_backoff) < 1.0: for g in opt.param_groups: g["lr"] = float(g["lr"]) * float(args.nan_backoff) raise ValueError("non-finite-loss") opt.zero_grad(set_to_none=True) if scaler.is_enabled(): scaler.scale(loss).backward() # Unscale for grad checks/clipping scaler.unscale_(opt) # Check for non-finite grads finite = True bad_param = None for n,p in model.named_parameters(): if p.grad is not None and not torch.isfinite(p.grad).all(): finite = False; bad_param = n; break if not finite: print("[NaN][train] non-finite grads; skipping step and backing off LR") if args.debug_nan: try: def tstat(t, name): if not isinstance(t, torch.Tensor): return finite = torch.isfinite(t) nan_n = (~finite).sum().item() print(f" {name}: shape={tuple(t.shape)}, min={t[finite].min().item() if finite.any() else 'n/a'}, max={t[finite].max().item() if finite.any() else 'n/a'}, nan/inf={nan_n}") print(f" sample: {meta.get('path','')}") tstat(img_t, 'img_t') tstat(heat_t, 'heat_t') tstat(neg_t if isinstance(neg_t, torch.Tensor) else torch.tensor([]), 'neg_t') tstat(raw_logits, 'raw_logits') tstat(logits, 'logits') prob_dbg = torch.sigmoid(logits) tstat(prob_dbg, 'prob') try: tstat(weight, 'weight') except Exception: pass try: tstat(base_pos_w, 'base_pos_w') except Exception: pass # Identify the first offending param and report grad stats if bad_param is not None: for n,p in model.named_parameters(): if n == bad_param and p.grad is not None: tstat(p.grad, f'grad[{n}]') break except Exception as _e: print(f" [debug_nan] failed to print grad stats: {_e}") opt.zero_grad(set_to_none=True) if args.nan_backoff and 0.0 < float(args.nan_backoff) < 1.0: for g in opt.param_groups: g["lr"] = float(g["lr"]) * float(args.nan_backoff) # Advance scaler state even when skipping step to avoid unscale_() error next iter scaler.update() continue else: if args.grad_clip and float(args.grad_clip) > 0: torch.nn.utils.clip_grad_norm_(model.parameters(), float(args.grad_clip)) scaler.step(opt) scaler.update() else: loss.backward() finite = True bad_param = None for n,p in model.named_parameters(): if p.grad is not None and not torch.isfinite(p.grad).all(): finite = False; bad_param = n; break if not finite: print("[NaN][train] non-finite grads; skipping step and backing off LR") if args.debug_nan: try: def tstat(t, name): if not isinstance(t, torch.Tensor): return finite = torch.isfinite(t) nan_n = (~finite).sum().item() print(f" {name}: shape={tuple(t.shape)}, min={t[finite].min().item() if finite.any() else 'n/a'}, max={t[finite].max().item() if finite.any() else 'n/a'}, nan/inf={nan_n}") print(f" sample: {meta.get('path','')}") tstat(img_t, 'img_t') tstat(heat_t, 'heat_t') tstat(neg_t if isinstance(neg_t, torch.Tensor) else torch.tensor([]), 'neg_t') tstat(raw_logits, 'raw_logits') tstat(logits, 'logits') prob_dbg = torch.sigmoid(logits) tstat(prob_dbg, 'prob') tstat(weight, 'weight') if bad_param is not None: for n,p in model.named_parameters(): if n == bad_param and p.grad is not None: tstat(p.grad, f'grad[{n}]') break except Exception as _e: print(f" [debug_nan] failed to print grad stats: {_e}") opt.zero_grad(set_to_none=True) if args.nan_backoff and 0.0 < float(args.nan_backoff) < 1.0: for g in opt.param_groups: g["lr"] = float(g["lr"]) * float(args.nan_backoff) # Advance scaler state even when skipping step to avoid unscale_() error next iter scaler.update() continue else: if args.grad_clip and float(args.grad_clip) > 0: torch.nn.utils.clip_grad_norm_(model.parameters(), float(args.grad_clip)) opt.step() except (torch.cuda.OutOfMemoryError, RuntimeError, ValueError) as e: msg=str(e).lower(); if "out of memory" in msg or "non-finite-loss" in msg: print("[OOM][train] skipping batch; clearing cache") import gc gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() continue else: raise epoch_loss += loss.item() # Validation (BCE loss + simple metrics @ 2,4,8 px) model.eval() val_metrics = {2: {"p":[], "r":[], "f1":[], "err":[]}, 4: {"p":[], "r":[], "f1":[], "err":[]}, 8: {"p":[], "r":[], "f1":[], "err":[]}} val_loss = 0.0 processed = 0 val_diag_pmax = [] val_diag_pmean = [] val_diag_gtprob = [] val_diag_npeaks = [] with torch.no_grad(): for batch in tqdm(val_loader, smoothing=0, desc="Validate"): if isinstance(batch, (list, tuple)) and len(batch) == 4: img_t, heat_t, _neg_t_unused, meta = batch else: img_t, heat_t, meta = batch img_t = img_t.to(device) try: with torch.amp.autocast('cuda', enabled=scaler.is_enabled()): logits = model(img_t)[0] # CxHxW if args.logit_clip and float(args.logit_clip) > 0: logits = torch.clamp(logits, min=-float(args.logit_clip), max=float(args.logit_clip)) # val BCE if pos_w_vec is not None: weight_v = 1.0 + (pos_w_vec - 1.0) * heat_t.to(device) else: weight_v = None val_bce = F.binary_cross_entropy_with_logits(logits.unsqueeze(0), heat_t.to(device), weight=weight_v) if torch.isfinite(val_bce): val_loss += float(val_bce) processed += 1 else: print("[NaN][val] non-finite loss; skipping sample") continue peaks = decode_peaks(logits, thresh=args.val_thresh, window=args.nms_window_val, per_channel_topk=args.val_topk, max_peaks=args.val_max_peaks, fallback_topk=getattr(args, 'val_fallback_topk', 0)) val_diag_npeaks.append(len(peaks)) prob = torch.sigmoid(logits) val_diag_pmax.append(float(prob.max().item())) val_diag_pmean.append(float(prob.mean().item())) # Convert preds to (x,y) pred_xy = [(x, y) for (x, y, c, s) in peaks] except (torch.cuda.OutOfMemoryError, RuntimeError, ValueError) as e: msg = str(e).lower() if "out of memory" in msg or "non-finite-loss" in msg: print("[OOM][val] skipping sample; clearing cache") import gc gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() continue else: raise # Ground truth from CSV transformed to scaled+padded coords; filter types based on training setup ipath = Path(meta["path"][0]) stem = ipath.stem csv_path = lbl_dir / f"{stem}_node_locations.csv" df_gt = pd.read_csv(csv_path) if "type" in df_gt.columns: tcol = df_gt["type"].astype(str).str.lower() if args.internal_only: # Keep internal + root (map root to internal at eval) keep = (tcol == "internal") | (tcol == "root") df_gt = df_gt[keep].reset_index(drop=True) df_gt.loc[:, "type"] = df_gt["type"].astype(str).str.lower().replace({"root":"internal"}) elif args.no_root_pred: # Predict tip/internal only; drop root df_gt = df_gt[tcol != "root"].reset_index(drop=True) # scale factors may be tensors from collate sx_t = meta["scale_x"] sy_t = meta["scale_y"] if isinstance(sx_t, torch.Tensor): sx = float(sx_t.view(-1)[0].item()) elif isinstance(sx_t, (list, tuple)): sx = float(sx_t[0]) else: sx = float(sx_t) if isinstance(sy_t, torch.Tensor): sy = float(sy_t.view(-1)[0].item()) elif isinstance(sy_t, (list, tuple)): sy = float(sy_t[0]) else: sy = float(sy_t) pads = meta["pad"] pad_left = pad_top = pad_right = pad_bottom = 0 if isinstance(pads, (list, tuple)): if len(pads) == 4 and all(isinstance(x, torch.Tensor) for x in pads): pad_left, pad_top, pad_right, pad_bottom = [int(x.view(-1)[0].item()) for x in pads] elif len(pads) == 1: p0 = pads[0] if isinstance(p0, (list, tuple)) and len(p0) == 4: pad_left, pad_top, pad_right, pad_bottom = [int(v) for v in p0] elif isinstance(p0, torch.Tensor) and p0.numel() >= 4: flat = p0.view(-1) pad_left, pad_top, pad_right, pad_bottom = [int(flat[i].item()) for i in range(4)] elif len(pads) == 4: try: pad_left, pad_top, pad_right, pad_bottom = [int(v) for v in pads] except Exception: pass elif isinstance(pads, torch.Tensor) and pads.numel() >= 4: flat = pads.view(-1) pad_left, pad_top, pad_right, pad_bottom = [int(flat[i].item()) for i in range(4)] gt_x = df_gt["x"].values.astype(float) * sx + pad_left gt_y = df_gt["y"].values.astype(float) * sy + pad_top gt = list(zip(gt_x.tolist(), gt_y.tolist())) # Sample predicted prob at GT points (max over channels) if 'prob' in locals(): prob_max = prob.max(dim=0).values vals = [] Ht, Wt = prob_max.shape for (gx, gy) in gt: xi = int(round(gx)); yi = int(round(gy)) if 0 <= yi < Ht and 0 <= xi < Wt: vals.append(float(prob_max[yi, xi].item())) if len(vals) > 0: val_diag_gtprob.append(float(np.mean(vals))) else: val_diag_gtprob.append(float('nan')) # Fast, vectorized greedy matching on requested device (default CPU to avoid GPU OOM) match_dev = torch.device(args.val_match_device if (args.val_match_device == 'cuda' and torch.cuda.is_available()) else 'cpu') for tau in [2,4,8]: p,r,f1,err = greedy_match_torch(pred_xy, gt, tau, device=match_dev, max_pairs=args.val_max_peaks) val_metrics[tau]["p"].append(p) val_metrics[tau]["r"].append(r) val_metrics[tau]["f1"].append(f1) val_metrics[tau]["err"].append(err) # Aggregate val_loss_avg = (val_loss/processed) if processed>0 else float('nan') log = {"epoch": epoch, "loss": epoch_loss/len(train_loader), "val_loss": val_loss_avg, "neg_lambda": float(neg_lambda_cur) if 'neg_lambda_cur' in locals() else float(getattr(args,'neg_lambda',1.0)), "hardneg_gamma": float(gamma_cur) if 'gamma_cur' in locals() else float(getattr(args,'hardneg_weight_gamma',0.5))} for tau in [2,4,8]: P = np.mean(val_metrics[tau]["p"]) R = np.mean(val_metrics[tau]["r"]) F1= np.mean(val_metrics[tau]["f1"]) E = np.mean([e for e in val_metrics[tau]["err"] if np.isfinite(e)]) if any(np.isfinite(val_metrics[tau]["err"])) else float('inf') log[f"val@{tau}px_P"] = P log[f"val@{tau}px_R"] = R log[f"val@{tau}px_F1"] = F1 log[f"val@{tau}px_err"] = E # Diagnostics: peak count and probability summaries if len(val_diag_npeaks) > 0: log["val_npeaks_mean"] = float(np.mean(val_diag_npeaks)) log["val_prob_max_mean"] = float(np.mean(val_diag_pmax)) log["val_prob_mean"] = float(np.mean(val_diag_pmean)) log["val_prob_at_gt_mean"] = float(np.nanmean(val_diag_gtprob)) logs.append(log) if args.write_csv: with open(metrics_csv_path, "a", newline="") as f: w = csv.writer(f) w.writerow([ log["epoch"], log.get("lr", cur_lr), log.get("lr_next", cur_lr), log["loss"], log["val_loss"], log.get("val@2px_P",0.0), log.get("val@2px_R",0.0), log.get("val@2px_F1",0.0), log.get("val@2px_err",float('inf')), log.get("val@4px_P",0.0), log.get("val@4px_R",0.0), log.get("val@4px_F1",0.0), log.get("val@4px_err",float('inf')), log.get("val@8px_P",0.0), log.get("val@8px_R",0.0), log.get("val@8px_F1",0.0), log.get("val@8px_err",float('inf')), ]) # Track best and early stop valid = np.isfinite(log["val_loss"]) and (log["val_loss"]>0) and processed>0 is_best = False if valid and (log["val_loss"] < (best_val - float(args.early_stop_min_delta))): best_val = log["val_loss"]; best_epoch = epoch; epochs_no_improve = 0 is_best = True else: if epoch > int(args.early_stop_burnin): epochs_no_improve += 1 # Scheduler step and record next LR lr_next = float(opt.param_groups[0]["lr"]) if "sched" in locals() and sched is not None: try: sched.step() lr_next = float(opt.param_groups[0]["lr"]) except Exception: pass log["lr_next"] = lr_next # Validate metrics and optionally abort on collapse valid = np.isfinite(log["val_loss"]) and (log["val_loss"]>0) and processed>0 # Also check P/R/F1 are finite and not all zero if valid: try: pr_ok = True for tau in [2,4,8]: P = log.get(f"val@{tau}px_P", np.nan) R = log.get(f"val@{tau}px_R", np.nan) F1= log.get(f"val@{tau}px_F1", np.nan) if not (np.isfinite(P) and np.isfinite(R) and np.isfinite(F1)): pr_ok = False; break valid = valid and pr_ok except Exception: valid = False if not valid: print("[early_stop] Invalid/degenerate validation metrics detected; stopping.") break # Save best whenever it improves (valid epoch) if valid and (log["val_loss"] < (best_val - float(args.early_stop_min_delta))): best_val = log["val_loss"]; best_epoch = epoch; epochs_no_improve = 0 torch.save({"model": model.state_dict(), "args": vars(args), "epoch": epoch}, out_dir / "best.pt") else: if epoch > int(args.early_stop_burnin): epochs_no_improve += 1 # Early stopping when patience exceeded after burn-in if epoch > int(args.early_stop_burnin) and epochs_no_improve >= int(args.early_stop_patience): print(f"[early_stop] No val_loss improvement for {epochs_no_improve} epochs since epoch {best_epoch}. Stopping.") break # Save periodic checkpoints and a few overlays if epoch % args.save_every == 0 or epoch == args.epochs: save_path = model_path if args.save_epoch_suffix: save_path = out_dir / f"{args.save_name}_e{epoch:03d}.pt" torch.save({"model": model.state_dict(), "args": vars(args), "epoch": epoch}, save_path) # Always save last and optionally best try: torch.save({"model": model.state_dict(), "args": vars(args), "epoch": epoch}, model_path) if is_best: torch.save({"model": model.state_dict(), "args": vars(args), "epoch": epoch, "val_loss": log["val_loss"]}, out_dir / "best.pt") except Exception: pass # Save 2 sample overlays (pred vs GT) from TEST set; fall back to VAL if test is empty src_ds = test_ds if len(test_ds) > 0 else val_ds save_n = min(2, len(src_ds)) for i in range(save_n): item = src_ds[i] if isinstance(item, (list, tuple)) and len(item) == 4: img_t, heat_t, _neg_unused, meta = item else: img_t, heat_t, meta = item img_t = img_t.unsqueeze(0).to(device) with torch.no_grad(): raw_logits = model(img_t) # Mask non-finite and clamp like training logits_s = torch.where(torch.isfinite(raw_logits), raw_logits, torch.zeros_like(raw_logits)) if args.logit_clip and float(args.logit_clip) > 0: logits_s = torch.clamp(logits_s, min=-float(args.logit_clip), max=float(args.logit_clip))[0] else: logits_s = logits_s[0] peaks = decode_peaks(logits_s, thresh=args.sample_thresh, window=args.nms_window_sample, per_channel_topk=args.sample_topk, max_peaks=max(args.sample_topk*3, args.val_max_peaks)) # Load GT from CSV and transform ipath = Path(meta["path"]) stem = ipath.stem csv_path = lbl_dir / f"{stem}_node_locations.csv" df_gt = pd.read_csv(csv_path) sx = float(meta["scale_x"]); sy = float(meta["scale_y"]) pad_left, pad_top, pad_right, pad_bottom = meta["pad"] gt_types = df_gt["type"].astype(str).tolist() if "type" in df_gt.columns else ["internal"]*len(df_gt) gt_pts = list(zip((df_gt["x"].values * sx + pad_left).tolist(), (df_gt["y"].values * sy + pad_top).tolist(), gt_types)) # Build overlay: GT in green, preds in red img_np = (img_t[0,0].detach().cpu().numpy()*255).astype(np.uint8) overlay = np.stack([img_np]*3, axis=-1) # Colors: tip=blue, internal=red, root=green def color_for_type(t): t = str(t).lower() if t == "tip": return (0,0,255) if t == "internal": return (255,0,0) if t == "root": return (0,255,0) return (255,255,0) # GT for (x,y,tlabel) in gt_pts: xi, yi = int(round(x)), int(round(y)) r = 2 y0 = max(0, yi-r); y1 = min(overlay.shape[0], yi+r+1) x0 = max(0, xi-r); x1 = min(overlay.shape[1], xi+r+1) col = color_for_type(tlabel) overlay[y0:y1, x0:x1, 0] = col[0] overlay[y0:y1, x0:x1, 1] = col[1] overlay[y0:y1, x0:x1, 2] = col[2] # Preds (use predicted channel); mark leftmost as root (green) left_idx = None if len(peaks) > 0: left_idx = min(range(len(peaks)), key=lambda k: peaks[k][0]) for idx,(x,y,c,s) in enumerate(peaks): xi, yi = int(round(x)), int(round(y)) r = 2 y0 = max(0, yi-r); y1 = min(overlay.shape[0], yi+r+1) x0 = max(0, xi-r); x1 = min(overlay.shape[1], xi+r+1) if left_idx is not None and idx == left_idx: col = (0,255,0) else: # tip channel assumed 0, internal 1 when no_root_pred col = (0,0,255) if c==0 else (255,0,0) overlay[y0:y1, x0:x1, 0] = col[0] overlay[y0:y1, x0:x1, 1] = col[1] overlay[y0:y1, x0:x1, 2] = col[2] out_path = samples_dir / f"epoch{epoch:03d}_sample{i}.png" try: Image.fromarray(overlay).save(out_path) print(f"[samples] saved {out_path}") except Exception as e: print(f"[samples] failed to save overlay: {e}") # Print per-epoch summary line print(json.dumps(log, indent=None)) # Final: optional plot of loss curves if args.plot: try: import matplotlib matplotlib.use('Agg') import matplotlib.pyplot as plt epochs = [d["epoch"] for d in logs] tr = [d["loss"] for d in logs] vl = [d.get("val_loss", None) for d in logs] plt.figure(figsize=(6,4)) plt.plot(epochs, tr, label="train_loss") if all(v is not None for v in vl): plt.plot(epochs, vl, label="val_loss") plt.xlabel("epoch"); plt.ylabel("loss"); plt.legend(); plt.tight_layout() plt.savefig(out_dir / "loss_plot.png", dpi=150) except Exception as e: with open(out_dir / "plot_error.txt", "w") as f: f.write(str(e)) print(str(model_path)) print(str(samples_dir)) if __name__ == "__main__": main()