VisionLanguageGroup/MicroscopyMatching
0
1import gradio as gr2from gradio_bbox_annotator import BBoxAnnotator3from PIL import Image4import numpy as np5import torch6import os7import shutil8import time9import json10import uuid11from pathlib import Path12import tempfile13import zipfile14import html15from skimage import measure16from matplotlib import cm17from glob import glob18from natsort import natsorted19from huggingface_hub import HfApi, upload_file20import spaces21 22from inference_seg import load_model as load_seg_model, run as run_seg23from inference_count import load_model as load_count_model, run as run_count24from inference_track import load_model as load_track_model, run as run_track25from _utils.image_io import (26 standardize_image, inspect_image, image_size, array_nbytes, pixel_stats, bit_depth,27 RECOMMENDED_SIZE, WARN_SIZE, MAX_SIZE, MIN_SIZE, MAX_READ_BYTES, TIFF_EXTENSIONS,28)29 30HF_TOKEN = os.getenv("HF_TOKEN")31DATASET_REPO = "VisionLanguageGroup/feedback"32 33 34print("===== clearing cache =====")35cache_path = os.path.expanduser("~/.cache/huggingface/gradio")36if os.path.exists(cache_path):37 try:38 shutil.rmtree(cache_path)39 print("✅ Deleted ~/.cache/huggingface/gradio")40 except:41 pass42 43SEG_MODEL = None44SEG_DEVICE = torch.device("cpu")45 46COUNT_MODEL = None47COUNT_DEVICE = torch.device("cpu")48 49TRACK_MODEL = None50TRACK_DEVICE = torch.device("cpu")51 52def load_all_models():53 global SEG_MODEL, SEG_DEVICE54 global COUNT_MODEL, COUNT_DEVICE55 global TRACK_MODEL, TRACK_DEVICE56 57 print("\n" + "="*60)58 print("📦 Loading Segmentation Model")59 print("="*60)60 SEG_MODEL, SEG_DEVICE = load_seg_model(use_box=False)61 62 print("\n" + "="*60)63 print("📦 Loading Counting Model")64 print("="*60)65 COUNT_MODEL, COUNT_DEVICE = load_count_model(use_box=False)66 67 print("\n" + "="*60)68 print("📦 Loading Tracking Model")69 print("="*60)70 TRACK_MODEL, TRACK_DEVICE = load_track_model(use_box=False)71 72 print("\n" + "="*60)73 print("✅ All Models Loaded Successfully")74 print("="*60)75 76load_all_models()77 78DATASET_DIR = Path("solver_cache")79DATASET_DIR.mkdir(parents=True, exist_ok=True)80 81def save_feedback_to_hf(query_id, feedback_type, feedback_text=None, img_path=None, bboxes=None):82 """Save feedback to Hugging Face Dataset"""83 84 if not HF_TOKEN:85 print("⚠️ No HF_TOKEN found, using local storage")86 save_feedback(query_id, feedback_type, feedback_text, img_path, bboxes)87 return88 89 # Shared by the record and its image so the two pair up; query_id alone is not unique (a session can submit more than once).90 stem = f"feedback_{query_id}_{int(time.time())}"91 92 try:93 api = HfApi()94 95 image_in_repo = None96 if img_path and os.path.exists(img_path):97 try:98 image_in_repo = f"images/{stem}.png"99 api.upload_file(100 path_or_fileobj=img_path,101 path_in_repo=image_in_repo,102 repo_id=DATASET_REPO,103 repo_type="dataset",104 token=HF_TOKEN105 )106 except Exception as e:107 print(f"⚠️ Failed to upload image: {e}")108 image_in_repo = None 109 110 feedback_data = {111 "query_id": query_id,112 "feedback_type": feedback_type,113 "feedback_text": feedback_text,114 "image_path": image_in_repo, 115 "bboxes": str(bboxes), 116 "datetime": time.strftime("%Y-%m-%d %H:%M:%S"),117 "timestamp": time.time()118 }119 120 filename = f"{stem}.json"121 122 with open(filename, 'w', encoding='utf-8') as f:123 json.dump(feedback_data, f, indent=2, ensure_ascii=False)124 125 api.upload_file(126 path_or_fileobj=filename,127 path_in_repo=f"data/{filename}",128 repo_id=DATASET_REPO,129 repo_type="dataset",130 token=HF_TOKEN131 )132 133 os.remove(filename)134 135 print(f"✅ Feedback saved to HF Dataset: {DATASET_REPO} ({stem})")136 137 except Exception as e:138 print(f"⚠️ Failed to save to HF Dataset: {e}")139 save_feedback(query_id, feedback_type, feedback_text, img_path, bboxes)140 141 142def save_feedback(query_id, feedback_type, feedback_text=None, img_path=None, bboxes=None):143 """Save feedback to local JSON file"""144 feedback_data = {145 "query_id": query_id,146 "feedback_type": feedback_type,147 "feedback_text": feedback_text,148 "image": img_path,149 "bboxes": bboxes,150 "datetime": time.strftime("%Y%m%d_%H%M%S")151 }152 feedback_file = DATASET_DIR / query_id / "feedback.json"153 feedback_file.parent.mkdir(parents=True, exist_ok=True)154 155 if feedback_file.exists():156 with feedback_file.open("r") as f:157 existing = json.load(f)158 if not isinstance(existing, list):159 existing = [existing]160 existing.append(feedback_data)161 feedback_data = existing162 else:163 feedback_data = [feedback_data]164 165 with feedback_file.open("w") as f:166 json.dump(feedback_data, f, indent=4, ensure_ascii=False)167 168def parse_first_bbox(bboxes):169 """Parse the first bounding box from the annotation input, supports dict or list format"""170 if not bboxes:171 return None172 b = bboxes[0]173 if isinstance(b, dict):174 x, y = float(b.get("x", 0)), float(b.get("y", 0))175 w, h = float(b.get("width", 0)), float(b.get("height", 0))176 return x, y, x + w, y + h177 if isinstance(b, (list, tuple)) and len(b) >= 4:178 return float(b[0]), float(b[1]), float(b[2]), float(b[3])179 return None180 181def parse_bboxes(bboxes):182 """Parse all bounding boxes from the annotation input"""183 if not bboxes:184 return None185 186 result = []187 for b in bboxes:188 if isinstance(b, dict):189 x, y = float(b.get("x", 0)), float(b.get("y", 0))190 w, h = float(b.get("width", 0)), float(b.get("height", 0))191 result.append([x, y, x + w, y + h])192 elif isinstance(b, (list, tuple)) and len(b) >= 4:193 result.append([float(b[0]), float(b[1]), float(b[2]), float(b[3])])194 195 return result196 197def colorize_mask(mask: np.ndarray, num_colors: int = 512) -> np.ndarray:198 """Convert a 2D mask of instance IDs to a color image for visualization."""199 def hsv_to_rgb(h, s, v):200 i = int(h * 6.0)201 f = h * 6.0 - i202 i = i % 6203 p = v * (1 - s)204 q = v * (1 - f * s)205 t = v * (1 - (1 - f) * s)206 if i == 0: r, g, b = v, t, p207 elif i == 1: r, g, b = q, v, p208 elif i == 2: r, g, b = p, v, t209 elif i == 3: r, g, b = p, q, v210 elif i == 4: r, g, b = t, p, v211 else: r, g, b = v, p, q212 return int(r * 255), int(g * 255), int(b * 255)213 214 palette = [(0, 0, 0)]215 for i in range(1, num_colors):216 h = (i % num_colors) / float(num_colors)217 palette.append(hsv_to_rgb(h, 1.0, 0.95))218 219 palette_arr = np.array(palette, dtype=np.uint8)220 color_idx = mask % num_colors221 return palette_arr[color_idx]222 223 224def render_seg_overlay(img_np, inst_mask, overlay_alpha):225 """Render segmentation overlay from cached image/mask."""226 if img_np is None or inst_mask is None:227 return None228 229 overlay = img_np.copy()230 alpha = float(np.clip(overlay_alpha, 0.0, 1.0))231 232 for inst_id in np.unique(inst_mask):233 if inst_id == 0:234 continue235 binary_mask = (inst_mask == inst_id).astype(np.uint8)236 color = get_well_spaced_color(inst_id)237 overlay[binary_mask == 1] = (1 - alpha) * overlay[binary_mask == 1] + alpha * color238 239 contours = measure.find_contours(binary_mask, 0.5)240 for contour in contours:241 contour = contour.astype(np.int32)242 valid_y = np.clip(contour[:, 0], 0, overlay.shape[0] - 1)243 valid_x = np.clip(contour[:, 1], 0, overlay.shape[1] - 1)244 overlay[valid_y, valid_x] = [1.0, 1.0, 0.0]245 246 overlay = np.clip(overlay * 255.0, 0, 255).astype(np.uint8)247 return Image.fromarray(overlay)248 249 250def render_count_overlay(img_np, density_normalized, overlay_alpha):251 """Render counting heatmap overlay from cached image/density."""252 if img_np is None or density_normalized is None:253 return None254 255 alpha = float(np.clip(overlay_alpha, 0.0, 1.0))256 cmap = cm.get_cmap("jet")257 density_colored = cmap(density_normalized)[:, :, :3]258 259 overlay = img_np.copy()260 threshold = 0.01261 significant_mask = density_normalized > threshold262 overlay[significant_mask] = (1 - alpha) * overlay[significant_mask] + alpha * density_colored[significant_mask]263 overlay = np.clip(overlay * 255.0, 0, 255).astype(np.uint8)264 return Image.fromarray(overlay)265 266 267def update_seg_overlay_alpha(overlay_alpha, seg_vis_cache):268 """Live update segmentation visualization without rerunning inference."""269 if not seg_vis_cache:270 return None271 return render_seg_overlay(seg_vis_cache.get("img_np"), seg_vis_cache.get("inst_mask"), overlay_alpha)272 273 274def update_count_overlay_alpha(overlay_alpha, count_vis_cache):275 """Live update counting visualization without rerunning inference."""276 if not count_vis_cache:277 return None278 return render_count_overlay(count_vis_cache.get("img_np"), count_vis_cache.get("density_normalized"), overlay_alpha)279 280 281def update_tracking_overlay_alpha(overlay_alpha, track_vis_cache):282 """Regenerate tracking visualization at new opacity using cached outputs."""283 if not track_vis_cache:284 return None285 286 tif_dir = track_vis_cache.get("tif_dir")287 output_dir = track_vis_cache.get("output_dir")288 valid_tif_files = track_vis_cache.get("valid_tif_files")289 if not tif_dir or not output_dir or not valid_tif_files:290 return None291 292 try:293 return create_tracking_visualization(294 tif_dir=tif_dir,295 output_dir=output_dir,296 valid_tif_files=valid_tif_files,297 overlay_alpha=overlay_alpha298 )299 except Exception as e:300 print(f"⚠️ Failed to update tracking opacity: {e}")301 return None302 303 304def cleanup_tracking_cache(track_vis_cache):305 """Delete cached tracking temp directories from the previous run."""306 if not track_vis_cache:307 return308 for key in ["input_temp_dir", "output_dir"]:309 path = track_vis_cache.get(key)310 if path and os.path.isdir(path):311 try:312 shutil.rmtree(path)313 except Exception:314 pass315 316 317# Per session. Refused rather than evicting the oldest, since there is no way to318# remove a single entry.319MAX_GALLERY_UPLOADS = 10320 321_GALLERY_RAW_SOURCE = {}322_RAW_SOURCE_CAP = 500323 324 325def _remember_gallery_source(thumb_path, raw_path):326 """Map a gallery thumbnail back to the original file it was made from.327 """328 _GALLERY_RAW_SOURCE[thumb_path] = raw_path329 while len(_GALLERY_RAW_SOURCE) > _RAW_SOURCE_CAP: # dicts keep insertion order330 _GALLERY_RAW_SOURCE.pop(next(iter(_GALLERY_RAW_SOURCE)))331 332 333def _annot_path(annot_value):334 """Extract the image path from a BBoxAnnotator value (path or (path, boxes))."""335 if not annot_value:336 return None337 if isinstance(annot_value, (list, tuple)):338 return annot_value[0] if len(annot_value) > 0 else None339 return annot_value340 341 342def _human_bytes(n):343 """Format a byte count as MB or GB, whichever reads better."""344 gb = n / 1024 ** 3345 return f"{gb:.1f} GB" if gb >= 1 else f"{n / 1024 ** 2:.0f} MB"346 347 348def check_image(img_path):349 """Validate an uploaded file, cheapest checks first.350 351 Returns a rejection message if the file cannot be used, otherwise None.352 Reject: empty file; unreadable / unsupported format; dimensions over353 MAX_SIZE; full array over MAX_READ_BYTES (a small time-lapse can still be354 huge); corrupt pixels; non-finite (NaN/inf) pixels. Advisory ``gr.Warning``355 (does not reject): larger than WARN_SIZE; smaller than MIN_SIZE; blank356 (uniform) image; very narrow dynamic range.357 358 Only the header is touched until the size guards pass; pixel checks decode359 image data afterwards, when doing so is bounded by the memory guard.360 """361 if os.path.exists(img_path) and os.path.getsize(img_path) == 0:362 return "This file is empty (0 bytes). Please upload a valid image file."363 364 w, h = image_size(img_path)365 if w <= 0 or h <= 0:366 # Neither reader could parse the header, so say so rather than failing367 # silently further down.368 ext = os.path.splitext(img_path)[1].lower() or "(no extension)"369 return (f"Could not read {ext} as an image. Supported formats: "370 f"TIFF / OME-TIFF (8/16/32-bit, stacks, multi-channel), PNG, JPG, and other formats supported by tifffile. "371 f"Please export to TIFF/PNG/JPG (e.g. from ImageJ/Fiji) first.")372 373 if w > MAX_SIZE or h > MAX_SIZE:374 return (f"Image is {w}×{h}, larger than the {MAX_SIZE}×{MAX_SIZE} limit. "375 f"Please crop or downsample it before uploading.")376 377 nbytes = array_nbytes(img_path)378 if nbytes > MAX_READ_BYTES:379 info = inspect_image(img_path)380 detail = f"{w}×{h}"381 if info.frames > 1:382 detail += f" × {info.frames} frames"383 if info.channels > 1:384 detail += f" × {info.channels} channels"385 return (f"This file needs {_human_bytes(nbytes)} of memory to open "386 f"({detail}), over the {_human_bytes(MAX_READ_BYTES)} limit. "387 f"Please crop it, or split out the frames you need.")388 389 if w > WARN_SIZE or h > WARN_SIZE:390 gr.Warning(f"Image is {w}×{h} and will be resized to "391 f"{RECOMMENDED_SIZE}×{RECOMMENDED_SIZE}, so fine detail could be lost. "392 f"You can try cropping the image for better results.",393 duration=None, title="⚠️ Large image")394 395 if 0 < min(w, h) < MIN_SIZE:396 gr.Warning(f"Image is only {w}×{h} and will be upscaled to "397 f"{RECOMMENDED_SIZE}×{RECOMMENDED_SIZE}, so results may be unreliable. "398 f"A larger image may work better.",399 duration=None, title="⚠️ Very small image")400 401 # Size guards passed, so pixel decoding is bounded. Check problems that402 # the header cannot reveal.403 rep = pixel_stats(img_path)404 if not rep.decoded:405 return ("This image appears to be corrupted or truncated and could not be read. "406 "Please re-export it (e.g. from ImageJ/Fiji) and try again.")407 if not rep.finite:408 return ("This image contains invalid pixel values (NaN or infinity). "409 "Please clean or re-export it before uploading.")410 if rep.vmin == rep.vmax:411 if rep.vmax == 0:412 kind = "completely black"413 elif rep.dtype_max and rep.vmax >= rep.dtype_max:414 kind = "completely white"415 else:416 kind = "a single uniform value"417 gr.Warning(f"This image is {kind}, it has no visible content, so results will not be meaningful.",418 duration=None, title="⚠️ Blank image")419 elif rep.low_range:420 gr.Warning("This image may use a very narrow intensity range (low contrast). Signal may be too weak for reliable results; consider adjusting acquisition or contrast before uploading.",421 duration=None, title="⚠️ Low dynamic range")422 return None423 424 425def _describe_read(info):426 """Plain-language summary of how a file's dimensions were interpreted.427 428 Deliberately avoids array jargon ("axis", "size-4"): the reader thinks in429 channels / z-slices / timepoints, not numpy axes.430 """431 # lines = [f"Detected shape: {info.shape}"]432 if info.frames >1:433 lines = [f"Detected {info.channels}-channel {info.frames}-frame image stack of size {info.width}×{info.height}."]434 else:435 lines = [f"Detected {info.channels}-channel image of size {info.width}×{info.height}."]436 if info.channels >= 3:437 lines.append(f"Loaded as 3-channel RGB image.")438 elif info.channels == 2:439 lines.append(f"Loaded as 2-channel image (red, green).")440 else:441 lines.append(f"Loaded as single-channel grayscale image.")442 443 if info.frames > 1:444 lines.append(f"Showing first frame by default (you can use the stack frame slider to pick another).")445 if info.sub_frames > 1:446 lines.append(f"Showing the first of {info.sub_frames} {info.sub_label}s by default (use the {info.sub_label} slider to pick another).")447 448 lines.append(449 "If wrong, please convert your image to a 8-bit RGB/Grayscale image file before uploading (there are available tools such as ImageJ/Fiji)."450 )451 return "<br>".join(lines)452 453 454def prepare_uploaded_image(annot_value):455 """On annotator upload: standardize frame 0 for preview and configure the456 stack sliders.457 458 Browsers cannot render 16-bit / float / multi-page TIFFs, so we standardize459 to an 8-bit RGB PNG. If the file is a stack, notify the user and reveal a460 frame slider (default frame 1); a file with a second stack axis (e.g. Z in a461 time+Z series) also gets a z-plane slider (0 = max-intensity projection).462 Sliders that don't apply stay hidden.463 464 Returns (annotator_value, raw_path, frame_slider_update, zplane_slider_update).465 """466 hidden = gr.update(visible=False, value=1)467 img_path = _annot_path(annot_value)468 if not img_path:469 return annot_value, None, hidden, hidden470 471 # Clear the annotator on refusal, so a rejected file cannot reach inference.472 rejected = check_image(img_path)473 if rejected:474 gr.Warning(rejected, duration=None, title="❌ Cannot use this file")475 return None, None, hidden, hidden476 477 info = inspect_image(img_path)478 display = standardize_image(img_path, frame=0,479 sub_frame=0 if info.sub_frames > 1 else None)480 481 # Stay quiet unless something non-obvious happened; a plain RGB or grayscale482 # image needs no explanation.483 if info.guessed or info.frames > 1 or info.channels > 3:484 gr.Info(_describe_read(info), duration=None, title="📚 Image Loading Info")485 486 slider = (gr.update(visible=True, maximum=info.frames, value=1) if info.frames > 1487 else hidden)488 zslider = (gr.update(visible=True, maximum=info.sub_frames, value=1,489 label=f"🔬 {info.sub_label.capitalize()} (choose plane to use)")490 if info.sub_frames > 1 else hidden)491 return display, img_path, slider, zslider492 493 494def select_frame(raw_path, frame_num, zplane=1):495 """Re-render the annotator preview for the chosen frame and z-plane (both496 1-based; z-plane is ignored for files with only one stack axis)."""497 if not raw_path:498 return gr.update()499 return standardize_image(raw_path, frame=int(frame_num) - 1, sub_frame=int(zplane) - 1)500 501 502@spaces.GPU503def segment_with_choice(use_box_choice, annot_value, overlay_alpha):504 """Segmentation handler - supports bounding box, returns colorized overlay and original mask path"""505 if annot_value is None or len(annot_value) < 1:506 print("❌ No annotation input")507 return None, None, {}508 509 img_path = annot_value[0]510 bboxes = annot_value[1] if len(annot_value) > 1 else []511 512 print(f"🖼️ Image path: {img_path}")513 # One standardized image for both the model and the overlay (WYSIWYG).514 img_path = standardize_image(img_path)515 print(f"🧪 Standardized image: {img_path}")516 box_array = None517 if use_box_choice == "Yes" and bboxes:518 box = parse_bboxes(bboxes)519 if box:520 box_array = box521 print(f"📦 Using bounding boxes: {box_array}")522 523 524 try:525 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")526 mask = run_seg(SEG_MODEL, img_path, box=box_array, device=device)527 print("📏 mask shape:", mask.shape, "dtype:", mask.dtype)528 except Exception as e:529 print(f"❌ Inference failed: {str(e)}")530 return None, None, {}531 532 temp_mask_file = tempfile.NamedTemporaryFile(delete=False, suffix=".tif")533 mask_img = Image.fromarray(mask.astype(np.uint16))534 mask_img.save(temp_mask_file.name)535 print(f"💾 Original mask saved to: {temp_mask_file.name}")536 537 try:538 img = Image.open(img_path)539 print("📷 Image mode:", img.mode, "size:", img.size)540 except Exception as e:541 print(f"❌ Failed to open image: {e}")542 return None, None, {}543 544 try:545 img_rgb = img.convert("RGB").resize(mask.shape[::-1], resample=Image.BILINEAR)546 img_np = np.array(img_rgb, dtype=np.float32)547 if img_np.max() > 1.5:548 img_np = img_np / 255.0549 except Exception as e:550 print(f"❌ Error in image conversion/resizing: {e}")551 return None, None, {}552 553 mask_np = np.array(mask)554 inst_mask = mask_np.astype(np.int32)555 unique_ids = np.unique(inst_mask)556 num_instances = len(unique_ids[unique_ids != 0])557 if num_instances == 0:558 print("⚠️ No instance found, returning dummy red image")559 return Image.new("RGB", mask.shape[::-1], (255, 0, 0)), None, {}560 561 overlay_img = render_seg_overlay(img_np, inst_mask, overlay_alpha)562 seg_vis_cache = {"img_np": img_np, "inst_mask": inst_mask}563 return overlay_img, temp_mask_file.name, seg_vis_cache564 565 566@spaces.GPU567def count_cells_handler(use_box_choice, annot_value, overlay_alpha):568 """Counting handler - supports bounding box, returns only density map"""569 if annot_value is None or len(annot_value) < 1:570 return None, None, "⚠️ Please provide an image.", {}571 572 image_path = annot_value[0]573 bboxes = annot_value[1] if len(annot_value) > 1 else []574 575 print(f"🖼️ Image path: {image_path}")576 # One standardized image for both the model and the overlay (WYSIWYG).577 image_path = standardize_image(image_path)578 print(f"🧪 Standardized image: {image_path}")579 box_array = None580 if use_box_choice == "Yes" and bboxes:581 box = parse_bboxes(bboxes)582 if box:583 box_array = box584 print(f"📦 Using bounding boxes: {box_array}")585 586 try:587 print(f"🔢 Counting - Image: {image_path}")588 589 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")590 result = run_count(591 COUNT_MODEL,592 image_path,593 box=box_array,594 device=device,595 visualize=True596 )597 598 if 'error' in result:599 return None, None, f"❌ Counting failed: {result['error']}", {}600 601 count = result['count']602 density_map = result['density_map']603 temp_density_file = tempfile.NamedTemporaryFile(delete=False, suffix=".npy")604 np.save(temp_density_file.name, density_map)605 print(f"💾 Density map saved to {temp_density_file.name}")606 607 608 try:609 img = Image.open(image_path)610 print("📷 Image mode:", img.mode, "size:", img.size)611 except Exception as e:612 print(f"❌ Failed to open image: {e}")613 return None, None, f"❌ Failed to open image: {str(e)}", {}614 615 try:616 img_rgb = img.convert("RGB").resize(density_map.shape[::-1], resample=Image.BILINEAR)617 img_np = np.array(img_rgb, dtype=np.float32)618 img_np = (img_np - img_np.min()) / (img_np.max() - img_np.min() + 1e-8)619 if img_np.max() > 1.5:620 img_np = img_np / 255.0621 except Exception as e:622 print(f"❌ Error in image conversion/resizing: {e}")623 return None, None, f"❌ Error in image conversion/resizing: {str(e)}", {}624 625 626 density_normalized = density_map.copy()627 if density_normalized.max() > 0:628 density_normalized = (density_normalized - density_normalized.min()) / (density_normalized.max() - density_normalized.min())629 630 overlay_img = render_count_overlay(img_np, density_normalized, overlay_alpha)631 result_text = f"✅ Detected {round(count)} objects"632 if use_box_choice == "Yes" and box_array:633 result_text += f"\n📦 Using bounding box: {box_array}"634 635 636 print(f"✅ Counting done - Count: {count:.1f}")637 638 count_vis_cache = {"img_np": img_np, "density_normalized": density_normalized}639 return overlay_img, temp_density_file.name, result_text, count_vis_cache640 641 642 except Exception as e:643 print(f"❌ Counting error: {e}")644 import traceback645 traceback.print_exc()646 return None, None, f"❌ Counting failed: {str(e)}", {}647 648 649def find_tif_dir(root_dir):650 """Recursively find the first directory containing .tif files"""651 for dirpath, _, filenames in os.walk(root_dir):652 if '__MACOSX' in dirpath:653 continue654 if any(f.lower().endswith('.tif') for f in filenames):655 return dirpath656 return None657 658def is_valid_tiff(filepath):659 """Check if a file is a valid TIFF image"""660 try:661 with Image.open(filepath) as img:662 img.verify()663 return True664 except Exception as e:665 return False666 667def find_valid_tif_dir(root_dir):668 """Recursively find the first directory containing valid .tif files"""669 for dirpath, dirnames, filenames in os.walk(root_dir):670 if '__MACOSX' in dirpath:671 continue672 673 potential_tifs = [674 os.path.join(dirpath, f) 675 for f in filenames 676 if f.lower().endswith(('.tif', '.tiff')) and not f.startswith('._')677 ]678 679 if not potential_tifs:680 continue681 682 valid_tifs = [f for f in potential_tifs if is_valid_tiff(f)]683 684 if valid_tifs:685 print(f"✅ Found {len(valid_tifs)} valid TIFF files in: {dirpath}")686 return dirpath687 688 return None689 690def create_ctc_results_zip(output_dir):691 """692 Create a ZIP file with CTC format results693 694 Parameters:695 -----------696 output_dir : str697 Directory containing tracking results (res_track.txt, etc.)698 699 Returns:700 --------701 zip_path : str702 Path to created ZIP file703 """704 # Create temp directory for ZIP705 temp_zip_dir = tempfile.mkdtemp()706 zip_filename = f"tracking_results_{time.strftime('%Y%m%d_%H%M%S')}.zip"707 zip_path = os.path.join(temp_zip_dir, zip_filename)708 709 print(f"📦 Creating results ZIP: {zip_path}")710 711 # Create ZIP with all tracking results712 with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:713 # Add all files from output directory714 for root, dirs, files in os.walk(output_dir):715 for file in files:716 file_path = os.path.join(root, file)717 arcname = os.path.relpath(file_path, output_dir)718 zipf.write(file_path, arcname)719 print(f" 📄 Added: {arcname}")720 721 # Add a README with summary722 readme_content = f"""Tracking Results Summary723 ========================724 725 Generated: {time.strftime('%Y-%m-%d %H:%M:%S')}726 727 Files:728 ------729 - res_track.txt: CTC format tracking data730 Format: track_id start_frame end_frame parent_id731 732 - Segmentation masks733 734 For more information on CTC format:735 http://celltrackingchallenge.net/736 """737 zipf.writestr("README.txt", readme_content)738 739 print(f"✅ ZIP created: {zip_path} ({os.path.getsize(zip_path) / 1024:.1f} KB)")740 return zip_path741 742 743def get_well_spaced_color(track_id, num_colors=256):744 """Generate well-spaced colors, using contrasting colors for adjacent IDs"""745 746 golden_ratio = 0.618033988749895747 hue = (track_id * golden_ratio) % 1.0748 749 import colorsys750 rgb = colorsys.hsv_to_rgb(hue, 0.9, 0.95)751 return np.array(rgb)752 753 754def extract_first_frame(tif_dir):755 """756 Extract the first frame from a directory of TIF files757 758 Returns:759 --------760 first_frame_path : str761 Path to the first TIF frame762 """763 tif_files = natsorted(glob(os.path.join(tif_dir, "*.tif")) + 764 glob(os.path.join(tif_dir, "*.tiff")))765 valid_tif_files = [f for f in tif_files 766 if not os.path.basename(f).startswith('._') and is_valid_tiff(f)]767 768 if valid_tif_files:769 return valid_tif_files[0]770 return None771 772def create_tracking_visualization(tif_dir, output_dir, valid_tif_files, overlay_alpha=0.3):773 """774 Create an animated GIF/video showing tracked objects with consistent colors775 776 Parameters:777 -----------778 tif_dir : str779 Directory containing input TIF frames780 output_dir : str781 Directory containing tracking results (masks)782 valid_tif_files : list783 List of valid TIF file paths784 785 Returns:786 --------787 video_path : str788 Path to generated visualization (GIF or first frame)789 """790 import numpy as np791 from matplotlib import colormaps792 from skimage import measure793 import tifffile794 795 # Look for tracking mask files in output directory796 # Common CTC formats: man_track*.tif, mask*.tif, or numbered masks797 mask_files = natsorted(glob(os.path.join(output_dir, "mask*.tif")) + 798 glob(os.path.join(output_dir, "man_track*.tif")) +799 glob(os.path.join(output_dir, "*.tif")))800 801 if not mask_files:802 print("⚠️ No mask files found in output directory")803 # Return first frame as fallback804 return valid_tif_files[0]805 806 print(f"📊 Found {len(mask_files)} mask files")807 808 809 frames = []810 alpha = float(np.clip(overlay_alpha, 0.0, 1.0)) # Transparency for overlay811 812 # Process each frame813 num_frames = min(len(valid_tif_files), len(mask_files))814 for i in range(num_frames):815 try:816 # Load original image using tifffile (handles ZSTD compression)817 try:818 img_np = tifffile.imread(valid_tif_files[i])819 820 # Normalize to [0, 1] range based on actual data type and values821 if img_np.dtype == np.uint8:822 img_np = img_np.astype(np.float32) / 255.0823 elif img_np.dtype == np.uint16:824 # Normalize uint16 to [0, 1] using actual min/max825 img_min, img_max = img_np.min(), img_np.max()826 if img_max > img_min:827 img_np = (img_np.astype(np.float32) - img_min) / (img_max - img_min)828 else:829 img_np = img_np.astype(np.float32) / 65535.0830 else:831 # For float or other types, normalize based on actual range832 img_np = img_np.astype(np.float32)833 img_min, img_max = img_np.min(), img_np.max()834 if img_max > img_min:835 img_np = (img_np - img_min) / (img_max - img_min)836 else:837 img_np = np.clip(img_np, 0, 1)838 839 # Convert to RGB if grayscale840 if img_np.ndim == 2:841 img_np = np.stack([img_np]*3, axis=-1)842 img_np = img_np.astype(np.float32)843 if img_np.max() > 1.5:844 img_np = img_np / 255.0845 except Exception as e:846 print(f"⚠️ Error loading image frame {i}: {e}")847 # Fallback to PIL848 img = Image.open(valid_tif_files[i]).convert("RGB")849 img_np = np.array(img, dtype=np.float32) / 255.0850 851 # Load tracking mask using tifffile (handles ZSTD compression)852 try:853 mask = tifffile.imread(mask_files[i])854 except Exception as e:855 print(f"⚠️ Error loading mask frame {i}: {e}")856 # Fallback to PIL857 mask = np.array(Image.open(mask_files[i]))858 859 # Resize mask to match image if needed860 if mask.shape[:2] != img_np.shape[:2]:861 from scipy.ndimage import zoom862 zoom_factors = [img_np.shape[0] / mask.shape[0], img_np.shape[1] / mask.shape[1]]863 mask = zoom(mask, zoom_factors, order=0).astype(mask.dtype)864 865 # Create overlay866 overlay = img_np.copy()867 868 # Get unique track IDs (excluding background 0)869 track_ids = np.unique(mask)870 track_ids = track_ids[track_ids != 0]871 872 # Color each tracked object873 for track_id in track_ids:874 # Create binary mask for this track875 binary_mask = (mask == track_id)876 877 # Get consistent color for this track ID878 # color = np.array(cmap(int(track_id) % 256)[:3])879 color = get_well_spaced_color(int(track_id))880 881 # Blend color onto image882 overlay[binary_mask] = (1 - alpha) * overlay[binary_mask] + alpha * color883 884 # Draw contours (optional, adds yellow boundaries)885 try:886 contours = measure.find_contours(binary_mask.astype(np.uint8), 0.5)887 for contour in contours:888 contour = contour.astype(np.int32)889 valid_y = np.clip(contour[:, 0], 0, overlay.shape[0] - 1)890 valid_x = np.clip(contour[:, 1], 0, overlay.shape[1] - 1)891 overlay[valid_y, valid_x] = [1.0, 1.0, 0.0] # Yellow contour892 except:893 pass # Skip contours if they fail894 895 # Convert to uint8896 overlay_uint8 = np.clip(overlay * 255.0, 0, 255).astype(np.uint8)897 frames.append(Image.fromarray(overlay_uint8))898 899 if i % 10 == 0 or i == num_frames - 1:900 print(f" 📸 Processed frame {i+1}/{num_frames}")901 902 except Exception as e:903 print(f"⚠️ Error processing frame {i}: {e}")904 import traceback905 traceback.print_exc()906 continue907 908 if not frames:909 print("⚠️ No frames were processed successfully")910 return valid_tif_files[0]911 912 # Save as animated GIF913 try:914 temp_gif = tempfile.NamedTemporaryFile(delete=False, suffix=".gif")915 frames[0].save(916 temp_gif.name,917 save_all=True,918 append_images=frames[1:],919 duration=200, # 200ms per frame = 5fps920 loop=0921 )922 temp_gif.close() # Close the file handle923 print(f"✅ Created tracking visualization GIF: {temp_gif.name}")924 print(f" Size: {os.path.getsize(temp_gif.name)} bytes, Frames: {len(frames)}")925 return temp_gif.name926 except Exception as e:927 print(f"⚠️ Failed to create GIF: {e}")928 import traceback929 traceback.print_exc()930 # Return first frame as static image fallback931 try:932 temp_img = tempfile.NamedTemporaryFile(delete=False, suffix=".png")933 frames[0].save(temp_img.name)934 temp_img.close()935 return temp_img.name936 except:937 return valid_tif_files[0]938 939@spaces.GPU940def track_video_handler(use_box_choice, first_frame_annot, zip_file_obj, overlay_alpha, prev_track_vis_cache):941 """942 Tracking handler - processes a ZIP of TIF frames, supports bounding box, returns visualization and results ZIP943 944 Parameters:945 -----------946 use_box_choice : str947 "Yes" or "No" - whether to use bounding box annotation for tracking948 first_frame_annot : tuple or None949 (image_path, bboxes) from BBoxAnnotator, only used if user annotated first frame950 zip_file_obj : File951 Uploaded ZIP file containing TIF sequence952 """953 if zip_file_obj is None:954 return None, "⚠️ Please upload a ZIP file containing video frames (.zip)", None, None, {}955 956 cleanup_tracking_cache(prev_track_vis_cache)957 temp_dir = None958 output_temp_dir = None959 960 try:961 # Parse bounding box if provided962 box_array = None963 if use_box_choice == "Yes" and first_frame_annot is not None:964 if isinstance(first_frame_annot, (list, tuple)) and len(first_frame_annot) > 1:965 bboxes = first_frame_annot[1]966 if bboxes:967 box = parse_bboxes(bboxes)968 if box:969 box_array = box970 print(f"📦 Using bounding boxes: {box_array}")971 972 # Extract input ZIP973 temp_dir = tempfile.mkdtemp()974 print(f"\n📦 Extracting to temporary directory: {temp_dir}")975 976 with zipfile.ZipFile(zip_file_obj.name, 'r') as zip_ref:977 extracted_count = 0978 skipped_count = 0979 980 for member in zip_ref.namelist():981 basename = os.path.basename(member)982 983 if ('__MACOSX' in member or 984 basename.startswith('._') or 985 basename.startswith('.DS_Store') or986 member.endswith('/')):987 skipped_count += 1988 continue989 990 try:991 zip_ref.extract(member, temp_dir)992 extracted_count += 1993 if basename.lower().endswith(('.tif', '.tiff')):994 print(f"📄 Extracted TIFF: {basename}")995 except Exception as e:996 print(f"⚠️ Failed to extract {member}: {e}")997 998 print(f"\n📊 Extracted: {extracted_count} files, Skipped: {skipped_count} files")999 1000 # Find valid TIFF directory1001 tif_dir = find_valid_tif_dir(temp_dir)1002 1003 if tif_dir is None:1004 return None, "❌ Did not find valid TIF directory", None, None, {}1005 1006 # Validate TIFF files1007 tif_files = natsorted(glob(os.path.join(tif_dir, "*.tif")) + 1008 glob(os.path.join(tif_dir, "*.tiff")))1009 valid_tif_files = [f for f in tif_files 1010 if not os.path.basename(f).startswith('._') and is_valid_tiff(f)]1011 1012 if len(valid_tif_files) == 0:1013 return None, "❌ Did not find valid TIF files", None, None, {}1014 1015 print(f"📈 Using {len(valid_tif_files)} TIF files")1016 1017 # Store paths for later visualization1018 first_frame_path = valid_tif_files[0]1019 1020 # Create temporary output directory for CTC results1021 output_temp_dir = tempfile.mkdtemp()1022 print(f"💾 CTC-format results will be saved to: {output_temp_dir}")1023 1024 # Run tracking with optional bounding box1025 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")1026 result = run_track(1027 TRACK_MODEL,1028 video_dir=tif_dir,1029 box=box_array, # Pass bounding box if specified1030 device=device,1031 output_dir=output_temp_dir1032 )1033 1034 if 'error' in result:1035 return None, f"❌ Tracking failed: {result['error']}", None, None, {}1036 1037 # Create visualization video of tracked objects1038 print("\n🎬 Creating tracking visualization...")1039 try:1040 tracking_video = create_tracking_visualization(1041 tif_dir, 1042 output_temp_dir, 1043 valid_tif_files,1044 overlay_alpha=overlay_alpha1045 )1046 except Exception as e:1047 print(f"⚠️ Failed to create visualization: {e}")1048 import traceback1049 traceback.print_exc()1050 # Fallback to first frame if visualization fails1051 try:1052 tracking_video = Image.open(first_frame_path)1053 except:1054 tracking_video = None1055 1056 # Create downloadable ZIP with results1057 try:1058 results_zip = create_ctc_results_zip(output_temp_dir)1059 except Exception as e:1060 print(f"⚠️ Failed to create ZIP: {e}")1061 results_zip = None1062 1063 bbox_info = ""1064 if box_array:1065 bbox_info = f"\n🔲 Using bounding box: [{box_array[0][0]}, {box_array[0][1]}, {box_array[0][2]}, {box_array[0][3]}]"1066 1067 result_text = (1068 f"✅ Tracking completed!\n"1069 f"\n"1070 f"🖼️ Processed frames: {len(valid_tif_files)}{bbox_info}\n"1071 f"\n"1072 f"📥 Use the \"Download (.zip)\" button above to get the results, which include:\n"1073 f"- res_track.txt (CTC-format tracking data)\n"1074 f"- Other tracking-related files\n"1075 f"- README.txt (Results description)"1076 )1077 1078 if use_box_choice == "Yes" and box_array:1079 result_text += f"\n📦 Using bounding box: {box_array}"1080 1081 print(f"\n✅ Tracking completed")1082 1083 track_vis_cache = {1084 "tif_dir": tif_dir,1085 "valid_tif_files": valid_tif_files,1086 "output_dir": output_temp_dir,1087 "input_temp_dir": temp_dir,1088 }1089 1090 return results_zip, result_text, gr.update(visible=True), tracking_video, track_vis_cache1091 1092 except zipfile.BadZipFile:1093 return None, "❌ Not a valid ZIP file", None, None, {}1094 except Exception as e:1095 import traceback1096 traceback.print_exc()1097 1098 # Clean up on error1099 for d in [temp_dir, output_temp_dir]:1100 if d:1101 try:1102 shutil.rmtree(d)1103 except:1104 pass1105 return None, f"❌ Tracking failed: {str(e)}", None, None, {}1106 1107 1108 1109# ===== Example Images =====1110example_images_seg = [f for f in glob("example_imgs/seg/*")]1111example_images_cnt = [f for f in glob("example_imgs/cnt/*")]1112example_tracking_zips = [f for f in glob("example_imgs/tra/*.zip")]1113 1114_CHANNEL_NOTE = {1: "grayscale", 3: "RGB", 4: "RGBA"}1115 1116 1117def describe_image_info(raw_path, frame=1):1118 """Metadata line under the result: file, size, channels, bit depth, frames.1119 1120 Reads the *original* upload (not the standardized preview) so the numbers1121 describe what the user actually provided. Metadata only - no pixel decode.1122 """1123 if not raw_path or not os.path.exists(raw_path):1124 return gr.update(value="", visible=False)1125 1126 info = inspect_image(raw_path)1127 bits = bit_depth(raw_path)1128 1129 fields = [f"<span class='ii-k'>File:</span> {html.escape(os.path.basename(raw_path))}"]1130 if info.width and info.height:1131 fields.append(f"<span class='ii-k'>Size:</span> {info.width} × {info.height} px")1132 note = _CHANNEL_NOTE.get(info.channels)1133 fields.append(f"<span class='ii-k'>Channels:</span> {info.channels}"1134 + (f" ({note})" if note else ""))1135 if bits:1136 fields.append(f"<span class='ii-k'>Bit depth:</span> {bits}-bit")1137 if info.frames > 1:1138 fields.append(f"<span class='ii-k'>Frames:</span> {int(frame)} / {info.frames}")1139 if info.sub_frames > 1:1140 fields.append(f"<span class='ii-k'>{info.sub_label.capitalize()}s:</span> {info.sub_frames}")1141 1142 body = "".join(f"<span class='ii-f'>{f}</span>" for f in fields)1143 return gr.update(1144 value=f"<div class='ii-wrap'><span class='ii-title'>ⓘ Image information</span>{body}</div>",1145 visible=True,1146 )1147 1148 1149# ===== Gradio UI =====1150CSS = """1151/* ── Layout ──────────────────────────────────────────── */1152.gradio-container {1153 max-width: 1320px !important;1154 margin: 0 auto !important;1155 font-family: 'Inter', 'Segoe UI', system-ui, sans-serif !important;1156 background: #fff !important;1157}1158 1159/* ── Header markdown polish ───────────────────────────── */1160.gradio-container .prose h1 {1161 font-size: 1.4rem !important;1162 font-weight: 700 !important;1163 color: #1c1c1a !important;1164 letter-spacing: -0.01em !important;1165 margin-bottom: 2px !important;1166}1167.gradio-container .prose h3 {1168 font-size: 1rem !important;1169 font-weight: 600 !important;1170 color: #bf1e2e !important;1171 margin-top: 14px !important;1172 margin-bottom: 4px !important;1173}1174.gradio-container .prose p {1175 margin-top: 4px !important;1176 margin-bottom: 6px !important;1177 color: #5c5a52 !important;1178 line-height: 1.7 !important;1179}1180.gradio-container .prose ul,1181.gradio-container .prose ol {1182 margin-top: 4px !important;1183 margin-bottom: 6px !important;1184}1185.gradio-container .prose li {1186 color: #5c5a52 !important;1187 line-height: 1.7 !important;1188}1189 1190/* ── Top-level header section: slim, no gradient card ─── */1191.gradio-container > .gap > .prose:first-child {1192 background: transparent !important;1193 border: none !important;1194 border-radius: 0 !important;1195 padding: 18px 4px 10px !important;1196 margin-bottom: 8px !important;1197 box-shadow: none !important;1198}1199 1200/* ── Mode selector: two-line task cards ──────────────── */