"""Tracking dataset exporters"""
from __future__ import annotations
import shutil
from collections import deque
from pathlib import Path
from typing import Any
import networkx as nx
import numpy as np
import pandas as pd
import tifffile
from tqdm.auto import tqdm
from acia.base import BaseImage, Contour, ImageSequenceSource, Overlay
from acia.segm.rasterize import frame_label_mask
from acia.tracking import TrackingSource
from acia.utils import largest_polygon, multi_mask_to_polygons
[docs]
class CellTrackingChallengeDatasetGT:
"""
Utility class to create segmentation/tracking output in form of the Cell Tracking Challenge (CTC)
"""
[docs]
def __init__(self):
self.sources = []
[docs]
def add(self, content: tuple[ImageSequenceSource, TrackingSource]):
self.sources.append(content)
[docs]
def write(self, base_folder: str | Path = "data", offset=0):
"""Exports a ctc ground truth dataset
Args:
base_folder (str | Path, optional): _description_. Defaults to "data".
"""
base_folder = Path(base_folder)
base_folder.mkdir(exist_ok=True, parents=True)
for i, (image_source, tracking_source) in enumerate(self.sources):
mode = "GT"
image_dir = base_folder / f"{i + offset:02}"
ann_dir = base_folder / f"{i + offset:02}_{mode}"
if image_dir.exists():
shutil.rmtree(image_dir)
if ann_dir.exists():
shutil.rmtree(ann_dir)
image_dir.mkdir()
ann_dir.mkdir()
height, width = -1, -1
# save the images
for t, image in enumerate(image_source):
if isinstance(image, BaseImage):
image = image.raw
height, width = image.shape[:2]
tifffile.imwrite(str(image_dir / f"t{t:04}.tif"), image[..., 0])
tracking_helper = CTCTrackingHelper(
tracking_source.overlay, tracking_source.tracking_graph, height, width
)
ctc_masks, ctc_tracking_format = tracking_helper.to_ctc_format()
# store segmentation masks
seg_dir = ann_dir / "SEG"
seg_dir.mkdir(exist_ok=True)
for t, mask in enumerate(ctc_masks):
tifffile.imwrite(str(seg_dir / f"man_seg{t:04}.tif"), mask)
# store tracking masks
track_dir = ann_dir / "TRA"
track_dir.mkdir(exist_ok=True)
for t, mask in enumerate(ctc_masks):
tifffile.imwrite(str(track_dir / f"man_track{t:04}.tif"), mask)
with open(
str(track_dir / "man_track.txt"), "w", encoding="utf-8"
) as output_file:
output_file.write("\n".join(ctc_tracking_format))
[docs]
class CellTrackingDatasetResult:
"""
Utility class to create segmentation/tracking output in form of the Cell Tracking Challenge (CTC)
"""
[docs]
def __init__(self):
self.sources = []
[docs]
def add(self, content: tuple[TrackingSource, tuple[int, int]]):
self.sources.append(content)
[docs]
def write(self, base_folder: str | Path = "data", offset=0):
base_folder = Path(base_folder)
base_folder.mkdir(exist_ok=True, parents=True)
for i, (tracking_source, (height, width)) in enumerate(self.sources):
mode = "RES"
ann_dir = base_folder / f"{i + offset:02}_{mode}"
if ann_dir.exists():
shutil.rmtree(ann_dir)
ann_dir.mkdir()
tracking_helper = CTCTrackingHelper(
tracking_source.overlay, tracking_source.tracking_graph, height, width
)
ctc_masks, ctc_tracking_format = tracking_helper.to_ctc_format()
for t, mask in enumerate(ctc_masks):
tifffile.imwrite(str(ann_dir / f"mask{t:04}.tif"), mask)
with open(
str(ann_dir / "res_track.txt"), "w", encoding="utf-8"
) as output_file:
output_file.write("\n".join(ctc_tracking_format))
[docs]
class CTCTrackingHelper:
"""Helper class for the CTC format generation"""
[docs]
def __init__(
self, overlay: Overlay, tracking_graph: nx.DiGraph, height: int, width: int
):
"""Create a new tracking helper object
Args:
tracking_graph (nx.DiGraph): tracking graph consisting of contour nodes
"""
self.contour_lookup = {cont.id: cont for cont in overlay}
self.tracking_graph = tracking_graph
# compute the life-cycles of individual cells
self.life_cycles = CTCTrackingHelper.compute_life_cycles(self.tracking_graph)
# create lookup (cont id --> life cycle index)
self.life_cycle_lookup = CTCTrackingHelper.create_life_cycle_lookup(
self.life_cycles # type: ignore[arg-type]
)
self.overlay = overlay
self.height = height
self.width = width
[docs]
@staticmethod
def compute_life_cycles(tracking_graph: nx.DiGraph) -> list[list[str]]:
"""Track life cycles of contour observations (from birth to division).
Args:
tracking_graph (nx.DiGraph): tracking graph
Returns:
List[List[str]]: List of life cycles (consisting of a list of contour ids)
"""
start_nodes = deque(
filter(lambda n: tracking_graph.in_degree(n) == 0, tracking_graph.nodes)
)
life_cycles = []
while len(start_nodes) > 0:
node = start_nodes.pop()
life_cycle = [node]
while tracking_graph.out_degree(node) == 1:
node = next(tracking_graph.successors(node))
life_cycle.append(node)
life_cycles.append(life_cycle)
if tracking_graph.out_degree(node) > 1:
start_nodes += list(tracking_graph.successors(node))
return life_cycles
[docs]
@staticmethod
def create_life_cycle_lookup(life_cycles: list[list[Contour]]):
"""Computes a mapping from contour ids to life cycle ids
Args:
life_cycles (List[List[Contour]]): _description_
Returns:
_type_: _description_
"""
contour_life_cycle_lookup = {}
for i, life_cycle in enumerate(life_cycles):
for cont_id in life_cycle:
contour_life_cycle_lookup[cont_id] = (
i # life cycle enumeration starts with 1
)
return contour_life_cycle_lookup
[docs]
@staticmethod
def convert_overlay_to_ctc_mask(
overlay: Overlay,
contour_life_cycle_lookup: dict[str, int],
height: int,
width: int,
) -> np.ndarray:
"""Creates a ctc mask with correct numbering for a frame overlay
Args:
overlay (Overlay): overlay containing the contours
contour_life_cycle_lookup (Dict[str, int]): lookup for the life cycle
Returns:
np.ndarray: the ctc mask for the frame
"""
assert height > 0 and width > 0
# numbered by life cycle rather than by cont.label, so the labels are
# passed in explicitly; exact_polygons keeps the pixel-centre rule this
# exported mask has always had
contours = list(overlay)
life_cycle_ids = [
contour_life_cycle_lookup[cont.id] + 1 # lifecycle ids must start with 1
for cont in contours
]
return frame_label_mask(
contours,
height=height,
width=width,
labels=life_cycle_ids,
exact_polygons=True,
).astype(np.uint16)
@staticmethod
def __load_masks(mask_path: Path) -> Overlay:
"""Load CTC masks into an overlay
Args:
mask_path (Path): Path to the folder containing the masks
Returns:
Overlay: Overlay containing all cell detections
"""
mask_path = Path(mask_path)
# load masks
mask_files = sorted(mask_path.glob("man_seg*.tif"))
masks = [tifffile.imread(file) for file in mask_files]
all_polygons = map(
multi_mask_to_polygons, tqdm(masks, desc="Convert masks to polygon...")
)
contours = []
for frame, frame_polygons in enumerate(
tqdm(all_polygons, desc="Convert to overlay...")
):
for id, poly in frame_polygons:
# a label whose mask has disconnected components has no single
# outline; its largest part represents the object
poly = largest_polygon(poly)
if poly is None:
continue
cc = np.stack(poly.exterior.coords.xy, axis=-1)
contours.append(Contour(cc, -1, frame, f"{frame}_{id}"))
return Overlay(contours, frames=list(range(0, len(masks))))
@staticmethod
def __load_tracking(tracking_path: Path) -> nx.DiGraph:
"""Load tracking from CTC txt file
Args:
tracking_path (Path): path to the folder containing the man_track.txt file
Returns:
nx.DiGraph: tracking graph with nodes and edges of every cell detection
"""
track_annotation = pd.read_csv(
tracking_path / "man_track.txt",
delimiter=" ",
names=["track_id", "start_frame", "end_frame", "parent_id"],
header=None,
)
tracking_graph = nx.DiGraph()
lc_lookup = {}
for _, row in track_annotation.iterrows():
start_frame = row["start_frame"]
end_frame = row["end_frame"]
id = row["track_id"]
parent_id = row["parent_id"]
node_items = [
f"{frame}_{id}" for frame in range(start_frame, end_frame + 1)
]
# add trajectory nodes
for frame in range(start_frame, end_frame + 1):
lc_lookup[f"{frame}_{id}"] = id
tracking_graph.add_node(f"{frame}_{id}", frame=frame)
# add trajectory edges
for a, b in zip(node_items, node_items[1:], strict=False):
tracking_graph.add_edge(a, b)
# add edge to parent (0 means no parent)
if parent_id != 0:
tracking_graph.add_edge(
f"{start_frame - 1}_{parent_id}", f"{start_frame}_{id}"
)
return tracking_graph