# %% [code] 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 # Import Zymatica Normalization try: from zymatica_integration.cuneiform_normalization import CuneiformScaler except ImportError: # Inline fallback if folder is missing 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} # ===================================================================== # CONFIGURATION & PARAMETERS (Score Push Preset) # ===================================================================== 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" # Global scales & tracking limits VOXEL_SCALE_UM = np.array([1.625, 0.40625, 0.40625]) # Z, Y, X µm/voxel scaler = CuneiformScaler(scale_factor=255.0) # Load preset values from "score_push" 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 # Graph filters 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 # Classical DoG parameters (Fallback mode) 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 # ===================================================================== # SPATIAL & NUMERICAL STABILITY MODULES (Cuneiform Normalization) # ===================================================================== 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) # Perform math in normalized range diff = norm_a - norm_b norm_dist = np.linalg.norm(diff) # Scale back to physical space 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)) # ===================================================================== # GEOM & IMAGE HELPERS # ===================================================================== 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]) # ===================================================================== # CLASSICAL FALLBACK: DETECTION & LINKING # ===================================================================== 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) # ===================================================================== # POST-PROCESSING GRAPH FILTERS (Hungarian relinking, gap-close, repairs) # ===================================================================== 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 = {} # node_id -> list of past position_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, []) # Acceleration-Aware Motion Priors (State Estimation) 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, }) # Update history list (capped at 5 past coordinates) 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 # Insert midpoint node 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, ) # Refine by intensity if raw volume is available 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] # Import TrajectoryCompressor with inline fallback 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): # 1. Build Union-Find to group nodes into trajectories 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) # Threshold chosen via cross-validation proxy metrics if mse < 1.8: keep_nodes.update(members) # Convert back to voxel scale 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.") # Update coordinates of nodes with smoothed coordinates to denoise jitter 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) # Single parent repair 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()) # Dynamic track length calibration based on average cells 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}") # Prune division geometry check (wide-angle check) 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 # It's a division candidate! Check biological symmetry 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 they move in similar direction (angle < 90 degrees), drop division if cos_theta > 0.0: valid_div = False if valid_div: filtered_edges.extend(ranked[:2]) else: # Keep only the closer daughter filtered_edges.append(ranked[0]) edges = filtered_edges nodes_by_id, edges = close_single_frame_gaps(nodes_by_id, edges, stats, dataset) # Filter short track components using dynamic min_track_len 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] # Trajectory SVD/DCT High-Frequency Noise Filtering if edges and len(nodes_by_id) > 10: nodes_by_id, edges = filter_trajectories_by_svd_dct(nodes_by_id, edges) # Prune isolated nodes 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 # ===================================================================== # MAIN RUN LOOP # ===================================================================== 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 # Define a mock Float16 class class Float16: pass # 1. Patch the in-memory sys.modules cache 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) # 2. Patch the on-disk file for subprocesses 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 for tracksdata compatibility patch_polars_globally() # Copy/extract repository and weights directory 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 # Try running the neural model pipeline 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: # Load predictions from geffs 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) # Accumulate output format 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: # Classical fallback pipeline 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()