import threading
from pathlib import Path
from typing import Union, Optional, Any, List, Dict, Tuple, Callable
import numpy as np
import pandas as pd
from skimage.measure import regionprops_table, label
from skimage.transform import resize
from tqdm import tqdm
from celldetective.utils.image_loaders import load_frames
from scipy.ndimage import binary_fill_holes
from scipy.ndimage import find_objects
import concurrent.futures
[docs]
def fill_label_holes(lbl_img: np.ndarray, **kwargs: Any) -> np.ndarray:
"""
Fill small holes in label image.
from https://github.com/stardist/stardist/blob/main/stardist/utils.py
Parameters
----------
lbl_img : ndarray
Label image.
**kwargs : dict
Additional arguments for `scipy.ndimage.binary_fill_holes`.
Returns
-------
ndarray
Label image with filled holes.
"""
# TODO: refactor 'fill_label_holes' and 'edt_prob' to share code
def grow(
sl: Tuple[slice, ...], interior: List[Tuple[bool, bool]]
) -> Tuple[slice, ...]:
"""
Grow slice.
Parameters
----------
sl : tuple
Slice tuple.
interior : list
List of interior flags.
"""
return tuple(
slice(s.start - int(w[0]), s.stop + int(w[1])) for s, w in zip(sl, interior)
)
def shrink(interior: List[Tuple[bool, bool]]) -> Tuple[slice, ...]:
"""
Shrink slice.
Parameters
----------
interior : list
List of interior flags.
"""
return tuple(slice(int(w[0]), (-1 if w[1] else None)) for w in interior)
objects = find_objects(lbl_img)
lbl_img_filled = np.zeros_like(lbl_img)
for i, sl in enumerate(objects, 1):
if sl is None:
continue
interior = [(s.start > 0, s.stop < sz) for s, sz in zip(sl, lbl_img.shape)]
shrink_slice = shrink(interior)
grown_mask = lbl_img[grow(sl, interior)] == i
mask_filled = binary_fill_holes(grown_mask, **kwargs)[shrink_slice]
lbl_img_filled[sl][mask_filled] = i
if lbl_img.min() < 0:
# preserve (and fill holes in) negative labels ('find_objects' ignores these)
lbl_neg_filled = -fill_label_holes(-np.minimum(lbl_img, 0))
mask = lbl_neg_filled < 0
lbl_img_filled[mask] = lbl_neg_filled[mask]
return lbl_img_filled
def _check_label_dims(
lbl: np.ndarray,
file: Optional[Union[str, Path]] = None,
template: Optional[np.ndarray] = None,
) -> np.ndarray:
"""
Check and resize label image to match template dimensions.
Parameters
----------
lbl : ndarray
Label image.
file : str, optional
Path to the file to load as template. Default is None.
template : ndarray, optional
Template image. Default is None.
Returns
-------
ndarray
Resized label image.
"""
if file is not None:
template = load_frames(0, file, scale=1, normalize_input=False)
elif template is not None:
template = template
else:
return lbl
if lbl.shape != template.shape[:2]:
lbl = resize(lbl, template.shape[:2], order=0)
return lbl
[docs]
def auto_correct_masks(
masks: np.ndarray,
bbox_factor: float = 1.75,
min_area: int = 9,
fill_labels: bool = False,
) -> np.ndarray:
"""
Correct segmentation masks to ensure consistency and remove anomalies.
This function processes a labeled mask image to correct anomalies and reassign labels.
It performs the following operations:
1. Corrects negative mask values by taking their absolute values.
2. Identifies and corrects segmented objects with a bounding box area that is disproportionately
larger than the actual object area. This indicates potential segmentation errors where separate objects
share the same label.
3. Removes small objects that are considered noise (default threshold is an area of less than 9 pixels).
4. Reorders the labels so they are consecutive from 1 up to the number of remaining objects (to avoid encoding errors).
Parameters
----------
masks : np.ndarray
A 2D array representing the segmented mask image with labeled regions. Each unique value
in the array represents a different object or cell.
bbox_factor : float, optional
A factor on cell area that is compared directly to the bounding box area of the cell, to detect remote cells
sharing a same label value. The default is `1.75`.
min_area : int, optional
Discard cells that have an area smaller than this minimum area (px²). The default is `9` (3x3 pixels).
fill_labels : bool, optional
Fill holes within cell masks automatically. The default is `False`.
Returns
-------
clean_labels : np.ndarray
A corrected version of the input mask, with anomalies corrected, small objects removed,
and labels reordered to be consecutive integers.
Notes
-----
- This function is useful for post-processing segmentation outputs to ensure high-quality
object detection, particularly in applications such as cell segmentation in microscopy images.
- The function assumes that the input masks contain integer labels and that the background
is represented by 0.
Examples
--------
>>> masks = np.array([[0, 0, 1, 1], [0, 2, 2, 1], [0, 2, 0, 0]])
>>> corrected_masks = auto_correct_masks(masks)
>>> corrected_masks
array([[0, 0, 1, 1],
[0, 2, 2, 1],
[0, 2, 0, 0]])
"""
assert masks.ndim == 2, "`masks` should be a 2D numpy array..."
# Avoid negative mask values
masks[masks < 0] = np.abs(masks[masks < 0])
props = pd.DataFrame(
regionprops_table(masks, properties=("label", "area", "area_bbox"))
)
max_lbl = props["label"].max()
corrected_lbl = masks.copy() # .astype(int)
for cell in props["label"].unique():
bbox_area = props.loc[props["label"] == cell, "area_bbox"].values
area = props.loc[props["label"] == cell, "area"].values
if bbox_area > bbox_factor * area: # condition for anomaly
lbl = masks == cell
lbl = lbl.astype(int)
relabelled = label(lbl, connectivity=2)
relabelled += max_lbl
relabelled[np.where(lbl == 0)] = 0
corrected_lbl[np.where(relabelled != 0)] = relabelled[
np.where(relabelled != 0)
]
max_lbl = np.amax(corrected_lbl)
# Second routine to eliminate objects too small
props2 = pd.DataFrame(
regionprops_table(corrected_lbl, properties=("label", "area", "area_bbox"))
)
for cell in props2["label"].unique():
area = props2.loc[props2["label"] == cell, "area"].values
lbl = corrected_lbl == cell
if area < min_area:
corrected_lbl[lbl] = 0
# Additionnal routine to reorder labels from 1 to number of cells
label_ids = np.unique(corrected_lbl)[1:]
clean_labels = corrected_lbl.copy()
for k, lbl in enumerate(label_ids):
clean_labels[corrected_lbl == lbl] = k + 1
clean_labels = clean_labels.astype(int)
if fill_labels:
clean_labels = fill_label_holes(clean_labels)
return clean_labels
[docs]
def relabel_segmentation(
labels: np.ndarray,
df: pd.DataFrame,
exclude_nans: bool = True,
column_labels: Dict[str, str] = {
"track": "TRACK_ID",
"frame": "FRAME",
"y": "POSITION_Y",
"x": "POSITION_X",
"label": "class_id",
},
threads: int = 1,
progress_callback: Optional[Callable[[float], bool]] = None,
) -> Optional[np.ndarray]:
"""
Relabel the segmentation labels with the tracking IDs from the tracks.
The function reassigns the mask value for each cell with the associated `TRACK_ID`, if it exists
in the trajectory table (`df`). If no track uses the cell mask, a new track with a single point
is generated on the fly (max of `TRACK_ID` values + i, for i=0 to N such cells). It supports
multithreaded processing for faster execution on large datasets.
Parameters
----------
labels : np.ndarray
A (TYX) array where each frame contains a 2D segmentation mask. Each unique
non-zero integer represents a labeled object.
df : pandas.DataFrame
A DataFrame containing tracking information with columns
specified in `column_labels`. Must include at least frame, track ID, and object ID.
exclude_nans : bool, optional
Whether to exclude rows in `df` with NaN values in the column specified by
`column_labels['label']`. Default is `True`.
column_labels : dict, optional
A dictionary specifying the column names in `df`. Default is:
- `'track'`: Track ID column name (`"TRACK_ID"`)
- `'frame'`: Frame column name (`"FRAME"`)
- `'y'`: Y-coordinate column name (`"POSITION_Y"`)
- `'x'`: X-coordinate column name (`"POSITION_X"`)
- `'label'`: Object ID column name (`"class_id"`)
threads : int, optional
Number of threads to use for multithreaded processing. Default is `1`.
progress_callback : callable, optional
A function to report progress. Should accept a single float argument (0-100) and return True to continue or False to cancel. Default is None.
Returns
-------
np.ndarray
A new (TYX) array with the same shape as `labels`, where objects are relabeled
according to their tracking identity in `df`.
Notes
-----
- For frames where labeled objects in `labels` do not match any entries in the `df`,
new track IDs are generated for the unmatched labels.
- The relabeling process maintains synchronization across threads using a shared
counter for generating unique track IDs.
Examples
--------
Relabel segmentation using tracking data:
>>> labels = np.random.randint(0, 5, (10, 100, 100))
>>> df = pd.DataFrame({
... "TRACK_ID": [1, 2, 1, 2],
... "FRAME": [0, 0, 1, 1],
... "class_id": [1, 2, 1, 2],
... })
>>> new_labels = relabel_segmentation(labels, df, threads=2)
Done.
Use custom column labels and exclude rows with NaNs:
>>> column_labels = {
... 'track': "track_id",
... 'frame': "time",
... 'label': "object_id"
... }
>>> new_labels = relabel_segmentation(labels, df, column_labels=column_labels, exclude_nans=True)
Done.
"""
n_threads = threads
df = df.sort_values(by=[column_labels["track"], column_labels["frame"]])
if exclude_nans:
df = df.dropna(subset=column_labels["label"])
new_labels = np.zeros_like(labels)
shared_data = {"s": 0}
# Progress tracking
shared_progress = {"val": 0, "lock": threading.Lock()}
total_frames = len(df[column_labels["frame"]].dropna().unique())
def rewrite_labels(indices: List[int]) -> None:
"""
Rewrite labels for a batch of frames.
Parameters
----------
indices : list
List of frame indices to process.
"""
all_track_ids = df[column_labels["track"]].dropna().unique()
# Check for cancellation
if progress_callback:
with shared_progress["lock"]:
if shared_progress.get("cancelled", False):
return
disable_tqdm = progress_callback is not None
for t in tqdm(indices, disable=disable_tqdm):
# Cancellation check inside loop
if progress_callback:
with shared_progress["lock"]:
if shared_progress.get("cancelled", False):
return
shared_progress["val"] += 1
p = int((shared_progress["val"] / total_frames) * 100)
if not progress_callback(p):
with shared_progress["lock"]:
shared_progress["cancelled"] = True
return
f = int(t)
cells = df.loc[
df[column_labels["frame"]] == f,
[column_labels["track"], column_labels["label"]],
].to_numpy()
tracks_at_t = list(cells[:, 0])
identities = list(cells[:, 1])
labels_at_t = list(np.unique(labels[f]))
if 0 in labels_at_t:
labels_at_t.remove(0)
labels_not_in_df = [lbl for lbl in labels_at_t if lbl not in identities]
for lbl in labels_not_in_df:
with threading.Lock(): # Synchronize access to `shared_data["s"]`
track_id = max(all_track_ids) + shared_data["s"]
shared_data["s"] += 1
tracks_at_t.append(track_id)
identities.append(lbl)
# exclude NaN
tracks_at_t = np.array(tracks_at_t)
identities = np.array(identities)
tracks_at_t = tracks_at_t[identities == identities]
identities = identities[identities == identities]
for k in range(len(identities)):
# need routine to check values from labels not in class_id of this frame and add new track id
loc_i, loc_j = np.where(labels[f] == identities[k])
track_id = tracks_at_t[k]
if track_id == track_id:
new_labels[f, loc_i, loc_j] = round(track_id)
# Multithreading
indices = list(df[column_labels["frame"]].dropna().unique())
chunks = np.array_split(indices, n_threads)
with concurrent.futures.ThreadPoolExecutor(max_workers=threads) as executor:
results = executor.map(
rewrite_labels, chunks
) # list(map(lambda x: executor.submit(self.parallel_job, x), chunks))
try:
for i, return_value in enumerate(results):
# print(f"Thread {i} output check: ", return_value)
pass
except Exception as e:
print("Exception: ", e)
if shared_progress.get("cancelled", False):
print("Relabeling cancelled.")
return None
print("\nDone.")
return new_labels