| |
| 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 |
| |
| try: |
| import hardneg_utils as hardneg |
| except Exception: |
| try: |
| from . import hardneg_utils as hardneg |
| except Exception: |
| import scripts.hardneg_utils as hardneg |
|
|
| 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 |
|
|
| |
| |
| |
|
|
| 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) |
|
|
| def draw_gaussians(h, w, points, sigma, type_to_channel): |
| C = len(type_to_channel) |
| 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: |
| |
| if not np.isfinite(x) or not np.isfinite(y): |
| continue |
| c = type_to_channel.get(t, type_to_channel.get("internal", 1)) |
| |
| 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): |
| |
| 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): |
| |
| return torch.from_numpy(arr.astype(np.float32)) |
|
|
|
|
| def to_tensor_negmask(arr: np.ndarray): |
| |
| 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) |
| |
| arr[y0:y0+t, x0:x1] = val |
| arr[y1-t:y1, x0:x1] = val |
| |
| 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 |
| |
| 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: |
| 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 |
| |
| if random.random() < 0.6: |
| arr = overlay_random_lines(arr, vertical=True, n_range=(1, 8), thickness=(1, 3), intensity=(0, 40)) |
| |
| if random.random() < 0.3: |
| arr = overlay_random_lines(arr, vertical=False, n_range=(1, 4), thickness=(1, 2), intensity=(0, 40)) |
| |
| if random.random() < 0.4: |
| arr = overlay_random_rectangles(arr, n_range=(0, 2), thickness=(1, 3), intensity=(0, 40)) |
| |
| if random.random() < 0.5: |
| arr = overlay_quadrilateral_shift(arr, delta_range=(-35, 35)) |
| |
| if random.random() < 0.6: |
| arr = overlay_node_occlusions(arr, points, max_frac=0.25, size_px=(3, 10)) |
| return arr |
|
|
| |
| |
| |
|
|
| 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 |
| |
| 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 |
| |
| 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) |
|
|
| |
| img_arr, sx, sy = resize_keep_aspect(img_arr, self.max_side) |
| |
| 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] |
|
|
| |
| df = pd.read_csv(csv_name) |
| |
| 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 |
| |
| 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" |
| |
| 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)) |
|
|
| |
| if self.augment and self.overlay_aug_prob > 0.0: |
| img_arr = apply_random_overlays(img_arr, points, prob=self.overlay_aug_prob) |
|
|
| |
| 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) |
|
|
| |
| heat = draw_gaussians(H, W, points, self.sigma, self.type_to_channel) |
|
|
| |
| img_t = to_tensor_image(img_arr) |
| heat_t = to_tensor_heatmap(heat) |
| neg_t = to_tensor_negmask(neg_mask) |
|
|
| 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 |
|
|
| |
| |
| |
|
|
| 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) |
| x2 = self.down1(x1) |
| x3 = self.down2(x2) |
| x4 = self.down3(x3) |
| x5 = self.down4(x4) |
|
|
| 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) |
|
|
| |
| |
| |
|
|
| 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) |
| |
| 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" |
|
|
| 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}") |
|
|
| |
| 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] |
|
|
| |
| stages = [ |
| [self.enc.conv1, self.enc.bn1], |
| [self.enc.layer1], |
| [self.enc.layer2], |
| [self.enc.layer3], |
| ] |
| |
| for i in range(min(freeze_stages, len(stages))): |
| for m in stages[i]: |
| for p in m.parameters(): |
| p.requires_grad = False |
|
|
| |
| |
| self.up4 = nn.ConvTranspose2d(ch["l4"], ch["l3"], 2, stride=2) |
| self.dec4 = DoubleConv(ch["l3"] + ch["l3"], ch["l3"]) |
| |
| self.up3 = nn.ConvTranspose2d(ch["l3"], ch["l2"], 2, stride=2) |
| self.dec3 = DoubleConv(ch["l2"] + ch["l2"], ch["l2"]) |
| |
| self.up2 = nn.ConvTranspose2d(ch["l2"], ch["l1"], 2, stride=2) |
| self.dec2 = DoubleConv(ch["l1"] + ch["l1"], ch["l1"]) |
| |
| self.up1 = nn.ConvTranspose2d(ch["l1"], ch["x0"], 2, stride=2) |
| self.dec1 = DoubleConv(ch["x0"] + ch["x0"], ch["x0"]) |
| |
| 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): |
| |
| 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): |
| |
| x0 = self.enc.conv1(x) |
| x0 = self.enc.bn1(x0) |
| x0 = self.enc.relu(x0) |
| x1 = self.enc.maxpool(x0) |
| if self.grad_ckpt: |
| l1 = self._ckpt(self.enc.layer1, x1) |
| l2 = self._ckpt(self.enc.layer2, l1) |
| l3 = self._ckpt(self.enc.layer3, l2) |
| l4 = self._ckpt(self.enc.layer4, l3) |
| else: |
| l1 = self.enc.layer1(x1) |
| l2 = self.enc.layer2(l1) |
| l3 = self.enc.layer3(l2) |
| l4 = self.enc.layer4(l3) |
|
|
| 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) |
|
|
| |
| |
| |
|
|
| 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, :, :] |
| |
| maxm = F.max_pool2d(p.unsqueeze(0), kernel_size=window, stride=1, padding=window//2).squeeze(0) |
| keep = (p == maxm) & (p > thresh) |
| idx = torch.nonzero(keep[0], as_tuple=False) |
| if idx.numel() == 0 and fallback_topk and fallback_topk > 0: |
| |
| 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]] |
| 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: |
| |
| 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) |
| G = torch.tensor(gt_pts, dtype=torch.float32, device=dev) |
| N, M = P.shape[0], G.shape[0] |
| if N == 0 or M == 0: |
| return 0.0, 0.0, 0.0, float('inf') |
|
|
| |
| D = torch.cdist(P.unsqueeze(0), G.unsqueeze(0)).squeeze(0) |
| |
| 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()) |
| |
| 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 |
|
|
| |
| |
| |
|
|
| 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)") |
| |
| 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") |
| |
| 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") |
| |
| 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") |
| |
| 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") |
| |
| 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") |
| |
| ap.add_argument("--prox_lambda", type=float, default=0.1, help="Weight for inter-class proximity penalty (tip/internal within 2px)") |
| |
| 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") |
| |
| 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)") |
| |
| 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"))] |
| ) |
|
|
| |
| 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 |
|
|
| |
| 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.") |
|
|
| |
| 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:] |
|
|
| |
| 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 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) |
|
|
| |
| 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) |
|
|
| |
| if args.channels_last: |
| model = model.to(memory_format=torch.channels_last) |
|
|
| |
| 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}") |
|
|
| |
| 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 |
|
|
| |
| 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"]) |
|
|
| |
| 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)) |
|
|
| |
| for epoch in tqdm(range(1, args.epochs + 1), smoothing=0, desc="Epochs"): |
| model.train() |
| if args.freeze_all_bn: |
| _set_bn_eval(model) |
| |
| 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": "<unknown>"} |
| img_t = img_t.to(device, non_blocking=True) |
| if args.channels_last: |
| img_t = img_t.contiguous(memory_format=torch.channels_last) |
| |
| 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) |
| |
| 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)) |
| |
| 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) |
|
|
| |
| 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)): |
| |
| pos_mask = heat_t.float() |
| pos_w = base_pos_w.float() |
| pos_denom = pos_mask.sum().clamp_min(1e-6) |
| pos_loss = (per_elem * pos_w * pos_mask).sum() / pos_denom |
|
|
| |
| 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) |
| neg_loss = (per_elem * neg_w * neg_mask).sum() / neg_denom |
| |
| bce = pos_loss + float(neg_lambda_cur) * neg_loss |
| else: |
| |
| 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_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)) |
| mean_gt = heat_t.mean(dim=(2,3)) |
| 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 |
|
|
| |
| if not torch.isfinite(loss): |
| print("[NaN][train] non-finite loss; skipping batch") |
| 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','<unknown>')}") |
| 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() |
| |
| scaler.unscale_(opt) |
| |
| 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','<unknown>')}") |
| 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 |
| |
| 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) |
| |
| 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','<unknown>')}") |
| 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) |
| |
| 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() |
|
|
| |
| 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] |
| if args.logit_clip and float(args.logit_clip) > 0: |
| logits = torch.clamp(logits, min=-float(args.logit_clip), max=float(args.logit_clip)) |
| |
| 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())) |
| |
| 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 |
|
|
| |
| 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 = (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: |
| |
| df_gt = df_gt[tcol != "root"].reset_index(drop=True) |
| |
| 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())) |
| |
| 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')) |
|
|
| |
| 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) |
|
|
| |
| 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 |
| |
| 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')), |
| ]) |
|
|
| |
| 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 |
|
|
| |
| 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 |
|
|
| |
| valid = np.isfinite(log["val_loss"]) and (log["val_loss"]>0) and processed>0 |
| |
| 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 |
|
|
| |
| 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 |
| |
| 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 |
|
|
| |
| 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) |
|
|
| |
| 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 |
|
|
| |
| 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) |
| |
| 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)) |
| |
| 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)) |
|
|
| |
| img_np = (img_t[0,0].detach().cpu().numpy()*255).astype(np.uint8) |
| overlay = np.stack([img_np]*3, axis=-1) |
| |
| 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) |
|
|
| |
| 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] |
| |
| 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: |
| |
| 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(json.dumps(log, indent=None)) |
|
|
| |
| 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() |
|
|
