Upload 7 files

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+standard-ilc.tif filter=lfs diff=lfs merge=lfs -text

HNE2cell_all_patch73_jit.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cd428608461bf295e4abfa646fb834c2d24acedca0155d43fdb817da42936593
+size 5138307858

inference.py

@@ -0,0 +1,316 @@
+"""
+HNE2Cell — Step 3: Cell Detection & Classification Inference
+
+Run the HNE2Cell model on extracted patches to detect and classify cells.
+Outputs: per-patch cell masks (PNG) and centroid CSVs with cell type annotations.
+
+Usage:
+    python inference.py \
+        --input_dir /path/to/patch_folders \
+        --output_dir /path/to/results \
+        --model_path ./HNE2cell_all_patch73_jit.pt \
+        --magnification 40 \
+        --batch_size 32
+
+Cell Types (16 classes):
+    0: Background      4: B             8: DC            12: Epithelial
+    1: Malignant        5: Plasma        9: Fibroblast    13: Immune_Other
+    2: CD4T             6: Macrophage   10: Endothelial   14: Stromal_Other
+    3: CD8T             7: Myeloid      11: Pericyte      15: Dead
+"""
+
+import os
+import argparse
+import glob
+
+import cv2
+import numpy as np
+import pandas as pd
+import torch
+from PIL import Image
+from torch.cuda.amp import autocast
+from torch.utils.data import DataLoader, Dataset
+from torchvision import transforms
+from tqdm import tqdm
+
+from post_processing import DetectionCellPostProcessor
+
+# ========================== Constants ======================================
+
+CELL_TYPES = {
+    0: "Background",
+    1: "Malignant",
+    2: "CD4T",
+    3: "CD8T",
+    4: "B",
+    5: "Plasma",
+    6: "Macrophage",
+    7: "Myeloid",
+    8: "DC",
+    9: "Fibroblast",
+    10: "Endothelial",
+    11: "Pericyte",
+    12: "Epithelial",
+    13: "Immune_Other",
+    14: "Stromal_Other",
+    15: "Dead",
+}
+
+# RGBA colors for mask visualization
+CELL_COLORS = {
+    0: [0, 0, 0, 0],
+    1: [255, 0, 0, 255],
+    2: [30, 144, 255, 255],
+    3: [65, 105, 225, 255],
+    4: [0, 0, 255, 255],
+    5: [100, 149, 237, 255],
+    6: [176, 224, 230, 255],
+    7: [70, 130, 180, 255],
+    8: [0, 191, 255, 255],
+    9: [34, 139, 34, 255],
+    10: [60, 179, 113, 255],
+    11: [50, 205, 50, 255],
+    12: [255, 140, 0, 255],
+    13: [176, 224, 230, 255],
+    14: [107, 142, 35, 255],
+    15: [128, 128, 128, 255],
+}
+
+# ImageNet-style normalization fitted to H&E data
+TRANSFORM = transforms.Compose(
+    [
+        transforms.Resize((224, 224)),
+        transforms.ToTensor(),
+        transforms.Normalize(
+            mean=[0.707223, 0.578729, 0.703617],
+            std=[0.211883, 0.230117, 0.177517],
+        ),
+    ]
+)
+
+
+# ========================== Dataset ========================================
+
+
+class PatchDataset(Dataset):
+    def __init__(self, file_paths, transform=None):
+        self.file_paths = file_paths
+        self.transform = transform
+
+    def __len__(self):
+        return len(self.file_paths)
+
+    def __getitem__(self, idx):
+        fpath = self.file_paths[idx]
+        img = cv2.imread(fpath)
+        if img is None:
+            print(f"[WARN] Failed to load: {fpath}")
+            img = np.zeros((256, 256, 3), dtype=np.uint8)
+        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+        if self.transform:
+            img = self.transform(Image.fromarray(img))
+        return fpath, img
+
+
+# ========================== Inference ======================================
+
+
+def process_batch(
+    batch,
+    model,
+    device,
+    mask_output_dir,
+    centroid_records,
+    magnification=40,
+):
+    """Run inference on one batch and save masks + centroid info."""
+    file_paths, images = batch
+
+    with torch.no_grad():
+        with autocast():
+            outputs = model(images.to(device, non_blocking=True))
+
+    for i, fpath in enumerate(file_paths):
+        slide_id = os.path.splitext(os.path.basename(fpath))[0]
+
+        # Extract per-sample predictions
+        cell_type_map = outputs["cell_type_map"][i].float().detach().cpu()
+        nuclei_binary_map = outputs["nuclei_binary_map"][i].float().detach().cpu()
+        hv_map = outputs["hv_map"][i].float().detach().cpu()
+        tissue_type_map = outputs["tissue_type_map"][i].float().detach().cpu()
+
+        # Build prediction map [H, W, 5]
+        pred_map = np.concatenate(
+            [
+                torch.argmax(tissue_type_map, dim=0)[..., None].numpy(),
+                torch.argmax(cell_type_map, dim=0)[..., None].numpy(),
+                torch.argmax(nuclei_binary_map, dim=0)[..., None].numpy(),
+                hv_map.permute(1, 2, 0).numpy(),
+            ],
+            axis=-1,
+        )
+
+        # Post-processing
+        post_processor = DetectionCellPostProcessor(
+            nr_types=cell_type_map.shape[0],
+            magnification=magnification,
+            gt=False,
+        )
+        _, type_pred = post_processor.post_process_cell_segmentation(pred_map)
+
+        # Create mask image
+        mask = np.ones((256, 256, 3), dtype=np.uint8) * 255
+        for cell in type_pred.values():
+            ctype = cell["type"]
+            rgba = CELL_COLORS.get(ctype, [255, 255, 255, 255])
+            bgr = [rgba[2], rgba[1], rgba[0]]
+            cv2.fillPoly(mask, [cell["contour"]], bgr)
+
+            centroid_records.append(
+                {
+                    "slide_id": slide_id,
+                    "x": cell["centroid"][0],
+                    "y": cell["centroid"][1],
+                    "celltype": ctype,
+                    "celltype_name": CELL_TYPES.get(ctype, "Unknown"),
+                }
+            )
+
+        # Save mask only if non-trivial
+        if not np.all(mask == 255):
+            cv2.imwrite(
+                os.path.join(mask_output_dir, f"{slide_id}_mask.png"), mask
+            )
+
+
+def run_inference(
+    patch_folders: list[str],
+    model,
+    device,
+    output_dir: str,
+    magnification: int = 40,
+    batch_size: int = 32,
+    num_workers: int = 4,
+):
+    """Run inference over a list of patch folders."""
+    model.to(device).eval()
+
+    for folder in patch_folders:
+        folder_name = os.path.basename(folder)
+        png_files = sorted(glob.glob(os.path.join(folder, "*.png")))
+
+        if not png_files:
+            print(f"[SKIP] {folder}: no PNG patches found")
+            continue
+
+        mask_dir = os.path.join(output_dir, "mask_patches", f"{folder_name}")
+        centroid_path = os.path.join(output_dir, "centroid", f"{folder_name}_centroid.csv")
+        os.makedirs(mask_dir, exist_ok=True)
+        os.makedirs(os.path.dirname(centroid_path), exist_ok=True)
+
+        dataset = PatchDataset(png_files, transform=TRANSFORM)
+        loader = DataLoader(
+            dataset,
+            batch_size=batch_size,
+            num_workers=num_workers,
+            pin_memory=True,
+            shuffle=False,
+            prefetch_factor=2,
+            persistent_workers=True,
+        )
+
+        centroids = []
+        for batch in tqdm(loader, desc=f"Inference: {folder_name}"):
+            process_batch(
+                batch, model, device, mask_dir, centroids, magnification
+            )
+
+        df = pd.DataFrame(centroids)
+        df.to_csv(centroid_path, index=False)
+        print(f"[DONE] {folder_name} → {centroid_path} ({len(df)} cells)")
+
+
+# =============================== CLI =======================================
+
+
+def main():
+    parser = argparse.ArgumentParser(description="HNE2Cell inference")
+    parser.add_argument(
+        "--input_dir",
+        type=str,
+        required=True,
+        help="Directory containing patch folders (each with *.png)",
+    )
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        required=True,
+        help="Output directory for masks and centroid CSVs",
+    )
+    parser.add_argument(
+        "--model_path",
+        type=str,
+        required=True,
+        help="Path to the TorchScript JIT model (.pt)",
+    )
+    parser.add_argument(
+        "--magnification",
+        type=int,
+        default=40,
+        choices=[20, 40],
+        help="Magnification of input patches. 40x recommended. (default: 40)",
+    )
+    parser.add_argument("--batch_size", type=int, default=32)
+    parser.add_argument("--num_workers", type=int, default=4)
+    parser.add_argument(
+        "--device",
+        type=str,
+        default="auto",
+        help="Device: 'cuda', 'cpu', or 'auto' (default: auto)",
+    )
+
+    args = parser.parse_args()
+
+    # Device
+    if args.device == "auto":
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    else:
+        device = torch.device(args.device)
+    print(f"Using device: {device}")
+
+    # Load model
+    print(f"Loading model: {args.model_path}")
+    model = torch.jit.load(args.model_path, map_location=device)
+    model.eval()
+
+    if args.magnification == 20:
+        print(
+            "⚠️  Running at 20x. Results are usable but 40x is recommended "
+            "for best accuracy, especially for small immune cells."
+        )
+
+    # Collect patch folders
+    patch_folders = sorted(
+        p
+        for p in glob.glob(os.path.join(args.input_dir, "*"))
+        if os.path.isdir(p)
+    )
+    # Also check if input_dir itself contains patches
+    if not patch_folders and glob.glob(os.path.join(args.input_dir, "*.png")):
+        patch_folders = [args.input_dir]
+
+    print(f"Found {len(patch_folders)} patch folder(s)")
+
+    run_inference(
+        patch_folders,
+        model,
+        device,
+        args.output_dir,
+        magnification=args.magnification,
+        batch_size=args.batch_size,
+        num_workers=args.num_workers,
+    )
+
+
+if __name__ == "__main__":
+    main()

normalize.py

@@ -0,0 +1,277 @@
+"""
+HNE2Cell — Step 1: Reinhard Color Normalization
+
+Normalize H&E stained whole-slide images (WSI) to a reference color distribution
+using the Reinhard method in LAB color space.
+
+Supported input formats: .svs, .tif, .tiff, .ndpi
+Output: Aligned-hne.tif (full-resolution normalized), Aligned-hne.jpg (4x downsampled preview)
+
+Usage:
+    python normalize.py \
+        --input_dir /path/to/slides \
+        --target /path/to/standard-ilc.tif \
+        --patch_size 128 \
+        --saturation_threshold 0.1
+"""
+
+import os
+import argparse
+import glob
+
+import numpy as np
+import tifffile as tiff
+from PIL import Image
+from skimage import color
+
+Image.MAX_IMAGE_PIXELS = None
+os.environ["OPENCV_IO_MAX_IMAGE_PIXELS"] = str(pow(2, 40))
+
+# ---------------------------------------------------------------------------
+# Optional: openslide (only needed for .svs / .ndpi)
+# ---------------------------------------------------------------------------
+try:
+    import openslide
+
+    OPENSLIDE_AVAILABLE = True
+except ImportError:
+    OPENSLIDE_AVAILABLE = False
+
+
+# ============================= I/O helpers =================================
+
+
+def load_image(image_path: str, level: int = 0) -> np.ndarray:
+    """Load a whole-slide image as an RGB numpy array.
+
+    Supports .svs/.ndpi (via OpenSlide) and .tif/.tiff (via tifffile).
+    """
+    ext = os.path.splitext(image_path)[1].lower()
+
+    if ext in (".svs", ".ndpi"):
+        if not OPENSLIDE_AVAILABLE:
+            raise ImportError(
+                "openslide-python is required to read .svs/.ndpi files. "
+                "Install it with: pip install openslide-python"
+            )
+        slide = openslide.OpenSlide(image_path)
+        image = slide.read_region((0, 0), level, slide.level_dimensions[level])
+        image = image.convert("RGB")
+        slide.close()
+        return np.array(image)
+
+    if ext in (".tif", ".tiff"):
+        image = tiff.imread(image_path)
+        if image.ndim == 2:
+            image = np.stack((image,) * 3, axis=-1)
+        elif image.ndim == 4 and image.shape[0] == 1:
+            image = image[0]
+        # Ensure RGB uint8
+        if image.dtype != np.uint8:
+            image = np.clip(image, 0, 255).astype(np.uint8)
+        return image
+
+    raise ValueError(f"Unsupported file format: {ext}")
+
+
+# ======================== Saturation filtering =============================
+
+
+def calculate_saturation(patch: Image.Image) -> float:
+    hsv = patch.convert("HSV")
+    return np.mean(np.array(hsv)[:, :, 1] / 255.0)
+
+
+def extract_high_saturation_patches(
+    image: np.ndarray, patch_size: int, saturation_threshold: float
+) -> list:
+    """Return list of ((x0, y0), patch_array) for patches above the saturation threshold."""
+    pil_img = Image.fromarray(image)
+    width, height = pil_img.size
+
+    patches = []
+    for i in range(width // patch_size):
+        for j in range(height // patch_size):
+            x0, y0 = i * patch_size, j * patch_size
+            patch = pil_img.crop((x0, y0, x0 + patch_size, y0 + patch_size))
+            if calculate_saturation(patch) >= saturation_threshold:
+                patches.append(((x0, y0), np.array(patch)))
+    return patches
+
+
+def reconstruct_from_patches(
+    width: int, height: int, patch_size: int, patches: list
+) -> np.ndarray:
+    """Place high-saturation patches back into a blank canvas (background = black)."""
+    canvas = np.zeros((height, width, 3), dtype=np.uint8)
+    for (x0, y0), arr in patches:
+        if arr.shape == (patch_size, patch_size, 3):
+            canvas[y0 : y0 + patch_size, x0 : x0 + patch_size, :] = arr
+    return canvas
+
+
+# =================== Reinhard color normalization ==========================
+
+
+def _color_convert_chunked(image, func, chunk_size=16384):
+    """Apply color conversion function in spatial chunks to limit memory."""
+    h, w, _ = image.shape
+    out = np.zeros_like(image, dtype=np.float32)
+    for i in range(0, h, chunk_size):
+        for j in range(0, w, chunk_size):
+            out[i : min(i + chunk_size, h), j : min(j + chunk_size, w), :] = func(
+                image[i : min(i + chunk_size, h), j : min(j + chunk_size, w), :]
+            )
+    return out
+
+
+def reinhard_normalize(source: np.ndarray, target: np.ndarray) -> np.ndarray:
+    """Reinhard color normalization in LAB space.
+
+    Only non-zero (tissue) pixels are used for statistics.
+    Returns float64 image in [0, 1] range.
+    """
+    src_lab = _color_convert_chunked(source, color.rgb2lab)
+    tgt_lab = color.rgb2lab(target)
+
+    for ch in range(3):
+        src_ch = src_lab[:, :, ch]
+        tgt_ch = tgt_lab[:, :, ch]
+
+        src_vals = src_ch[src_ch != 0]
+        tgt_vals = tgt_ch[tgt_ch != 0]
+
+        if len(src_vals) == 0 or len(tgt_vals) == 0:
+            continue
+
+        src_mean, src_std = src_vals.mean(), src_vals.std()
+        tgt_mean, tgt_std = tgt_vals.mean(), tgt_vals.std()
+
+        if src_std < 1e-6:
+            continue
+
+        src_lab[:, :, ch] = np.where(
+            src_ch != 0,
+            (src_ch - src_mean) * (tgt_std / src_std) + tgt_mean,
+            0,
+        )
+
+    return _color_convert_chunked(src_lab, color.lab2rgb)
+
+
+# ============================= Main pipeline ===============================
+
+
+def normalize_slide(
+    slide_path: str,
+    target_image: np.ndarray,
+    patch_size: int = 128,
+    saturation_threshold: float = 0.1,
+    output_dir: str | None = None,
+    skip_existing: bool = True,
+):
+    """Full normalization pipeline for a single slide."""
+
+    if output_dir is None:
+        output_dir = os.path.dirname(slide_path)
+
+    output_tif = os.path.join(output_dir, "Aligned-hne.tif")
+    output_jpg = os.path.join(output_dir, "Aligned-hne.jpg")
+
+    if skip_existing and os.path.exists(output_tif):
+        print(f"[SKIP] {slide_path} — Aligned-hne.tif already exists.")
+        return
+
+    print(f"[LOAD] {slide_path}")
+    raw = load_image(slide_path)
+    h, w = raw.shape[:2]
+
+    # 1. Saturation-based tissue detection
+    patches = extract_high_saturation_patches(
+        raw, patch_size, saturation_threshold
+    )
+    reconstructed = reconstruct_from_patches(w, h, patch_size, patches)
+
+    # 2. (Optional) save intermediate reconstruction
+    recon_path = os.path.join(output_dir, "recon.tif")
+    bigtiff = reconstructed.nbytes > 4 * 1024**3
+    tiff.imwrite(recon_path, reconstructed, bigtiff=bigtiff)
+
+    # 3. Reinhard normalization
+    normalized = reinhard_normalize(reconstructed, target_image)
+    normalized_u8 = (normalized * 255).astype(np.uint8)
+
+    # 4. Save outputs
+    tiff.imwrite(output_tif, normalized_u8, bigtiff=bigtiff)
+
+    resized = Image.fromarray(normalized_u8).resize(
+        (w // 4, h // 4), Image.LANCZOS
+    )
+    resized.save(output_jpg, quality=90)
+
+    print(f"[DONE] {slide_path} → {output_tif}")
+
+
+# =============================== CLI =======================================
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Reinhard color normalization for H&E WSIs"
+    )
+    parser.add_argument(
+        "--input_dir",
+        type=str,
+        required=True,
+        help="Root directory to search for slide files (.svs, .tif, .tiff, .ndpi)",
+    )
+    parser.add_argument(
+        "--target",
+        type=str,
+        required=True,
+        help="Path to the reference/target image (.tif)",
+    )
+    parser.add_argument("--patch_size", type=int, default=128)
+    parser.add_argument("--saturation_threshold", type=float, default=0.1)
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default=None,
+        help="If set, all outputs go here. Otherwise, outputs are saved next to each slide.",
+    )
+
+    args = parser.parse_args()
+
+    # Load target image once
+    target_image = load_image(args.target)
+
+    # Collect slides
+    extensions = ("*.svs", "*.tif", "*.tiff", "*.ndpi")
+    slides = []
+    for ext in extensions:
+        slides.extend(glob.glob(os.path.join(args.input_dir, "**", ext), recursive=True))
+
+    # Exclude files that are already outputs
+    slides = [
+        s
+        for s in slides
+        if os.path.basename(s) not in ("Aligned-hne.tif", "Aligned-hne.tiff", "recon.tif")
+    ]
+
+    print(f"Found {len(slides)} slide(s) in {args.input_dir}")
+
+    for slide_path in slides:
+        try:
+            normalize_slide(
+                slide_path,
+                target_image,
+                patch_size=args.patch_size,
+                saturation_threshold=args.saturation_threshold,
+                output_dir=args.output_dir,
+            )
+        except Exception as e:
+            print(f"[ERROR] {slide_path}: {e}")
+
+
+if __name__ == "__main__":
+    main()

patchify.py

@@ -0,0 +1,254 @@
+"""
+HNE2Cell — Step 2: Patch Extraction
+
+Extract overlapping patches from color-normalized H&E images for cell detection.
+Supports both 20x and 40x magnification (40x recommended for best results).
+
+Usage:
+    # 40x (recommended)
+    python patchify.py \
+        --input_dir /path/to/slides \
+        --patch_size 256 \
+        --overlap 64 \
+        --magnification 40 \
+        --workers 8
+
+    # 20x (supported but 40x preferred)
+    python patchify.py \
+        --input_dir /path/to/slides \
+        --patch_size 256 \
+        --overlap 64 \
+        --magnification 20 \
+        --workers 8
+
+Notes:
+    - 40x magnification is recommended for optimal cell detection accuracy.
+    - 20x is supported and functional, but fine-grained cell boundaries
+      (especially small immune cells) may be less precise.
+    - Input: Aligned-hne.tif (output of normalize.py)
+    - Output: <section>/patches_<mag>x_p<patch>_o<overlap>/<name>_<x>_<y>.png
+"""
+
+import os
+import argparse
+import glob
+import time
+from multiprocessing import Pool
+
+import numpy as np
+from PIL import Image
+from tqdm import tqdm
+
+Image.MAX_IMAGE_PIXELS = None
+
+
+# ========================== Utility functions ==============================
+
+
+def black_to_white(pil_img: Image.Image) -> Image.Image:
+    """Replace pure-black (0,0,0) pixels with white — avoids dark-border artifacts."""
+    arr = np.array(pil_img)
+    if arr.ndim == 3 and arr.shape[2] >= 3:
+        mask = (arr[..., :3] == 0).all(axis=-1)
+        arr[mask] = 255
+    return Image.fromarray(arr)
+
+
+def make_start_positions(length: int, patch_size: int, stride: int) -> list[int]:
+    """Generate start positions so the last patch always reaches the edge."""
+    if length < patch_size:
+        return [0]
+    starts = list(range(0, length - patch_size + 1, stride))
+    last = length - patch_size
+    if starts[-1] != last:
+        starts.append(last)
+    return starts
+
+
+# ========================== Core patching ==================================
+
+
+def extract_patches(
+    image_path: str,
+    output_dir: str,
+    patch_size: int = 256,
+    overlap: int = 64,
+    prefix: str = "patch",
+) -> int:
+    """Crop overlapping patches from a single image and save as PNG.
+
+    Returns the number of patches saved.
+    """
+    os.makedirs(output_dir, exist_ok=True)
+
+    stride = patch_size - overlap
+    assert stride > 0, f"overlap ({overlap}) must be < patch_size ({patch_size})"
+
+    img = Image.open(image_path).convert("RGB")
+    width, height = img.size
+
+    xs = make_start_positions(width, patch_size, stride)
+    ys = make_start_positions(height, patch_size, stride)
+
+    count = 0
+    with tqdm(total=len(xs) * len(ys), desc=prefix, unit="patch", leave=False) as pbar:
+        for x0 in xs:
+            for y0 in ys:
+                patch = img.crop((x0, y0, x0 + patch_size, y0 + patch_size))
+                patch = black_to_white(patch)
+                patch.save(
+                    os.path.join(output_dir, f"{prefix}_{x0}_{y0}.png"),
+                    format="PNG",
+                )
+                count += 1
+                pbar.update(1)
+    return count
+
+
+# =================== Per-section processing (for Pool) =====================
+
+# These will be set once in main() before the pool is created
+_ARGS = {}
+
+
+def _process_section(section_dir: str) -> str:
+    """Process a single section directory. Designed for multiprocessing.Pool."""
+
+    patch_size = _ARGS["patch_size"]
+    overlap = _ARGS["overlap"]
+    magnification = _ARGS["magnification"]
+    input_filename = _ARGS["input_filename"]
+
+    # Locate input file
+    candidates = [
+        os.path.join(section_dir, f"{input_filename}.tif"),
+        os.path.join(section_dir, f"{input_filename}.tiff"),
+    ]
+    image_path = next((p for p in candidates if os.path.exists(p)), None)
+
+    if image_path is None:
+        return f"[SKIP] {section_dir}: {input_filename}.tif not found"
+
+    stride = patch_size - overlap
+    out_dir = os.path.join(
+        section_dir,
+        f"patches_{magnification}x_p{patch_size}_o{overlap}",
+    )
+
+    section_name = os.path.basename(section_dir)
+
+    t0 = time.time()
+    n = extract_patches(
+        image_path=image_path,
+        output_dir=out_dir,
+        patch_size=patch_size,
+        overlap=overlap,
+        prefix=section_name,
+    )
+    dt = time.time() - t0
+
+    return (
+        f"[OK] {section_name} | {magnification}x | "
+        f"stride={stride} | {n} patches | {dt:.1f}s → {out_dir}"
+    )
+
+
+# =============================== CLI =======================================
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Extract overlapping patches from normalized H&E images"
+    )
+    parser.add_argument(
+        "--input_dir",
+        type=str,
+        required=True,
+        help="Root directory containing section folders with Aligned-hne.tif files",
+    )
+    parser.add_argument(
+        "--input_filename",
+        type=str,
+        default="Aligned-hne",
+        help="Base filename of the normalized image (default: Aligned-hne)",
+    )
+    parser.add_argument(
+        "--patch_size", type=int, default=256, help="Patch size in pixels (default: 256)"
+    )
+    parser.add_argument(
+        "--overlap", type=int, default=64, help="Overlap in pixels (default: 64)"
+    )
+    parser.add_argument(
+        "--magnification",
+        type=int,
+        default=40,
+        choices=[20, 40],
+        help="Slide magnification. 40x recommended; 20x supported. (default: 40)",
+    )
+    parser.add_argument(
+        "--pattern",
+        type=str,
+        default="*",
+        help="Glob pattern to match section folders (default: '*')",
+    )
+    parser.add_argument(
+        "--workers", type=int, default=8, help="Number of parallel workers (default: 8)"
+    )
+
+    args = parser.parse_args()
+
+    if args.magnification == 20:
+        print(
+            "⚠️  20x magnification is supported but 40x is recommended for best "
+            "cell detection accuracy (especially small immune cells)."
+        )
+
+    # Collect section directories
+    section_dirs = sorted(
+        p
+        for p in glob.glob(os.path.join(args.input_dir, args.pattern))
+        if os.path.isdir(p)
+    )
+
+    if not section_dirs:
+        # Maybe input_dir itself contains the image directly
+        candidates = [
+            os.path.join(args.input_dir, f"{args.input_filename}.tif"),
+            os.path.join(args.input_dir, f"{args.input_filename}.tiff"),
+        ]
+        if any(os.path.exists(c) for c in candidates):
+            section_dirs = [args.input_dir]
+        else:
+            raise SystemExit(
+                f"No section folders matching '{args.pattern}' found in {args.input_dir}"
+            )
+
+    print(f"Found {len(section_dirs)} section(s) | {args.magnification}x | "
+          f"patch={args.patch_size} overlap={args.overlap}")
+
+    # Set global args for worker processes
+    global _ARGS
+    _ARGS = {
+        "patch_size": args.patch_size,
+        "overlap": args.overlap,
+        "magnification": args.magnification,
+        "input_filename": args.input_filename,
+    }
+
+    if args.workers <= 1 or len(section_dirs) == 1:
+        results = [_process_section(d) for d in tqdm(section_dirs, desc="Sections")]
+    else:
+        with Pool(processes=min(args.workers, len(section_dirs))) as pool:
+            results = list(
+                tqdm(
+                    pool.imap_unordered(_process_section, section_dirs),
+                    total=len(section_dirs),
+                    desc="Sections",
+                )
+            )
+
+    print("\n".join(results))
+
+
+if __name__ == "__main__":
+    main()

post_processing.py

@@ -0,0 +1,348 @@
+# -*- coding: utf-8 -*-
+# PostProcessing Pipeline
+#
+# Adapted from HoverNet
+# HoverNet Network (https://doi.org/10.1016/j.media.2019.101563)
+# Code Snippet adapted from HoverNet implementation (https://github.com/vqdang/hover_net)
+#
+# @ Fabian Hörst, fabian.hoerst@uk-essen.de
+# Institute for Artifical Intelligence in Medicine,
+# University Medicine Essen
+
+
+import warnings
+from typing import Tuple, Literal
+
+import cv2
+import numpy as np
+from scipy.ndimage import measurements
+from scipy.ndimage.morphology import binary_fill_holes
+from skimage.segmentation import watershed
+import torch
+
+# import sys
+# sys.path.append("/home01/k123a01/CellViTR/cell_segmentation/utils/")
+from tools import get_bounding_box, remove_small_objects
+
+
+def noop(*args, **kargs):
+    pass
+
+
+warnings.warn = noop
+
+
+class DetectionCellPostProcessor:
+    def __init__(
+        self,
+        nr_types: int = None,
+        magnification: Literal[20, 40] = 40,
+        gt: bool = False,
+    ) -> None:
+        """DetectionCellPostProcessor for postprocessing prediction maps and get detected cells
+
+        Args:
+            nr_types (int, optional): Number of cell types, including background (background = 0). Defaults to None.
+            magnification (Literal[20, 40], optional): Which magnification the data has. Defaults to 40.
+            gt (bool, optional): If this is gt data (used that we do not suppress tiny cells that may be noise in a prediction map).
+                Defaults to False.
+
+        Raises:
+            NotImplementedError: Unknown magnification
+        """
+        self.nr_types = nr_types
+        self.magnification = magnification
+        self.gt = gt
+
+        if magnification == 40:
+            self.object_size = 10
+            self.k_size = 21
+        elif magnification == 20:
+            self.object_size = 3  # 3 or 40, we used 5
+            self.k_size = 11  # 11 or 41, we used 13
+        else:
+            raise NotImplementedError("Unknown magnification")
+        if gt:  # to not supress something in gt!
+            self.object_size = 100
+            self.k_size = 21
+
+    def post_process_cell_segmentation(
+        self,
+        pred_map: np.ndarray,
+    ) -> Tuple[np.ndarray, dict]:
+        """Post processing of one image tile
+
+        Args:
+            pred_map (np.ndarray): Combined output of tp, np and hv branches, in the same order. Shape: (H, W, 4)
+
+        Returns:
+            Tuple[np.ndarray, dict]:
+                np.ndarray: Instance map for one image. Each nuclei has own integer. Shape: (H, W)
+                dict: Instance dictionary. Main Key is the nuclei instance number (int), with a dict as value.
+                    For each instance, the dictionary contains the keys: bbox (bounding box), centroid (centroid coordinates),
+                    contour, type_prob (probability), type (nuclei type)
+        """
+        if self.nr_types is not None:
+            pred_type = pred_map[..., 1:2]
+            pred_inst = pred_map[..., 2:]
+            pred_type = pred_type.astype(np.int32)
+            # print('pred_type',pred_type)
+        else:
+            pred_inst = pred_map
+
+        pred_inst = np.squeeze(pred_inst)
+        pred_inst = self.__proc_np_hv(
+            pred_inst, object_size=self.object_size, ksize=self.k_size
+        )
+        # print('pred_inst',pred_inst)
+        inst_id_list = np.unique(pred_inst)[1:]  # exlcude background
+        # print('inst_id_list',inst_id_list)
+        inst_info_dict = {}
+        for inst_id in inst_id_list:
+            inst_map = pred_inst == inst_id
+            rmin, rmax, cmin, cmax = get_bounding_box(inst_map)
+            inst_bbox = np.array([[rmin, cmin], [rmax, cmax]])
+            inst_map = inst_map[
+                inst_bbox[0][0] : inst_bbox[1][0], inst_bbox[0][1] : inst_bbox[1][1]
+            ]
+            inst_map = inst_map.astype(np.uint8)
+            inst_moment = cv2.moments(inst_map)
+            inst_contour = cv2.findContours(
+                inst_map, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE
+            )
+            # * opencv protocol format may break
+            inst_contour = np.squeeze(inst_contour[0][0].astype("int32"))
+            # < 3 points dont make a contour, so skip, likely artifact too
+            # as the contours obtained via approximation => too small or sthg
+            if inst_contour.shape[0] < 3:
+                continue
+            if len(inst_contour.shape) != 2:
+                continue  # ! check for trickery shape
+            inst_centroid = [
+                (inst_moment["m10"] / inst_moment["m00"]),
+                (inst_moment["m01"] / inst_moment["m00"]),
+            ]
+            inst_centroid = np.array(inst_centroid)
+            inst_contour[:, 0] += inst_bbox[0][1]  # X
+            inst_contour[:, 1] += inst_bbox[0][0]  # Y
+            inst_centroid[0] += inst_bbox[0][1]  # X
+            inst_centroid[1] += inst_bbox[0][0]  # Y
+            inst_info_dict[inst_id] = {  # inst_id should start at 1
+                "bbox": inst_bbox,
+                "centroid": inst_centroid,
+                "contour": inst_contour,
+                "type_prob": None,
+                "type": None,
+                "all_type_prob": []
+            }
+
+        #### * Get class of each instance id, stored at index id-1 (inst_id = number of deteced nucleus)
+        for inst_id in list(inst_info_dict.keys()):
+            rmin, cmin, rmax, cmax = (inst_info_dict[inst_id]["bbox"]).flatten()
+            inst_map_crop = pred_inst[rmin:rmax, cmin:cmax]
+            inst_type_crop = pred_type[rmin:rmax, cmin:cmax]
+            inst_map_crop = inst_map_crop == inst_id
+            inst_type = inst_type_crop[inst_map_crop]
+            type_list, type_pixels = np.unique(inst_type, return_counts=True)
+            type_list = list(zip(type_list, type_pixels))
+
+            type_probs = {}  # 각 인스턴스에 대한 cell type의 확률
+            total_pixels = np.sum(inst_map_crop) + 1.0e-6  # 0으로 나누지 않게 하기 위한 작은 값 추가
+            for cell_type, pixel_count in type_list:
+                type_probs[cell_type] = float(pixel_count / total_pixels)
+
+            all_type_prob = [type_probs.get(i, 0.0) for i in range(self.nr_types)]  # 전체 cell type에 대한 확률
+            inst_info_dict[inst_id]["all_type_prob"] = all_type_prob
+            #print("inst all probs :", inst_info_dict[inst_id]["all_type_prob"])
+
+            type_list = sorted(type_list, key=lambda x: x[1], reverse=True)
+            inst_type = type_list[0][0]
+            if inst_type == 0:  # ! pick the 2nd most dominant if exist
+                if len(type_list) > 1:
+                    inst_type = type_list[1][0]
+            type_dict = {v[0]: v[1] for v in type_list}
+            type_prob = type_dict[inst_type] / (np.sum(inst_map_crop) + 1.0e-6)
+            inst_info_dict[inst_id]["type"] = int(inst_type)
+            inst_info_dict[inst_id]["type_prob"] = float(type_prob)
+
+        return pred_inst, inst_info_dict
+
+    def __proc_np_hv(
+        self, pred: np.ndarray, object_size: int = 10, ksize: int = 21
+    ) -> np.ndarray:
+        """Process Nuclei Prediction with XY Coordinate Map and generate instance map (each instance has unique integer)
+
+        Separate Instances (also overlapping ones) from binary nuclei map and hv map by using morphological operations and watershed
+
+        Args:
+            pred (np.ndarray): Prediction output, assuming. Shape: (H, W, 3)
+                * channel 0 contain probability map of nuclei
+                * channel 1 containing the regressed X-map
+                * channel 2 containing the regressed Y-map
+            object_size (int, optional): Smallest oject size for filtering. Defaults to 10
+            k_size (int, optional): Sobel Kernel size. Defaults to 21
+        Returns:
+            np.ndarray: Instance map for one image. Each nuclei has own integer. Shape: (H, W)
+        """
+        pred = np.array(pred, dtype=np.float32)
+
+        blb_raw = pred[..., 0]
+        h_dir_raw = pred[..., 1]
+        v_dir_raw = pred[..., 2]
+
+        # processing
+        blb = np.array(blb_raw >= 0.5, dtype=np.int32)
+
+        blb = measurements.label(blb)[0]  # ndimage.label(blb)[0]
+        blb = remove_small_objects(blb, min_size=10)  # 10
+        blb[blb > 0] = 1  # background is 0 already
+
+        h_dir = cv2.normalize(
+            h_dir_raw,
+            None,
+            alpha=0,
+            beta=1,
+            norm_type=cv2.NORM_MINMAX,
+            dtype=cv2.CV_32F,
+        )
+        v_dir = cv2.normalize(
+            v_dir_raw,
+            None,
+            alpha=0,
+            beta=1,
+            norm_type=cv2.NORM_MINMAX,
+            dtype=cv2.CV_32F,
+        )
+
+        # ksize = int((20 * scale_factor) + 1) # 21 vs 41
+        # obj_size = math.ceil(10 * (scale_factor**2)) #10 vs 40
+
+        sobelh = cv2.Sobel(h_dir, cv2.CV_64F, 1, 0, ksize=ksize)
+        sobelv = cv2.Sobel(v_dir, cv2.CV_64F, 0, 1, ksize=ksize)
+
+        sobelh = 1 - (
+            cv2.normalize(
+                sobelh,
+                None,
+                alpha=0,
+                beta=1,
+                norm_type=cv2.NORM_MINMAX,
+                dtype=cv2.CV_32F,
+            )
+        )
+        sobelv = 1 - (
+            cv2.normalize(
+                sobelv,
+                None,
+                alpha=0,
+                beta=1,
+                norm_type=cv2.NORM_MINMAX,
+                dtype=cv2.CV_32F,
+            )
+        )
+
+        overall = np.maximum(sobelh, sobelv)
+        overall = overall - (1 - blb)
+        overall[overall < 0] = 0
+
+        dist = (1.0 - overall) * blb
+        ## nuclei values form mountains so inverse to get basins
+        dist = -cv2.GaussianBlur(dist, (3, 3), 0)
+
+        overall = np.array(overall >= 0.4, dtype=np.int32)
+
+        marker = blb - overall
+        marker[marker < 0] = 0
+        marker = binary_fill_holes(marker).astype("uint8")
+        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
+        marker = cv2.morphologyEx(marker, cv2.MORPH_OPEN, kernel)
+        marker = measurements.label(marker)[0]
+        marker = remove_small_objects(marker, min_size=object_size)
+
+        proced_pred = watershed(dist, markers=marker, mask=blb)
+
+        return proced_pred
+
+
+def calculate_instances(
+    pred_types: torch.Tensor, pred_insts: torch.Tensor
+) -> list[dict]:
+    """Best used for GT
+
+    Args:
+        pred_types (torch.Tensor): Binary or type map ground-truth.
+             Shape must be (B, C, H, W) with C=1 for binary or num_nuclei_types for multi-class.
+        pred_insts (torch.Tensor): Ground-Truth instance map with shape (B, H, W)
+
+    Returns:
+        list[dict]: Dictionary with nuclei informations, output similar to post_process_cell_segmentation
+    """
+    type_preds = []
+    pred_types = pred_types.permute(0, 2, 3, 1)
+    for i in range(pred_types.shape[0]):
+        pred_type = torch.argmax(pred_types, dim=-1)[i].detach().cpu().numpy()
+        pred_inst = pred_insts[i].detach().cpu().numpy()
+        inst_id_list = np.unique(pred_inst)[1:]  # exlcude background
+        inst_info_dict = {}
+        for inst_id in inst_id_list:
+            inst_map = pred_inst == inst_id
+            rmin, rmax, cmin, cmax = get_bounding_box(inst_map)
+            inst_bbox = np.array([[rmin, cmin], [rmax, cmax]])
+            inst_map = inst_map[
+                inst_bbox[0][0] : inst_bbox[1][0], inst_bbox[0][1] : inst_bbox[1][1]
+            ]
+            inst_map = inst_map.astype(np.uint8)
+            inst_moment = cv2.moments(inst_map)
+            inst_contour = cv2.findContours(
+                inst_map, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE
+            )
+            # * opencv protocol format may break
+            inst_contour = np.squeeze(inst_contour[0][0].astype("int32"))
+            # < 3 points dont make a contour, so skip, likely artifact too
+            # as the contours obtained via approximation => too small or sthg
+            if inst_contour.shape[0] < 3:
+                continue
+            if len(inst_contour.shape) != 2:
+                continue  # ! check for trickery shape
+            inst_centroid = [
+                (inst_moment["m10"] / inst_moment["m00"]),
+                (inst_moment["m01"] / inst_moment["m00"]),
+            ]
+            inst_centroid = np.array(inst_centroid)
+            inst_contour[:, 0] += inst_bbox[0][1]  # X
+            inst_contour[:, 1] += inst_bbox[0][0]  # Y
+            inst_centroid[0] += inst_bbox[0][1]  # X
+            inst_centroid[1] += inst_bbox[0][0]  # Y
+            inst_info_dict[inst_id] = {  # inst_id should start at 1
+                "bbox": inst_bbox,
+                "centroid": inst_centroid,
+                "contour": inst_contour,
+                "type_prob": None,
+                "type": None,
+            }
+        #### * Get class of each instance id, stored at index id-1 (inst_id = number of deteced nucleus)
+        for inst_id in list(inst_info_dict.keys()):
+            rmin, cmin, rmax, cmax = (inst_info_dict[inst_id]["bbox"]).flatten()
+            inst_map_crop = pred_inst[rmin:rmax, cmin:cmax]
+            inst_type_crop = pred_type[rmin:rmax, cmin:cmax]
+            inst_map_crop = inst_map_crop == inst_id
+            inst_type = inst_type_crop[inst_map_crop]
+            type_list, type_pixels = np.unique(inst_type, return_counts=True)
+            type_list = list(zip(type_list, type_pixels))
+            type_list = sorted(type_list, key=lambda x: x[1], reverse=True)
+            inst_type = type_list[0][0]
+            if inst_type == 0:  # ! pick the 2nd most dominant if exist
+                if len(type_list) > 1:
+                    inst_type = type_list[1][0]
+            type_dict = {v[0]: v[1] for v in type_list}
+            type_prob = type_dict[inst_type] / (np.sum(inst_map_crop) + 1.0e-6)
+            inst_info_dict[inst_id]["type"] = int(inst_type)
+            inst_info_dict[inst_id]["type_prob"] = float(type_prob)
+        type_preds.append(inst_info_dict)
+
+    return type_preds
+
+
+
+
+

tools.py

@@ -0,0 +1,400 @@
+# -*- coding: utf-8 -*-
+# Helpful functions Pipeline
+#
+# Adapted from HoverNet
+# HoverNet Network (https://doi.org/10.1016/j.media.2019.101563)
+# Code Snippet adapted from HoverNet implementation (https://github.com/vqdang/hover_net)
+#
+# @ Fabian Hörst, fabian.hoerst@uk-essen.de
+# Institute for Artifical Intelligence in Medicine,
+# University Medicine Essen
+
+
+import math
+from typing import Tuple
+
+import numpy as np
+import scipy
+from numba import njit, prange
+from scipy import ndimage
+from scipy.optimize import linear_sum_assignment
+from skimage.draw import polygon
+
+
+def get_bounding_box(img):
+    """Get bounding box coordinate information."""
+    rows = np.any(img, axis=1)
+    cols = np.any(img, axis=0)
+    rmin, rmax = np.where(rows)[0][[0, -1]]
+    cmin, cmax = np.where(cols)[0][[0, -1]]
+    # due to python indexing, need to add 1 to max
+    # else accessing will be 1px in the box, not out
+    rmax += 1
+    cmax += 1
+    return [rmin, rmax, cmin, cmax]
+
+
+@njit
+def cropping_center(x, crop_shape, batch=False):
+    """Crop an input image at the centre.
+
+    Args:
+        x: input array
+        crop_shape: dimensions of cropped array
+
+    Returns:
+        x: cropped array
+
+    """
+    orig_shape = x.shape
+    if not batch:
+        h0 = int((orig_shape[0] - crop_shape[0]) * 0.5)
+        w0 = int((orig_shape[1] - crop_shape[1]) * 0.5)
+        x = x[h0 : h0 + crop_shape[0], w0 : w0 + crop_shape[1], ...]
+    else:
+        h0 = int((orig_shape[1] - crop_shape[0]) * 0.5)
+        w0 = int((orig_shape[2] - crop_shape[1]) * 0.5)
+        x = x[:, h0 : h0 + crop_shape[0], w0 : w0 + crop_shape[1], ...]
+    return x
+
+
+def remove_small_objects(pred, min_size=64, connectivity=1):
+    """Remove connected components smaller than the specified size.
+
+    This function is taken from skimage.morphology.remove_small_objects, but the warning
+    is removed when a single label is provided.
+
+    Args:
+        pred: input labelled array
+        min_size: minimum size of instance in output array
+        connectivity: The connectivity defining the neighborhood of a pixel.
+
+    Returns:
+        out: output array with instances removed under min_size
+
+    """
+    out = pred
+
+    if min_size == 0:  # shortcut for efficiency
+        return out
+
+    if out.dtype == bool:
+        selem = ndimage.generate_binary_structure(pred.ndim, connectivity)
+        ccs = np.zeros_like(pred, dtype=np.int32)
+        ndimage.label(pred, selem, output=ccs)
+    else:
+        ccs = out
+
+    try:
+        component_sizes = np.bincount(ccs.ravel())
+    except ValueError:
+        raise ValueError(
+            "Negative value labels are not supported. Try "
+            "relabeling the input with `scipy.ndimage.label` or "
+            "`skimage.morphology.label`."
+        )
+
+    too_small = component_sizes < min_size
+    too_small_mask = too_small[ccs]
+    out[too_small_mask] = 0
+
+    return out
+
+
+def pair_coordinates(
+    setA: np.ndarray, setB: np.ndarray, radius: float
+) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """Use the Munkres or Kuhn-Munkres algorithm to find the most optimal
+    unique pairing (largest possible match) when pairing points in set B
+    against points in set A, using distance as cost function.
+
+    Args:
+        setA (np.ndarray): np.array (float32) of size Nx2 contains the of XY coordinate
+                    of N different points
+        setB (np.ndarray): np.array (float32) of size Nx2 contains the of XY coordinate
+                    of N different points
+        radius (float): valid area around a point in setA to consider
+                a given coordinate in setB a candidate for match
+
+    Returns:
+        Tuple[np.ndarray, np.ndarray, np.ndarray]:
+            pairing: pairing is an array of indices
+                where point at index pairing[0] in set A paired with point
+                in set B at index pairing[1]
+            unparedA: remaining point in set A unpaired
+            unparedB: remaining point in set B unpaired
+    """
+    # * Euclidean distance as the cost matrix
+    pair_distance = scipy.spatial.distance.cdist(setA, setB, metric="euclidean")
+
+    # * Munkres pairing with scipy library
+    # the algorithm return (row indices, matched column indices)
+    # if there is multiple same cost in a row, index of first occurence
+    # is return, thus the unique pairing is ensured
+    indicesA, paired_indicesB = linear_sum_assignment(pair_distance)
+
+    # extract the paired cost and remove instances
+    # outside of designated radius
+    pair_cost = pair_distance[indicesA, paired_indicesB]
+
+    pairedA = indicesA[pair_cost <= radius]
+    pairedB = paired_indicesB[pair_cost <= radius]
+
+    pairing = np.concatenate([pairedA[:, None], pairedB[:, None]], axis=-1)
+    unpairedA = np.delete(np.arange(setA.shape[0]), pairedA)
+    unpairedB = np.delete(np.arange(setB.shape[0]), pairedB)
+
+    return pairing, unpairedA, unpairedB
+
+
+def fix_duplicates(inst_map: np.ndarray) -> np.ndarray:
+    """Re-label duplicated instances in an instance labelled mask.
+
+    Parameters
+    ----------
+        inst_map : np.ndarray
+            Instance labelled mask. Shape (H, W).
+
+    Returns
+    -------
+        np.ndarray:
+            The instance labelled mask without duplicated indices.
+            Shape (H, W).
+    """
+    current_max_id = np.amax(inst_map)
+    inst_list = list(np.unique(inst_map))
+    if 0 in inst_list:
+        inst_list.remove(0)
+
+    for inst_id in inst_list:
+        inst = np.array(inst_map == inst_id, np.uint8)
+        remapped_ids = ndimage.label(inst)[0]
+        remapped_ids[remapped_ids > 1] += current_max_id
+        inst_map[remapped_ids > 1] = remapped_ids[remapped_ids > 1]
+        current_max_id = np.amax(inst_map)
+
+    return inst_map
+
+
+def polygons_to_label_coord(
+    coord: np.ndarray, shape: Tuple[int, int], labels: np.ndarray = None
+) -> np.ndarray:
+    """Render polygons to image given a shape.
+
+    Parameters
+    ----------
+        coord.shape : np.ndarray
+            Shape: (n_polys, n_rays)
+        shape : Tuple[int, int]
+            Shape of the output mask.
+        labels : np.ndarray, optional
+            Sorted indices of the centroids.
+
+    Returns
+    -------
+        np.ndarray:
+            Instance labelled mask. Shape: (H, W).
+    """
+    coord = np.asarray(coord)
+    if labels is None:
+        labels = np.arange(len(coord))
+
+    assert coord.ndim == 3 and coord.shape[1] == 2 and len(coord) == len(labels)
+
+    lbl = np.zeros(shape, np.int32)
+
+    for i, c in zip(labels, coord):
+        rr, cc = polygon(*c, shape)
+        lbl[rr, cc] = i + 1
+
+    return lbl
+
+
+def ray_angles(n_rays: int = 32):
+    """Get linearly spaced angles for rays."""
+    return np.linspace(0, 2 * np.pi, n_rays, endpoint=False)
+
+
+def dist_to_coord(
+    dist: np.ndarray, points: np.ndarray, scale_dist: Tuple[int, int] = (1, 1)
+) -> np.ndarray:
+    """Convert list of distances and centroids from polar to cartesian coordinates.
+
+    Parameters
+    ----------
+        dist : np.ndarray
+            The centerpoint pixels of the radial distance map. Shape (n_polys, n_rays).
+        points : np.ndarray
+            The centroids of the instances. Shape: (n_polys, 2).
+        scale_dist : Tuple[int, int], default=(1, 1)
+            Scaling factor.
+
+    Returns
+    -------
+        np.ndarray:
+            Cartesian cooridnates of the polygons. Shape (n_polys, 2, n_rays).
+    """
+    dist = np.asarray(dist)
+    points = np.asarray(points)
+    assert (
+        dist.ndim == 2
+        and points.ndim == 2
+        and len(dist) == len(points)
+        and points.shape[1] == 2
+        and len(scale_dist) == 2
+    )
+    n_rays = dist.shape[1]
+    phis = ray_angles(n_rays)
+    coord = (dist[:, np.newaxis] * np.array([np.sin(phis), np.cos(phis)])).astype(
+        np.float32
+    )
+    coord *= np.asarray(scale_dist).reshape(1, 2, 1)
+    coord += points[..., np.newaxis]
+    return coord
+
+
+def polygons_to_label(
+    dist: np.ndarray,
+    points: np.ndarray,
+    shape: Tuple[int, int],
+    prob: np.ndarray = None,
+    thresh: float = -np.inf,
+    scale_dist: Tuple[int, int] = (1, 1),
+) -> np.ndarray:
+    """Convert distances and center points to instance labelled mask.
+
+    Parameters
+    ----------
+        dist : np.ndarray
+            The centerpoint pixels of the radial distance map. Shape (n_polys, n_rays).
+        points : np.ndarray
+            The centroids of the instances. Shape: (n_polys, 2).
+        shape : Tuple[int, int]:
+            Shape of the output mask.
+        prob : np.ndarray, optional
+            The centerpoint pixels of the regressed distance transform.
+            Shape: (n_polys, n_rays).
+        thresh : float, default=-np.inf
+            Threshold for the regressed distance transform.
+        scale_dist : Tuple[int, int], default=(1, 1)
+            Scaling factor.
+
+    Returns
+    -------
+        np.ndarray:
+            Instance labelled mask. Shape (H, W).
+    """
+    dist = np.asarray(dist)
+    points = np.asarray(points)
+    prob = np.inf * np.ones(len(points)) if prob is None else np.asarray(prob)
+
+    assert dist.ndim == 2 and points.ndim == 2 and len(dist) == len(points)
+    assert len(points) == len(prob) and points.shape[1] == 2 and prob.ndim == 1
+
+    ind = prob > thresh
+    points = points[ind]
+    dist = dist[ind]
+    prob = prob[ind]
+
+    ind = np.argsort(prob, kind="stable")
+    points = points[ind]
+    dist = dist[ind]
+
+    coord = dist_to_coord(dist, points, scale_dist=scale_dist)
+
+    return polygons_to_label_coord(coord, shape=shape, labels=ind)
+
+
+@njit(cache=True, fastmath=True)
+def intersection(boxA: np.ndarray, boxB: np.ndarray):
+    """Compute area of intersection of two boxes.
+
+    Parameters
+    ----------
+        boxA : np.ndarray
+            First boxes
+        boxB : np.ndarray
+            Second box
+
+    Returns
+    -------
+        float64:
+            Area of intersection
+    """
+    xA = max(boxA[..., 0], boxB[..., 0])
+    xB = min(boxA[..., 2], boxB[..., 2])
+    dx = xB - xA
+    if dx <= 0:
+        return 0.0
+
+    yA = max(boxA[..., 1], boxB[..., 1])
+    yB = min(boxA[..., 3], boxB[..., 3])
+    dy = yB - yA
+    if dy <= 0.0:
+        return 0.0
+
+    return dx * dy
+
+
+@njit(parallel=True)
+def get_bboxes(
+    dist: np.ndarray, points: np.ndarray
+) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, int]:
+    """Get bounding boxes from the non-zero pixels of the radial distance maps.
+
+    This is basically a translation from the stardist repo cpp code to python
+
+    NOTE: jit compiled and parallelized with numba.
+
+    Parameters
+    ----------
+        dist : np.ndarray
+            The non-zero values of the radial distance maps. Shape: (n_nonzero, n_rays).
+        points : np.ndarray
+            The yx-coordinates of the non-zero points. Shape (n_nonzero, 2).
+
+    Returns
+    -------
+    Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, int]:
+        Returns the x0, y0, x1, y1 bbox coordinates, bbox areas and the maximum
+        radial distance in the image.
+    """
+    n_polys = dist.shape[0]
+    n_rays = dist.shape[1]
+
+    bbox_x1 = np.zeros(n_polys)
+    bbox_x2 = np.zeros(n_polys)
+    bbox_y1 = np.zeros(n_polys)
+    bbox_y2 = np.zeros(n_polys)
+
+    areas = np.zeros(n_polys)
+    angle_pi = 2 * math.pi / n_rays
+    max_dist = 0
+
+    for i in prange(n_polys):
+        max_radius_outer = 0
+        py = points[i, 0]
+        px = points[i, 1]
+
+        for k in range(n_rays):
+            d = dist[i, k]
+            y = py + d * np.sin(angle_pi * k)
+            x = px + d * np.cos(angle_pi * k)
+
+            if k == 0:
+                bbox_x1[i] = x
+                bbox_x2[i] = x
+                bbox_y1[i] = y
+                bbox_y2[i] = y
+            else:
+                bbox_x1[i] = min(x, bbox_x1[i])
+                bbox_x2[i] = max(x, bbox_x2[i])
+                bbox_y1[i] = min(y, bbox_y1[i])
+                bbox_y2[i] = max(y, bbox_y2[i])
+
+            max_radius_outer = max(d, max_radius_outer)
+
+        areas[i] = (bbox_x2[i] - bbox_x1[i]) * (bbox_y2[i] - bbox_y1[i])
+        max_dist = max(max_dist, max_radius_outer)
+
+    return bbox_x1, bbox_y1, bbox_x2, bbox_y2, areas, max_dist