|
|
| import os
|
| import sys
|
| import csv
|
| import json
|
| import math
|
| import time
|
| import glob
|
| import shutil
|
| import zipfile
|
| import subprocess
|
| from pathlib import Path
|
| import numpy as np
|
| import pandas as pd
|
| from scipy.ndimage import gaussian_filter, maximum_filter
|
| from scipy.optimize import linear_sum_assignment
|
| from scipy.spatial import cKDTree
|
|
|
|
|
| try:
|
| from zymatica_integration.cuneiform_normalization import CuneiformScaler
|
| except ImportError:
|
|
|
| class CuneiformScaler:
|
| def __init__(self, scale_factor=255.0):
|
| self.scale_factor = float(scale_factor)
|
| def normalize(self, coords):
|
| return coords / self.scale_factor
|
| def denormalize(self, coords_norm):
|
| return coords_norm * self.scale_factor
|
| def check_float16_safety(self, coords):
|
| max_val = np.max(np.abs(coords))
|
| return {"max_coordinate_value": float(max_val), "is_float16_safe": max_val**2 < 65504.0}
|
|
|
|
|
|
|
|
|
| COMPETITION = "biohub-cell-tracking-during-development"
|
| COMP_DIR_CANDIDATES = [
|
| Path(f"/kaggle/input/competitions/{COMPETITION}"),
|
| Path(f"/kaggle/input/{COMPETITION}"),
|
| Path("."),
|
| ]
|
| COMP_DIR = next((path for path in COMP_DIR_CANDIDATES if path.exists()), COMP_DIR_CANDIDATES[-1])
|
| TEST_DIR = COMP_DIR / "test" if (COMP_DIR / "test").exists() else COMP_DIR
|
|
|
| WORKING_DIR = Path(".")
|
| REPO_DIR = WORKING_DIR / "tracking_repo"
|
| SUBMISSION_PATH = WORKING_DIR / "submission.csv"
|
|
|
|
|
| VOXEL_SCALE_UM = np.array([1.625, 0.40625, 0.40625])
|
| scaler = CuneiformScaler(scale_factor=255.0)
|
|
|
|
|
| DET_THRESHOLD = 0.99
|
| UNET_BATCH_SIZE = 4
|
| USE_ILP = True
|
| ILP_EDGE_WEIGHT = -1.0
|
| ILP_APPEARANCE_WEIGHT = 0.1
|
| ILP_DISAPPEARANCE_WEIGHT = 0.1
|
| ILP_DIVISION_WEIGHT = 1.0
|
|
|
|
|
| OUTPUT_EDGE_MAX_UM = 14.5
|
| OUTPUT_ENFORCE_NEXT_FRAME = True
|
| OUTPUT_SINGLE_PARENT_REPAIR = True
|
| OUTPUT_SINGLE_CHILD_REPAIR = False
|
| OUTPUT_PRUNE_ISOLATED = True
|
| OUTPUT_MOTION_RELINK = True
|
|
|
| MOTION_RELINK_TIGHT_UM = 6.2
|
| MOTION_RELINK_RELAXED_UM = 10.4
|
| MOTION_RELINK_VELOCITY_WEIGHT = 0.52
|
| MOTION_RELINK_LEARNED_BONUS = 0.78
|
| MOTION_RELINK_MAX_FRAME_NODES = 2800
|
|
|
| OUTPUT_GAP_CLOSE = True
|
| GAP_CLOSE_MAX_GAP = 1
|
| GAP_CLOSE_UM = 6.2
|
| GAP_CLOSE_REUSE_EXISTING = True
|
| GAP_CLOSE_REUSE_UM = 3.4
|
| GAP_CLOSE_MAX_ADDED_FRAC = 0.052
|
| GAP_CLOSE_MAX_ADDED_ABS = 2200
|
|
|
| GAP_REFINE_SYNTHETIC = True
|
| GAP_REFINE_WIN_Z = 1
|
| GAP_REFINE_WIN_YX = 3
|
| GAP_REFINE_MAX_SHIFT_UM = 3.1
|
|
|
| OUTPUT_FILTER_SHORT_TRACKS = False
|
| OUTPUT_MIN_TRACK_LEN = 4
|
| OUTPUT_KEEP_DIVISION_COMPONENTS = True
|
|
|
| OUTPUT_LINEFIT_SMOOTH = True
|
| OUTPUT_LINEFIT_WEIGHT = 0.72
|
| OUTPUT_LINEFIT_WINDOW = 2
|
|
|
| OUTPUT_GAP2_RECOVERY = True
|
| GAP2_MAX_TOTAL_UM = 9.7
|
| GAP2_MAX_STEP_UM = 4.05
|
| GAP2_MAX_LINKS_FRAC = 0.0032
|
| GAP2_MAX_LINKS_ABS = 140
|
| GAP2_REQUIRE_CONTEXT = True
|
| GAP2_FRAME_FRAC_CAP = 0.0045
|
|
|
| OUTPUT_SAFE_DIVISIONS = True
|
| SAFE_DIV_MAX_UM = 4.8
|
| SAFE_DIV_SISTER_MAX_UM = 7.0
|
| SAFE_DIV_EXISTING_CHILD_MAX_UM = 7.6
|
| SAFE_DIV_FRAME_FRAC_CAP = 0.008
|
| SAFE_DIV_GLOBAL_FRAC_CAP = 0.0042
|
|
|
|
|
| XY_DS = 4
|
| MIN_PEAK_DIST = 2
|
| NMS_RADIUS_UM = 4.0
|
| REFINE_RZ, REFINE_RYX = 2, 5
|
| DOG_SIGMAS = (1.0, 1.8, 3.0)
|
| DOG_K = 1.6
|
| DOG_THR_PCT = 80.0
|
| GENEROUS_DOG_PCT = 55.0
|
|
|
|
|
|
|
|
|
| def _scale_distance_um(a: np.ndarray, b: np.ndarray) -> float:
|
| """
|
| Computes Euclidean distance in physical space (µm).
|
| Applies Cuneiform Normalization internally to prevent FP16 overflows.
|
| """
|
| norm_a = scaler.normalize(a * VOXEL_SCALE_UM)
|
| norm_b = scaler.normalize(b * VOXEL_SCALE_UM)
|
|
|
|
|
| diff = norm_a - norm_b
|
| norm_dist = np.linalg.norm(diff)
|
|
|
|
|
| return float(scaler.denormalize(norm_dist))
|
|
|
| def edge_distance_um(source: dict, target: dict) -> float:
|
| pos_s = np.array([float(source["z"]), float(source["y"]), float(source["x"])])
|
| pos_t = np.array([float(target["z"]), float(target["y"]), float(target["x"])])
|
| return _scale_distance_um(pos_s, pos_t)
|
|
|
| def point_distance_um(a: tuple, b: tuple) -> float:
|
| return _scale_distance_um(np.array(a), np.array(b))
|
|
|
|
|
|
|
|
|
| def _read_meta(zarr_path: Path) -> tuple[tuple[int, ...], np.dtype]:
|
| meta = json.loads((zarr_path / "0" / "zarr.json").read_text())
|
| return tuple(int(v) for v in meta["shape"]), np.dtype(meta["data_type"])
|
|
|
| def _read_volume_frame(zarr_path: Path, t: int, shape: tuple, dtype: np.dtype) -> np.ndarray:
|
| chunk_path = zarr_path / "0" / "c" / str(t) / "0" / "0" / "0"
|
| try:
|
| import blosc2
|
| raw = chunk_path.read_bytes()
|
| arr = np.frombuffer(blosc2.decompress(raw), dtype=dtype)
|
| if arr.size == int(np.prod(shape[1:])):
|
| return arr.reshape(shape[1:]).copy()
|
| except Exception:
|
| pass
|
| import zarr
|
| return np.asarray(zarr.open(zarr_path / "0", mode="r")[t])
|
|
|
|
|
|
|
|
|
| def _pool(vol, f):
|
| if f <= 1: return vol.astype(np.float32)
|
| Z, Y, X = vol.shape; Y2, X2 = (Y // f) * f, (X // f) * f
|
| return vol[:, :Y2, :X2].astype(np.float32).reshape(Z, Y2 // f, f, X2 // f, f).mean(axis=(2, 4))
|
|
|
| def _peaks(sm, thr, d):
|
| mx = maximum_filter(sm, size=2 * int(d) + 1, mode='nearest')
|
| return np.argwhere((sm >= mx) & (sm > thr)).astype(np.int32)
|
|
|
| def _refine(vol, zyx):
|
| Z, Y, X = vol.shape; z, y, x = (int(round(v)) for v in zyx)
|
| z0, z1 = max(0, z - REFINE_RZ), min(Z, z + REFINE_RZ + 1)
|
| y0, y1 = max(0, y - REFINE_RYX), min(Y, y + REFINE_RYX + 1)
|
| x0, x1 = max(0, x - REFINE_RYX), min(X, x + REFINE_RYX + 1)
|
| crop = vol[z0:z1, y0:y1, x0:x1].astype(np.float32); bg = float(crop.min())
|
| w = np.clip(crop - bg, 0, None); s = float(w.sum())
|
| if s <= 0: return np.array([z, y, x], float), 0.0
|
| zz, yy, xx = np.mgrid[z0:z1, y0:y1, x0:x1]
|
| return np.array([(zz * w).sum(), (yy * w).sum(), (xx * w).sum()]) / s, float(crop.max() - bg)
|
|
|
| def _nms(coords, scores, radius_um):
|
| if len(coords) <= 1: return coords, scores
|
| pts = coords * VOXEL_SCALE_UM[None, :]; order = np.argsort(-scores)
|
| tree = cKDTree(pts); killed = np.zeros(len(coords), bool); keep = []
|
| for i in order:
|
| if killed[i]: continue
|
| keep.append(int(i)); killed[tree.query_ball_point(pts[i], r=radius_um)] = True
|
| keep = np.array(keep); return coords[keep], scores[keep]
|
|
|
| def _scale_back(pk):
|
| full = pk.astype(float)
|
| full[:, 1] = full[:, 1] * XY_DS + (XY_DS - 1) / 2
|
| full[:, 2] = full[:, 2] * XY_DS + (XY_DS - 1) / 2
|
| return full
|
|
|
| def detect_cells_classical(vol, pct=DOG_THR_PCT):
|
| pooled = _pool(vol, XY_DS)
|
| coords, scores = [], []
|
| for sg in DOG_SIGMAS:
|
| dog = gaussian_filter(pooled, sg) - gaussian_filter(pooled, sg * DOG_K)
|
| posv = dog[dog > 0]
|
| if posv.size == 0: continue
|
| pk = _peaks(dog, float(np.percentile(posv, pct)), MIN_PEAK_DIST)
|
| if len(pk) == 0: continue
|
| resp = dog[pk[:, 0], pk[:, 1], pk[:, 2]].astype(float)
|
| resp = resp / max(resp.max(), 1e-6)
|
| for p, r in zip(_scale_back(pk), resp):
|
| c, _ = _refine(vol, p); coords.append(c); scores.append(float(r))
|
| if not coords: return np.zeros((0, 3)), np.zeros(0)
|
| return _nms(np.array(coords), np.array(scores), NMS_RADIUS_UM)
|
|
|
|
|
|
|
|
|
| def motion_relink_edges(nodes_by_id, stats, learned_edge_probs=None):
|
| if not OUTPUT_MOTION_RELINK or not nodes_by_id:
|
| return []
|
|
|
| learned_edge_probs = learned_edge_probs or {}
|
|
|
| ids_by_t = {}
|
| for node_id, node in nodes_by_id.items():
|
| ids_by_t.setdefault(int(node["t"]), []).append(node_id)
|
|
|
| position_um = {
|
| node_id: np.array([float(node["z"]), float(node["y"]), float(node["x"])]) * VOXEL_SCALE_UM
|
| for node_id, node in nodes_by_id.items()
|
| }
|
|
|
| history_positions_um = {}
|
| selected_edges = []
|
|
|
| def assign_pass(src_ids, tgt_ids, gate_um):
|
| if not src_ids or not tgt_ids: return []
|
| big = gate_um * 1000.0 + 1.0
|
| cost = np.full((len(src_ids), len(tgt_ids)), big, dtype=np.float64)
|
| raw_dist = np.full_like(cost, np.inf)
|
| motion_dist = np.full_like(cost, np.inf)
|
| for i, src_id in enumerate(src_ids):
|
| src_pos = position_um[src_id]
|
| history = history_positions_um.get(src_id, [])
|
|
|
|
|
| if len(history) >= 2:
|
| prev_pos_1 = history[-1]
|
| prev_pos_2 = history[-2]
|
| v1 = src_pos - prev_pos_1
|
| v2 = prev_pos_1 - prev_pos_2
|
| a = v1 - v2
|
| predicted = src_pos + v1 + 0.5 * a
|
| elif len(history) == 1:
|
| prev_pos = history[-1]
|
| v = src_pos - prev_pos
|
| predicted = src_pos + v
|
| else:
|
| predicted = src_pos
|
|
|
| for j, tgt_id in enumerate(tgt_ids):
|
| tgt_pos = position_um[tgt_id]
|
| raw = _scale_distance_um(src_pos / VOXEL_SCALE_UM, tgt_pos / VOXEL_SCALE_UM)
|
| if raw > gate_um: continue
|
| motion = _scale_distance_um(predicted / VOXEL_SCALE_UM, tgt_pos / VOXEL_SCALE_UM)
|
| prob = learned_edge_probs.get((src_id, tgt_id), 0.0)
|
| raw_dist[i, j] = raw
|
| motion_dist[i, j] = motion
|
| cost[i, j] = motion + 0.05 * raw - MOTION_RELINK_LEARNED_BONUS * prob
|
| ri, rc = linear_sum_assignment(cost)
|
| return [(src_ids[r], tgt_ids[c], raw_dist[r, c], motion_dist[r, c], learned_edge_probs.get((src_ids[r], tgt_ids[c]), 0.0))
|
| for r, c in zip(ri, rc) if cost[r, c] < big]
|
|
|
| for t in sorted(ids_by_t):
|
| src_ids = ids_by_t.get(t, [])
|
| tgt_ids = ids_by_t.get(t + 1, [])
|
| if not src_ids or not tgt_ids: continue
|
| unmatched_src = set(src_ids)
|
| unmatched_tgt = set(tgt_ids)
|
| frame_matches = []
|
| for pass_name, gate in (("tight", MOTION_RELINK_TIGHT_UM), ("relaxed", MOTION_RELINK_RELAXED_UM)):
|
| p_src = [n for n in src_ids if n in unmatched_src]
|
| p_tgt = [n for n in tgt_ids if n in unmatched_tgt]
|
| matches = assign_pass(p_src, p_tgt, gate)
|
| for s, tg, raw, motion, pr in matches:
|
| if s not in unmatched_src or tg not in unmatched_tgt: continue
|
| unmatched_src.remove(s)
|
| unmatched_tgt.remove(tg)
|
| frame_matches.append((s, tg, raw, motion, pass_name, pr))
|
|
|
| for s, tg, raw, motion, pass_name, pr in frame_matches:
|
| selected_edges.append({
|
| "source_id": s,
|
| "target_id": tg,
|
| "edge_prob": pr,
|
| "distance_um": raw,
|
| "motion_distance_um": motion,
|
| "motion_relinked": 1,
|
| })
|
|
|
| s_history = history_positions_um.get(s, [])
|
| history_positions_um[tg] = (s_history + [position_um[s]])[-5:]
|
| return selected_edges
|
|
|
| def refine_centroid_by_intensity(vol, init_zyx, window_radius=(1, 3, 3)):
|
| z_c, y_c, x_c = int(round(init_zyx[0])), int(round(init_zyx[1])), int(round(init_zyx[2]))
|
| H, W, D = vol.shape
|
|
|
| z_start = max(0, z_c - window_radius[0])
|
| z_end = min(H, z_c + window_radius[0] + 1)
|
| y_start = max(0, y_c - window_radius[1])
|
| y_end = min(W, y_c + window_radius[1] + 1)
|
| x_start = max(0, x_c - window_radius[2])
|
| x_end = min(D, x_c + window_radius[2] + 1)
|
|
|
| sub_vol = vol[z_start:z_end, y_start:y_end, x_start:x_end]
|
| if sub_vol.size == 0 or np.sum(sub_vol) == 0:
|
| return init_zyx
|
|
|
| zs, ys, xs = np.mgrid[z_start:z_end, y_start:y_end, x_start:x_end]
|
| sum_int = np.sum(sub_vol)
|
| expected_z = np.sum(zs * sub_vol) / sum_int
|
| expected_y = np.sum(ys * sub_vol) / sum_int
|
| expected_x = np.sum(xs * sub_vol) / sum_int
|
|
|
| refined = np.array([expected_z, expected_y, expected_x])
|
| shift_um = (refined - init_zyx) * VOXEL_SCALE_UM
|
| shift_dist = np.linalg.norm(shift_um)
|
| if shift_dist > 1.5:
|
| refined = init_zyx + (shift_um / shift_dist * 1.5) / VOXEL_SCALE_UM
|
| return refined
|
|
|
| def close_single_frame_gaps(nodes_by_id, edges, stats, dataset=None):
|
| if not OUTPUT_GAP_CLOSE or not edges: return nodes_by_id, edges
|
| outgoing = {int(e["source_id"]) for e in edges}
|
| incoming = {int(e["target_id"]) for e in edges}
|
| incident = outgoing | incoming
|
|
|
| ends = {}
|
| starts = {}
|
| for nid, node in nodes_by_id.items():
|
| t = int(node["t"])
|
| if nid not in outgoing: ends.setdefault(t, []).append(nid)
|
| if nid not in incoming: starts.setdefault(t, []).append(nid)
|
|
|
| next_id = max(nodes_by_id) + 1 if nodes_by_id else 1
|
| new_edges = []
|
| used_starts = set()
|
|
|
| for t, end_ids in sorted(ends.items()):
|
| start_ids = [sid for sid in starts.get(t + 2, []) if sid not in used_starts]
|
| if not end_ids or not start_ids: continue
|
|
|
| d = np.zeros((len(end_ids), len(start_ids)))
|
| for i, eid in enumerate(end_ids):
|
| for j, sid in enumerate(start_ids):
|
| d[i, j] = point_distance_um(
|
| (nodes_by_id[eid]["z"], nodes_by_id[eid]["y"], nodes_by_id[eid]["x"]),
|
| (nodes_by_id[sid]["z"], nodes_by_id[sid]["y"], nodes_by_id[sid]["x"])
|
| )
|
| threshold = GAP_CLOSE_UM * 2
|
| big = threshold * 1000.0 + 1.0
|
| cost = np.where(d <= threshold, d, big)
|
| ri, rc = linear_sum_assignment(cost)
|
|
|
| for r, c in zip(ri, rc):
|
| if d[r, c] > threshold: continue
|
| src_id = end_ids[r]
|
| tgt_id = start_ids[c]
|
| if src_id in outgoing or tgt_id in used_starts: continue
|
|
|
|
|
| mid_t = t + 1
|
| src_node = nodes_by_id[src_id]
|
| tgt_node = nodes_by_id[tgt_id]
|
| mid_pos = (
|
| (float(src_node["z"]) + float(tgt_node["z"])) / 2.0,
|
| (float(src_node["y"]) + float(tgt_node["y"])) / 2.0,
|
| (float(src_node["x"]) + float(tgt_node["x"])) / 2.0,
|
| )
|
|
|
|
|
| refined_pos = mid_pos
|
| if dataset:
|
| zarr_path = TEST_DIR / f"{dataset}.zarr"
|
| if zarr_path.exists():
|
| try:
|
| shape, dtype = _read_meta(zarr_path)
|
| vol = _read_volume_frame(zarr_path, mid_t, shape, dtype)
|
| refined_pos = refine_centroid_by_intensity(vol, np.array(mid_pos))
|
| except Exception as ex:
|
| print(f" [Warning] Centroid refinement failed for t={mid_t}: {ex}")
|
|
|
| nodes_by_id[next_id] = {
|
| "node_id": next_id,
|
| "t": mid_t,
|
| "z": refined_pos[0],
|
| "y": refined_pos[1],
|
| "x": refined_pos[2],
|
| }
|
| new_edges.append({
|
| "source_id": src_id,
|
| "target_id": next_id,
|
| "edge_prob": None,
|
| "distance_um": edge_distance_um(src_node, nodes_by_id[next_id])
|
| })
|
| new_edges.append({
|
| "source_id": next_id,
|
| "target_id": tgt_id,
|
| "edge_prob": None,
|
| "distance_um": edge_distance_um(nodes_by_id[next_id], tgt_node)
|
| })
|
| outgoing.add(src_id)
|
| incoming.add(next_id)
|
| outgoing.add(next_id)
|
| incoming.add(tgt_id)
|
| used_starts.add(tgt_id)
|
| next_id += 1
|
|
|
| return nodes_by_id, [*edges, *new_edges]
|
|
|
|
|
| try:
|
| from zymatica_integration.svd_dct_compression import TrajectoryCompressor
|
| except ImportError:
|
| class TrajectoryCompressor:
|
| def __init__(self, rank=2, k_coef=8):
|
| self.rank = rank
|
| self.k_coef = k_coef
|
| def compress(self, trajectory):
|
| T, D = trajectory.shape
|
| mean_vector = np.mean(trajectory, axis=0)
|
| centered = trajectory - mean_vector
|
| U, S, Vh = np.linalg.svd(centered, full_matrices=False)
|
| r = min(self.rank, D)
|
| U_scaled = U[:, :r] * np.sqrt(S[:r])
|
| V_scaled = Vh[:r, :].T * np.sqrt(S[:r])
|
| U_dct = np.zeros((self.k_coef, r))
|
| from scipy.fft import dct
|
| for col in range(r):
|
| c_dct = dct(U_scaled[:, col], norm='ortho')
|
| k_eff = min(self.k_coef, T)
|
| U_dct[:k_eff, col] = c_dct[:k_eff]
|
| return {"mean": mean_vector, "U_dct_coefs": U_dct, "V_scaled": V_scaled, "original_shape": (T, D)}
|
| def decompress(self, compressed_dict):
|
| mean_vector = compressed_dict["mean"]
|
| U_dct_coefs = compressed_dict["U_dct_coefs"]
|
| V_scaled = compressed_dict["V_scaled"]
|
| T, D = compressed_dict["original_shape"]
|
| r = U_dct_coefs.shape[1]
|
| from scipy.fft import idct
|
| U_recon = np.zeros((T, r))
|
| for col in range(r):
|
| full_dct = np.zeros(T)
|
| k_eff = min(self.k_coef, T)
|
| full_dct[:k_eff] = U_dct_coefs[:k_eff, col]
|
| U_recon[:, col] = idct(full_dct, norm='ortho')
|
| return np.dot(U_recon, V_scaled.T) + mean_vector
|
|
|
| def filter_trajectories_by_svd_dct(nodes_by_id, edges):
|
|
|
| parent = {nid: nid for nid in nodes_by_id}
|
| def find(n):
|
| while parent[n] != n:
|
| parent[n] = parent[parent[n]]
|
| n = parent[n]
|
| return n
|
| def union(a, b):
|
| ra, rb = find(a), find(b)
|
| if ra != rb: parent[ra] = rb
|
|
|
| for e in edges:
|
| union(int(e["source_id"]), int(e["target_id"]))
|
|
|
| tracks = {}
|
| for nid in nodes_by_id:
|
| tracks.setdefault(find(nid), []).append(nid)
|
|
|
| keep_nodes = set()
|
| pruned_count = 0
|
| smoothed_coords = {}
|
|
|
| compressor = TrajectoryCompressor(rank=2, k_coef=6)
|
|
|
| for root, members in tracks.items():
|
| if len(members) < 10:
|
| keep_nodes.update(members)
|
| continue
|
|
|
| members_sorted = sorted(members, key=lambda nid: int(nodes_by_id[nid]["t"]))
|
| coords = np.array([
|
| [float(nodes_by_id[nid]["z"]), float(nodes_by_id[nid]["y"]), float(nodes_by_id[nid]["x"])]
|
| for nid in members_sorted
|
| ])
|
| coords_um = coords * VOXEL_SCALE_UM
|
|
|
| try:
|
| compressed = compressor.compress(coords_um)
|
| recon_um = compressor.decompress(compressed)
|
| mse = np.mean((coords_um - recon_um) ** 2)
|
|
|
| if mse < 1.8:
|
| keep_nodes.update(members)
|
|
|
| recon_voxel = recon_um / VOXEL_SCALE_UM
|
| for idx, nid in enumerate(members_sorted):
|
| smoothed_coords[nid] = recon_voxel[idx]
|
| else:
|
| pruned_count += 1
|
| except Exception:
|
| keep_nodes.update(members)
|
|
|
| if pruned_count > 0:
|
| print(f" [SVD/DCT Trajectory Filter] Pruned {pruned_count} noisy trajectories.")
|
|
|
|
|
| for nid, coord in smoothed_coords.items():
|
| nodes_by_id[nid]["z"] = float(coord[0])
|
| nodes_by_id[nid]["y"] = float(coord[1])
|
| nodes_by_id[nid]["x"] = float(coord[2])
|
|
|
| filtered_nodes = {nid: n for nid, n in nodes_by_id.items() if nid in keep_nodes}
|
| filtered_edges = [
|
| e for e in edges
|
| if int(e["source_id"]) in filtered_nodes and int(e["target_id"]) in filtered_nodes
|
| ]
|
| return filtered_nodes, filtered_edges
|
|
|
| def filter_output_graph(nodes_by_id, raw_edges, dataset=None, avg_cells=600):
|
| stats = {}
|
| edges = []
|
| for edge in raw_edges:
|
| source = nodes_by_id.get(int(edge["source_id"]))
|
| target = nodes_by_id.get(int(edge["target_id"]))
|
| if source is None or target is None: continue
|
| if OUTPUT_ENFORCE_NEXT_FRAME and int(target["t"]) != int(source["t"]) + 1: continue
|
| dist = edge_distance_um(source, target)
|
| edge["distance_um"] = dist
|
| if dist > OUTPUT_EDGE_MAX_UM: continue
|
| edges.append(edge)
|
|
|
| if OUTPUT_MOTION_RELINK:
|
| edges = motion_relink_edges(nodes_by_id, stats)
|
|
|
|
|
| if OUTPUT_SINGLE_PARENT_REPAIR and edges:
|
| best_by_target = {}
|
| for edge in edges:
|
| tgt = int(edge["target_id"])
|
| if tgt not in best_by_target or float(edge.get("edge_prob") or 0) > float(best_by_target[tgt].get("edge_prob") or 0):
|
| best_by_target[tgt] = edge
|
| edges = list(best_by_target.values())
|
|
|
|
|
| if avg_cells < 500:
|
| min_track_len = 4
|
| elif avg_cells < 1200:
|
| min_track_len = 6
|
| else:
|
| min_track_len = 8
|
| print(f" [Dynamic Calibration] avg_cells={avg_cells:.1f} -> min_track_len={min_track_len}")
|
|
|
|
|
| if edges:
|
| by_source = {}
|
| for edge in edges:
|
| by_source.setdefault(int(edge["source_id"]), []).append(edge)
|
|
|
| filtered_edges = []
|
| for src_id, src_edges in by_source.items():
|
| if len(src_edges) <= 1:
|
| filtered_edges.extend(src_edges)
|
| continue
|
|
|
|
|
| source = nodes_by_id[src_id]
|
| s_pos = np.array([float(source["z"]), float(source["y"]), float(source["x"])]) * VOXEL_SCALE_UM
|
|
|
| ranked = sorted(src_edges, key=lambda e: (float(e.get("edge_prob") or 0.0), -float(e["distance_um"])), reverse=True)
|
| t1 = nodes_by_id[int(ranked[0]["target_id"])]
|
| t2 = nodes_by_id[int(ranked[1]["target_id"])]
|
|
|
| p1 = np.array([float(t1["z"]), float(t1["y"]), float(t1["x"])]) * VOXEL_SCALE_UM
|
| p2 = np.array([float(t2["z"]), float(t2["y"]), float(t2["x"])]) * VOXEL_SCALE_UM
|
|
|
| v1 = p1 - s_pos
|
| v2 = p2 - s_pos
|
|
|
| norm_v1 = np.linalg.norm(v1)
|
| norm_v2 = np.linalg.norm(v2)
|
|
|
| valid_div = True
|
| if norm_v1 > 1e-5 and norm_v2 > 1e-5:
|
| cos_theta = np.dot(v1, v2) / (norm_v1 * norm_v2)
|
|
|
| if cos_theta > 0.0:
|
| valid_div = False
|
|
|
| if valid_div:
|
| filtered_edges.extend(ranked[:2])
|
| else:
|
|
|
| filtered_edges.append(ranked[0])
|
| edges = filtered_edges
|
|
|
| nodes_by_id, edges = close_single_frame_gaps(nodes_by_id, edges, stats, dataset)
|
|
|
|
|
| if min_track_len > 1 and edges:
|
| parent = {nid: nid for nid in nodes_by_id}
|
| def find(n):
|
| while parent[n] != n:
|
| parent[n] = parent[parent[n]]
|
| n = parent[n]
|
| return n
|
| def union(a, b):
|
| ra, rb = find(a), find(b)
|
| if ra != rb: parent[ra] = rb
|
|
|
| for edge in edges:
|
| union(int(edge["source_id"]), int(edge["target_id"]))
|
|
|
| components = {}
|
| for nid in nodes_by_id:
|
| components.setdefault(find(nid), []).append(nid)
|
|
|
| keep = set()
|
| for members in components.values():
|
| if len(members) >= min_track_len:
|
| keep.update(members)
|
|
|
| if keep:
|
| nodes_by_id = {nid: n for nid, n in nodes_by_id.items() if nid in keep}
|
| edges = [e for e in edges if int(e["source_id"]) in nodes_by_id and int(e["target_id"]) in nodes_by_id]
|
|
|
|
|
| if edges and len(nodes_by_id) > 10:
|
| nodes_by_id, edges = filter_trajectories_by_svd_dct(nodes_by_id, edges)
|
|
|
|
|
| if OUTPUT_PRUNE_ISOLATED:
|
| incident = {int(e["source_id"]) for e in edges} | {int(e["target_id"]) for e in edges}
|
| nodes_by_id = {nid: n for nid, n in nodes_by_id.items() if nid in incident}
|
| edges = [e for e in edges if int(e["source_id"]) in nodes_by_id and int(e["target_id"]) in nodes_by_id]
|
|
|
| return nodes_by_id, edges, stats
|
|
|
|
|
|
|
|
|
| def materialize_folder_or_zip(src_dir, src_zip, dst_dir):
|
| if dst_dir.exists():
|
| if dst_dir.is_file() or dst_dir.is_symlink():
|
| dst_dir.unlink()
|
| else:
|
| shutil.rmtree(dst_dir)
|
|
|
| if src_dir.exists() and src_dir.is_dir():
|
| print(f"Copying directory {src_dir} to {dst_dir}...")
|
| shutil.copytree(src_dir, dst_dir)
|
| elif src_zip.exists() and src_zip.is_file():
|
| print(f"Extracting {src_zip} to {dst_dir}...")
|
| dst_dir.mkdir(parents=True, exist_ok=True)
|
| with zipfile.ZipFile(src_zip) as zf:
|
| zf.extractall(dst_dir)
|
| else:
|
| raise FileNotFoundError(f"Could not find folder {src_dir} or zip {src_zip}")
|
|
|
| def patch_polars_globally():
|
| try:
|
| import polars as pl
|
|
|
|
|
| class Float16:
|
| pass
|
|
|
|
|
| import sys
|
| for name in list(sys.modules):
|
| if name == "polars" or name.startswith("polars."):
|
| mod = sys.modules[name]
|
| if mod is not None:
|
| setattr(mod, "Float16", Float16)
|
|
|
|
|
| init_file = Path(pl.__file__)
|
| if init_file.exists():
|
| content = init_file.read_text()
|
| if "class Float16" not in content:
|
| patch = "\n\n# Patch for tracksdata compatibility\nclass Float16:\n pass\n"
|
| init_file.write_text(content + patch)
|
| print("Successfully monkeypatched polars package in site-packages and memory.")
|
| except Exception as e:
|
| print(f"Failed to patch polars globally: {e}")
|
|
|
| def setup_neural_environment():
|
| slug = "biohub-tracking-support-pack-50ep-v1"
|
| candidates = [
|
| Path(f"/kaggle/input/datasets/pilkwang/{slug}"),
|
| Path(f"/kaggle/input/{slug}"),
|
| Path(f"/kaggle/input/{slug}/{slug}"),
|
| Path(f"PublicNotebook/{slug}"),
|
| ]
|
| artifacts_dir = None
|
| for cand in candidates:
|
| if (cand / "repo.zip").exists() or (cand / "repo").exists():
|
| artifacts_dir = cand
|
| break
|
|
|
| if artifacts_dir is None:
|
| raise FileNotFoundError(f"Could not find support pack model artifacts in candidates.")
|
|
|
| print(f"Found support pack at: {artifacts_dir}")
|
|
|
| wheels_dir = artifacts_dir / "wheels"
|
| if wheels_dir.exists():
|
| print("Installing offline dependencies from wheels...")
|
| specs = [
|
| "tracksdata", "zarr>=3.0.10,<4", "pyscipopt", "geff>=1.1.3.1.1",
|
| "geff-spec<1.2", "ilpy>=0.5.1", "blosc2", "donfig",
|
| "numcodecs", "bidict", "psygnal", "rustworkx"
|
| ]
|
| cmd = [sys.executable, "-m", "pip", "install", "--no-index", "--no-deps", "--find-links", str(wheels_dir)] + specs
|
| subprocess.run(cmd, check=True)
|
| print("Dependency installation completed.")
|
|
|
|
|
| patch_polars_globally()
|
|
|
|
|
| materialize_folder_or_zip(artifacts_dir / "repo", artifacts_dir / "repo.zip", REPO_DIR)
|
| materialize_folder_or_zip(artifacts_dir / "weights", artifacts_dir / "weights.zip", REPO_DIR / "weights")
|
|
|
| print("Neural environment setup completed successfully.")
|
| return artifacts_dir
|
|
|
| def run_neural_predictions(test_stems):
|
| splits_path = REPO_DIR / "kaggle_test_splits_50ep.json"
|
| splits_path.write_text(json.dumps([{"split": 0, "train": [], "test": test_stems}], indent=2))
|
|
|
| predict_cmd = [
|
| sys.executable,
|
| "scripts/predict_unet_transformer.py",
|
| "--data-dir", str(TEST_DIR),
|
| "--splits", str(splits_path.name),
|
| "--split", "0",
|
| "--weights", f"weights/unet_transformer/split_0/edge_predictor_best.pth",
|
| "--unet-batch-size", str(UNET_BATCH_SIZE),
|
| "--det-threshold", str(DET_THRESHOLD),
|
| "--ilp-edge-weight", str(ILP_EDGE_WEIGHT),
|
| "--ilp-appearance-weight", str(ILP_APPEARANCE_WEIGHT),
|
| "--ilp-disappearance-weight", str(ILP_DISAPPEARANCE_WEIGHT),
|
| "--ilp-division-weight", str(ILP_DIVISION_WEIGHT),
|
| ]
|
| if USE_ILP:
|
| predict_cmd.append("--use-ilp")
|
|
|
| print(f"Running predictions: {' '.join(predict_cmd)}")
|
| subprocess.run(predict_cmd, cwd=REPO_DIR, env={**os.environ, "PYTHONPATH": "src"}, check=True)
|
|
|
| def load_graph_from_geff(geff_path):
|
| import tracksdata as td
|
| graph = td.graph.IndexedRXGraph.from_geff(geff_path)
|
| g = graph[0] if isinstance(graph, tuple) else graph
|
|
|
| nodes_by_id = {}
|
| for row in g.node_attrs().iter_rows(named=True):
|
| node_id = int(row["node_id"])
|
| nodes_by_id[node_id] = {
|
| "node_id": node_id,
|
| "t": int(row["t"]),
|
| "z": float(row["z"]),
|
| "y": float(row["y"]),
|
| "x": float(row["x"]),
|
| }
|
|
|
| raw_edges = []
|
| for row in g.edge_attrs().iter_rows(named=True):
|
| edge_prob = row.get("edge_prob") if hasattr(row, "get") else None
|
| raw_edges.append({
|
| "source_id": int(row["source_id"]),
|
| "target_id": int(row["target_id"]),
|
| "edge_prob": None if edge_prob is None else float(edge_prob),
|
| })
|
|
|
| return nodes_by_id, raw_edges
|
|
|
| def run_pipeline():
|
| print("Initializing Language U Microscopy Pipeline...")
|
| zarr_files = sorted(TEST_DIR.glob("*.zarr"))
|
| if not zarr_files:
|
| print(f"No .zarr files found in {TEST_DIR}. Creating a dry-run test trajectory instead.")
|
| return
|
|
|
| print(f"Found {len(zarr_files)} test volumes.")
|
| test_stems = [p.name[:-5] for p in zarr_files]
|
|
|
| all_nodes = []
|
| all_edges = []
|
| row_counter = 0
|
|
|
|
|
| use_neural = False
|
| try:
|
| setup_neural_environment()
|
| run_neural_predictions(test_stems)
|
| use_neural = True
|
| print("Neural prediction run completed successfully. Processing prediction graphs...")
|
| except Exception as e:
|
| print(f"\n[Hybrid Setup] Neural environment/inference failed: {e}")
|
| print("Falling back to refined classical DoG tracking pipeline...\n")
|
| use_neural = False
|
|
|
| if use_neural:
|
|
|
| geffs = sorted((REPO_DIR / "predictions").glob("*/unet_transformer/split_0/*.geff"))
|
| for geff_path in geffs:
|
| dataset = geff_path.name[:-5]
|
| print(f"Post-processing prediction graph for {dataset}...")
|
| nodes_by_id, raw_edges = load_graph_from_geff(geff_path)
|
|
|
| times_with_cells = set(n["t"] for n in nodes_by_id.values())
|
| avg_cells_per_frame = len(nodes_by_id) / max(len(times_with_cells), 1)
|
|
|
| nodes_by_id, edges, _ = filter_output_graph(nodes_by_id, raw_edges, dataset, avg_cells_per_frame)
|
|
|
|
|
| for nid, node in sorted(nodes_by_id.items()):
|
| all_nodes.append({
|
| "id": row_counter,
|
| "dataset": dataset,
|
| "row_type": "node",
|
| "node_id": int(node["node_id"]),
|
| "t": int(node["t"]),
|
| "z": int(round(float(node["z"]))),
|
| "y": int(round(float(node["y"]))),
|
| "x": int(round(float(node["x"]))),
|
| "source_id": -1,
|
| "target_id": -1
|
| })
|
| row_counter += 1
|
|
|
| for edge in edges:
|
| all_edges.append({
|
| "id": row_counter,
|
| "dataset": dataset,
|
| "row_type": "edge",
|
| "node_id": -1,
|
| "t": -1,
|
| "z": -1,
|
| "y": -1,
|
| "x": -1,
|
| "source_id": int(edge["source_id"]),
|
| "target_id": int(edge["target_id"])
|
| })
|
| row_counter += 1
|
| else:
|
|
|
| for zarr_path in zarr_files:
|
| dataset = zarr_path.name[:-5]
|
| print(f"Processing {dataset}...")
|
| shape, dtype = _read_meta(zarr_path)
|
| T = shape[0]
|
|
|
| nodes_by_id = {}
|
| raw_edges = []
|
| node_idx = 1
|
|
|
| prev_ids = []
|
| prev_coords = np.zeros((0, 3))
|
|
|
| for t in range(T):
|
| vol = _read_volume_frame(zarr_path, t, shape, dtype)
|
| coords, scores = detect_cells_classical(vol)
|
| ids = list(range(node_idx, node_idx + len(coords)))
|
| node_idx += len(coords)
|
|
|
| for i, c in zip(ids, coords):
|
| nodes_by_id[i] = {"node_id": i, "t": t, "z": c[0], "y": c[1], "x": c[2]}
|
|
|
| if t > 0 and prev_ids:
|
| if len(prev_coords) > 0 and len(coords) > 0:
|
| d = np.zeros((len(prev_coords), len(coords)))
|
| for r, pc in enumerate(prev_coords):
|
| for c_col, cc in enumerate(coords):
|
| d[r, c_col] = _scale_distance_um(pc, cc)
|
| cost = np.where(d <= GAP_CLOSE_UM, d, 1e9)
|
| ri, rc = linear_sum_assignment(cost)
|
| for r, col in zip(ri, rc):
|
| if cost[r, col] < 1e9:
|
| raw_edges.append({
|
| "source_id": prev_ids[r],
|
| "target_id": ids[col],
|
| "edge_prob": 1.0 - (cost[r, col] / GAP_CLOSE_UM)
|
| })
|
| prev_ids = ids
|
| prev_coords = coords
|
|
|
| times_with_cells = set(n["t"] for n in nodes_by_id.values())
|
| avg_cells_per_frame = len(nodes_by_id) / max(len(times_with_cells), 1)
|
|
|
| nodes_by_id, edges, _ = filter_output_graph(nodes_by_id, raw_edges, dataset, avg_cells_per_frame)
|
|
|
| for nid, node in sorted(nodes_by_id.items()):
|
| all_nodes.append({
|
| "id": row_counter,
|
| "dataset": dataset,
|
| "row_type": "node",
|
| "node_id": int(node["node_id"]),
|
| "t": int(node["t"]),
|
| "z": int(round(float(node["z"]))),
|
| "y": int(round(float(node["y"]))),
|
| "x": int(round(float(node["x"]))),
|
| "source_id": -1,
|
| "target_id": -1
|
| })
|
| row_counter += 1
|
|
|
| for edge in edges:
|
| all_edges.append({
|
| "id": row_counter,
|
| "dataset": dataset,
|
| "row_type": "edge",
|
| "node_id": -1,
|
| "t": -1,
|
| "z": -1,
|
| "y": -1,
|
| "x": -1,
|
| "source_id": int(edge["source_id"]),
|
| "target_id": int(edge["target_id"])
|
| })
|
| row_counter += 1
|
|
|
| pd.DataFrame(all_nodes + all_edges).to_csv(SUBMISSION_PATH, index=False)
|
| print(f"Submission saved to {SUBMISSION_PATH} with {len(all_nodes)} nodes and {len(all_edges)} edges.")
|
|
|
| if __name__ == "__main__":
|
| run_pipeline()
|
|
|
