diff --git a/src/napari_iohub/_view_tracks.py b/src/napari_iohub/_view_tracks.py index 2fd6598..622b23f 100644 --- a/src/napari_iohub/_view_tracks.py +++ b/src/napari_iohub/_view_tracks.py @@ -5,10 +5,11 @@ import typing from typing import TYPE_CHECKING -import numpy as np import pandas as pd from iohub.ngff import open_ome_zarr from magicgui import magic_factory +from ultrack.reader.napari_reader import read_csv +from xarray import open_zarr from napari_iohub._reader import fov_to_layers @@ -31,10 +32,9 @@ def _zarr_modes(label: str) -> dict[str, str]: def open_image_and_tracks( images_dataset: pathlib.Path, tracks_dataset: pathlib.Path, - features_dir: pathlib.Path, + features_dataset: pathlib.Path, fov_name: str, features_type: str, - features_name: str, expand_z_for_tracking_labels: bool = True, ) -> typing.List[napari.types.LayerDataTuple]: """ @@ -49,15 +49,12 @@ def open_image_and_tracks( tracks_dataset : pathlib.Path Path to the tracking labels dataset (HCS OME-Zarr). Potentially with a singleton Z dimension. - features_dir : pathlib.Path - Path to the directory containing the features. - Has the same directory structure as the HCS stores. + features_dataset : pathlib.Path + Path to the predicted embeddings. fov_name : str Name of the FOV to load, e.g. `"A/12/2"`. features_type : str Name of the subdirectory containing the feature files. - features_name : str - Name of the `.npy` feature file. expand_z_for_tracking_labels : bool Whether to expand the tracking labels to the Z dimension of the images. @@ -71,6 +68,12 @@ def open_image_and_tracks( image_plate = open_ome_zarr(images_dataset) image_fov = image_plate[fov_name] image_layers = fov_to_layers(image_fov) + _logger.info(f"Loading {features_type} from {str(features_dataset)}") + features = ( + open_zarr(features_dataset) + .set_index(sample=["fov_name", "track_id"])[features_type] + .sel(fov_name=fov_name) + ) _logger.info(f"Loading tracking labels from {tracks_dataset}") tracks_plate = open_ome_zarr(tracks_dataset) tracks_fov = tracks_plate[fov_name] @@ -81,17 +84,17 @@ def open_image_and_tracks( labels_layer[0][0] = labels_layer[0][0].repeat(image_z, axis=1) tracks_csv = next((tracks_dataset / fov_name).glob("*.csv")) _logger.info(f"Loading tracks from {str(tracks_csv)}") + tracks_layer = read_csv(tracks_csv) + image_layers.append(tracks_layer) df = pd.read_csv(tracks_csv) - tracks = df[["track_id", "t"]].rename( - columns={"track_id": "label", "t": "frame"} - ) - features_match = features_dir / fov_name / features_type / features_name - _logger.info(f"Loading features from {str(features_match)}") - features_data = np.load(features_match) - features = pd.DataFrame( - features_data, - columns=[f"feature_{i}" for i in range(features_data.shape[1])], + tracks = ( + df[["track_id", "t"]] + .rename(columns={"track_id": "label", "t": "frame"}) + .set_index("track_id") ) - labels_layer[1]["features"] = pd.concat([tracks, features], axis=1) + tracks = tracks[features["track_id"]].reset_index() + labels_layer[1]["features"] = pd.DataFrame( + index=tracks, data=features.values + ).reset_index() image_layers.append(labels_layer) return image_layers