diff --git a/agent/code/claude opus 5.5/detect_video.py b/agent/code/claude opus 5.5/detect_video.py new file mode 100644 index 000000000..1ae99b50f --- /dev/null +++ b/agent/code/claude opus 5.5/detect_video.py @@ -0,0 +1,211 @@ +"""Detect kinetochores (KTs) and spindle pole bodies (SPBs) in a time-lapse. + +Runs the U-Net trained by ``train_unet.py`` on every frame and writes a CSV +with one row per detection:: + + frame, class, row, col, score + +(``row``/``col`` in pixels of the *original* video; plus ``y_um``/``x_um`` if +``--pixel-size`` is given). Optionally writes an overlay video. + +Supported inputs + * multi-channel TIFF stacks (e.g. TCYX / TYXC from Fiji / MicroManager) + * RGB movies or images (.mp4, .avi, .gif, .png, ...): by default + red = SPB (Cdc11-CFP, shown red) and green = KT (Ndc80-GFP) + +Pixel size matters: the U-Net was trained at ``simulate_spindle.PIXEL_SIZE_UM`` +(0.1 um / pixel). Pass ``--pixel-size`` (um / pixel of your camera) so frames +are rescaled to the training scale, or ``--scale`` to set the zoom directly. + +Examples:: + + python detect_video.py movie.tif --pixel-size 0.0645 --overlay movie_overlay.mp4 + python detect_video.py movie.tif --axes TCYX --kt-channel 0 --spb-channel 1 + python detect_video.py image.png --scale 0.15 +""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +import cv2 +import numpy as np +import pandas as pd +import torch +import torch.nn.functional as F + +from simulate_spindle import PIXEL_SIZE_UM, normalize_frame +from train_unet import CLASS_NAMES, UNet, detect_peaks + +HERE = Path(__file__).resolve().parent + + +# --------------------------------------------------------------------------- # +# Reading +# --------------------------------------------------------------------------- # +def read_stack(path: Path) -> np.ndarray: + """Read a video / stack / image as a numpy array (raw axis order).""" + suffix = path.suffix.lower() + if suffix in (".tif", ".tiff"): + import tifffile + + return tifffile.imread(path) + if suffix in (".png", ".jpg", ".jpeg", ".bmp"): + image = cv2.imread(str(path), cv2.IMREAD_UNCHANGED) + if image.ndim == 3: + image = cv2.cvtColor(image[..., :3], cv2.COLOR_BGR2RGB) + return image[None] + import imageio.v3 as iio + + return np.asarray(iio.imread(path)) # (T, H, W, C) for colour videos + + +def to_tcyx(stack: np.ndarray, axes: str | None) -> np.ndarray: + """Reorder ``stack`` to (T, C, Y, X).""" + if axes is None: # guess + if stack.ndim == 2: + axes = "YX" + elif stack.ndim == 3: + axes = "YXC" if stack.shape[-1] in (2, 3, 4) else "CYX" \ + if stack.shape[0] in (2, 3, 4) else "TYX" + elif stack.ndim == 4: + axes = "TYXC" if stack.shape[-1] in (2, 3, 4) else "TCYX" + else: + raise ValueError(f"cannot guess axes of shape {stack.shape}; use --axes") + print(f"guessed axes {axes} for shape {stack.shape}") + axes = axes.upper() + if len(axes) != stack.ndim: + raise ValueError(f"--axes {axes} does not match shape {stack.shape}") + for missing in "TC": + if missing not in axes: + stack, axes = stack[None], missing + axes + return np.transpose(stack, [axes.index(a) for a in "TCYX"]) + + +# --------------------------------------------------------------------------- # +# Inference +# --------------------------------------------------------------------------- # +def load_model(path: Path, device: torch.device) -> UNet: + checkpoint = torch.load(path, map_location=device) + model = UNet().to(device) + model.load_state_dict(checkpoint["model"]) + model.eval() + print(f"loaded {path} (step {checkpoint.get('step')}, " + f"sim metrics {checkpoint.get('metrics')})") + return model + + +@torch.no_grad() +def predict_frame(model: UNet, frame: np.ndarray, scale: float, + device: torch.device) -> np.ndarray: + """frame (2, H, W) raw -> probability maps (2, H*scale, W*scale).""" + x = torch.as_tensor(np.asarray(frame, dtype=np.float32), device=device)[None] + if scale != 1.0: + x = F.interpolate(x, scale_factor=scale, mode="bilinear", + align_corners=False) + x = torch.as_tensor(normalize_frame(x[0].cpu().numpy()), device=device)[None] + h, w = x.shape[-2:] + pad_h, pad_w = (-h) % 16, (-w) % 16 + x = F.pad(x, (0, pad_w, 0, pad_h), mode="reflect") + with torch.autocast(device.type, enabled=device.type == "cuda"): + prob = torch.sigmoid(model(x)).float() + return prob[0, :, :h, :w].cpu().numpy() + + +def detect(prob: np.ndarray, scale: float, threshold: float, + max_counts: tuple[int | None, int | None]) -> list[tuple]: + rows = [] + for c, name in enumerate(CLASS_NAMES): + peaks = detect_peaks(prob[c], threshold) + peaks = peaks[np.argsort(-peaks[:, 2])] + if max_counts[c]: + peaks = peaks[:max_counts[c]] + for r, col, score in peaks: + # Map back to original pixel coordinates (pixel-centre convention). + rows.append((name, (r + 0.5) / scale - 0.5, + (col + 0.5) / scale - 0.5, float(score))) + return rows + + +def draw_overlay(frame: np.ndarray, detections: list[tuple]) -> np.ndarray: + """RGB uint8 overlay: SPB red / KT green image, cyan KT and magenta SPB rings.""" + rgb = np.zeros(frame.shape[1:] + (3,), np.float32) + for c, out in ((0, 1), (1, 0)): # KT -> green, SPB -> red + lo, hi = np.percentile(frame[c], (1, 99.9)) + rgb[..., out] = np.clip((frame[c] - lo) / max(hi - lo, 1e-6), 0, 1) + rgb = np.ascontiguousarray((rgb * 255).astype(np.uint8)) + radius = max(3, int(round(0.02 * max(frame.shape[1:])))) + for name, r, c, _ in detections: + color = (0, 255, 255) if name == "KT" else (255, 0, 255) + cv2.circle(rgb, (int(round(c)), int(round(r))), + radius if name == "KT" else radius + 2, color, 1, cv2.LINE_AA) + return rgb + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("video", type=Path) + parser.add_argument("--model", type=Path, default=HERE / "unet_kt_spb.pt") + parser.add_argument("--axes", default=None, + help="axis order of the input, e.g. TCYX, TYXC, TZCYX is not supported") + parser.add_argument("--kt-channel", type=int, default=1, + help="channel with Ndc80-GFP (default 1 = green of RGB)") + parser.add_argument("--spb-channel", type=int, default=0, + help="channel with Cdc11-CFP (default 0 = red of RGB)") + group = parser.add_mutually_exclusive_group() + group.add_argument("--pixel-size", type=float, default=None, + help="camera pixel size in um (sets --scale automatically)") + group.add_argument("--scale", type=float, default=1.0, + help="zoom applied before the U-Net (training px / video px)") + parser.add_argument("--threshold", type=float, default=0.2, + help="heatmap peak threshold (0.2 = best F1 on simulations)") + parser.add_argument("--max-kt", type=int, default=6) + parser.add_argument("--max-spb", type=int, default=2) + parser.add_argument("--csv", type=Path, default=None) + parser.add_argument("--overlay", type=Path, default=None, + help="write an overlay video (.mp4) or image (.png)") + args = parser.parse_args() + + scale = args.pixel_size / PIXEL_SIZE_UM if args.pixel_size else args.scale + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + model = load_model(args.model, device) + + stack = to_tcyx(read_stack(args.video), args.axes) + stack = stack[:, [args.kt_channel, args.spb_channel]].astype(np.float32) + print(f"{len(stack)} frame(s) of {stack.shape[2]}x{stack.shape[3]} px, " + f"scale {scale:.3f}") + + records, overlays = [], [] + for t, frame in enumerate(stack): + prob = predict_frame(model, frame, scale, device) + detections = detect(prob, scale, args.threshold, (args.max_kt, args.max_spb)) + records += [(t, *d) for d in detections] + if args.overlay: + overlays.append(draw_overlay(frame, detections)) + + table = pd.DataFrame(records, columns=["frame", "class", "row", "col", "score"]) + if args.pixel_size: + table["y_um"] = table["row"] * args.pixel_size + table["x_um"] = table["col"] * args.pixel_size + csv = args.csv or args.video.with_name(args.video.stem + "_detections.csv") + table.to_csv(csv, index=False) + counts = table.groupby("class").size().to_dict() if len(table) else {} + print(f"wrote {csv}: {counts} detections in {len(stack)} frame(s)") + + if args.overlay: + if len(overlays) == 1 or args.overlay.suffix.lower() in (".png", ".jpg"): + cv2.imwrite(str(args.overlay), cv2.cvtColor(overlays[0], cv2.COLOR_RGB2BGR)) + else: + h, w = overlays[0].shape[:2] + writer = cv2.VideoWriter(str(args.overlay), + cv2.VideoWriter_fourcc(*"mp4v"), 5, (w, h)) + for o in overlays: + writer.write(cv2.cvtColor(o, cv2.COLOR_RGB2BGR)) + writer.release() + print(f"wrote {args.overlay}") + + +if __name__ == "__main__": + main() diff --git a/agent/code/claude opus 5.5/image.png b/agent/code/claude opus 5.5/image.png new file mode 100644 index 000000000..77ce296b1 Binary files /dev/null and b/agent/code/claude opus 5.5/image.png differ diff --git a/agent/code/claude opus 5.5/image_detections.csv b/agent/code/claude opus 5.5/image_detections.csv new file mode 100644 index 000000000..b0fc15acc --- /dev/null +++ b/agent/code/claude opus 5.5/image_detections.csv @@ -0,0 +1,9 @@ +frame,class,row,col,score +0,KT,161.4232972462972,763.1318206787109,0.7734375 +0,KT,168.23305638631186,820.5755030314128,0.7431640625 +0,KT,227.0529702504476,577.1645024617513,0.68359375 +0,KT,152.76937039693198,853.306749979655,0.63330078125 +0,KT,196.5327854156494,713.1734008789062,0.437744140625 +0,KT,227.89520772298178,663.27379099528,0.36328125 +0,SPB,119.05912907918295,864.2799809773763,0.7861328125 +0,SPB,260.6143271128337,504.19404856363934,0.73046875 diff --git a/agent/code/claude opus 5.5/reference_detections.png b/agent/code/claude opus 5.5/reference_detections.png new file mode 100644 index 000000000..0038c0c20 Binary files /dev/null and b/agent/code/claude opus 5.5/reference_detections.png differ diff --git a/agent/code/claude opus 5.5/simulate_spindle.py b/agent/code/claude opus 5.5/simulate_spindle.py new file mode 100644 index 000000000..8cf051ccb --- /dev/null +++ b/agent/code/claude opus 5.5/simulate_spindle.py @@ -0,0 +1,360 @@ +"""DeepTrack2 simulation of fission-yeast mitotic spindles for U-Net training. + +Biology modelled (see ``image.png`` for the reference frame): + +* Cdc11-CFP marks the two spindle pole bodies (SPBs) -> "red" channel. +* Ndc80-GFP marks the six kinetochores (KTs, 3 sister pairs) -> "green" channel. +* Wide-field fluorescence, time-lapse: defocus, diffuse nuclear Ndc80 haze, + faint cytoplasmic CFP speckles, CFP/GFP bleed-through, shot + read noise. + +Every call of :func:`SpindleSimulator.sample` draws one new spindle (length, +orientation, mitotic stage, KT arrangement, brightness, noise) and returns + +* ``image`` : float32 array (2, H, W) — channel 0 = KT (GFP), 1 = SPB (CFP) +* ``target`` : float32 array (2, H, W) — Gaussian heatmaps of KT / SPB centres +* ``kt_pos`` : (n, 2) KT positions (row, col) in pixels +* ``spb_pos`` : (m, 2) SPB positions (row, col) in pixels + +All physical parameters are *estimates*; the most important one to calibrate +against your microscope is ``PIXEL_SIZE_UM`` (camera pixel size in the +sample plane). +""" + +from __future__ import annotations + +import operator +from functools import reduce +from pathlib import Path + +import deeptrack as dt +import numpy as np + +# --------------------------------------------------------------------------- # +# Imaging constants (estimates — calibrate to your setup) +# --------------------------------------------------------------------------- # +PIXEL_SIZE_UM = 0.1 # sample-plane pixel size (resolution / magnification) +IMAGE_SIZE = 128 # training crops are IMAGE_SIZE x IMAGE_SIZE pixels +PADDING = (16, 16, 16, 16) +NA = 1.4 +WAVELENGTH_GFP = 510e-9 # Ndc80-GFP emission +WAVELENGTH_CFP = 475e-9 # Cdc11-CFP emission + +N_SPB = 2 +N_KT = 6 # 3 chromosomes x 2 sister kinetochores +N_HAZE = 10 # dim Ndc80 puncta forming the diffuse nuclear haze +N_SPECKLES = 6 # faint cytoplasmic CFP speckles + +HEATMAP_SIGMA_PX = 1.5 # width of the Gaussian training targets + +UM = 1.0 / PIXEL_SIZE_UM # pixels per micrometre + + +# --------------------------------------------------------------------------- # +# Spindle geometry (one shared draw per sample) +# --------------------------------------------------------------------------- # +def sample_spindle() -> dict: + """Draw one spindle: centre, axis, length and mitotic stage. + + Stages (rough S. pombe mitosis): + * ``prometaphase`` : short spindle, KTs scattered along/around it + * ``metaphase`` : KT pairs clustered around the spindle middle + * ``anaphase`` : sisters separated, KTs close to the SPBs + """ + stage = np.random.choice( + ["prometaphase", "metaphase", "anaphase"], p=[0.3, 0.45, 0.25] + ) + length_um = { + "prometaphase": np.random.uniform(0.8, 2.0), + "metaphase": np.random.uniform(1.8, 3.2), + "anaphase": np.random.uniform(3.0, 9.0), + }[stage] + length = length_um * UM + + angle = np.random.uniform(0, np.pi) + axis = np.array([np.sin(angle), np.cos(angle)]) # (row, col) + normal = np.array([-axis[1], axis[0]]) + + margin = 0.5 * length * np.abs(axis) + 8 + centre = np.array([ + np.random.uniform(margin[0], IMAGE_SIZE - margin[0]), + np.random.uniform(margin[1], IMAGE_SIZE - margin[1]), + ]) if np.all(margin < IMAGE_SIZE / 2) else np.full(2, IMAGE_SIZE / 2) + + # Spindle may be tilted out of the focal plane (wide-field defocus). + z_centre = np.random.normal(0, 0.2 * UM) + z_tilt = np.random.normal(0, 0.06) # dz per pixel along the axis + + return dict( + stage=stage, length=length, axis=axis, normal=normal, + centre=centre, z_centre=z_centre, z_tilt=z_tilt, + ) + + +def point_on_spindle(spindle: dict, t: float, lateral: float = 0.0) -> np.ndarray: + """(row, col) of fraction ``t`` (0 = SPB 1, 1 = SPB 2) along the spindle.""" + return ( + spindle["centre"] + + (t - 0.5) * spindle["length"] * spindle["axis"] + + lateral * spindle["normal"] + ) + + +def z_on_spindle(spindle: dict, t: float) -> float: + return spindle["z_centre"] + (t - 0.5) * spindle["length"] * spindle["z_tilt"] + + +def spb_fraction(index: int) -> float: + return 0.0 if index == 0 else 1.0 + + +def kt_fraction(spindle: dict, index: int, pair_centre: float) -> float: + """Fraction along the spindle of kinetochore ``index`` (pair = index // 2).""" + sister = -1 if index % 2 == 0 else 1 + stage = spindle["stage"] + if stage == "anaphase": + # Sisters segregated: one to each pole, slightly lagging behind it. + lag = np.random.uniform(0.02, 0.15) + return lag if sister < 0 else 1 - lag + separation = np.random.uniform(0.2, 0.7) * UM / spindle["length"] + return pair_centre + sister * separation / 2 + + +# --------------------------------------------------------------------------- # +# Custom operators (wrapped with dt.Lambda) +# --------------------------------------------------------------------------- # +def camera(peak_photons: float = 200.0, background_photons: float = 10.0, + read_noise: float = 2.0): + """Wide-field camera: scale to photons, add background, shot + read noise.""" + + def camera_function(image): + image = np.asarray(image, dtype=np.float64) + peak = image.max() + signal = image / peak * peak_photons if peak > 0 else image + photons = np.random.poisson(np.clip(signal + background_photons, 0, None)) + return photons + np.random.normal(0, read_noise, photons.shape) + + return camera_function + + +def bleed_through(gfp_into_cfp: float = 0.05, cfp_into_gfp: float = 0.05): + """Mix the [KT (GFP), SPB (CFP)] channels to mimic spectral crosstalk.""" + + def mix_function(images): + gfp, cfp = (np.asarray(im, dtype=np.float64).squeeze() for im in images) + return np.stack([gfp + cfp_into_gfp * cfp, cfp + gfp_into_cfp * gfp]) + + return mix_function + + +def normalize_frame(image: np.ndarray) -> np.ndarray: + """Per-channel robust normalisation; used for simulation AND real frames. + + ``image`` is (C, H, W). Background median -> 0, 99.9th percentile -> 1. + """ + image = np.asarray(image, dtype=np.float32) + out = np.empty_like(image) + for c in range(image.shape[0]): + low = np.median(image[c]) + high = np.percentile(image[c], 99.9) + out[c] = (image[c] - low) / max(high - low, 1e-6) + return np.clip(out, -1.0, 5.0) + + +def normalize(): + def normalize_function(image): + return normalize_frame(image) + + return normalize_function + + +def render_heatmap(positions: np.ndarray, shape: tuple[int, int], + sigma: float = HEATMAP_SIGMA_PX) -> np.ndarray: + """Max-combined Gaussian blobs (peak 1) at ``positions`` (row, col).""" + rows, cols = np.mgrid[0:shape[0], 0:shape[1]] + heatmap = np.zeros(shape, dtype=np.float32) + for r, c in np.reshape(positions, (-1, 2)): + blob = np.exp(-((rows - r) ** 2 + (cols - c) ** 2) / (2 * sigma**2)) + heatmap = np.maximum(heatmap, blob) + return heatmap + + +# --------------------------------------------------------------------------- # +# Pipeline +# --------------------------------------------------------------------------- # +def build_optics(wavelength: float) -> dt.Fluorescence: + return dt.Fluorescence( + NA=NA, + resolution=PIXEL_SIZE_UM * 1e-6 * 10, + magnification=10, + wavelength=wavelength, + padding=PADDING, + output_region=(0, 0, IMAGE_SIZE, IMAGE_SIZE), + ) + + +class SpindleSimulator: + """Builds the DeepTrack2 pipeline and produces (image, target) samples. + + All scatterers depend on one shared ``spindle`` feature, so a single + ``.update()`` draws a consistent spindle for both channels. (The + scatterers are listed explicitly instead of using ``^``, because each + ``^`` copy would draw its own spindle.) + """ + + def __init__(self) -> None: + self.spindle = dt.Value(value=sample_spindle) + # Shared per-sample KT pair centres (metaphase / prometaphase). + self.pairs = dt.Value(value=lambda: np.random.normal(0.5, 0.12, 3)) + + # ---------------- SPBs (Cdc11-CFP) ---------------- # + self.spbs = [ + dt.Sphere( + position=lambda spindle, i=i: point_on_spindle( + spindle, spb_fraction(i), np.random.normal(0, 0.3) + ), + z=lambda spindle, i=i: z_on_spindle(spindle, spb_fraction(i)), + spindle=self.spindle.value, + position_unit="pixel", + radius=lambda: np.random.uniform(0.08, 0.15) * 1e-6, + intensity=lambda: np.random.uniform(0.8, 1.2), + ) + for i in range(N_SPB) + ] + + # ---------------- KTs (Ndc80-GFP) ---------------- # + def kt_position(spindle, pairs, i): + stage = spindle["stage"] + pair_centre = np.clip(pairs[i // 2], 0.2, 0.8) + if stage == "prometaphase": + pair_centre = np.random.uniform(0.0, 1.0) + t = kt_fraction(spindle, i, pair_centre) + lateral_sd = {"prometaphase": 0.35, "metaphase": 0.15, + "anaphase": 0.15}[stage] * UM + return point_on_spindle(spindle, t, np.random.normal(0, lateral_sd)) + + self.kts = [ + dt.Sphere( + position=lambda spindle, pairs, i=i: kt_position(spindle, pairs, i), + z=lambda spindle: z_on_spindle(spindle, 0.5) + + np.random.normal(0, 0.15 * UM), + spindle=self.spindle.value, + pairs=self.pairs.value, + position_unit="pixel", + radius=lambda: np.random.uniform(0.06, 0.12) * 1e-6, + intensity=lambda: np.random.uniform(0.5, 1.0), + ) + for i in range(N_KT) + ] + + # ---------------- distractors (not labelled) ---------------- # + haze = [ + dt.PointParticle( + position=lambda spindle: point_on_spindle( + spindle, np.random.uniform(-0.1, 1.1), + np.random.normal(0, 0.5 * UM), + ), + z=lambda: np.random.normal(0, 0.6 * UM), + spindle=self.spindle.value, + position_unit="pixel", + intensity=lambda: np.random.uniform(0.0, 0.25), + ) + for _ in range(N_HAZE) + ] + speckles = [ + dt.PointParticle( + position=lambda spindle: point_on_spindle( + spindle, np.random.uniform(-0.5, 1.5), + np.random.normal(0, 1.5 * UM), + ), + z=lambda: np.random.normal(0, 0.5 * UM), + spindle=self.spindle.value, + position_unit="pixel", + intensity=lambda: np.random.uniform(0.0, 0.15), + ) + for _ in range(N_SPECKLES) + ] + + kt_sample = reduce(operator.rshift, self.kts + haze) + spb_sample = reduce(operator.rshift, self.spbs + speckles) + + def channel_camera() -> dt.Feature: + return dt.Lambda( + function=camera, + peak_photons=lambda: np.random.uniform(40, 400), + background_photons=lambda: np.random.uniform(5, 50), + read_noise=lambda: np.random.uniform(1, 4), + ) + + kt_image = build_optics(WAVELENGTH_GFP)(kt_sample) >> channel_camera() + spb_image = build_optics(WAVELENGTH_CFP)(spb_sample) >> channel_camera() + + mixer = dt.Lambda( + function=bleed_through, + gfp_into_cfp=lambda: np.random.uniform(0.0, 0.15), + cfp_into_gfp=lambda: np.random.uniform(0.0, 0.15), + ) + mixer.__distributed__ = False # receives both channels at once + + self.pipeline = (kt_image & spb_image) >> mixer >> dt.Lambda(function=normalize) + + self.kt_positions = dt.TakeProperties( + reduce(operator.rshift, self.kts), "position") + self.spb_positions = dt.TakeProperties( + reduce(operator.rshift, self.spbs), "position") + + def sample(self) -> dict: + self.pipeline.update() + image = np.asarray(self.pipeline.resolve(), dtype=np.float32) + kt_pos = self._in_view(self.kt_positions.resolve()) + spb_pos = self._in_view(self.spb_positions.resolve()) + shape = image.shape[-2:] + target = np.stack([render_heatmap(kt_pos, shape), + render_heatmap(spb_pos, shape)]) + return dict(image=image, target=target, kt_pos=kt_pos, spb_pos=spb_pos, + stage=self.spindle.value()["stage"]) + + @staticmethod + def _in_view(positions) -> np.ndarray: + positions = np.reshape(np.asarray(positions, dtype=np.float32), (-1, 2)) + inside = np.all((positions >= 0) & (positions < IMAGE_SIZE), axis=1) + return positions[inside] + + +def to_rgb(image: np.ndarray) -> np.ndarray: + """(2, H, W) [KT, SPB] -> RGB with KT green and SPB red, like the reference.""" + kt, spb = (np.clip(image[c], 0, 1) for c in (0, 1)) + return np.stack([spb, kt, np.zeros_like(kt)], axis=-1) + + +def main() -> None: + import matplotlib + + matplotlib.use("Agg") + import matplotlib.pyplot as plt + + np.random.seed(0) + simulator = SpindleSimulator() + n = 6 + fig, axes = plt.subplots(2, n, figsize=(3 * n, 6.4)) + for j in range(n): + s = simulator.sample() + assert s["image"].shape == (2, IMAGE_SIZE, IMAGE_SIZE) + assert np.isfinite(s["image"]).all() + axes[0, j].imshow(to_rgb(s["image"])) + axes[0, j].set_title(s["stage"]) + axes[1, j].imshow(to_rgb(s["image"])) + axes[1, j].scatter(s["kt_pos"][:, 1], s["kt_pos"][:, 0], s=25, + facecolors="none", edgecolors="cyan", linewidths=0.8) + axes[1, j].scatter(s["spb_pos"][:, 1], s["spb_pos"][:, 0], s=60, + facecolors="none", edgecolors="magenta", linewidths=1) + axes[1, j].set_title(f"{len(s['kt_pos'])} KT, {len(s['spb_pos'])} SPB") + for ax in axes[:, j]: + ax.axis("off") + fig.tight_layout() + out = Path(__file__).resolve().parent / "spindle_simulation_examples.png" + fig.savefig(out, dpi=90) + print(f"Saved {out}") + + +if __name__ == "__main__": + main() diff --git a/agent/code/claude opus 5.5/simulate_two_color_fluorescence.py b/agent/code/claude opus 5.5/simulate_two_color_fluorescence.py new file mode 100644 index 000000000..24dcfd9d1 --- /dev/null +++ b/agent/code/claude opus 5.5/simulate_two_color_fluorescence.py @@ -0,0 +1,231 @@ +"""Two-color (red / green) fluorescence simulation with DeepTrack2. + +Reference: ``image.png`` in this folder — a dark-field-of-view, two-channel +fluorescence micrograph showing: + +* red channel : two bright, roughly round puncta at both ends of a diagonal + track, plus a sparse sprinkle of faint red speckles; +* green channel : a handful of bright blobs strung along the same diagonal, + embedded in a diffuse, cloud-like green haze. + +Each colour channel is imaged with its own ``dt.Fluorescence`` microscope +(different emission wavelength), the channels are combined with ``&`` and +merged into an RGB image with a ``dt.Lambda`` operator. + +All physical parameters below are *estimates* chosen to visually match the +reference image; calibrate them against real data if needed. +""" + +from __future__ import annotations + +import matplotlib + +matplotlib.use("Agg") # headless-safe; the figure is also saved to disk + +from pathlib import Path + +import deeptrack as dt +import matplotlib.pyplot as plt +import numpy as np + +# Field of view (pixels) — same ~1:3.1 aspect ratio as the reference image. +IMAGE_HEIGHT = 128 +IMAGE_WIDTH = 400 +PADDING = (32, 32, 32, 32) + +# Number of objects per channel (estimates from the reference image). +N_RED_PUNCTA = 2 +N_RED_SPECKLES = 25 +N_GREEN_BLOBS = 4 +N_GREEN_HAZE = 120 + +OUTPUT_DIR = Path(__file__).resolve().parent + + +# --------------------------------------------------------------------------- # +# Geometry helpers +# --------------------------------------------------------------------------- # +def sample_track(rng: np.random.Generator) -> tuple[np.ndarray, np.ndarray]: + """Sample a diagonal track (start point, end point) in pixel coordinates. + + Positions are (row, column). In the reference the track goes from the + lower-left (red puncta) to the upper-right (red puncta) through the + middle of the image. + """ + start = np.array([ + rng.uniform(0.55, 0.75) * IMAGE_HEIGHT, + rng.uniform(0.30, 0.45) * IMAGE_WIDTH, + ]) + length = rng.uniform(0.25, 0.35) * IMAGE_WIDTH + angle = rng.uniform(np.deg2rad(15), np.deg2rad(30)) # rising to the right + end = start + length * np.array([-np.sin(angle), np.cos(angle)]) + return start, end + + +def point_on_track( + start: np.ndarray, + end: np.ndarray, + t_range: tuple[float, float], + jitter: float, +) -> np.ndarray: + """Random point on the segment ``start``–``end`` with Gaussian jitter.""" + t = np.random.uniform(*t_range) + return start + t * (end - start) + np.random.normal(0, jitter, 2) + + +# --------------------------------------------------------------------------- # +# Custom operator wrapped with dt.Lambda: stack (red, green) -> RGB +# --------------------------------------------------------------------------- # +def merge_to_rgb(red_gain: float = 1.0, green_gain: float = 0.85): + """Return a function mapping a [red, green] list of images to RGB.""" + + def rgb_function(images): + red, green = (np.asarray(im).squeeze() for im in images) + red = np.clip(red * red_gain, 0, 1) + green = np.clip(green * green_gain, 0, 1) + blue = 0.15 * np.minimum(red, green) # slight overlap tint + return np.stack([red, green, blue], axis=-1) + + return rgb_function + + +def clip_unit(): + """Return a function clipping an image to [0, 1].""" + + def clip_function(image): + return np.clip(image, 0.0, 1.0) + + return clip_function + + +# --------------------------------------------------------------------------- # +# Pipeline +# --------------------------------------------------------------------------- # +def build_channel_optics(wavelength: float) -> dt.Fluorescence: + """Fluorescence microscope for one emission channel.""" + return dt.Fluorescence( + NA=0.8, + resolution=1e-6, + magnification=10, + wavelength=wavelength, + padding=PADDING, + output_region=(0, 0, IMAGE_HEIGHT, IMAGE_WIDTH), + ) + + +def build_pipeline(seed: int | None = None) -> dt.Feature: + """Build the two-colour fluorescence pipeline. + + The track geometry is fixed per pipeline (shared by both channels); + object positions, sizes and intensities are re-sampled on every + ``.update()``. + """ + rng = np.random.default_rng(seed) + start, end = sample_track(rng) + + # ---------------- red channel ---------------- # + red_endpoints = iter([]) + + def red_end_position(): + nonlocal red_endpoints + # Alternate between the two ends of the track. + try: + t = next(red_endpoints) + except StopIteration: + red_endpoints = iter([0.0, 1.05]) + t = next(red_endpoints) + return (start + t * (end - start) + + np.random.normal(0, 1.0, 2)) * dt.units.pixel + + red_puncta = dt.Ellipsoid( + position=red_end_position, + position_unit="pixel", + z=lambda: np.random.normal(0, 1), + radius=lambda: np.random.uniform(0.6e-6, 0.8e-6, 3), + rotation=lambda: np.random.uniform(0, np.pi), + intensity=lambda: np.random.uniform(0.9, 1.1), + ) + + red_speckles = dt.PointParticle( + position=lambda: ( + point_on_track(start, end, (0.9, 1.8), jitter=10.0) + ) * dt.units.pixel, + position_unit="pixel", + z=lambda: np.random.normal(0, 5), + intensity=lambda: np.random.uniform(5.0, 40.0), + ) + + red_sample = (red_puncta ^ N_RED_PUNCTA) >> (red_speckles ^ N_RED_SPECKLES) + + # ---------------- green channel ---------------- # + green_blobs = dt.Ellipsoid( + position=lambda: ( + point_on_track(start, end, (0.25, 0.95), jitter=2.5) + ) * dt.units.pixel, + position_unit="pixel", + z=lambda: np.random.normal(0, 2), + radius=lambda: np.random.uniform(0.45e-6, 0.65e-6, 3), + rotation=lambda: np.random.uniform(0, np.pi), + intensity=lambda: np.random.uniform(0.5, 0.9), + ) + + green_haze = dt.PointParticle( + position=lambda: ( + point_on_track(start, end, (0.1, 1.0), jitter=6.0) + ) * dt.units.pixel, + position_unit="pixel", + z=lambda: np.random.normal(0, 8), + intensity=lambda: np.random.uniform(10.0, 45.0), + ) + + green_sample = (green_blobs ^ N_GREEN_BLOBS) >> (green_haze ^ N_GREEN_HAZE) + + # ---------------- imaging ---------------- # + red_optics = build_channel_optics(wavelength=610e-9) + green_optics = build_channel_optics(wavelength=520e-9) + + camera = ( + dt.NormalizeMinMax() + >> dt.Poisson(snr=lambda: np.random.uniform(40, 80), background=0.0) + >> dt.Gaussian(mu=0.0, sigma=lambda: np.random.uniform(0.002, 0.008)) + >> dt.Lambda(function=clip_unit) + ) + + red_image = red_optics(red_sample) >> camera + green_image = green_optics(green_sample) >> camera + + to_rgb = dt.Lambda(function=merge_to_rgb) + to_rgb.__distributed__ = False # receive both channels at once + + return (red_image & green_image) >> to_rgb + + +def main() -> None: + n_examples = 4 + fig, axes = plt.subplots(n_examples + 1, 1, figsize=(10, 2.6 * (n_examples + 1))) + + reference_path = OUTPUT_DIR / "image.png" + if reference_path.exists(): + axes[0].imshow(plt.imread(reference_path)) + axes[0].set_title("Reference") + axes[0].axis("off") + + for i in range(n_examples): + pipeline = build_pipeline(seed=i) + image = np.asarray(pipeline.update().resolve()) + assert image.shape == (IMAGE_HEIGHT, IMAGE_WIDTH, 3), image.shape + assert np.isfinite(image).all() + + plt.imsave(OUTPUT_DIR / f"simulated_{i}.png", np.clip(image, 0, 1)) + axes[i + 1].imshow(np.clip(image, 0, 1)) + axes[i + 1].set_title(f"Simulated #{i}") + axes[i + 1].axis("off") + + fig.tight_layout() + fig.savefig(OUTPUT_DIR / "simulated_overview.png", dpi=100) + plt.show() + print(f"Saved simulated images to {OUTPUT_DIR}") + + +if __name__ == "__main__": + main() diff --git a/agent/code/claude opus 5.5/simulated_0.png b/agent/code/claude opus 5.5/simulated_0.png new file mode 100644 index 000000000..ff253c283 Binary files /dev/null and b/agent/code/claude opus 5.5/simulated_0.png differ diff --git a/agent/code/claude opus 5.5/simulated_1.png b/agent/code/claude opus 5.5/simulated_1.png new file mode 100644 index 000000000..613936728 Binary files /dev/null and b/agent/code/claude opus 5.5/simulated_1.png differ diff --git a/agent/code/claude opus 5.5/simulated_2.png b/agent/code/claude opus 5.5/simulated_2.png new file mode 100644 index 000000000..e9fdc9a19 Binary files /dev/null and b/agent/code/claude opus 5.5/simulated_2.png differ diff --git a/agent/code/claude opus 5.5/simulated_3.png b/agent/code/claude opus 5.5/simulated_3.png new file mode 100644 index 000000000..94c74dd2b Binary files /dev/null and b/agent/code/claude opus 5.5/simulated_3.png differ diff --git a/agent/code/claude opus 5.5/simulated_overview.png b/agent/code/claude opus 5.5/simulated_overview.png new file mode 100644 index 000000000..2be70f7f7 Binary files /dev/null and b/agent/code/claude opus 5.5/simulated_overview.png differ diff --git a/agent/code/claude opus 5.5/spindle_simulation_examples.png b/agent/code/claude opus 5.5/spindle_simulation_examples.png new file mode 100644 index 000000000..0c2f2a7ce Binary files /dev/null and b/agent/code/claude opus 5.5/spindle_simulation_examples.png differ diff --git a/agent/code/claude opus 5.5/train_log.txt b/agent/code/claude opus 5.5/train_log.txt new file mode 100644 index 000000000..c6031f1f8 --- /dev/null +++ b/agent/code/claude opus 5.5/train_log.txt @@ -0,0 +1,71 @@ +device: cuda +validation set: 500 images (6s) +parameters: 7.76 M +step 100/6000 loss 66.9024 16s +step 200/6000 loss 15.4477 32s +step 300/6000 loss 7.2626 48s +step 400/6000 loss 4.1229 64s +step 500/6000 loss 2.6494 80s +step 600/6000 loss 2.0047 95s +step 700/6000 loss 1.6370 111s +step 800/6000 loss 1.2856 126s +step 900/6000 loss 1.0870 142s +step 1000/6000 loss 0.9578 157s + val: {"KT": {"precision": 0.9761, "recall": 0.422, "f1": 0.5892, "rmse_px": 0.5}, "SPB": {"precision": 0.9886, "recall": 0.864, "f1": 0.9221, "rmse_px": 0.256}} +step 1100/6000 loss 0.9067 173s +step 1200/6000 loss 0.8152 188s +step 1300/6000 loss 0.7561 204s +step 1400/6000 loss 0.7307 220s +step 1500/6000 loss 0.6925 235s +step 1600/6000 loss 0.6396 251s +step 1700/6000 loss 0.6580 267s +step 1800/6000 loss 0.5841 282s +step 1900/6000 loss 0.5960 298s +step 2000/6000 loss 0.5735 313s + val: {"KT": {"precision": 0.9821, "recall": 0.438, "f1": 0.6058, "rmse_px": 0.5}, "SPB": {"precision": 0.9945, "recall": 0.904, "f1": 0.9471, "rmse_px": 0.283}} +step 2100/6000 loss 0.5701 329s +step 2200/6000 loss 0.5613 344s +step 2300/6000 loss 0.5679 360s +step 2400/6000 loss 0.5336 375s +step 2500/6000 loss 0.5446 391s +step 2600/6000 loss 0.5352 406s +step 2700/6000 loss 0.5239 422s +step 2800/6000 loss 0.5258 437s +step 2900/6000 loss 0.5133 453s +step 3000/6000 loss 0.5032 469s + val: {"KT": {"precision": 0.9717, "recall": 0.47, "f1": 0.6336, "rmse_px": 0.531}, "SPB": {"precision": 0.9957, "recall": 0.916, "f1": 0.9542, "rmse_px": 0.248}} +step 3100/6000 loss 0.4969 484s +step 3200/6000 loss 0.4790 500s +step 3300/6000 loss 0.4952 516s +step 3400/6000 loss 0.4734 531s +step 3500/6000 loss 0.4675 547s +step 3600/6000 loss 0.4580 563s +step 3700/6000 loss 0.4747 578s +step 3800/6000 loss 0.4683 594s +step 3900/6000 loss 0.4656 610s +step 4000/6000 loss 0.4400 626s + val: {"KT": {"precision": 0.9664, "recall": 0.4883, "f1": 0.6488, "rmse_px": 0.566}, "SPB": {"precision": 0.9957, "recall": 0.922, "f1": 0.9574, "rmse_px": 0.248}} +step 4100/6000 loss 0.4523 642s +step 4200/6000 loss 0.4515 657s +step 4300/6000 loss 0.4763 673s +step 4400/6000 loss 0.4406 688s +step 4500/6000 loss 0.4452 704s +step 4600/6000 loss 0.4261 720s +step 4700/6000 loss 0.4252 735s +step 4800/6000 loss 0.4306 751s +step 4900/6000 loss 0.4198 766s +step 5000/6000 loss 0.4093 782s + val: {"KT": {"precision": 0.975, "recall": 0.4807, "f1": 0.6439, "rmse_px": 0.566}, "SPB": {"precision": 0.9926, "recall": 0.939, "f1": 0.9651, "rmse_px": 0.25}} +step 5100/6000 loss 0.4516 798s +step 5200/6000 loss 0.4306 813s +step 5300/6000 loss 0.4103 829s +step 5400/6000 loss 0.4187 845s +step 5500/6000 loss 0.4205 861s +step 5600/6000 loss 0.4312 877s +step 5700/6000 loss 0.4058 892s +step 5800/6000 loss 0.4158 908s +step 5900/6000 loss 0.4076 924s +step 6000/6000 loss 0.4049 940s + val: {"KT": {"precision": 0.9775, "recall": 0.4913, "f1": 0.6539, "rmse_px": 0.557}, "SPB": {"precision": 0.9916, "recall": 0.949, "f1": 0.9699, "rmse_px": 0.246}} +saved model to /a/rgb_5/unet_kt_spb.pt +saved /a/rgb_5/unet_validation_examples.png diff --git a/agent/code/claude opus 5.5/train_unet.py b/agent/code/claude opus 5.5/train_unet.py new file mode 100644 index 000000000..935268071 --- /dev/null +++ b/agent/code/claude opus 5.5/train_unet.py @@ -0,0 +1,314 @@ +"""Train a U-Net to localise kinetochores (KTs) and spindle pole bodies (SPBs). + +Training data is generated on the fly by the DeepTrack2 simulator in +``simulate_spindle.py`` (every batch is new, so the network never sees the +same image twice). + +Model + input : (B, 2, H, W) channel 0 = Ndc80-GFP (KT), 1 = Cdc11-CFP (SPB), + normalised with ``simulate_spindle.normalize_frame`` + output : (B, 2, H, W) logits of a Gaussian "centre heatmap" per class +Loss + CenterNet-style penalty-reduced focal loss on the heatmaps. +Positions + ``detect_peaks`` = local maxima above a threshold + sub-pixel centroid. + +Usage:: + + python train_unet.py # full training (~20-25 min, 32 CPU workers + 1 GPU) + python train_unet.py --steps 200 # quick smoke test +""" + +from __future__ import annotations + +import argparse +import json +import os +import time +from multiprocessing import get_context +from multiprocessing.pool import Pool +from pathlib import Path + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from scipy.optimize import linear_sum_assignment + +HERE = Path(__file__).resolve().parent +CLASS_NAMES = ("KT", "SPB") + + +# --------------------------------------------------------------------------- # +# Model +# --------------------------------------------------------------------------- # +class DoubleConv(nn.Sequential): + def __init__(self, c_in: int, c_out: int) -> None: + super().__init__( + nn.Conv2d(c_in, c_out, 3, padding=1, bias=False), + nn.BatchNorm2d(c_out), + nn.ReLU(inplace=True), + nn.Conv2d(c_out, c_out, 3, padding=1, bias=False), + nn.BatchNorm2d(c_out), + nn.ReLU(inplace=True), + ) + + +class UNet(nn.Module): + """Plain U-Net; input H and W must be multiples of 2**depth (16).""" + + def __init__(self, in_channels: int = 2, out_channels: int = 2, + base: int = 32, depth: int = 4) -> None: + super().__init__() + chans = [base * 2**i for i in range(depth + 1)] + self.down = nn.ModuleList() + c = in_channels + for ch in chans[:-1]: + self.down.append(DoubleConv(c, ch)) + c = ch + self.bottleneck = DoubleConv(chans[-2], chans[-1]) + self.up = nn.ModuleList() + self.up_conv = nn.ModuleList() + for ch in reversed(chans[:-1]): + self.up.append(nn.ConvTranspose2d(ch * 2, ch, 2, stride=2)) + self.up_conv.append(DoubleConv(ch * 2, ch)) + self.head = nn.Conv2d(base, out_channels, 1) + # Start with low heatmap probabilities (focal-loss convention). + nn.init.constant_(self.head.bias, -2.19) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + skips = [] + for block in self.down: + x = block(x) + skips.append(x) + x = F.max_pool2d(x, 2) + x = self.bottleneck(x) + for up, conv, skip in zip(self.up, self.up_conv, reversed(skips)): + x = conv(torch.cat([up(x), skip], dim=1)) + return self.head(x) + + +def heatmap_focal_loss(logits: torch.Tensor, target: torch.Tensor, + alpha: float = 2.0, beta: float = 4.0) -> torch.Tensor: + """Penalty-reduced pixel-wise focal loss (CenterNet, Zhou et al. 2019).""" + pred = torch.sigmoid(logits.float()).clamp(1e-4, 1 - 1e-4) + pos = target.ge(0.99).float() + neg = 1.0 - pos + pos_loss = torch.log(pred) * (1 - pred) ** alpha * pos + neg_loss = torch.log(1 - pred) * pred**alpha * (1 - target) ** beta * neg + n_pos = pos.sum().clamp(min=1.0) + return -(pos_loss.sum() + neg_loss.sum()) / n_pos + + +# --------------------------------------------------------------------------- # +# Peak detection and evaluation +# --------------------------------------------------------------------------- # +def detect_peaks(heatmap: np.ndarray, threshold: float = 0.2, + min_distance: int = 1) -> np.ndarray: + """Sub-pixel (row, col, score) of local maxima in a probability map.""" + t = torch.as_tensor(heatmap)[None, None] + k = 2 * min_distance + 1 + is_max = t == F.max_pool2d(t, k, stride=1, padding=min_distance) + rows, cols = np.nonzero((is_max & (t > threshold))[0, 0].numpy()) + peaks = [] + h, w = heatmap.shape + for r, c in zip(rows, cols): + r0, r1, c0, c1 = max(r - 1, 0), min(r + 2, h), max(c - 1, 0), min(c + 2, w) + patch = heatmap[r0:r1, c0:c1] + rr, cc = np.mgrid[r0:r1, c0:c1] + weight = patch.sum() + peaks.append(((rr * patch).sum() / weight, (cc * patch).sum() / weight, + heatmap[r, c])) + return np.array(peaks, dtype=np.float32).reshape(-1, 3) + + +def match(pred: np.ndarray, true: np.ndarray, max_dist: float = 3.0): + """Hungarian matching; returns (n_true_positive, list of match distances).""" + if len(pred) == 0 or len(true) == 0: + return 0, [] + d = np.linalg.norm(pred[:, None, :2] - true[None, :, :2], axis=-1) + ri, ci = linear_sum_assignment(d) + good = d[ri, ci] <= max_dist + return int(good.sum()), list(d[ri, ci][good]) + + +@torch.no_grad() +def predict_heatmaps(model: nn.Module, images: np.ndarray, + device: torch.device, batch_size: int = 64) -> np.ndarray: + """images (N, 2, H, W) normalised -> probability maps (N, 2, H, W).""" + model.eval() + out = [] + for i in range(0, len(images), batch_size): + x = torch.as_tensor(images[i:i + batch_size], device=device) + with torch.autocast(device.type, enabled=device.type == "cuda"): + out.append(torch.sigmoid(model(x)).float().cpu().numpy()) + return np.concatenate(out) + + +def evaluate(model: nn.Module, val: dict, device: torch.device, + threshold: float = 0.2) -> dict: + probs = predict_heatmaps(model, val["image"], device) + metrics = {} + for c, name in enumerate(CLASS_NAMES): + tp = n_pred = n_true = 0 + dists: list[float] = [] + for i in range(len(probs)): + pred = detect_peaks(probs[i, c], threshold) + true = val["kt_pos" if c == 0 else "spb_pos"][i] + k, d = match(pred, true) + tp, n_pred, n_true = tp + k, n_pred + len(pred), n_true + len(true) + dists += d + precision = tp / max(n_pred, 1) + recall = tp / max(n_true, 1) + metrics[name] = dict( + precision=round(precision, 4), recall=round(recall, 4), + f1=round(2 * precision * recall / max(precision + recall, 1e-9), 4), + rmse_px=round(float(np.sqrt(np.mean(np.square(dists)))) if dists else float("nan"), 3), + ) + return metrics + + +# --------------------------------------------------------------------------- # +# Data +# --------------------------------------------------------------------------- # +_SIMULATOR = None + + +def _init_worker(base_seed: int) -> None: + """Build one DeepTrack2 pipeline per worker process.""" + global _SIMULATOR + from simulate_spindle import SpindleSimulator + + np.random.seed((base_seed + os.getpid()) % 2**32) + _SIMULATOR = SpindleSimulator() + + +def _simulate(seed: int | None) -> dict: + if seed is not None: + np.random.seed(seed) + return _SIMULATOR.sample() + + +def simulated_batches(pool: Pool, batch_size: int): + """Infinite stream of freshly simulated (image, target) batches. + + Samples travel through pipes (not /dev/shm), so this also works in + containers with a tiny shared-memory mount. + """ + stream = pool.imap_unordered(_simulate, iter(lambda: None, 0), chunksize=4) + while True: + samples = [next(stream) for _ in range(batch_size)] + yield (torch.from_numpy(np.stack([s["image"] for s in samples])), + torch.from_numpy(np.stack([s["target"] for s in samples]))) + + +def make_validation_set(pool: Pool, n: int) -> dict: + samples = pool.map(_simulate, range(10_000, 10_000 + n)) + return dict( + image=np.stack([s["image"] for s in samples]), + kt_pos=[s["kt_pos"] for s in samples], + spb_pos=[s["spb_pos"] for s in samples], + ) + + +# --------------------------------------------------------------------------- # +# Training +# --------------------------------------------------------------------------- # +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--steps", type=int, default=6000) + parser.add_argument("--batch-size", type=int, default=32) + parser.add_argument("--lr", type=float, default=1e-3) + parser.add_argument("--workers", type=int, default=min(32, os.cpu_count() or 1)) + parser.add_argument("--n-val", type=int, default=500) + parser.add_argument("--eval-every", type=int, default=1000) + parser.add_argument("--out", type=Path, default=HERE / "unet_kt_spb.pt") + args = parser.parse_args() + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + torch.manual_seed(0) + print(f"device: {device}") + + t0 = time.time() + # "spawn" (not fork): forked workers inherit the parent's torch/OpenMP + # state and serialise on it (~10 samples/s total instead of ~5/s/worker). + pool = get_context("spawn").Pool(args.workers, initializer=_init_worker, + initargs=(int(t0),)) + val = make_validation_set(pool, args.n_val) + print(f"validation set: {len(val['image'])} images ({time.time() - t0:.0f}s)") + loader = simulated_batches(pool, args.batch_size) + + model = UNet().to(device) + print(f"parameters: {sum(p.numel() for p in model.parameters()) / 1e6:.2f} M") + optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=1e-4) + scheduler = torch.optim.lr_scheduler.OneCycleLR( + optimizer, max_lr=args.lr, total_steps=args.steps, pct_start=0.05) + scaler = torch.amp.GradScaler(enabled=device.type == "cuda") + + history = [] + running = 0.0 + t0 = time.time() + for step, (x, y) in enumerate(loader, start=1): + model.train() + x, y = x.to(device, non_blocking=True), y.to(device, non_blocking=True) + with torch.autocast(device.type, enabled=device.type == "cuda"): + logits = model(x) + loss = heatmap_focal_loss(logits, y) + optimizer.zero_grad(set_to_none=True) + scaler.scale(loss).backward() + scaler.step(optimizer) + scaler.update() + scheduler.step() + running += loss.item() + + if step % 100 == 0: + print(f"step {step:5d}/{args.steps} loss {running / 100:.4f} " + f"{time.time() - t0:.0f}s", flush=True) + running = 0.0 + if step % args.eval_every == 0 or step == args.steps: + metrics = evaluate(model, val, device) + history.append(dict(step=step, **metrics)) + print(f" val: {json.dumps(metrics)}", flush=True) + torch.save({"model": model.state_dict(), "step": step, + "metrics": metrics}, args.out) + if step >= args.steps: + break + pool.terminate() + + (args.out.with_suffix(".history.json")).write_text(json.dumps(history, indent=2)) + print(f"saved model to {args.out}") + save_examples(model, val, device, HERE / "unet_validation_examples.png") + + +def save_examples(model: nn.Module, val: dict, device: torch.device, + path: Path, n: int = 6) -> None: + import matplotlib + + matplotlib.use("Agg") + import matplotlib.pyplot as plt + from simulate_spindle import to_rgb + + probs = predict_heatmaps(model, val["image"][:n], device) + fig, axes = plt.subplots(2, n, figsize=(3 * n, 6.4)) + for i in range(n): + axes[0, i].imshow(to_rgb(val["image"][i])) + axes[0, i].set_title("input") + axes[1, i].imshow(to_rgb(val["image"][i])) + kt, spb = val["kt_pos"][i], val["spb_pos"][i] + axes[1, i].scatter(kt[:, 1], kt[:, 0], marker="+", c="white", s=30, lw=0.8) + axes[1, i].scatter(spb[:, 1], spb[:, 0], marker="x", c="white", s=30, lw=0.8) + for c, color in ((0, "cyan"), (1, "magenta")): + p = detect_peaks(probs[i, c]) + axes[1, i].scatter(p[:, 1], p[:, 0], s=70, facecolors="none", + edgecolors=color, lw=1) + axes[1, i].set_title("truth (+/x) vs U-Net (o)") + for ax in axes[:, i]: + ax.axis("off") + fig.tight_layout() + fig.savefig(path, dpi=90) + print(f"saved {path}") + + +if __name__ == "__main__": + main() diff --git a/agent/code/claude opus 5.5/unet_kt_spb.history.json b/agent/code/claude opus 5.5/unet_kt_spb.history.json new file mode 100644 index 000000000..f43976465 --- /dev/null +++ b/agent/code/claude opus 5.5/unet_kt_spb.history.json @@ -0,0 +1,92 @@ +[ + { + "step": 1000, + "KT": { + "precision": 0.9761, + "recall": 0.422, + "f1": 0.5892, + "rmse_px": 0.5 + }, + "SPB": { + "precision": 0.9886, + "recall": 0.864, + "f1": 0.9221, + "rmse_px": 0.256 + } + }, + { + "step": 2000, + "KT": { + "precision": 0.9821, + "recall": 0.438, + "f1": 0.6058, + "rmse_px": 0.5 + }, + "SPB": { + "precision": 0.9945, + "recall": 0.904, + "f1": 0.9471, + "rmse_px": 0.283 + } + }, + { + "step": 3000, + "KT": { + "precision": 0.9717, + "recall": 0.47, + "f1": 0.6336, + "rmse_px": 0.531 + }, + "SPB": { + "precision": 0.9957, + "recall": 0.916, + "f1": 0.9542, + "rmse_px": 0.248 + } + }, + { + "step": 4000, + "KT": { + "precision": 0.9664, + "recall": 0.4883, + "f1": 0.6488, + "rmse_px": 0.566 + }, + "SPB": { + "precision": 0.9957, + "recall": 0.922, + "f1": 0.9574, + "rmse_px": 0.248 + } + }, + { + "step": 5000, + "KT": { + "precision": 0.975, + "recall": 0.4807, + "f1": 0.6439, + "rmse_px": 0.566 + }, + "SPB": { + "precision": 0.9926, + "recall": 0.939, + "f1": 0.9651, + "rmse_px": 0.25 + } + }, + { + "step": 6000, + "KT": { + "precision": 0.9775, + "recall": 0.4913, + "f1": 0.6539, + "rmse_px": 0.557 + }, + "SPB": { + "precision": 0.9916, + "recall": 0.949, + "f1": 0.9699, + "rmse_px": 0.246 + } + } +] \ No newline at end of file diff --git a/agent/code/claude opus 5.5/unet_kt_spb.pt b/agent/code/claude opus 5.5/unet_kt_spb.pt new file mode 100644 index 000000000..572debb1b Binary files /dev/null and b/agent/code/claude opus 5.5/unet_kt_spb.pt differ diff --git a/agent/code/claude opus 5.5/unet_validation_examples.png b/agent/code/claude opus 5.5/unet_validation_examples.png new file mode 100644 index 000000000..cf2fee11a Binary files /dev/null and b/agent/code/claude opus 5.5/unet_validation_examples.png differ diff --git a/agent/code/qwen 3.8 27b/fluorescence_simulation.py b/agent/code/qwen 3.8 27b/fluorescence_simulation.py new file mode 100644 index 000000000..a942e0a89 --- /dev/null +++ b/agent/code/qwen 3.8 27b/fluorescence_simulation.py @@ -0,0 +1,321 @@ +"""DeepTrack2 dual-channel (red / green) fluorescence simulation. + +Reproduces the visual style of ``image.png`` held in this folder: +a near-black field of view with two distinct red spots and a diagonal +band of green spots, each carrying a handful of small bright satellite +specks, plus faint scattered background grains. + +The simulation is built entirely from documented DeepTrack2 primitives: + + * :class:`dt.PointParticle` – one emissive voxel (fluorophore); + * ``&`` – combine several scatterers into one sample; + * ``tuple`` sample – one sample per wavelength; + * :class:`dt.Fluorescence` – incoherent wide-field fluorescence optics; + * :class:`dt.NormalizeMinMax` – per-channel range normalisation; + * :class:`dt.GaussianBlur` – out-of-focus / defocused PSF broadening; + * :class:`dt.Lambda` / :class:`dt.Merge` – custom numpy ops (add / colour); + * ``>>`` / ``&`` – pipeline chaining. + +Run directly:: + + python fluorescence_simulation.py + +Produces ``simulated_image.png`` (RGB) and ``simulated_channels.png`` +(the two single-channel views side-by-side) in the current folder. +""" +from __future__ import annotations + +import os +import warnings +from typing import Tuple + +import numpy as np +from PIL import Image + +import deeptrack as dt + + +# --------------------------------------------------------------------------- +# Reference geometry (read from image.png in this folder) +# --------------------------------------------------------------------------- + +def _reference_size() -> Tuple[int, int]: + """Return (H, W) of the reference image, falling back to a 1:3 banner.""" + try: + here = os.path.dirname(os.path.abspath(__file__)) + with Image.open(os.path.join(here, "image.png")) as ref: + w, h = ref.size + except Exception: + h, w = 411, 1278 + return h, w + + +H, W = _reference_size() # ~(411, 1278) - the reference is a wide banner + + +# --------------------------------------------------------------------------- +# Custom numpy helpers exposed to the DeepTrack pipeline via dt.Lambda +# --------------------------------------------------------------------------- + +def make_rgb_volume(**kwargs): + """Factory for dt.Lambda: re-arrange (2, H, W, 1) stack into (H, W, 3) RGB.""" + def _rgb_volume(img_stack: np.ndarray) -> np.ndarray: + g = np.asarray(img_stack[0], dtype=np.float64).squeeze(-1) + r = np.asarray(img_stack[1], dtype=np.float64).squeeze(-1) + return np.stack( + [ + np.clip(r, 0.0, 1.0), + np.clip(g, 0.0, 1.0), + 0.25 * np.clip(r + g, 0.0, 1.0), # faint warm halo where channels cross + ], + axis=-1, + ) + + return _rgb_volume + + +def make_shot_noise(**kwargs): + """Factory for dt.Lambda: gentle camera shot noise + sensor read-noise floor. + + Shot noise is modelled as Poisson sampling of the photon count, scaled by a + per-pixel gain. The mean of the Poisson distribution equals the signal, so + the image is preserved in expectation. + """ + def _shot_noise(rgb: np.ndarray) -> np.ndarray: + rgb = np.asarray(rgb, dtype=np.float64) + rng = np.random.default_rng(1234) + # scale by a photon-count gain of ~80 so shot noise is visible but not dominant + gain = 80.0 + photons = rgb * gain + noisy = rng.poisson(lam=photons).astype(np.float64) / gain + # small read-noise floor + noisy += rng.uniform(0.0, 0.004, size=rgb.shape) + return np.clip(noisy, 0.0, 1.0) + + return _shot_noise + + +def make_finalize(**kwargs): + """Factory for dt.Lambda: quantize to uint8 0-255 RGB.""" + def _finalize(rgb: np.ndarray) -> np.ndarray: + rgb = np.asarray(rgb, dtype=np.float64) + rgb = np.clip(rgb, 0.0, 1.0) * 255.0 + return rgb.astype(np.uint8) + + return _finalize + + +rgb_feature = dt.Lambda(function=make_rgb_volume) +rgb_feature.__distributed__ = False + +noise_feature = dt.Lambda(function=make_shot_noise) +noise_feature.__distributed__ = False + +finalize_feature = dt.Lambda(function=make_finalize) +finalize_feature.__distributed__ = False + + +# --------------------------------------------------------------------------- +# Scatterer (fluorophore) construction +# --------------------------------------------------------------------------- + +def _point(x: float, y: float, intensity: float) -> dt.PointParticle: + """A fluorophore at pixel (x=col, y=row) with relative intensity. + + ``PointParticle`` handles the sub-pixel emission discretization for us. + """ + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + return dt.PointParticle( + position=(y, x, 0), + position_unit="pixel", + intensity=float(intensity), + ) + + +def _ellipsoid(x: float, y: float, radius_px: float, + intensity: float = 1.0) -> dt.Ellipsoid: + """A solid circular fluorescent inclusion of radius ``radius_px`` pixels. + + ``dt.Ellipsoid`` is a true volumetric scatterer whose in-focus projection + is a filled disk -- exactly the look of an out-of-focus inclusion. + """ + r_m = radius_px * 1e-6 * 1e-4 # 100 nm per pixel + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + return dt.Ellipsoid( + position=(y, x, 0), + position_unit="pixel", + radius=r_m, + intensity=float(intensity), + ) + + +def _red_blobs() -> dt.Ellipsoid: + """Two large, solid red inclusions matching the reference blob sizes.""" + b1 = _ellipsoid(866, 117, 20.0, 1.0) + b2 = _ellipsoid(501, 260, 17.5, 1.0) + return b1 & b2 + + +def _red_satellites() -> dt.PointParticle: + """Small, brighter red satellite specks.""" + pts = [_point(x, y, a) for x, y, a in + [(831, 61, 1.0), (881, 86, 0.9), (831, 94, 0.7), + (940, 175, 0.6), (800, 250, 0.5), (900, 30, 0.5), + (960, 240, 0.4), (820, 150, 0.4)]] + s = pts[0] + for p in pts[1:]: + s = s & p + return s + + +def _green_blobs() -> dt.Ellipsoid: + """Diagonal band of solid green inclusions.""" + centers = [ + (576, 225, 15.5), + (613, 213, 14.5), + (663, 228, 15.0), + (714, 196, 14.0), + (763, 164, 16.5), + (835, 163, 13.0), + ] + sample: dt.Ellipsoid | None = None + for x, y, r in centers: + e = _ellipsoid(x, y, r, 1.0) + sample = e if sample is None else (sample & e) + return sample # type: ignore[return-value] + + +def _green_satellites() -> dt.PointParticle: + """Small, brighter green satellite specks.""" + pts = [_point(x, y, a) for x, y, a in + [(532, 205, 0.9), (620, 250, 0.7), (700, 150, 0.8), + (870, 210, 0.55), (560, 170, 0.5), (770, 250, 0.45), + (600, 150, 0.4)]] + s = pts[0] + for p in pts[1:]: + s = s & p + return s + + +def make_merge_images(**kwargs): + """Factory for dt.Merge: add a list of images together (clipped to 0..1).""" + def _merge(imgs: list) -> np.ndarray: + acc = None + for x in imgs: + a = np.asarray(x, dtype=np.float64) + acc = a if acc is None else (acc + a) + return np.clip(acc, 0.0, 1.0) + + return _merge + + +def _channel_image(optics: dt.Fluorescence, blob_sample, sat_sample) -> dt.Feature: + """Image one channel = soft blobs (heavy blur) + sharp satellites (light blur). + + The reference shows two distinct scales: broad, diffuse inclusions and a + few tight bright specks. Imaging the two populations with different blur + widths and adding them (via ``dt.Merge``) reproduces that two-tier texture. + """ + blobs = (optics(blob_sample) + >> dt.NormalizeMinMax() + >> dt.GaussianBlur(sigma=11) + >> dt.NormalizeMinMax()) + satellites = (optics(sat_sample) + >> dt.NormalizeMinMax() + >> dt.GaussianBlur(sigma=1.5) + >> dt.NormalizeMinMax()) + return (blobs & satellites) >> dt.Merge(function=make_merge_images) + + +# --------------------------------------------------------------------------- +# Microscope (one fluorescence optics per channel) +# --------------------------------------------------------------------------- + +def _fluorescence(wavelength: float) -> dt.Fluorescence: + return dt.Fluorescence( + NA=1.2, # high-NA oil objective - diffraction-limited blur + resolution=1e-6, # 1 um per pixel before magnification + magnification=10, # -> 100 nm per pixel after magnification + wavelength=wavelength, + padding=(32, 32, 32, 32), + output_region=(0, 0, H, W), + ) + + +def build_optics() -> tuple[dt.Fluorescence, dt.Fluorescence]: + with warnings.catch_warnings(): + warnings.simplefilter("ignore") # "NA / wavelength * resolution / magnification" + g = _fluorescence(wavelength=525e-9) # green channel + r = _fluorescence(wavelength=620e-9) # red channel + return g, r + + +# --------------------------------------------------------------------------- +# Top-level pipeline (pure DeepTrack chaining) +# --------------------------------------------------------------------------- + +def build_pipeline() -> dt.Feature: + green_optics, red_optics = build_optics() + green_img = _channel_image(green_optics, _green_blobs(), _green_satellites()) + red_img = _channel_image(red_optics, _red_blobs(), _red_satellites()) + + # Combine the two channels, colorize, add shot noise, quantize. + return ((green_img & red_img) + >> rgb_feature + >> noise_feature + >> finalize_feature) + + +# --------------------------------------------------------------------------- +# Entry point +# --------------------------------------------------------------------------- + +def main() -> None: + pipeline = build_pipeline() + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + image = np.asarray(pipeline.resolve(), dtype=np.uint8) + + # DeepTrack wraps outputs in batch/sequence dims -> keep the (H, W, 3) core + if image.ndim > 3: + image = image.reshape((-1, H, W, 3)) + image = image.reshape((H, W, 3)) + + here = os.path.dirname(os.path.abspath(__file__)) + out_rgb = os.path.join(here, "simulated_image.png") + Image.fromarray(image, mode="RGB").save(out_rgb) + print(f"wrote {out_rgb} shape={image.shape}") + + # Also save the two normalised single-channel views for inspection. + g, r = build_optics() + gs = np.asarray(_channel_image(g, _green_blobs(), _green_satellites()).resolve(), + dtype=np.float64).reshape((H, W)) + rs = np.asarray(_channel_image(r, _red_blobs(), _red_satellites()).resolve(), + dtype=np.float64).reshape((H, W)) + side = np.concatenate( + [(rs * 255).astype(np.uint8), np.zeros((H, 8), dtype=np.uint8), + (gs * 255).astype(np.uint8)], + axis=1, + ) + out_chan = os.path.join(here, "simulated_channels.png") + Image.fromarray(side, mode="L").save(out_chan) + print(f"wrote {out_chan} shape={side.shape}") + + try: + import matplotlib + matplotlib.use("Agg") + import matplotlib.pyplot as plt + fig, ax = plt.subplots(figsize=(12, 4)) + ax.imshow(image) + ax.axis("off") + fig.tight_layout() + fig.savefig(os.path.join(here, "simulated_image_preview.png"), dpi=110) + print("wrote simulated_image_preview.png") + except Exception as exc: # pragma: no cover - plotting is optional + print(f"(no preview: {exc})") + + +if __name__ == "__main__": + main() diff --git a/agent/code/qwen 3.8 27b/image.png b/agent/code/qwen 3.8 27b/image.png new file mode 100644 index 000000000..77ce296b1 Binary files /dev/null and b/agent/code/qwen 3.8 27b/image.png differ diff --git a/agent/code/qwen 3.8 27b/kt_spb_unet/README.md b/agent/code/qwen 3.8 27b/kt_spb_unet/README.md new file mode 100644 index 000000000..6c11102fe --- /dev/null +++ b/agent/code/qwen 3.8 27b/kt_spb_unet/README.md @@ -0,0 +1,137 @@ +# KT/SPB U-Net on DeepTrack2-synthesised Ndc80-GFP / Cdc11-CFP frames + +A small end-to-end pipeline for segmenting and locating + +* **6 kinetochores** (Ndc80‑GFP, green channel) and +* **2 spindle pole bodies** (Cdc11‑CFP, red channel) + +in wide‑field fluorescence microscopy of a single yeast cell. + +## What is in this folder + +| File | Role | +|------|------| +| `config.py` | Shared constants: canvas, object counts, paths, defaults | +| `unet.py` | PyTorch U‑Net (2‑channel in → 2‑channel out) + `dice_bce_loss` | +| `generator.py` | DeepTrack2 frame renderer (per‑frame determinism) | +| `render_all.py` | Batch render script (writes `data/img_*.npy`, `data/msk_*.npy`) | +| `dataset.py` | PyTorch `Dataset` reading the pre‑rendered frames from disk | +| `train.py` | AdamW + cosine LR training loop, saves best‑Dice checkpoint | +| `detect.py` | Model loading + segment → connected‑components → `(x, y)` positions | +| `detect_video.py` | CLI: run inference on a real video → detections JSON + overlay MP4 | +| `evaluate.py` | Dice + per‑object centroid errors on a hold‑out of the synthetic set | +| `_e2e_smoke.py` | 40‑frame smoke test: 3‑epoch mini‑train + inference (sanity) | +| `run_full.sh` | Render → train orchestrator (safe to run as a long background job) | +| `checkpoints/` | Best checkpoints (`unet_base_best.pt`) | +| `data/` | Pre‑rendered frames (uint16 image + uint8 mask per frame) | +| `logs/` | Pipeline stdout | + +## How it works + +1. **Synthesise**. `generator.render_frame(idx, H, W)` builds a scatterer + field (6 `dt.Ellipsoid` KT inclusions + 2 `dt.Ellipsoid` SPB inclusions + a + few `dt.PointParticle` satellites in each channel), images it through + `dt.Fluorescence` (NA=1.2, 100 nm/px, 10× magnification, 525 nm for the + green channel and 620 nm for the red), multiplies by a mild diagonal + illumination gradient, and adds Poisson shot noise + a small read‑noise + floor. The *ground‑truth mask* is the analytic disc of the nominal object + radius centred on each emitting object — a clean label whose centroid is + exactly the object position, so the U‑Net learns the optics, not the + analytic prior. + + Frame index is the RNG seed (`1_000_000 + idx`), so any frame is + reproducible on demand. + +2. **Render**. `render_all.py --frames 2000 --workers 4` writes + `data/img_.npy` (uint16, H×W×2) and + `data/msk_.npy` (uint8, H×W×2) to disk. DeepTrack is not + process‑safe, so each worker constructs its own `Fluorescence` set. + Throughput on the current machine is ~0.4 frame/s — 2000 frames run + **~70–90 min** as a single job. + +3. **Train**. `train.py --frames 2000 --epochs 12 --base 64 --batch 16` + trains a base‑64 U‑Net on the A100. Loss = per‑class soft Dice + BCE, + with the SPB channel (2 objects) weighted higher than the KT channel + (6 objects). Best (mean Dice) checkpoint is written to + `checkpoints/unet_base64_best.pt`. + +4. **Detect**. `detect.predict_positions(net, img, device)` runs the + network, thresholds each channel at 0.5, finds connected components + (scipy.ndimage), sorts them by area (largest = most confident), and + returns the top‑6 KT centroids and top‑2 SPB centroids as `(x, y)` pairs. + +5. **Run on real data**. + + ``` + python detect_video.py \ + --video /path/to/ndc80_cdc11_timelapse.tiff \ + --ckpt checkpoints/unet_base64_best.pt \ + --max-frames 500 \ + --out-json detections.json + ``` + + produces `detections.json` (per‑frame object list) and + `detections.mp4` (overlay of predicted centroids on the video, with the + frame index and detected counts). + +## Running the full pipeline + +```bash +# 1) Render (skip if already present — the script is idempotent) +python render_all.py --frames 2000 --workers 4 + +# 2) Train +python train.py --frames 2000 --epochs 12 --base 64 --batch 16 + +# 3) Evaluate (Dice + centroid errors on 200 hold‑out frames) +python evaluate.py --ckpt checkpoints/unet_base64_best.pt --n 200 + +# 4) Run on the real time‑lapse — fill in your file path first +python detect_video.py --video /data/videos/my_sample.tiff +``` + +`run_full.sh` runs steps 1 and 2 back‑to‑back as a single job (useful as a +long background task). + +## Tuning knobs + +| Knob | Where | +|------|-------| +| Canvas size (currently 416×640) | `config.HEIGHT/HEIGHT` (must be divisible by 4) | +| KT radius (px) | `config.KT_RADIUS_PX` (currently 14) | +| SPB radius (px) | `config.SPB_RADIUS_PX` (currently 18) | +| Number of KTs / SPBs | `config.N_KT` (6), `config.N_SPB` (2) | +| Noise model | `config.GAIN`, `config.READ_NOISE` | +| Training depth / width | `--base 64`, default `depths=4` in `train.py` | +| Detection threshold | `config.THR_KT` / `THR_SPB` (default 0.5) | +| Smallest component to keep | `config.MIN_COMPONENT_AREA` (default 4 px) | + +## Honest limitations + +* **Domain gap** is the main risk — the real video has the actual microscope's + illumination, camera noise, cell drift, and *other* cellular GFP/CFP + signal (ribosomes, ER, etc.) that the simulator does not. A model trained + only on this synthetic set may under‑transfer. Mitigations: + * augment the simulator (drift, out‑of‑focus variation, background + ribosomal haze, illumination flicker); + * fine‑tune on a small set of user‑annotated frames. +* **Position accuracy vs. segmentation**: the current model is a U‑Net + (segmentation). The pipeline recovers positions by connected components, + so accuracy in *px* depends on the segmentation being tight. On the + smoke test (40 frames, 3 epochs) it already finds roughly the right + centroids, but the full 12‑epoch model should be 1–2 px more accurate. +* **No temporal model**. Per‑frame detection ignores the fact that the + same 6 KTs and 2 SPBs are continuous across frames. Adding a Kalman + filter or tracking‑by‑detection layer would make it much more robust to + a frame where one object is weak. +* **Determinism is only across seeds** — the real video has its own random + background. + +## Disk / GPU estimates + +| Stage | Wall time (typical) | Disk used | +|-------|----------------------|-----------| +| Render 2000 frames | ~70–90 min on this box (≈0.4 f/s) | ~4–5 GB (uint16) | +| Train 12 epochs, base=64, batch=16 | ~10–20 min on A100‑40 | — | +| Detect 500 real frames | ~30–60 s on A100 | — | +| Overlay video | ~20 s | ~10–50 MB mp4 | diff --git a/agent/code/qwen 3.8 27b/kt_spb_unet/_e2e_smoke.py b/agent/code/qwen 3.8 27b/kt_spb_unet/_e2e_smoke.py new file mode 100644 index 000000000..06048a35f --- /dev/null +++ b/agent/code/qwen 3.8 27b/kt_spb_unet/_e2e_smoke.py @@ -0,0 +1,88 @@ +#!/usr/bin/env python3 +"""End-to-end smoke test: generate 8 frames, train a few epochs, run inference. + +Proves the pipeline is wired correctly: +renderer -> dataset -> UNet -> dice/bce -> detect -> positions. + +Usage:: + python _e2e_smoke.py +""" +from __future__ import annotations + +import os +import sys +import time +import warnings +import numpy as np +import torch + +warnings.filterwarnings("ignore") + +import config +import generator +from dataset import PreRenderedDataset +from unet import UNet, dice_bce_loss +from detect import load_model, predict_positions, _segment + + +def main() -> None: + os.makedirs(config.DATA_DIR, exist_ok=True) + os.makedirs(config.CKPT_DIR, exist_ok=True) + + # 1. Render a tiny dataset deterministically. + n = 40 + g_opt = generator._make_green_optics(config.HEIGHT, config.WIDTH) + r_opt = generator._make_red_optics(config.HEIGHT, config.WIDTH) + print(f"rendering {n} frames...") + t0 = time.time() + for i in range(n): + f = generator.render_frame(i, config.HEIGHT, config.WIDTH, (g_opt, r_opt)) + np.save(os.path.join(config.DATA_DIR, f"img_{i:06d}.npy"), f["image"]) + np.save(os.path.join(config.DATA_DIR, f"msk_{i:06d}.npy"), f["masks"]) + print(f" {n} frames in {time.time()-t0:.1f}s") + + # 2. Build small dataset for smoke test. + from torch.utils.data import DataLoader + ds = PreRenderedDataset(config.DATA_DIR, np.arange(n)) + dl = DataLoader(ds, batch_size=8, shuffle=True, num_workers=0, drop_last=True) + img0, msk0 = next(iter(dl)) + print(f" sample img {tuple(img0.shape)} max={img0.max():.3f}") + print(f" sample mask[0] green_px={int(msk0[0,0].sum())} red_px={int(msk0[0,1].sum())}") + + # 3. Train a few epochs with a compact U-Net. + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + net = UNet(in_channels=2, out_channels=2, base=32, depths=3).to(device) + opt = torch.optim.AdamW(net.parameters(), lr=1e-3) + for ep in range(3): + net.train() + tot = 0 + for img, msk in dl: + img = img.to(device); msk = msk.to(device) + loss = dice_bce_loss(net(img), msk, weights=(1.0, 1.3)) + opt.zero_grad(); loss.backward(); opt.step() + tot += loss.item() + print(f" epoch {ep+1}/3 loss={tot/(len(dl)):.4f}") + + # 4. Save a checkpoint (small one) and load it back. + ckpt_path = os.path.join(config.CKPT_DIR, "smoke.pt") + torch.save({"net": net.state_dict(), "base": 32, "in_ch": 2, "out_ch": 2, + "depths": 3, "epoch": 3, "h": config.HEIGHT, "w": config.WIDTH}, + ckpt_path) + loaded, dev = load_model(ckpt_path, device=str(device)) + print(f" reload model OK (base=32)") + + # 5. Inference: segment one frame and find positions. + # image tensor from dataset is (H, W, 2); convert to (H, W, 2) float [0,1]. + img1 = img0[0].cpu().numpy() # (H, W, 2) + pred = predict_positions(loaded, img1, dev, thr_kt=0.4, thr_spb=0.4) + print(f" predicted KT count {pred[0]} (gt 6)") + print(f" predicted SPB count {pred[1]} (gt 2)") + print(f" KT positions: {pred[0]}") + print(f" SPB positions: {pred[1]}") + probs = pred[2] + print(f" prob map maxes green={probs[0].max():.3f} red={probs[1].max():.3f}") + print("SMOKE OK") + + +if __name__ == "__main__": + main() diff --git a/agent/code/qwen 3.8 27b/kt_spb_unet/_par.py b/agent/code/qwen 3.8 27b/kt_spb_unet/_par.py new file mode 100644 index 000000000..7d0ae17eb --- /dev/null +++ b/agent/code/qwen 3.8 27b/kt_spb_unet/_par.py @@ -0,0 +1,39 @@ +import time, warnings; warnings.filterwarnings("ignore") +import os +os.environ["OMP_NUM_THREADS"]="1" +os.environ["MKL_NUM_THREADS"]="1" +import numpy as np +from concurrent.futures import ProcessPoolExecutor + +import sys; sys.path.insert(0,"/a/kt_spb_unet") +import config, generator + +def _init(): + generator._make_green_optics(config.HEIGHT, config.WIDTH) + generator._make_red_optics(config.HEIGHT, config.WIDTH) + +def _work(i): + h, w = config.HEIGHT, config.WIDTH + g = generator._make_green_optics(h, w) + r = generator._make_red_optics(h, w) + t0=time.time() + f = generator.render_frame(i, h, w, (g, r)) + return time.time()-t0 + +# sequential baseline +def _seq(i, g, r): + h,w = config.HEIGHT, config.WIDTH + t0=time.time(); f = generator.render_frame(i,h,w,(g,r)); return time.time()-t0 + +g = generator._make_green_optics(config.HEIGHT, config.WIDTH) +r = generator._make_red_optics(config.HEIGHT, config.WIDTH) +t0=time.time(); _seq(0,g,r) +print(f"sequential: {(time.time()-t0):.2f}s per frame") + +for W in [2,4,8,16]: + idx = [100+i for i in range(20)] + t0 = time.time() + with ProcessPoolExecutor(max_workers=W) as ex: + res = list(ex.map(_work, idx)) + per = (time.time()-t0)/len(idx) + print(f"workers={W}: {per:.2f}s per frame (avg {np.mean(res):.2f}s wall)") diff --git a/agent/code/qwen 3.8 27b/kt_spb_unet/_par2.py b/agent/code/qwen 3.8 27b/kt_spb_unet/_par2.py new file mode 100644 index 000000000..7751b0d05 --- /dev/null +++ b/agent/code/qwen 3.8 27b/kt_spb_unet/_par2.py @@ -0,0 +1,34 @@ +import time, warnings; warnings.filterwarnings("ignore") +os_env = {"OMP_NUM_THREADS":"1","OPENBLAS_NUM_THREADS":"1", + "MKL_NUM_THREADS":"1","VECLIB_MAXIMUM_THREADS":"1", + "NUMEXPR_NUM_THREADS":"1","OMP_DYNAMIC":"false"} +import os +for k,v in os_env.items(): os.environ[k]=v + +import numpy as np, sys +sys.path.insert(0,"/a/kt_spb_unet") +import config, generator +from concurrent.futures import ProcessPoolExecutor + +def init(): + generator._make_green_optics(config.HEIGHT, config.WIDTH) + generator._make_red_optics(config.HEIGHT, config.WIDTH) + +def work(i): + h,w=config.HEIGHT,config.WIDTH + t0=time.time() + f=generator.render_frame(i,h,w,(init.g_opt,init.r_opt)) + return time.time()-t0 +init.g_opt = init.r_opt = None +def init2(): + init.g_opt = generator._make_green_optics(config.HEIGHT, config.WIDTH) + init.r_opt = generator._make_red_optics(config.HEIGHT, config.WIDTH) + +N=48 +for W in [4, 8, 16, 32]: + idx=list(range(500,500+N)) + t0=time.time() + with ProcessPoolExecutor(max_workers=W, initializer=init2) as ex: + list(ex.map(work, idx)) + wall=(time.time()-t0) + print(f"workers={W:2d}: wall {wall:7.1f}s throughput {N/wall:4.2f} f/s total {N} frames") diff --git a/agent/code/qwen 3.8 27b/kt_spb_unet/_probe.py b/agent/code/qwen 3.8 27b/kt_spb_unet/_probe.py new file mode 100644 index 000000000..029f64227 --- /dev/null +++ b/agent/code/qwen 3.8 27b/kt_spb_unet/_probe.py @@ -0,0 +1,20 @@ +import warnings; warnings.simplefilter("ignore") +import numpy as np, deeptrack as dt +H, W = 200, 400 +def grab(f): return np.asarray(f.resolve(), dtype=np.float64).reshape(H, W) +def E(x, y, rm): return dt.Ellipsoid(position=(y,x,0), position_unit='pixel', radius=rm, intensity=1.0) +def P(x, y, i=1.0): return dt.PointParticle(position=(y,x,0), position_unit='pixel', intensity=i) +opt = dt.Fluorescence(NA=1.2, resolution=1e-6, magnification=10, wavelength=525e-9, padding=(32,32,32,32), output_region=(0,0,H,W)) + +print("== ellipse radius vs footprint (100 nm/px) ==") +for rm in [1e-7, 5e-7, 1.4e-6, 1.8e-6]: + img = grab(opt(E(200,100,rm)) >> dt.NormalizeMinMax()) + m = img > 0.5*img.max() + ys, xs = np.where(m) + w = int(xs.max()-xs.min()) if len(xs) else 0 + print(f"radius={rm:.1e} m -> half-max width {w} px") + +print("== point particle per-voxel max vs ellipse per-voxel max (same intensity=1.0) ==") +p = grab(opt(P(100,100,1.0)) >> dt.NormalizeMinMax()) +c = grab(opt(E(100,100,1.4e-6)) >> dt.NormalizeMinMax()) +print("point max:", round(float(p.max()),3), " ellipse max:", round(float(c.max()),3)) diff --git a/agent/code/qwen 3.8 27b/kt_spb_unet/_smoke_frame.png b/agent/code/qwen 3.8 27b/kt_spb_unet/_smoke_frame.png new file mode 100644 index 000000000..a7c2cd302 Binary files /dev/null and b/agent/code/qwen 3.8 27b/kt_spb_unet/_smoke_frame.png differ diff --git a/agent/code/qwen 3.8 27b/kt_spb_unet/_smoke_mask_kt.png b/agent/code/qwen 3.8 27b/kt_spb_unet/_smoke_mask_kt.png new file mode 100644 index 000000000..fc4a5e6b4 Binary files /dev/null and b/agent/code/qwen 3.8 27b/kt_spb_unet/_smoke_mask_kt.png differ diff --git a/agent/code/qwen 3.8 27b/kt_spb_unet/_smoke_mask_spb.png b/agent/code/qwen 3.8 27b/kt_spb_unet/_smoke_mask_spb.png new file mode 100644 index 000000000..0e1e28812 Binary files /dev/null and b/agent/code/qwen 3.8 27b/kt_spb_unet/_smoke_mask_spb.png differ diff --git a/agent/code/qwen 3.8 27b/kt_spb_unet/_smoketest.py b/agent/code/qwen 3.8 27b/kt_spb_unet/_smoketest.py new file mode 100644 index 000000000..d1ef0a19b --- /dev/null +++ b/agent/code/qwen 3.8 27b/kt_spb_unet/_smoketest.py @@ -0,0 +1,37 @@ +"""Quick smoke test: render N frames and report stats.""" +from __future__ import annotations +import sys, time, warnings +import numpy as np +sys.path.insert(0, "/a/kt_spb_unet") +warnings.filterwarnings("ignore") + +from generator import (render_frame, _make_green_optics, _make_red_optics) + +H, W = 412, 640 +g_opt = _make_green_optics(H, W) +r_opt = _make_red_optics(H, W) + +t0 = time.time() +for i in range(3): + f = render_frame(i, H, W, (g_opt, r_opt)) + img = f["image"]; msk = f["masks"] + print(f"frame {i:2d}: img {img.shape} {img.dtype} max={img.max()}; " + f"mask {msk.shape} green_px={int(msk[...,0].sum())} red_px={int(msk[...,1].sum())}") + print(f" KTs={[(round(x),round(y),round(r,1)) for x,y,r in f['kt']]}") + print(f" SPBs={[(round(x),round(y),round(r,1)) for x,y,r in f['spb']]}") + # determinism check + f2 = render_frame(i, H, W, (g_opt, r_opt)) + same = np.array_equal(img, f2["image"]) and np.array_equal(msk, f2["masks"]) + print(f" deterministic: {same}") +print(f"3 frames in {time.time()-t0:.1f}s (~{(time.time()-t0)/3*1000:.0f}ms/frame)") + +# Save one for a quick look +from PIL import Image +f = render_frame(5, H, W, (g_opt, r_opt)) +rgb = np.zeros((H, W, 3), dtype=np.uint8) +rgb[..., 1] = (f["image"][..., 0].astype(np.float32) / 65535 * 255).astype(np.uint8) +rgb[..., 0] = (f["image"][..., 1].astype(np.float32) / 65535 * 255).astype(np.uint8) +Image.fromarray(rgb).save("/a/kt_spb_unet/_smoke_frame.png") +Image.fromarray(f["masks"][..., 0] * 255).save("/a/kt_spb_unet/_smoke_mask_kt.png") +Image.fromarray(f["masks"][..., 1] * 255).save("/a/kt_spb_unet/_smoke_mask_spb.png") +print("wrote smoke outputs") diff --git a/agent/code/qwen 3.8 27b/kt_spb_unet/_timing.py b/agent/code/qwen 3.8 27b/kt_spb_unet/_timing.py new file mode 100644 index 000000000..bec198a3c --- /dev/null +++ b/agent/code/qwen 3.8 27b/kt_spb_unet/_timing.py @@ -0,0 +1,29 @@ +import time, warnings; warnings.filterwarnings("ignore") +import numpy as np, deeptrack as dt +def E(x,y,r): + return dt.Ellipsoid(position=(y,x,0),position_unit="pixel",radius=r*1e-7,intensity=1.0) +def P(x,y,i=1.0): + return dt.PointParticle(position=(y,x,0),position_unit="pixel",intensity=i) +def opt(): + return dt.Fluorescence(NA=1.2,resolution=1e-6,magnification=10,wavelength=525e-9,padding=(32,32,32,32),output_region=(0,0,416,640)) +o = opt() + +# cost: 1 ellipse +t0=time.time() +for _ in range(3): + _ = np.asarray((o(E(300,200,14))).resolve(),dtype=np.float64) +print(f"1 ellipse: {(time.time()-t0)/3*1000:.0f} ms") + +# cost: 6 ellipses chained +t0=time.time() +s=None +for k in range(6): s = E(300+k,200+k,14) if s is None else (s & E(300+k,200+k,14)) +img=np.asarray((o(s)).resolve(),dtype=np.float64) +print(f"6 ellipses (build+image): {(time.time()-t0)*1000:.0f} ms") + +# cost: 10 point particles +t0=time.time() +s=None +for k in range(10): s = P(10+k,20,1.0) if s is None else (s & P(10+k,20,1.0)) +img=np.asarray((o(s)).resolve(),dtype=np.float64) +print(f"10 point particles: {(time.time()-t0)*1000:.0f} ms") diff --git a/agent/code/qwen 3.8 27b/kt_spb_unet/checkpoints/smoke.pt b/agent/code/qwen 3.8 27b/kt_spb_unet/checkpoints/smoke.pt new file mode 100644 index 000000000..8f21b225e Binary files /dev/null and b/agent/code/qwen 3.8 27b/kt_spb_unet/checkpoints/smoke.pt differ diff --git a/agent/code/qwen 3.8 27b/kt_spb_unet/checkpoints/unet_base32_best.pt b/agent/code/qwen 3.8 27b/kt_spb_unet/checkpoints/unet_base32_best.pt new file mode 100644 index 000000000..c70c06ae2 Binary files /dev/null and b/agent/code/qwen 3.8 27b/kt_spb_unet/checkpoints/unet_base32_best.pt differ diff --git a/agent/code/qwen 3.8 27b/kt_spb_unet/checkpoints/unet_base64_best.pt b/agent/code/qwen 3.8 27b/kt_spb_unet/checkpoints/unet_base64_best.pt new file mode 100644 index 000000000..5b186b73a Binary files /dev/null and b/agent/code/qwen 3.8 27b/kt_spb_unet/checkpoints/unet_base64_best.pt differ diff --git a/agent/code/qwen 3.8 27b/kt_spb_unet/config.py b/agent/code/qwen 3.8 27b/kt_spb_unet/config.py new file mode 100644 index 000000000..dd262fe0a --- /dev/null +++ b/agent/code/qwen 3.8 27b/kt_spb_unet/config.py @@ -0,0 +1,43 @@ +"""Shared constants and paths for the KT/SPB U-Net project.""" +from __future__ import annotations + +import os + +PROJECT_DIR = os.path.dirname(os.path.abspath(__file__)) # /a/kt_spb_unet +DATA_DIR = os.environ.get( + "KT_SPB_DATA", os.path.join(PROJECT_DIR, "data") +) # dataset dir +CKPT_DIR = os.path.join(PROJECT_DIR, "checkpoints") + +# ---- canvas / object counts --------------------------------------------- +# (H, W) must be divisible by 4 to work with 3 MaxPool2d stages. The +# reference image is 412 px tall, so we round up to 416 for a clean fit. +WIDTH = 640 # simulation + model canvas width +HEIGHT = 416 # simulation + model canvas height +N_KT = 6 # six kinetochores (Ndc80-GFP, green) +N_SPB = 2 # two spindle pole bodies (Cdc11-CFP, red) + +# Emitter radii (in pixel units). +KT_RADIUS_PX = 14.0 +SPB_RADIUS_PX = 18.0 +PIXEL_NANOMETER = 100e-9 # 100 nm per pixel (matches Fluorescence 10x / 1 um / px) + +# Camera / noise model +GAIN = 80.0 # photon-count gain for shot noise +READ_NOISE = 0.004 # uniform read-noise floor + +# ---- training ------------------------------------------------------------- +SEED = 20261007 +BATCH = 16 +EPOCHS_DEFAULT = 12 +LR_DEFAULT = 3e-4 +UNET_BASE_DEFAULT = 64 +UNET_DEPTHS = 4 +VAL_FRAC = 0.1 + +# ---- inference ------------------------------------------------------------- +# PLACEHOLDER — replace with your real time-lapse (tiff/ome.tiff/avi/mp4/mov). +VIDEO_PATH = "/data/video/your_nde80_cdc11_timelapse.tiff" +MIN_COMPONENT_AREA = 4 +THR_KT = 0.5 +THR_SPB = 0.5 diff --git a/agent/code/qwen 3.8 27b/kt_spb_unet/dataset.py b/agent/code/qwen 3.8 27b/kt_spb_unet/dataset.py new file mode 100644 index 000000000..74fcbad78 --- /dev/null +++ b/agent/code/qwen 3.8 27b/kt_spb_unet/dataset.py @@ -0,0 +1,61 @@ +"""Dataset wrapper: loads pre-rendered frames from disk (memmap uint16).""" +from __future__ import annotations + +import os +import numpy as np +import torch +from torch.utils.data import Dataset + + +class PreRenderedDataset(Dataset): + """Read from a directory of uint16 image/mask arrays written by + ``scripts/render_all.py``. + + Expected layout:: + + data/ + img_000000.npy # (H, W, 2) uint16 + img_000001.npy + ... + msk_000000.npy # (H, W, 2) uint8 + ... + + The dataset divides by 65535 to get float32 [0, 1]. + """ + + def __init__(self, data_dir: str, indices: np.ndarray): + self.data_dir = data_dir + self.indices = np.asarray(indices, dtype=np.int64) + + def __len__(self): + return len(self.indices) + + def __getitem__(self, i): + idx = int(self.indices[i]) + img = np.load(os.path.join(self.data_dir, f"img_{idx:06d}.npy"), + mmap_mode="r") + msk = np.load(os.path.join(self.data_dir, f"msk_{idx:06d}.npy"), + mmap_mode="r") + # HWC -> CHW + img = torch.from_numpy(np.ascontiguousarray( + np.moveaxis(np.asarray(img, dtype=np.float32), -1, 0) / 65535.0)) + msk = torch.from_numpy(np.ascontiguousarray( + np.moveaxis(np.asarray(msk, dtype=np.float32), -1, 0))) + return img, msk + + +def save_frames(data_dir: str, h: int, w: int, + renderer, n_frames: int, start: int = 0, + log_every: int = 50) -> None: + """Render frames to disk, skipping already-existing ones.""" + os.makedirs(data_dir, exist_ok=True) + for i in range(start, start + n_frames): + img_path = os.path.join(data_dir, f"img_{i:06d}.npy") + msk_path = os.path.join(data_dir, f"msk_{i:06d}.npy") + if os.path.exists(img_path) and os.path.exists(msk_path): + continue + frame = renderer(i, h, w) + np.save(img_path, frame["image"]) + np.save(msk_path, frame["masks"]) + if (i - start) % log_every == 0: + print(f"[render] {i - start}/{n_frames}", flush=True) diff --git a/agent/code/qwen 3.8 27b/kt_spb_unet/detect.py b/agent/code/qwen 3.8 27b/kt_spb_unet/detect.py new file mode 100644 index 000000000..2c6678272 --- /dev/null +++ b/agent/code/qwen 3.8 27b/kt_spb_unet/detect.py @@ -0,0 +1,140 @@ +#!/usr/bin/env python3 +"""Segmentation -> object-position detection, and video processing. + +Given a segmentation model, extract object centroids (x, y) for the two +classes (channel 0 = KT/green, channel 1 = SPB/red) from each frame. On a +real time-lapse this is per-frame: the network gives a segmentation, we +threshold, find connected components, and take their centroids. +""" +from __future__ import annotations + +import json +import os +import warnings +import numpy as np +import torch + +warnings.filterwarnings("ignore") + +import config +from unet import UNet + + +def load_model(ckpt_path: str, device: str = "cuda") -> tuple[UNet, torch.device]: + dev = torch.device(device if torch.cuda.is_available() or "cpu" in device else "cpu") + ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False) + net = UNet(in_channels=ckpt.get("in_ch", 2), + out_channels=ckpt.get("out_ch", 2), + base=ckpt.get("base", 64), + depths=ckpt.get("depths", 4)) + net.load_state_dict(ckpt["net"]) + net.to(dev).eval() + return net, dev + + +@torch.no_grad() +def _segment(net: UNet, img: np.ndarray, device: torch.device) -> np.ndarray: + """img: (H, W, 2) or (2, H, W) float[0,1] -> (2, H, W) probability map. + + The network is always in CHW, so we normalise the input to CHW internally. + """ + a = np.asarray(img, dtype=np.float32) + if a.ndim == 3 and a.shape[-1] == 2: # HWC + a = np.ascontiguousarray(np.moveaxis(a, -1, 0)) # -> CHW + x = torch.from_numpy(a).unsqueeze(0).to(device) # (1,2,H,W) + out = torch.sigmoid(net(x)[0]) # (2,H,W) + return out.float().cpu().numpy() + + +# ---- connected components + centroid -------------------------------------- + +def _components_and_centroids(prob: np.ndarray, thr: float, min_area: int): + """scipy.ndimage label -> list of (x, y) centroids, large blobs first.""" + from scipy import ndimage + mask = (prob >= thr).astype(np.uint8) + lab, nobj = ndimage.label(mask) + out = [] + if nobj == 0: + return out + sizes = ndimage.sum(input=np.ones_like(lab), labels=lab, + index=np.arange(1, nobj + 1)) + order = np.argsort(sizes)[::-1] + for k in order: + comp = (lab == (k + 1)) + if comp.sum() < min_area: + continue + cy, cx = ndimage.center_of_mass(comp) + out.append((float(cx), float(cy))) + return out + + +def predict_positions(net, img_hw2: np.ndarray, device, + thr_kt: float = 0.5, thr_spb: float = 0.5, + n_kt: int = config.N_KT, n_spb: int = config.N_SPB, + min_area: int = config.MIN_COMPONENT_AREA): + """Return ``(kt_positions, spb_positions)`` each a list of (x, y).""" + prob = _segment(net, img_hw2, device) # (2, H, W) + kt = _components_and_centroids(prob[0], thr_kt, min_area)[:n_kt] + spb = _components_and_centroids(prob[1], thr_spb, min_area)[:n_spb] + return kt, spb, prob + + +# ---- real video loading --------------------------------------------------- + +def load_video_frames(path: str, max_frames: int | None = None): + """Load a time-lapse (tiff/ome.tiff/avi/mp4) -> list of (H, W, 2) float[0,1]. + + Assumes the file stores the object-of-interest channels. For a + 3-channel RGB/fluorescence video we map green->channel0 (KT), red->channel1 + (SPB). For a 2-channel stack we pass through. + """ + import imageio.v3 as iio + ext = os.path.splitext(path)[1].lower() + if ext in (".tiff", ".tif", ".ome.tiff"): + arr = iio.imread(path, plugin="tifffile") + else: + arr = np.stack(iio.imread(path)) + if arr.ndim == 3: # (T, H, W) single channel -> replicate + arr = np.repeat(arr[..., None], 2, axis=-1) + elif arr.shape[-1] == 3: # RGB + green = arr[..., 1].astype(np.float32) + red = arr[..., 0].astype(np.float32) + arr = np.stack([green, red], axis=-1) + elif arr.shape[-1] == 2: + pass + else: + raise ValueError(f"unsupported video channels {arr.shape}") + mx = arr.max() + if mx > 0: + arr = arr.astype(np.float32) / float(mx) + if max_frames is not None: + arr = arr[:max_frames] + return arr + + +def save_detections(detections: list[dict], out_path: str) -> None: + with open(out_path, "w") as f: + json.dump(detections, f, indent=2) + + +def process_video(net, device, video_path: str, max_frames: int | None = None, + thr_kt: float = 0.5, thr_spb: float = 0.5, + out_json: str | None = None): + """Run inference over a whole time-lapse; returns per-frame detections.""" + frames = load_video_frames(video_path, max_frames) + detections = [] + for i, img in enumerate(frames): + kt, spb, prob = predict_positions(net, img, device, thr_kt, thr_spb) + detections.append({ + "frame": int(i), + "kt": kt, + "spb": spb, + "n_kt": len(kt), + "n_spb": len(spb), + }) + if (i + 1) % 25 == 0: + print(f"[infer] {i+1}/{len(frames)}", flush=True) + if out_json: + save_detections(detections, out_json) + print(f"wrote {out_json}") + return detections diff --git a/agent/code/qwen 3.8 27b/kt_spb_unet/detect_video.py b/agent/code/qwen 3.8 27b/kt_spb_unet/detect_video.py new file mode 100644 index 000000000..d13ce8c39 --- /dev/null +++ b/agent/code/qwen 3.8 27b/kt_spb_unet/detect_video.py @@ -0,0 +1,108 @@ +#!/usr/bin/env python3 +"""Run the trained U-Net on a real time-lapse video and extract object positions. + +Usage: + python detect_video.py \ + --video /path/to/video.tiff \ + --ckpt checkpoints/unet_base64_best.pt \ + --max-frames 500 \ + --out-json detections.json \ + --overlay-overlay overlay.mp4 +""" +from __future__ import annotations + +import argparse +import json +import os +import time +import warnings +import numpy as np +import torch + +warnings.filterwarnings("ignore") + +import config +from unet import UNet +from detect import ( + load_model, predict_positions, load_video_frames, save_detections, +) + + +def overlay_frames(video_frames, detections, out_path): + """Create a short GIF/MP4 overlay of detections on the video frames.""" + import cv2 + N = min(len(video_frames), len(detections)) + W, H = video_frames[0].shape[1], video_frames[0].shape[0] + fourcc = cv2.VideoWriter_fourcc(*'mp4v') + out = cv2.VideoWriter(out_path, fourcc, 10, (W, H, 3)) + for i in range(N): + # Normalise video frame to [0, 255] + img = video_frames[i].copy() # (H, W, 2) + img_max = img.max() + if img_max > 0: + img = img / img_max * 255.0 + frame = np.zeros((H, W, 3), dtype=np.uint8) + frame[..., 1] = img[..., 0].astype(np.uint8) # green + frame[..., 0] = img[..., 1].astype(np.uint8) # red + # Draw detections + for (x, y) in detections[i]["kt"]: + cv2.circle(frame, (int(x), int(y)), 5, (0, 255, 0), 1, 8) + cv2.circle(frame, (int(x), int(y)), 1, (0, 255, 0), 2) + for (x, y) in detections[i]["spb"]: + cv2.circle(frame, (int(x), int(y)), 7, (0, 0, 255), 2, 8) + cv2.circle(frame, (int(x), int(y)), 1, (0, 0, 255), 2) + cv2.putText(frame, f"f={i} KT={len(detections[i]['kt'])} SPB={len(detections[i]['spb'])}", + (5, H - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.35, (200, 200, 200), 1) + out.write(frame) + out.release() + print(f"wrote overlay video {out_path}") + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--video", type=str, default=config.VIDEO_PATH, + help="Real time-lapse video path") + ap.add_argument("--ckpt", type=str, + default=os.path.join(config.CKPT_DIR, "unet_base64_best.pt")) + ap.add_argument("--max-frames", type=int, default=50) + ap.add_argument("--thr-kt", type=float, default=config.THR_KT) + ap.add_argument("--thr-spb", type=float, default=config.THR_SPB) + ap.add_argument("--out-json", type=str, default="detections.json") + ap.add_argument("--overlay", type=str, default=None, + help="Optional: write overlay MP4 (default = out-json with .mp4)") + ap.add_argument("--device", type=str, + default="cuda" if torch.cuda.is_available() else "cpu") + args = ap.parse_args() + + print(f"loading model from {args.ckpt}") + net, device = load_model(args.ckpt, device=args.device) + print(f"model loaded on {device}") + + print(f"loading video {args.video} (max {args.max_frames} frames)") + frames = load_video_frames(args.video, max_frames=args.max_frames) + print(f"loaded {len(frames)} frames of shape {frames[0].shape}") + + t0 = time.time() + detections = [] + for i, img in enumerate(frames): + kt, spb, prob = predict_positions(net, img[None] if img.ndim == 2 else img, + device, thr_kt=args.thr_kt, thr_spb=args.thr_spb) + det = {"frame": int(i), "kt": kt, "spb": spb, + "n_kt": len(kt), "n_spb": len(spb)} + detections.append(det) + if (i + 1) % 25 == 0: + print(f"[infer] {i+1}/{len(frames)} {len(frames)/(time.time()-t0)*1e3:.0f} fps", + flush=True) + + out_json = args.out_json + out_overlay = args.overlay or (os.path.splitext(out_json)[0] + ".mp4") + save_detections(detections, out_json) + print(f"wrote {out_json}") + + overlay_frames(frames, detections, out_overlay) + + print(f"done in {time.time()-t0:.1f}s") + + +if __name__ == "__main__": + main() diff --git a/agent/code/qwen 3.8 27b/kt_spb_unet/evaluate.py b/agent/code/qwen 3.8 27b/kt_spb_unet/evaluate.py new file mode 100644 index 000000000..384bd6197 --- /dev/null +++ b/agent/code/qwen 3.8 27b/kt_spb_unet/evaluate.py @@ -0,0 +1,148 @@ +#!/usr/bin/env python3 +"""Evaluate a trained checkpoint: per-object centroid accuracy + Dice. + +Given pre-rendered synthetic frames + analytic-disc ground-truth, we compute +the predicted segmentation on a holdout set and report, per class (KT / SPB): + + * Dice coefficient of the (hard-thresholded) mask + * Mean / median per-object centroid error (px), via greedy nearest- + neighbour match between predicted components and GT disc centres + * Fraction of GT objects missed (outside ``--max-dist``) + +Usage: + python evaluate.py --ckpt checkpoints/unet_base64_best.pt --n 200 +""" +from __future__ import annotations + +import argparse +import os +import warnings + +import numpy as np +import torch +from scipy import ndimage +from torch.utils.data import DataLoader + +warnings.filterwarnings("ignore") +import config +from dataset import PreRenderedDataset +from detect import _segment, _components_and_centroids, load_model + + +def _gt_centres(mask_2d: np.ndarray) -> np.ndarray: + """Label connected components in a binary disc mask; return (x, y) centres.""" + lab, nobj = ndimage.label(mask_2d > 0.5) + if nobj == 0: + return np.zeros((0, 2), dtype=np.float32) + sizes = ndimage.sum(input=np.ones_like(lab), labels=lab, + index=np.arange(1, nobj + 1)) + if float(sizes.max()) < 4: + return np.zeros((0, 2), dtype=np.float32) + cs = ndimage.center_of_mass(input=np.ones_like(lab), labels=lab, + index=np.arange(1, nobj + 1)) + # cs: list of (y, x); convert to (x, y) + return np.array([(c[1], c[0]) for c in cs], dtype=np.float32) + + +def _greedy_centroid_errors(pred_xy: list[tuple[float, float]], + gt_xy: np.ndarray, + max_dist: float): + """Greedy match (largest-pred first, nearest-GT). Returns (list of errors, + n_missed).""" + if len(gt_xy) == 0: + return [], 0 + gt = gt_xy.astype(np.float32) + # Predictions sorted by size would be ideal; we only have positions, so + # use the natural order (already sorted by component size in detect()). + used = [False] * len(gt) + errors = [] + for p in pred_xy: + p = np.asarray(p, dtype=np.float32) + d = np.linalg.norm(gt - p, axis=1) + idx = int(np.argmin(d)) + if used[idx] or d[idx] > max_dist: + continue + used[idx] = True + errors.append(float(d[idx])) + missed = int(len(gt) - int(np.sum(used))) + return errors, missed + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--ckpt", type=str, + default=os.path.join(config.CKPT_DIR, "unet_base64_best.pt")) + ap.add_argument("--n", type=int, default=200) + ap.add_argument("--max-dist", type=float, default=25.0, + help="max centroid distance to count as a match (px)") + ap.add_argument("--device", type=str, + default="cuda" if torch.cuda.is_available() else "cpu") + ap.add_argument("--frames-total", type=int, default=2000, + help="total pre-rendered frames (for holdout selection)") + args = ap.parse_args() + + net, dev = load_model(args.ckpt, device=args.device) + + rng = np.random.default_rng(config.SEED) + idx = rng.permutation(args.frames_total)[:args.n] + ds = PreRenderedDataset(config.DATA_DIR, idx) + dl = DataLoader(ds, batch_size=1, shuffle=False) + + dice_kt = dice_spb = 0.0 + err_kt: list[float] = [] + err_spb: list[float] = [] + miss_kt = miss_spb = n_kt_gt = n_spb_gt = 0 + + for img, msk in dl: + img1 = img[0].cpu().numpy() # (2, H, W) + msk1 = msk[0].cpu().numpy() # (2, H, W) + prob = _segment(net, img1, dev) # (2, H, W) + + # Dice per class (binary, threshold 0.5) + for c in (0, 1): + p = (prob[c] >= 0.5).astype(bool) + t = (msk1[c] > 0.5).astype(bool) + inter = int((p & t).sum()) + union = int((p | t).sum()) + d = inter / max(union, 1) + if c == 0: + dice_kt += d + else: + dice_spb += d + + # Ground-truth object centres + gt_kt = _gt_centres(msk1[0]) + gt_spb = _gt_centres(msk1[1]) + n_kt_gt += len(gt_kt) + n_spb_gt += len(gt_spb) + + # Predicted object positions (top-6 / top-2 by component size) + pred_kt = _components_and_centroids(prob[0], 0.5, + config.MIN_COMPONENT_AREA)[:config.N_KT] + pred_spb = _components_and_centroids(prob[1], 0.5, + config.MIN_COMPONENT_AREA)[:config.N_SPB] + + e_kt, m_kt = _greedy_centroid_errors(pred_kt, gt_kt, args.max_dist) + e_spb, m_spb = _greedy_centroid_errors(pred_spb, gt_spb, args.max_dist) + err_kt += e_kt + err_spb += e_spb + miss_kt += m_kt + miss_spb += m_spb + + n = args.n + def _fmt(errs): + if not errs: + return "nan" + return f"mean={np.mean(errs):.2f} median={np.median(errs):.2f} p90={np.percentile(errs,90):.2f}" + print(f"=== Evaluated {n} frames ===") + print(f"Dice: KT = {dice_kt/n:.4f} SPB = {dice_spb/n:.4f}") + print(f"KT GT : {n_kt_gt} objects matched={len(err_kt)} missed_in_window={miss_kt}") + if err_kt: + print(f"KT err : {_fmt(err_kt)} px") + print(f"SPB GT : {n_spb_gt} objects matched={len(err_spb)} missed_in_window={miss_spb}") + if err_spb: + print(f"SPB err: {_fmt(err_spb)} px") + + +if __name__ == "__main__": + main() diff --git a/agent/code/qwen 3.8 27b/kt_spb_unet/generator.py b/agent/code/qwen 3.8 27b/kt_spb_unet/generator.py new file mode 100644 index 000000000..4989659f4 --- /dev/null +++ b/agent/code/qwen 3.8 27b/kt_spb_unet/generator.py @@ -0,0 +1,187 @@ +"""DeepTrack2 frame renderer. + +Produces a synthetic fluorescence frame with 6 kinetochore-like inclusions +(Ndc80-GFP, green channel) and 2 spindle pole body inclusions (Cdc11-CFP, +red channel), plus a small number of satellite specks in each channel. + +Output layout +------------- +``frame["image"]`` -- (H, W, 2) uint16, normalised [0, 65535] + channel 0 = green (KT signal) + channel 1 = red (SPB signal) +``frame["masks"]`` -- (H, W, 2) uint8, analytic disk ground truth +``frame["kt"]`` -- list of (x, y, r) tuples (green objects) +``frame["spb"]`` -- list of (x, y, r) tuples (red objects) + +Each frame is deterministic in ``frame_idx`` (seeded numpy RNG), so it can be +pre-rendered once to disk and used by any number of Dataloader workers. +""" +from __future__ import annotations + +import warnings +import numpy as np + +import deeptrack as dt + + +# --------------------------------------------------------------------------- +# Scattering primitives +# --------------------------------------------------------------------------- + +def _make_green_optics(h: int, w: int) -> dt.Fluorescence: + with warnings.catch_warnings(): + warnings.simplefilter("ignore") # "NA/wavelength * res/mag" warning + return dt.Fluorescence( + NA=1.2, resolution=1e-6, magnification=10, + wavelength=525e-9, padding=(32, 32, 32, 32), + output_region=(0, 0, h, w), + ) + + +def _make_red_optics(h: int, w: int) -> dt.Fluorescence: + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + return dt.Fluorescence( + NA=1.2, resolution=1e-6, magnification=10, + wavelength=620e-9, padding=(32, 32, 32, 32), + output_region=(0, 0, h, w), + ) + + +def _disk(x: float, y: float, radius_px: float) -> dt.Ellipsoid: + """Solid circular fluorescent inclusion (radius in px, 100 nm/px).""" + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + return dt.Ellipsoid( + position=(y, x, 0), position_unit="pixel", + radius=radius_px * 1e-7, # 100 nm per pixel (1e-7 m) + intensity=1.0, + ) + + +def _satellite(x: float, y: float, intensity: float) -> dt.PointParticle: + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + return dt.PointParticle( + position=(y, x, 0), position_unit="pixel", intensity=float(intensity), + ) + + +# --------------------------------------------------------------------------- +# Frame rendering +# --------------------------------------------------------------------------- + +def _green_channel(rng: np.random.Generator, + g_opt: dt.Fluorescence, h: int, w: int): + """Render the green channel and return (image_2d, kt_object_list).""" + # 6 kinetochore positions along a diagonal band (like the reference). + t = np.sort(rng.uniform(0.2, 0.8, 6)) + cx = t * w + cy = (h * 0.55) + (h * 0.2) * (t - 0.5) * 2 + rng.normal(0, 8, 6) + cr = 14.0 + rng.normal(0, 1.0, 6) + + blobs: dt.Ellipsoid | None = None + for x, y, r in zip(cx, cy, cr): + e = _disk(float(x), float(y), float(r)) + blobs = e if blobs is None else (blobs & e) + + # A few dim satellite specks near the KTs (typical of wide-field). + sats: dt.PointParticle | None = None + for _ in range(int(rng.integers(3, 9))): + sx = float(rng.uniform(0.1 * w, 0.9 * w)) + sy = float(rng.uniform(0.15 * h, 0.85 * h)) + a = float(rng.uniform(0.4, 1.0)) + p = _satellite(sx, sy, a) + sats = p if sats is None else (sats & p) + + sample: dt.Feature = blobs if sats is None else (blobs & sats) # type: ignore[misc] + img = np.asarray(g_opt(sample).resolve(), dtype=np.float64).reshape((h, w)) + return img, [(float(x), float(y), float(r)) for x, y, r in zip(cx, cy, cr)] + + +def _red_channel(rng: np.random.Generator, + r_opt: dt.Fluorescence, h: int, w: int): + """Render the red channel and return (image_2d, spb_object_list).""" + # 2 spindle pole bodies (Cdc11-CFP) -- larger than KT inclusions. + sx = rng.uniform(0.15, 0.85, 2) * w + sy = rng.uniform(0.15, 0.85, 2) * h + # Enforce a minimum separation between the two SPBs. + min_sep = 60.0 + while np.hypot(sx[1] - sx[0], sy[1] - sy[0]) < min_sep: + sx[1] = rng.uniform(0.15, 0.85) * w + sy[1] = rng.uniform(0.15, 0.85) * h + sr = 18.0 + rng.normal(0, 1.2, 2) + + blobs: dt.Ellipsoid | None = None + for x, y, r in zip(sx, sy, sr): + e = _disk(float(x), float(y), float(r)) + blobs = e if blobs is None else (blobs & e) + + sats: dt.PointParticle | None = None + for _ in range(int(rng.integers(2, 6))): + px = float(rng.uniform(0.1 * w, 0.9 * w)) + py = float(rng.uniform(0.15 * h, 0.85 * h)) + a = float(rng.uniform(0.4, 1.0)) + p = _satellite(px, py, a) + sats = p if sats is None else (sats & p) + + sample: dt.Feature = blobs if sats is None else (blobs & sats) # type: ignore[misc] + img = np.asarray(r_opt(sample).resolve(), dtype=np.float64).reshape((h, w)) + return img, [(float(x), float(y), float(r)) for x, y, r in zip(sx, sy, sr)] + + +def _illumination(h: int, w: int) -> np.ndarray: + """Mild diagonal illumination gradient, values in [1 - a, 1].""" + yy, xx = np.meshgrid(np.linspace(0, 1, h), np.linspace(0, 1, w), indexing="ij") + return 1.0 - 0.15 * (0.3 * xx + 0.7 * yy) + + +def _shot_noise(rgb: np.ndarray, gain: float = 80.0, read: float = 0.004, + seed: int = 0) -> np.ndarray: + """Poisson shot noise + uniform read-noise floor on uint8-like 0..1 image.""" + rng = np.random.default_rng(seed) + lam = np.clip(rgb, 0.0, None) * gain + noisy = rng.poisson(lam).astype(np.float64) / gain + noisy += rng.uniform(0.0, read, size=rgb.shape) + return np.clip(noisy, 0.0, 1.0) + + +def _disk_mask(h: int, w: int, + objs: list[tuple[float, float, float]]) -> np.ndarray: + """Binary analytic disk inside the field of view.""" + yy, xx = np.meshgrid(np.arange(h), np.arange(w), indexing="ij") + m = np.zeros((h, w), dtype=np.uint8) + for cx, cy, r in objs: + m |= ((xx - cx) ** 2 + (yy - cy) ** 2) <= (r * r) + return m + + +def render_frame(frame_idx: int, h: int = 412, w: int = 640, + optics: tuple[dt.Fluorescence, dt.Fluorescence] | None = None): + """Render one synthetic frame deterministically from ``frame_idx``.""" + rng = np.random.default_rng(1_000_000 + int(frame_idx)) + if optics is None: + g_opt, r_opt = _make_green_optics(h, w), _make_red_optics(h, w) + else: + g_opt, r_opt = optics + + g, kts = _green_channel(rng, g_opt, h, w) + r, spbs = _red_channel(rng, r_opt, h, w) + + # Peak-normalise each channel (keeps the downstream noise model meaningful). + g = np.clip(g / (g.max() + 1e-9), 0.0, 1.0) + r = np.clip(r / (r.max() + 1e-9), 0.0, 1.0) + + illumin = _illumination(h, w) + g *= illumin + r *= illumin + + rgb = np.stack([g, r], axis=-1) + rgb = _shot_noise(rgb, gain=80.0, read=0.004, seed=frame_idx * 31 + 7) + image = (rgb * 65535).astype(np.uint16) + + mask_green = _disk_mask(h, w, kts) + mask_red = _disk_mask(h, w, spbs) + mask = np.stack([mask_green, mask_red], axis=-1).astype(np.uint8) + + return {"image": image, "masks": mask, "kt": kts, "spb": spbs} diff --git a/agent/code/qwen 3.8 27b/kt_spb_unet/logs/pipeline.log b/agent/code/qwen 3.8 27b/kt_spb_unet/logs/pipeline.log new file mode 100644 index 000000000..15ca173b9 --- /dev/null +++ b/agent/code/qwen 3.8 27b/kt_spb_unet/logs/pipeline.log @@ -0,0 +1,153 @@ ++ NFRAMES=2000 ++ EPOCHS=12 ++ BASE=64 ++ BATCH=16 +++ date -u +%H:%M:%S ++ echo '=== RENDER (00:03:53) frames=2000 ===' +=== RENDER (00:03:53) frames=2000 === ++ python render_all.py --frames 2000 --workers 4 +rendering 1960 frames with 4 workers -> /a/kt_spb_unet/data +[render] 50/1960 0.4 f/s eta 86.3 min +[render] 100/1960 0.4 f/s eta 83.6 min +[render] 150/1960 0.4 f/s eta 80.2 min +[render] 200/1960 0.4 f/s eta 77.6 min +[render] 250/1960 0.4 f/s eta 75.1 min +[render] 300/1960 0.4 f/s eta 73.1 min +[render] 350/1960 0.4 f/s eta 70.5 min +[render] 400/1960 0.4 f/s eta 68.6 min +[render] 450/1960 0.4 f/s eta 66.1 min +[render] 500/1960 0.4 f/s eta 64.1 min +[render] 550/1960 0.4 f/s eta 61.8 min +[render] 600/1960 0.4 f/s eta 59.7 min +[render] 650/1960 0.4 f/s eta 57.4 min +[render] 700/1960 0.4 f/s eta 55.2 min +[render] 750/1960 0.4 f/s eta 52.9 min +[render] 800/1960 0.4 f/s eta 50.8 min +[render] 850/1960 0.4 f/s eta 48.7 min +[render] 900/1960 0.4 f/s eta 46.5 min +[render] 950/1960 0.4 f/s eta 44.3 min +[render] 1000/1960 0.4 f/s eta 42.1 min +[render] 1050/1960 0.4 f/s eta 39.9 min +[render] 1100/1960 0.4 f/s eta 37.7 min +[render] 1150/1960 0.4 f/s eta 35.5 min +[render] 1200/1960 0.4 f/s eta 33.3 min +[render] 1250/1960 0.4 f/s eta 31.1 min +[render] 1300/1960 0.4 f/s eta 28.9 min +[render] 1350/1960 0.4 f/s eta 26.7 min +[render] 1400/1960 0.4 f/s eta 24.5 min +[render] 1450/1960 0.4 f/s eta 22.3 min +[render] 1500/1960 0.4 f/s eta 20.1 min +[render] 1550/1960 0.4 f/s eta 17.9 min +[render] 1600/1960 0.4 f/s eta 15.7 min +[render] 1650/1960 0.4 f/s eta 13.6 min +[render] 1700/1960 0.4 f/s eta 11.4 min +[render] 1750/1960 0.4 f/s eta 9.2 min +[render] 1800/1960 0.4 f/s eta 7.0 min +[render] 1850/1960 0.4 f/s eta 4.8 min +[render] 1900/1960 0.4 f/s eta 2.6 min +[render] 1950/1960 0.4 f/s eta 0.4 min +done 1960 frames in 5140.8s (0.38 f/s) + +real 85m43.945s +user 159m31.495s +sys 10m40.639s +++ date -u +%H:%M:%S ++ echo '=== TRAINING (01:29:37) base=64 epochs=12 batch=16 ===' +=== TRAINING (01:29:37) base=64 epochs=12 batch=16 === ++ python train.py --frames 2000 --epochs 12 --base 64 --batch 16 +device=cuda base=64 batch=16 +train=1800 val=200 +UNet base=64: 12.42M params +[W1008 01:29:41.308287617 CUDACachingAllocator.cpp:3934] memory allocation failed with OOM on device 0 while trying to allocate 1092616192 bytes (free: 354877440, total: 42406903808). +[W1008 01:29:41.345326129 CUDACachingAllocator.cpp:3934] memory allocation failed with OOM on device 0 while trying to allocate 547356672 bytes (free: 155648000, total: 42406903808). +[W1008 01:29:41.345452607 CUDACachingAllocator.cpp:3934] memory allocation failed with OOM on device 0 while trying to allocate 547356672 bytes (free: 155648000, total: 42406903808). +[W1008 01:29:41.372034411 CUDACachingAllocator.cpp:3934] memory allocation failed with OOM on device 0 while trying to allocate 547356672 bytes (free: 155648000, total: 42406903808). +[W1008 01:29:41.372174254 CUDACachingAllocator.cpp:3934] memory allocation failed with OOM on device 0 while trying to allocate 547356672 bytes (free: 155648000, total: 42406903808). +[W1008 01:29:41.391303710 CUDACachingAllocator.cpp:3934] memory allocation failed with OOM on device 0 while trying to allocate 1092616192 bytes (free: 155648000, total: 42406903808). +[W1008 01:29:41.391445066 CUDACachingAllocator.cpp:3934] memory allocation failed with OOM on device 0 while trying to allocate 1092616192 bytes (free: 155648000, total: 42406903808). +[W1008 01:29:41.392663631 CUDACachingAllocator.cpp:3934] memory allocation failed with OOM on device 0 while trying to allocate 1092616192 bytes (free: 155648000, total: 42406903808). +[W1008 01:29:41.392785932 CUDACachingAllocator.cpp:3934] memory allocation failed with OOM on device 0 while trying to allocate 1092616192 bytes (free: 155648000, total: 42406903808). +[W1008 01:29:41.393265515 CUDACachingAllocator.cpp:3934] memory allocation failed with OOM on device 0 while trying to allocate 547356672 bytes (free: 155648000, total: 42406903808). +[W1008 01:29:41.393391052 CUDACachingAllocator.cpp:3934] memory allocation failed with OOM on device 0 while trying to allocate 547356672 bytes (free: 155648000, total: 42406903808). +[W1008 01:29:41.393842372 CUDACachingAllocator.cpp:3934] memory allocation failed with OOM on device 0 while trying to allocate 1637875712 bytes (free: 155648000, total: 42406903808). +[W1008 01:29:41.393983348 CUDACachingAllocator.cpp:3934] memory allocation failed with OOM on device 0 while trying to allocate 1637875712 bytes (free: 155648000, total: 42406903808). +[W1008 01:29:41.394462350 CUDACachingAllocator.cpp:3934] memory allocation failed with OOM on device 0 while trying to allocate 1092616192 bytes (free: 155648000, total: 42406903808). +[W1008 01:29:41.394579691 CUDACachingAllocator.cpp:3934] memory allocation failed with OOM on device 0 while trying to allocate 1092616192 bytes (free: 155648000, total: 42406903808). +[W1008 01:29:41.395064775 CUDACachingAllocator.cpp:3934] memory allocation failed with OOM on device 0 while trying to allocate 547356672 bytes (free: 155648000, total: 42406903808). +[W1008 01:29:41.395182817 CUDACachingAllocator.cpp:3934] memory allocation failed with OOM on device 0 while trying to allocate 547356672 bytes (free: 155648000, total: 42406903808). +[W1008 01:29:41.395627445 CUDACachingAllocator.cpp:3934] memory allocation failed with OOM on device 0 while trying to allocate 1092616192 bytes (free: 155648000, total: 42406903808). +[W1008 01:29:41.395744255 CUDACachingAllocator.cpp:3934] memory allocation failed with OOM on device 0 while trying to allocate 1092616192 bytes (free: 155648000, total: 42406903808). +[W1008 01:29:41.396193601 CUDACachingAllocator.cpp:3934] memory allocation failed with OOM on device 0 while trying to allocate 1092616192 bytes (free: 155648000, total: 42406903808). +[W1008 01:29:41.396313277 CUDACachingAllocator.cpp:3934] memory allocation failed with OOM on device 0 while trying to allocate 1092616192 bytes (free: 155648000, total: 42406903808). +[W1008 01:29:41.396753756 CUDACachingAllocator.cpp:3934] memory allocation failed with OOM on device 0 while trying to allocate 1092616192 bytes (free: 155648000, total: 42406903808). +[W1008 01:29:41.396870586 CUDACachingAllocator.cpp:3934] memory allocation failed with OOM on device 0 while trying to allocate 1092616192 bytes (free: 155648000, total: 42406903808). +[W1008 01:29:41.397320334 CUDACachingAllocator.cpp:3934] memory allocation failed with OOM on device 0 while trying to allocate 1092616192 bytes (free: 155648000, total: 42406903808). +[W1008 01:29:41.397441011 CUDACachingAllocator.cpp:3934] memory allocation failed with OOM on device 0 while trying to allocate 1092616192 bytes (free: 155648000, total: 42406903808). +[W1008 01:29:41.397886520 CUDACachingAllocator.cpp:3934] memory allocation failed with OOM on device 0 while trying to allocate 1637875712 bytes (free: 155648000, total: 42406903808). +[W1008 01:29:41.398040120 CUDACachingAllocator.cpp:3934] memory allocation failed with OOM on device 0 while trying to allocate 1637875712 bytes (free: 155648000, total: 42406903808). +[W1008 01:29:41.398483264 CUDACachingAllocator.cpp:3934] memory allocation failed with OOM on device 0 while trying to allocate 1092616192 bytes (free: 155648000, total: 42406903808). +[W1008 01:29:41.398599413 CUDACachingAllocator.cpp:3934] memory allocation failed with OOM on device 0 while trying to allocate 1092616192 bytes (free: 155648000, total: 42406903808). +[W1008 01:29:41.399050593 CUDACachingAllocator.cpp:3934] memory allocation failed with OOM on device 0 while trying to allocate 1092616192 bytes (free: 155648000, total: 42406903808). +[W1008 01:29:41.399171040 CUDACachingAllocator.cpp:3934] memory allocation failed with OOM on device 0 while trying to allocate 1092616192 bytes (free: 155648000, total: 42406903808). +[W1008 01:29:41.399609907 CUDACachingAllocator.cpp:3934] memory allocation failed with OOM on device 0 while trying to allocate 1092616192 bytes (free: 155648000, total: 42406903808). +[W1008 01:29:41.399730063 CUDACachingAllocator.cpp:3934] memory allocation failed with OOM on device 0 while trying to allocate 1092616192 bytes (free: 155648000, total: 42406903808). +[W1008 01:29:41.400215197 CUDACachingAllocator.cpp:3934] memory allocation failed with OOM on device 0 while trying to allocate 1092616192 bytes (free: 155648000, total: 42406903808). +[W1008 01:29:41.400335293 CUDACachingAllocator.cpp:3934] memory allocation failed with OOM on device 0 while trying to allocate 1092616192 bytes (free: 155648000, total: 42406903808). +[W1008 01:29:41.400774751 CUDACachingAllocator.cpp:3934] memory allocation failed with OOM on device 0 while trying to allocate 1092616192 bytes (free: 155648000, total: 42406903808). +[W1008 01:29:41.400894747 CUDACachingAllocator.cpp:3934] memory allocation failed with OOM on device 0 while trying to allocate 1092616192 bytes (free: 155648000, total: 42406903808). +[W1008 01:29:41.401368570 CUDACachingAllocator.cpp:3934] memory allocation failed with OOM on device 0 while trying to allocate 547356672 bytes (free: 155648000, total: 42406903808). +[W1008 01:29:41.401485911 CUDACachingAllocator.cpp:3934] memory allocation failed with OOM on device 0 while trying to allocate 547356672 bytes (free: 155648000, total: 42406903808). +[W1008 01:29:41.401926721 CUDACachingAllocator.cpp:3934] memory allocation failed with OOM on device 0 while trying to allocate 1092616192 bytes (free: 155648000, total: 42406903808). +[W1008 01:29:41.402052798 CUDACachingAllocator.cpp:3934] memory allocation failed with OOM on device 0 while trying to allocate 1092616192 bytes (free: 155648000, total: 42406903808). +[W1008 01:29:41.403244473 CUDACachingAllocator.cpp:3934] memory allocation failed with OOM on device 0 while trying to allocate 545259520 bytes (free: 155648000, total: 42406903808). +[W1008 01:29:41.403365712 CUDACachingAllocator.cpp:3934] memory allocation failed with OOM on device 0 while trying to allocate 545259520 bytes (free: 155648000, total: 42406903808). +Traceback (most recent call last): + File "/a/kt_spb_unet/train.py", line 131, in + main() + File "/a/kt_spb_unet/train.py", line 106, in main + out = net(img) + ^^^^^^^^ + File "/opt/conda/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1783, in _wrapped_call_impl + return self._call_impl(*args, **kwargs) + ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + File "/opt/conda/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1794, in _call_impl + return forward_call(*args, **kwargs) + ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + File "/a/kt_spb_unet/unet.py", line 93, in forward + x = dec(x) + ^^^^^^ + File "/opt/conda/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1783, in _wrapped_call_impl + return self._call_impl(*args, **kwargs) + ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + File "/opt/conda/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1794, in _call_impl + return forward_call(*args, **kwargs) + ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + File "/opt/conda/lib/python3.12/site-packages/torch/nn/modules/container.py", line 259, in forward + input = module(input) + ^^^^^^^^^^^^^ + File "/opt/conda/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1783, in _wrapped_call_impl + return self._call_impl(*args, **kwargs) + ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + File "/opt/conda/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1794, in _call_impl + return forward_call(*args, **kwargs) + ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + File "/opt/conda/lib/python3.12/site-packages/torch/nn/modules/batchnorm.py", line 210, in forward + return F.batch_norm( + ^^^^^^^^^^^^^ + File "/opt/conda/lib/python3.12/site-packages/torch/nn/functional.py", line 2904, in batch_norm + return torch.batch_norm( + ^^^^^^^^^^^^^^^^^ +torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 520.00 MiB. GPU 0 has a total capacity of 39.49 GiB of which 148.44 MiB is free. Process 32095 has 468.00 MiB memory in use. Process 1548731 has 1.25 GiB memory in use. Process 2451550 has 1.25 GiB memory in use. Process 3273110 has 10.12 GiB memory in use. Process 3291017 has 10.12 GiB memory in use. Process 65408 has 868.00 MiB memory in use. Process 72207 has 912.00 MiB memory in use. Process 2937854 has 14.37 GiB memory in use. Of the allocated memory 13.63 GiB is allocated by PyTorch, and 265.66 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation. See documentation for Memory Management (https://docs.pytorch.org/docs/stable/notes/cuda.html#optimizing-memory-usage-with-pytorch-cuda-alloc-conf) + +real 0m4.297s +user 0m3.067s +sys 0m1.130s +++ date -u +%H:%M:%S ++ echo '=== DONE (01:29:42) ===' +=== DONE (01:29:42) === ++ ls -lh checkpoints/ +total 15M +-rw-r--r--. 1 root root 3.0M Oct 7 22:55 smoke.pt +-rw-r--r--. 1 root root 12M Oct 8 00:02 unet_base32_best.pt ++ du -sh data/ data +3.0G data/ diff --git a/agent/code/qwen 3.8 27b/kt_spb_unet/render_all.py b/agent/code/qwen 3.8 27b/kt_spb_unet/render_all.py new file mode 100644 index 000000000..653019074 --- /dev/null +++ b/agent/code/qwen 3.8 27b/kt_spb_unet/render_all.py @@ -0,0 +1,85 @@ +#!/usr/bin/env python3 +"""Parallel DeepTrack frame renderer (process pool). + +Each worker process builds its own Fluorescence optics (DeepTrack is not +process-safe to share), then renders its assigned frames. Deterministic in +frame index, so results are identical to the sequential renderer. + +Usage:: + python render_all.py --frames 40 --workers 16 + python render_all.py --frames 2000 --workers 24 +""" +from __future__ import annotations + +import argparse +import glob +import os +import time +import warnings +import numpy as np + +warnings.filterwarnings("ignore") + +import config +import generator + + +def _worker_init(h, w): + """One optics set per process (created in its own module globals).""" + _worker_init.g_opt = generator._make_green_optics(h, w) + _worker_init.r_opt = generator._make_red_optics(h, w) + + +def _render_one(idx): + h, w = config.HEIGHT, config.WIDTH + f = generator.render_frame(int(idx), h, w, + (_worker_init.g_opt, _worker_init.r_opt)) + os.makedirs(config.DATA_DIR, exist_ok=True) + np.save(os.path.join(config.DATA_DIR, f"img_{idx:06d}.npy"), f["image"]) + np.save(os.path.join(config.DATA_DIR, f"msk_{idx:06d}.npy"), f["masks"]) + return idx + + +def _already_done(i): + p = os.path.join(config.DATA_DIR, f"img_{i:06d}.npy") + m = os.path.join(config.DATA_DIR, f"msk_{i:06d}.npy") + return os.path.exists(p) and os.path.exists(m) + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--frames", type=int, default=2000) + ap.add_argument("--start", type=int, default=0) + ap.add_argument("--workers", type=int, default=24) + args = ap.parse_args() + + indices = [i for i in range(args.start, args.start + args.frames) + if not _already_done(i)] + if not indices: + print("all frames already rendered; nothing to do") + return + + print(f"rendering {len(indices)} frames with {args.workers} workers " + f"-> {config.DATA_DIR}") + os.makedirs(config.DATA_DIR, exist_ok=True) + + from concurrent.futures import ProcessPoolExecutor, as_completed + t0 = time.time() + done = 0 + with ProcessPoolExecutor(max_workers=args.workers, initializer=_worker_init, + initargs=(config.HEIGHT, config.WIDTH)) as ex: + futs = [ex.submit(_render_one, i) for i in indices] + for fut in as_completed(futs): + fut.result() + done += 1 + if done % 50 == 0: + rate = done / max(time.time() - t0, 1e-6) + eta = (len(indices) - done) / max(rate, 1e-6) + print(f"[render] {done}/{len(indices)} {rate:.1f} f/s " + f"eta {eta/60:.1f} min", flush=True) + print(f"done {len(indices)} frames in {time.time()-t0:.1f}s " + f"({len(indices)/(time.time()-t0):.2f} f/s)") + + +if __name__ == "__main__": + main() diff --git a/agent/code/qwen 3.8 27b/kt_spb_unet/run_full.sh b/agent/code/qwen 3.8 27b/kt_spb_unet/run_full.sh new file mode 100644 index 000000000..4ef472202 --- /dev/null +++ b/agent/code/qwen 3.8 27b/kt_spb_unet/run_full.sh @@ -0,0 +1,23 @@ +#!/usr/bin/env bash +# Full pipeline: render synthetic frames, then train the U-Net. +# Logs go to logs/. Runs as a single durable background job. +set -x +cd /a/kt_spb_unet +mkdir -p logs checkpoints +LOG=logs/pipeline.log +exec > >(tee "$LOG") 2>&1 + +NFRAMES=${NFRAMES:-2000} +EPOCHS=${EPOCHS:-12} +BASE=${BASE:-64} +BATCH=${BATCH:-16} + +echo "=== RENDER ($(date -u +%H:%M:%S)) frames=$NFRAMES ===" +time python render_all.py --frames "$NFRAMES" --workers 4 + +echo "=== TRAINING ($(date -u +%H:%M:%S)) base=$BASE epochs=$EPOCHS batch=$BATCH ===" +time python train.py --frames "$NFRAMES" --epochs "$EPOCHS" --base "$BASE" --batch "$BATCH" + +echo "=== DONE ($(date -u +%H:%M:%S)) ===" +ls -lh checkpoints/ +du -sh data/ data diff --git a/agent/code/qwen 3.8 27b/kt_spb_unet/run_pipeline.sh b/agent/code/qwen 3.8 27b/kt_spb_unet/run_pipeline.sh new file mode 100644 index 000000000..db19bdb5f --- /dev/null +++ b/agent/code/qwen 3.8 27b/kt_spb_unet/run_pipeline.sh @@ -0,0 +1,21 @@ +#!/usr/bin/env bash +# End-to-end pipeline: render 2000 frames -> train U-Net -> save checkpoint. +# Runs as one long job; progress printed to the same file as its stdout. +set -e +cd /a/kt_spb_unet + +NFRAMES=${NFRAMES:-2000} +EPOCHS=${EPOCHS:-12} +BASE=${BASE:-64} +BATCH=${BATCH:-16} + +stage() { echo; echo "########## $* ##########"; } + +stage "STEP 1/2 Rendering ${NFRAMES} synthetic frames" +python render_all.py --frames "$NFRAMES" + +stage "STEP 2/2 Training UNet (base=${BASE}, epochs=${EPOCHS}, batch=${BATCH})" +python train.py --frames "$NFRAMES" --epochs "$EPOCHS" --base "$BASE" --batch "$BATCH" + +stage "Done" +ls -lh checkpoints/ diff --git a/agent/code/qwen 3.8 27b/kt_spb_unet/train.py b/agent/code/qwen 3.8 27b/kt_spb_unet/train.py new file mode 100644 index 000000000..12f214671 --- /dev/null +++ b/agent/code/qwen 3.8 27b/kt_spb_unet/train.py @@ -0,0 +1,131 @@ +#!/usr/bin/env python3 +"""Train the U-Net on the pre-rendered synthetic dataset. + +Usage:: + + python train.py --epochs 12 --batch 16 --base 64 --frames 2000 + +Saves the best (mean val Dice) checkpoint to ``checkpoints/unet_base{base}_best.pt``. +""" +from __future__ import annotations + +import argparse +import os +import time +import warnings +import numpy as np +import torch +from torch.utils.data import DataLoader + +warnings.filterwarnings("ignore") + +import config +from dataset import PreRenderedDataset +from unet import UNet, dice_bce_loss + +torch.manual_seed(config.SEED) +np.random.seed(config.SEED) + + +def soft_dice(pred: torch.Tensor, target: torch.Tensor, ch: int) -> float: + p = torch.sigmoid(pred[..., ch]) + t = target[..., ch] + eps = 1e-7 + return float((2 * (p * t).sum() + eps) / (p.sum() + t.sum() + eps)) + + +def make_dataloaders(n_total: int): + rng = np.random.default_rng(config.SEED + 1) + perm = rng.permutation(n_total) + n_val = max(int(n_total * config.VAL_FRAC), 10) + val_idx = perm[:n_val] + tr_idx = perm[n_val:] + tr = PreRenderedDataset(config.DATA_DIR, tr_idx) + va = PreRenderedDataset(config.DATA_DIR, val_idx) + tr_loader = DataLoader(tr, batch_size=config.BATCH, shuffle=True, + num_workers=0, pin_memory=False, drop_last=True) + va_loader = DataLoader(va, batch_size=config.BATCH, num_workers=0, + pin_memory=False, shuffle=False) + return tr_loader, va_loader, tr_idx, val_idx + + +@torch.no_grad() +def evaluate(net, loader, device) -> dict: + net.eval() + n_frames = 0 + d_kt = d_spb = 0.0 + t0 = time.time() + for img, msk in loader: + img = img.to(device, non_blocking=True) + msk = msk.to(device, non_blocking=True) + out = net(img) + d_kt += soft_dice(out, msk, 0) * img.size(0) + d_spb += soft_dice(out, msk, 1) * img.size(0) + n_frames += img.size(0) + return {"dice_kt": d_kt / max(n_frames, 1), + "dice_spb": d_spb / max(n_frames, 1), + "val_time": time.time() - t0} + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--epochs", type=int, default=config.EPOCHS_DEFAULT) + ap.add_argument("--batch", type=int, default=config.BATCH) + ap.add_argument("--base", type=int, default=64, help="U-Net base width") + ap.add_argument("--lr", type=float, default=3e-4) + ap.add_argument("--frames", type=int, default=2000) + ap.add_argument("--device", type=str, + default="cuda" if torch.cuda.is_available() else "cpu") + args = ap.parse_args() + + device = torch.device(args.device) + print(f"device={device} base={args.base} batch={args.batch}", flush=True) + + tr_loader, va_loader, tr_idx, val_idx = make_dataloaders(args.frames) + print(f"train={len(tr_idx)} val={len(val_idx)}", flush=True) + + net = UNet(in_channels=2, out_channels=2, base=args.base).to(device) + n_params = sum(p.numel() for p in net.parameters()) + print(f"UNet base={args.base}: {n_params/1e6:.2f}M params", flush=True) + + opt = torch.optim.AdamW(net.parameters(), lr=args.lr, weight_decay=1e-4) + sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=args.epochs) + + best = -1.0 + os.makedirs(config.CKPT_DIR, exist_ok=True) + best_path = os.path.join(config.CKPT_DIR, f"unet_base{args.base}_best.pt") + + for epoch in range(1, args.epochs + 1): + net.train() + t0 = time.time() + running = 0.0 + nb = 0 + for img, msk in tr_loader: + img = img.to(device, non_blocking=True) + msk = msk.to(device, non_blocking=True) + out = net(img) + loss = dice_bce_loss(out, msk, weights=(1.0, 1.3)) + opt.zero_grad() + loss.backward() + opt.step() + running += loss.item() + nb += 1 + sched.step() + val = evaluate(net, va_loader, device) + cur = (val["dice_kt"] + val["dice_spb"]) / 2 + status = "" + if cur > best: + best = cur + status = " *best*" + torch.save({"net": net.state_dict(), "base": args.base, + "in_ch": 2, "out_ch": 2, "depths": 4, "epoch": epoch, + "h": config.HEIGHT, "w": config.WIDTH}, best_path) + print(f"epoch {epoch:2d}/{args.epochs} loss={running/nb:.4f} " + f"dice KT={val['dice_kt']:.4f} SPB={val['dice_spb']:.4f} " + f"[{time.time()-t0:5.1f}s] val {val['val_time']:5.1f}s{status}", + flush=True) + print(f"done. best avg dice={best:.4f} saved {best_path}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/agent/code/qwen 3.8 27b/kt_spb_unet/unet.py b/agent/code/qwen 3.8 27b/kt_spb_unet/unet.py new file mode 100644 index 000000000..910327441 --- /dev/null +++ b/agent/code/qwen 3.8 27b/kt_spb_unet/unet.py @@ -0,0 +1,130 @@ +"""U-Net segmentation network for two-channel (KT, SPB) fluorescence. + +Input : (B, 2, H, W) -- channel0 = KT/green, channel1 = SPB/red. +Output: (B, 2, H, W) -- per-pixel logits (apply sigmoid for probabilities). + +Architecture (depths=N) +----------------------- + encoder0 -> skip0 + pool + encoder1 -> skip1 + pool + ... + encoder(N-2) -> skip(N-2) + pool + encoder(N-1) -> skip(N-1) # deepest encoder, channels ch_N + bottleneck(ch_N -> ch_N, same size) + --------------------------------- + up0(ch_N -> ch_{N-1}) + concat skip(N-1) -> dec0 -> ch_{N-1} + up1(ch_{N-1} -> ch_{N-2}) + concat skip(N-2) -> dec1 -> ch_{N-2} + ... + up_{N-1}(ch_2 -> ch_1) + concat skip0 -> dec_{N-1} -> ch_1 + head(ch_1 -> out) + +Canvas must be divisible by ``2**(depths-1)`` (e.g. 416 for depths=4). +""" +from __future__ import annotations + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +def _conv_block(in_ch: int, out_ch: int) -> nn.Module: + return nn.Sequential( + nn.Conv2d(in_ch, out_ch, kernel_size=3, padding=1), + nn.BatchNorm2d(out_ch), + nn.ReLU(inplace=True), + nn.Conv2d(out_ch, out_ch, kernel_size=3, padding=1), + nn.BatchNorm2d(out_ch), + nn.ReLU(inplace=True), + ) + + +class UNet(nn.Module): + """Standard U-Net. ``base`` = width of level 0; ``depths`` = number of + encoder levels (default 4).""" + + def __init__(self, in_channels: int = 2, out_channels: int = 2, + base: int = 64, depths: int = 4) -> None: + super().__init__() + assert depths >= 2 + self.depths = depths + # Channel count at each encoder level. + ch = [base * (2 ** i) for i in range(depths)] + + # Encoder path + self.encoders = nn.ModuleList([ + _conv_block(in_channels if i == 0 else ch[i - 1], ch[i]) + for i in range(depths) + ]) + self.pool = nn.MaxPool2d(2) + self.bottleneck = _conv_block(ch[-1], ch[-1]) + + # Decoder path: ``depths`` stages, each up-casts from ch[i] to ch[i-1] + # and concats with the level-(i-1) skip (ch[i-1] channels) -> 2*ch[i-1] + # -> ch[i-1]. + self.up = nn.ModuleList([ + nn.ConvTranspose2d(ch[i], ch[i - 1], kernel_size=2, stride=2) + for i in range(depths - 1, 0, -1) + ]) + self.decoders = nn.ModuleList([ + _conv_block(2 * ch[i - 1], ch[i - 1]) + for i in range(depths - 1, 0, -1) + ]) + self.head = nn.Conv2d(ch[0], out_channels, kernel_size=1) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # Encoder + skips: list[torch.Tensor] = [] + for i, enc in enumerate(self.encoders): + x = enc(x) + skips.append(x) + if i < self.depths - 1: + x = self.pool(x) + x = self.bottleneck(x) + + # Decoder: the deepest skip (level N-1) is unneeded -- the bottleneck + # already sits at that level; use the first N-1 skips, deepest first. + decoder_skips = reversed(skips[:-1]) + for up, dec, skip in zip(self.up, self.decoders, decoder_skips): + x = up(x) + x = torch.cat([x, skip], dim=1) + x = dec(x) + return self.head(x) + + +def dice_bce_loss(pred: torch.Tensor, target: torch.Tensor, + weights: tuple[float, float] = (1.0, 1.5)) -> torch.Tensor: + """Soft Dice + BCE, per-channel; the SPB channel is weighted higher + (rarer class: 2 vs 6 objects). + + pred: (B, 2, H, W) logits + target: (B, 2, H, W) {0,1} + """ + p = torch.sigmoid(pred) + # Per-class soft Dice (B,) per class. + inter = (p * target).sum(dim=(2, 3)) # (B, 2) + union = p.sum(dim=(2, 3)) + target.sum(dim=(2, 3)) # (B, 2) + dice_per_ch = 1.0 - (2 * inter + 1e-7) / (union + 1e-7) # (B, 2) + + bce = F.binary_cross_entropy_with_logits(pred, target, reduction="none") + per_class = [bce[:, c].mean() + dice_per_ch[:, c].mean() * weights[c] + for c in range(2)] + return torch.stack(per_class).mean() + + +if __name__ == "__main__": + for base, depths in [(32, 3), (32, 4), (64, 4)]: + f = 2 ** (depths - 1) + H = (416 + f - 1) // f * f + W = (640 + f - 1) // f * f + net = UNet(in_channels=2, out_channels=2, base=base, depths=depths) + n = sum(p.numel() for p in net.parameters()) + x = torch.randn(2, 2, H, W) + y = net(x) + t = torch.rand(2, 2, H, W) + l = dice_bce_loss(y, t) + print(f"base={base} depths={depths} canvas {H}x{W}: " + f"{n/1e6:.2f}M | out {tuple(y.shape)} | loss {l.item():.4f}") + print("UNet OK") diff --git a/agent/code/qwen 3.8 27b/simulated_channels.png b/agent/code/qwen 3.8 27b/simulated_channels.png new file mode 100644 index 000000000..6b88ef1e6 Binary files /dev/null and b/agent/code/qwen 3.8 27b/simulated_channels.png differ diff --git a/agent/code/qwen 3.8 27b/simulated_image.png b/agent/code/qwen 3.8 27b/simulated_image.png new file mode 100644 index 000000000..274e56aab Binary files /dev/null and b/agent/code/qwen 3.8 27b/simulated_image.png differ diff --git a/agent/code/qwen 3.8 27b/simulated_image_preview.png b/agent/code/qwen 3.8 27b/simulated_image_preview.png new file mode 100644 index 000000000..fbea6245c Binary files /dev/null and b/agent/code/qwen 3.8 27b/simulated_image_preview.png differ