datasets
dreem.datasets
¶
Data loading and preprocessing.
Modules:
| Name | Description |
|---|---|
base_dataset |
Module containing logic for loading datasets. |
cell_tracking_dataset |
Module containing cell tracking challenge dataset. |
data_utils |
Module containing helper functions for datasets. |
microscopy_dataset |
Module containing microscopy dataset. |
preprocessors |
Preprocessor classes for dataset preparation steps. |
sleap_dataset |
Module containing logic for loading sleap datasets. |
tracking_dataset |
Module containing Lightning module wrapper around all other datasets. |
Classes:
| Name | Description |
|---|---|
BaseDataset |
Base Dataset for microscopy and sleap datasets to override. |
CellTrackingDataset |
Dataset for loading cell tracking challenge data. |
MicroscopyDataset |
Dataset for loading Microscopy Data. |
SleapDataset |
Dataset for loading animal behavior data from sleap. |
TrackingDataset |
Lightning dataset used to load dataloaders for train, test and validation. |
BaseDataset
¶
Bases: Dataset
Base Dataset for microscopy and sleap datasets to override.
Methods:
| Name | Description |
|---|---|
__getitem__ |
Get an element of the dataset. |
__init__ |
Initialize Dataset. |
__len__ |
Get the size of the dataset. |
create_chunks_other |
Legacy chunking logic. Does not support unannotated segments. |
create_chunks_slp |
Get indexing for data. |
get_indices |
Retrieve label and frame indices given batch index. |
get_instances |
Build chunk of frames. |
no_batching_fn |
Collate function used to overwrite dataloader batching function. |
process_segments |
Process segments to stitch. |
Source code in dreem/datasets/base_dataset.py
class BaseDataset(Dataset):
"""Base Dataset for microscopy and sleap datasets to override."""
def __init__(
self,
label_files: list[str],
vid_files: list[str],
padding: int,
crop_size: Union[int, list[int]],
chunk: bool,
clip_length: int,
mode: str,
augmentations: dict | None = None,
n_chunks: int | float = 1.0,
seed: int | None = None,
gt_list: str | None = None,
):
"""Initialize Dataset.
Args:
label_files: a list of paths to label files. Should at least contain
detections for inference, detections + tracks for training.
vid_files: list of paths to video files.
padding: amount of padding around object crops
crop_size: the size of the object crops
chunk: whether or not to chunk the dataset into batches
clip_length: the number of frames in each chunk
mode: `train` or `val`. Determines whether this dataset is used for
training or validation. Currently doesn't affect dataset logic
augmentations: An optional dict mapping augmentations to parameters.
See subclasses for details.
n_chunks: Number of chunks to subsample from.
Can either a fraction of the dataset (ie (0,1.0]) or number of chunks
seed: set a seed for reproducibility
gt_list: An optional path to .txt file containing ground truth for
cell tracking challenge datasets.
"""
self.vid_files = vid_files
self.label_files = label_files
self.padding = padding
self.crop_size = crop_size
self.chunk = chunk
self.clip_length = clip_length
self.mode = mode
self.n_chunks = n_chunks
self.seed = seed
if self.seed is not None:
np.random.seed(self.seed)
if augmentations and self.mode == "train":
self.instance_dropout = augmentations.pop(
"InstanceDropout", {"p": 0.0, "n": 0}
)
self.node_dropout = data_utils.NodeDropout(
**augmentations.pop("NodeDropout", {"p": 0.0, "n": 0})
)
self.augmentations = data_utils.build_augmentations(augmentations)
else:
self.instance_dropout = {"p": 0.0, "n": 0}
self.node_dropout = data_utils.NodeDropout(p=0.0, n=0)
self.augmentations = None
# Initialize in subclasses
self.frame_idx = None
self.labels = None
self.gt_list = None
def process_segments(
self, i: int, segments_to_stitch: list[torch.Tensor], clip_length: int
) -> None:
"""Process segments to stitch.
Modifies state variables chunked_frame_idx and label_idx.
Args:
segments_to_stitch: list of segments to stitch
i: index of the video
clip_length: the number of frames in each chunk
Returns: None
"""
stitched_segment = torch.cat(segments_to_stitch)
frame_idx_split = torch.split(stitched_segment, clip_length)
self.chunked_frame_idx.extend(frame_idx_split)
self.label_idx.extend(len(frame_idx_split) * [i])
def create_chunks_slp(self) -> None:
"""Get indexing for data.
Creates both indexes for selecting dataset (label_idx) and frame in
dataset (chunked_frame_idx). If chunking is false, we index directly
using the frame ids. Setting chunking to true creates a list of lists
containing chunk frames for indexing. This is useful for computational
efficiency and data shuffling. To be called by subclass __init__()
"""
self.chunked_frame_idx, self.label_idx = [], []
# go through each slp file and create chunks that respect max_batching_gap
for i, slp_file in enumerate(self.label_files):
annotated_segments = self.annotated_segments[slp_file]
segments_to_stitch = []
prev_end = annotated_segments[0][1] # end of first segment
for start, end in annotated_segments:
# check if the start of current segment is within
# batching_max_gap of end of previous
if (
(int(start) - int(prev_end) < self.max_batching_gap)
or not self.chunk
): # also takes care of first segment as start < prev_end
segments_to_stitch.append(torch.arange(start, end + 1))
prev_end = end
else:
# stitch previous set of segments before creating a new chunk
self.process_segments(i, segments_to_stitch, self.clip_length)
# reset segments_to_stitch as we are starting a new chunk
segments_to_stitch = [torch.arange(start, end + 1)]
prev_end = end
if not self.chunk:
self.process_segments(
i, segments_to_stitch, self.labels[i].video.shape[0]
)
else:
# add last chunk after the loop
if segments_to_stitch:
self.process_segments(i, segments_to_stitch, self.clip_length)
if self.n_chunks > 0 and self.n_chunks <= 1.0:
n_chunks = int(self.n_chunks * len(self.chunked_frame_idx))
elif self.n_chunks <= len(self.chunked_frame_idx):
n_chunks = int(self.n_chunks)
else:
n_chunks = len(self.chunked_frame_idx)
if n_chunks > 0 and n_chunks < len(self.chunked_frame_idx):
sample_idx = np.random.choice(
np.arange(len(self.chunked_frame_idx)), n_chunks, replace=False
)
self.chunked_frame_idx = [self.chunked_frame_idx[i] for i in sample_idx]
self.label_idx = [self.label_idx[i] for i in sample_idx]
# workaround for empty batch bug (needs to be changed).
# Check for batch with with only 1/10 size of clip length.
# Arbitrary thresholds
remove_idx = []
for i, frame_chunk in enumerate(self.chunked_frame_idx):
if (
len(frame_chunk) <= min(int(self.clip_length / 10), 5)
# and frame_chunk[-1] % self.clip_length == 0
):
logger.warning(
f"Warning: Batch containing frames {frame_chunk} from video "
f"{self.vid_files[self.label_idx[i]]} has {len(frame_chunk)} frames. "
f"Removing to avoid empty batch possibility with failed frame loading"
)
remove_idx.append(i)
if len(remove_idx) > 0:
for i in sorted(remove_idx, reverse=True):
self.chunked_frame_idx.pop(i)
self.label_idx.pop(i)
def create_chunks_other(self) -> None:
"""Legacy chunking logic. Does not support unannotated segments.
Creates both indexes for selecting dataset (label_idx) and frame in
dataset (chunked_frame_idx). If chunking is false, we index directly
using the frame ids. Setting chunking to true creates a list of lists
containing chunk frames for indexing. This is useful for computational
efficiency and data shuffling. To be called by subclass __init__()
"""
if self.chunk:
self.chunked_frame_idx, self.label_idx = [], []
for i, frame_idx in enumerate(self.frame_idx):
frame_idx_split = torch.split(frame_idx, self.clip_length)
self.chunked_frame_idx.extend(frame_idx_split)
self.label_idx.extend(len(frame_idx_split) * [i])
if self.n_chunks > 0 and self.n_chunks <= 1.0:
n_chunks = int(self.n_chunks * len(self.chunked_frame_idx))
elif self.n_chunks <= len(self.chunked_frame_idx):
n_chunks = int(self.n_chunks)
else:
n_chunks = len(self.chunked_frame_idx)
if n_chunks > 0 and n_chunks < len(self.chunked_frame_idx):
sample_idx = np.random.choice(
np.arange(len(self.chunked_frame_idx)), n_chunks, replace=False
)
self.chunked_frame_idx = [self.chunked_frame_idx[i] for i in sample_idx]
self.label_idx = [self.label_idx[i] for i in sample_idx]
# workaround for empty batch bug (needs to be changed).
# Check for batch with with only 1/10 size of clip length.
# Arbitrary thresholds
remove_idx = []
for i, frame_chunk in enumerate(self.chunked_frame_idx):
if (
len(frame_chunk) <= min(int(self.clip_length / 10), 5)
# and frame_chunk[-1] % self.clip_length == 0
):
logger.warning(
f"Warning: Batch containing frames {frame_chunk} from video {self.vid_files[self.label_idx[i]]} has {len(frame_chunk)} frames. Removing to avoid empty batch possibility with failed frame loading"
)
remove_idx.append(i)
if len(remove_idx) > 0:
for i in sorted(remove_idx, reverse=True):
self.chunked_frame_idx.pop(i)
self.label_idx.pop(i)
else:
self.chunked_frame_idx = self.frame_idx
self.label_idx = [i for i in range(len(self.labels))]
def __len__(self) -> int:
"""Get the size of the dataset.
Returns:
the size or the number of chunks in the dataset
"""
return len(self.chunked_frame_idx)
def no_batching_fn(self, batch: list[Frame]) -> list[Frame]:
"""Collate function used to overwrite dataloader batching function.
Args:
batch: the chunk of frames to be returned
Returns:
The batch
"""
return batch
def __getitem__(self, idx: int) -> list[Frame]:
"""Get an element of the dataset.
Args:
idx: the index of the batch. Note this is not the index of the video
or the frame.
Returns:
A list of `Frame`s in the chunk containing the metadata + instance features.
"""
label_idx, frame_idx = self.get_indices(idx)
return self.get_instances(label_idx, frame_idx)
def get_indices(self, idx: int):
"""Retrieve label and frame indices given batch index.
This method should be implemented in any subclass of the BaseDataset.
Args:
idx: the index of the batch.
Raises:
NotImplementedError: If this method is not overridden in a subclass.
"""
raise NotImplementedError("Must be implemented in subclass")
def get_instances(self, label_idx: list[int], frame_idx: list[int]):
"""Build chunk of frames.
This method should be implemented in any subclass of the BaseDataset.
Args:
label_idx: The index of the labels.
frame_idx: The index of the frames.
Raises:
NotImplementedError: If this method is not overridden in a subclass.
"""
raise NotImplementedError("Must be implemented in subclass")
__getitem__(idx)
¶
Get an element of the dataset.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
idx
|
int
|
the index of the batch. Note this is not the index of the video or the frame. |
required |
Returns:
| Type | Description |
|---|---|
list[Frame]
|
A list of |
Source code in dreem/datasets/base_dataset.py
def __getitem__(self, idx: int) -> list[Frame]:
"""Get an element of the dataset.
Args:
idx: the index of the batch. Note this is not the index of the video
or the frame.
Returns:
A list of `Frame`s in the chunk containing the metadata + instance features.
"""
label_idx, frame_idx = self.get_indices(idx)
return self.get_instances(label_idx, frame_idx)
__init__(label_files, vid_files, padding, crop_size, chunk, clip_length, mode, augmentations=None, n_chunks=1.0, seed=None, gt_list=None)
¶
Initialize Dataset.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
label_files
|
list[str]
|
a list of paths to label files. Should at least contain detections for inference, detections + tracks for training. |
required |
vid_files
|
list[str]
|
list of paths to video files. |
required |
padding
|
int
|
amount of padding around object crops |
required |
crop_size
|
Union[int, list[int]]
|
the size of the object crops |
required |
chunk
|
bool
|
whether or not to chunk the dataset into batches |
required |
clip_length
|
int
|
the number of frames in each chunk |
required |
mode
|
str
|
|
required |
augmentations
|
dict | None
|
An optional dict mapping augmentations to parameters. See subclasses for details. |
None
|
n_chunks
|
int | float
|
Number of chunks to subsample from. Can either a fraction of the dataset (ie (0,1.0]) or number of chunks |
1.0
|
seed
|
int | None
|
set a seed for reproducibility |
None
|
gt_list
|
str | None
|
An optional path to .txt file containing ground truth for cell tracking challenge datasets. |
None
|
Source code in dreem/datasets/base_dataset.py
def __init__(
self,
label_files: list[str],
vid_files: list[str],
padding: int,
crop_size: Union[int, list[int]],
chunk: bool,
clip_length: int,
mode: str,
augmentations: dict | None = None,
n_chunks: int | float = 1.0,
seed: int | None = None,
gt_list: str | None = None,
):
"""Initialize Dataset.
Args:
label_files: a list of paths to label files. Should at least contain
detections for inference, detections + tracks for training.
vid_files: list of paths to video files.
padding: amount of padding around object crops
crop_size: the size of the object crops
chunk: whether or not to chunk the dataset into batches
clip_length: the number of frames in each chunk
mode: `train` or `val`. Determines whether this dataset is used for
training or validation. Currently doesn't affect dataset logic
augmentations: An optional dict mapping augmentations to parameters.
See subclasses for details.
n_chunks: Number of chunks to subsample from.
Can either a fraction of the dataset (ie (0,1.0]) or number of chunks
seed: set a seed for reproducibility
gt_list: An optional path to .txt file containing ground truth for
cell tracking challenge datasets.
"""
self.vid_files = vid_files
self.label_files = label_files
self.padding = padding
self.crop_size = crop_size
self.chunk = chunk
self.clip_length = clip_length
self.mode = mode
self.n_chunks = n_chunks
self.seed = seed
if self.seed is not None:
np.random.seed(self.seed)
if augmentations and self.mode == "train":
self.instance_dropout = augmentations.pop(
"InstanceDropout", {"p": 0.0, "n": 0}
)
self.node_dropout = data_utils.NodeDropout(
**augmentations.pop("NodeDropout", {"p": 0.0, "n": 0})
)
self.augmentations = data_utils.build_augmentations(augmentations)
else:
self.instance_dropout = {"p": 0.0, "n": 0}
self.node_dropout = data_utils.NodeDropout(p=0.0, n=0)
self.augmentations = None
# Initialize in subclasses
self.frame_idx = None
self.labels = None
self.gt_list = None
__len__()
¶
Get the size of the dataset.
Returns:
| Type | Description |
|---|---|
int
|
the size or the number of chunks in the dataset |
create_chunks_other()
¶
Legacy chunking logic. Does not support unannotated segments.
Creates both indexes for selecting dataset (label_idx) and frame in dataset (chunked_frame_idx). If chunking is false, we index directly using the frame ids. Setting chunking to true creates a list of lists containing chunk frames for indexing. This is useful for computational efficiency and data shuffling. To be called by subclass init()
Source code in dreem/datasets/base_dataset.py
def create_chunks_other(self) -> None:
"""Legacy chunking logic. Does not support unannotated segments.
Creates both indexes for selecting dataset (label_idx) and frame in
dataset (chunked_frame_idx). If chunking is false, we index directly
using the frame ids. Setting chunking to true creates a list of lists
containing chunk frames for indexing. This is useful for computational
efficiency and data shuffling. To be called by subclass __init__()
"""
if self.chunk:
self.chunked_frame_idx, self.label_idx = [], []
for i, frame_idx in enumerate(self.frame_idx):
frame_idx_split = torch.split(frame_idx, self.clip_length)
self.chunked_frame_idx.extend(frame_idx_split)
self.label_idx.extend(len(frame_idx_split) * [i])
if self.n_chunks > 0 and self.n_chunks <= 1.0:
n_chunks = int(self.n_chunks * len(self.chunked_frame_idx))
elif self.n_chunks <= len(self.chunked_frame_idx):
n_chunks = int(self.n_chunks)
else:
n_chunks = len(self.chunked_frame_idx)
if n_chunks > 0 and n_chunks < len(self.chunked_frame_idx):
sample_idx = np.random.choice(
np.arange(len(self.chunked_frame_idx)), n_chunks, replace=False
)
self.chunked_frame_idx = [self.chunked_frame_idx[i] for i in sample_idx]
self.label_idx = [self.label_idx[i] for i in sample_idx]
# workaround for empty batch bug (needs to be changed).
# Check for batch with with only 1/10 size of clip length.
# Arbitrary thresholds
remove_idx = []
for i, frame_chunk in enumerate(self.chunked_frame_idx):
if (
len(frame_chunk) <= min(int(self.clip_length / 10), 5)
# and frame_chunk[-1] % self.clip_length == 0
):
logger.warning(
f"Warning: Batch containing frames {frame_chunk} from video {self.vid_files[self.label_idx[i]]} has {len(frame_chunk)} frames. Removing to avoid empty batch possibility with failed frame loading"
)
remove_idx.append(i)
if len(remove_idx) > 0:
for i in sorted(remove_idx, reverse=True):
self.chunked_frame_idx.pop(i)
self.label_idx.pop(i)
else:
self.chunked_frame_idx = self.frame_idx
self.label_idx = [i for i in range(len(self.labels))]
create_chunks_slp()
¶
Get indexing for data.
Creates both indexes for selecting dataset (label_idx) and frame in dataset (chunked_frame_idx). If chunking is false, we index directly using the frame ids. Setting chunking to true creates a list of lists containing chunk frames for indexing. This is useful for computational efficiency and data shuffling. To be called by subclass init()
Source code in dreem/datasets/base_dataset.py
def create_chunks_slp(self) -> None:
"""Get indexing for data.
Creates both indexes for selecting dataset (label_idx) and frame in
dataset (chunked_frame_idx). If chunking is false, we index directly
using the frame ids. Setting chunking to true creates a list of lists
containing chunk frames for indexing. This is useful for computational
efficiency and data shuffling. To be called by subclass __init__()
"""
self.chunked_frame_idx, self.label_idx = [], []
# go through each slp file and create chunks that respect max_batching_gap
for i, slp_file in enumerate(self.label_files):
annotated_segments = self.annotated_segments[slp_file]
segments_to_stitch = []
prev_end = annotated_segments[0][1] # end of first segment
for start, end in annotated_segments:
# check if the start of current segment is within
# batching_max_gap of end of previous
if (
(int(start) - int(prev_end) < self.max_batching_gap)
or not self.chunk
): # also takes care of first segment as start < prev_end
segments_to_stitch.append(torch.arange(start, end + 1))
prev_end = end
else:
# stitch previous set of segments before creating a new chunk
self.process_segments(i, segments_to_stitch, self.clip_length)
# reset segments_to_stitch as we are starting a new chunk
segments_to_stitch = [torch.arange(start, end + 1)]
prev_end = end
if not self.chunk:
self.process_segments(
i, segments_to_stitch, self.labels[i].video.shape[0]
)
else:
# add last chunk after the loop
if segments_to_stitch:
self.process_segments(i, segments_to_stitch, self.clip_length)
if self.n_chunks > 0 and self.n_chunks <= 1.0:
n_chunks = int(self.n_chunks * len(self.chunked_frame_idx))
elif self.n_chunks <= len(self.chunked_frame_idx):
n_chunks = int(self.n_chunks)
else:
n_chunks = len(self.chunked_frame_idx)
if n_chunks > 0 and n_chunks < len(self.chunked_frame_idx):
sample_idx = np.random.choice(
np.arange(len(self.chunked_frame_idx)), n_chunks, replace=False
)
self.chunked_frame_idx = [self.chunked_frame_idx[i] for i in sample_idx]
self.label_idx = [self.label_idx[i] for i in sample_idx]
# workaround for empty batch bug (needs to be changed).
# Check for batch with with only 1/10 size of clip length.
# Arbitrary thresholds
remove_idx = []
for i, frame_chunk in enumerate(self.chunked_frame_idx):
if (
len(frame_chunk) <= min(int(self.clip_length / 10), 5)
# and frame_chunk[-1] % self.clip_length == 0
):
logger.warning(
f"Warning: Batch containing frames {frame_chunk} from video "
f"{self.vid_files[self.label_idx[i]]} has {len(frame_chunk)} frames. "
f"Removing to avoid empty batch possibility with failed frame loading"
)
remove_idx.append(i)
if len(remove_idx) > 0:
for i in sorted(remove_idx, reverse=True):
self.chunked_frame_idx.pop(i)
self.label_idx.pop(i)
get_indices(idx)
¶
Retrieve label and frame indices given batch index.
This method should be implemented in any subclass of the BaseDataset.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
idx
|
int
|
the index of the batch. |
required |
Raises:
| Type | Description |
|---|---|
NotImplementedError
|
If this method is not overridden in a subclass. |
Source code in dreem/datasets/base_dataset.py
def get_indices(self, idx: int):
"""Retrieve label and frame indices given batch index.
This method should be implemented in any subclass of the BaseDataset.
Args:
idx: the index of the batch.
Raises:
NotImplementedError: If this method is not overridden in a subclass.
"""
raise NotImplementedError("Must be implemented in subclass")
get_instances(label_idx, frame_idx)
¶
Build chunk of frames.
This method should be implemented in any subclass of the BaseDataset.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
label_idx
|
list[int]
|
The index of the labels. |
required |
frame_idx
|
list[int]
|
The index of the frames. |
required |
Raises:
| Type | Description |
|---|---|
NotImplementedError
|
If this method is not overridden in a subclass. |
Source code in dreem/datasets/base_dataset.py
def get_instances(self, label_idx: list[int], frame_idx: list[int]):
"""Build chunk of frames.
This method should be implemented in any subclass of the BaseDataset.
Args:
label_idx: The index of the labels.
frame_idx: The index of the frames.
Raises:
NotImplementedError: If this method is not overridden in a subclass.
"""
raise NotImplementedError("Must be implemented in subclass")
no_batching_fn(batch)
¶
process_segments(i, segments_to_stitch, clip_length)
¶
Process segments to stitch.
Modifies state variables chunked_frame_idx and label_idx.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
segments_to_stitch
|
list[Tensor]
|
list of segments to stitch |
required |
i
|
int
|
index of the video |
required |
clip_length
|
int
|
the number of frames in each chunk |
required |
Returns: None
Source code in dreem/datasets/base_dataset.py
def process_segments(
self, i: int, segments_to_stitch: list[torch.Tensor], clip_length: int
) -> None:
"""Process segments to stitch.
Modifies state variables chunked_frame_idx and label_idx.
Args:
segments_to_stitch: list of segments to stitch
i: index of the video
clip_length: the number of frames in each chunk
Returns: None
"""
stitched_segment = torch.cat(segments_to_stitch)
frame_idx_split = torch.split(stitched_segment, clip_length)
self.chunked_frame_idx.extend(frame_idx_split)
self.label_idx.extend(len(frame_idx_split) * [i])
CellTrackingDataset
¶
Bases: BaseDataset
Dataset for loading cell tracking challenge data.
Methods:
| Name | Description |
|---|---|
__init__ |
Initialize CellTrackingDataset. |
get_indices |
Retrieve label and frame indices given batch index. |
get_instances |
Get an element of the dataset. |
Source code in dreem/datasets/cell_tracking_dataset.py
class CellTrackingDataset(BaseDataset):
"""Dataset for loading cell tracking challenge data."""
def __init__(
self,
gt_list: list[list[str]],
raw_img_list: list[list[str]],
data_dirs: Optional[list[str]] = None,
padding: int = 5,
crop_size: int = 20,
chunk: bool = False,
clip_length: int = 10,
mode: str = "train",
augmentations: dict | None = None,
n_chunks: int | float = 1.0,
seed: int | None = None,
max_batching_gap: int = 15,
use_tight_bbox: bool = False,
ctc_track_meta: list[str] | None = None,
apply_mask_to_crop: bool = False,
**kwargs,
):
"""Initialize CellTrackingDataset.
Args:
gt_list: filepaths of gt label images in a list of lists (each list
corresponds to a dataset)
raw_img_list: filepaths of original tif images in a list of lists
(each list corresponds to a dataset)
data_dirs: paths to data directories
padding: amount of padding around object crops
crop_size: the size of the object crops. Can be either:
- An integer specifying a single crop size for all objects
- A list of integers specifying different crop sizes for
different data directories
chunk: whether or not to chunk the dataset into batches
clip_length: the number of frames in each chunk
mode: `train` or `val`. Determines whether this dataset is used for
training or validation. Currently doesn't affect dataset logic
augmentations: An optional dict mapping augmentations to parameters.
The keys
should map directly to augmentation classes in albumentations. Example:
augs = {
'Rotate': {'limit': [-90, 90]},
'GaussianBlur': {'blur_limit': (3, 7), 'sigma_limit': 0},
'RandomContrast': {'limit': 0.2}
}
n_chunks: Number of chunks to subsample from.
Can either a fraction of the dataset (ie (0,1.0]) or number of chunks
seed: set a seed for reproducibility
max_batching_gap: the max number of frames that can be unlabelled
before starting a new batch
use_tight_bbox: whether to use tight bounding box (around keypoints)
instead of the default square bounding box
ctc_track_meta: filepaths of man_track.txt files in a list of lists
(each list corresponds to a dataset)
apply_mask_to_crop: whether to apply the mask to the crop
**kwargs: Additional keyword arguments (unused but accepted for compatibility)
"""
super().__init__(
gt_list,
raw_img_list,
padding,
crop_size,
chunk,
clip_length,
mode,
augmentations,
n_chunks,
seed,
ctc_track_meta,
)
self.raw_img_list = raw_img_list
self.gt_list = gt_list
self.ctc_track_meta = ctc_track_meta
self.data_dirs = data_dirs
self.chunk = chunk
self.clip_length = clip_length
self.crop_size = crop_size
self.padding = padding
self.mode = mode.lower()
self.n_chunks = n_chunks
self.seed = seed
self.max_batching_gap = max_batching_gap
self.use_tight_bbox = use_tight_bbox
self.skeleton = sio.Skeleton(nodes=["centroid"])
self.apply_mask_to_crop = apply_mask_to_crop
if not isinstance(self.data_dirs, list):
self.data_dirs = [self.data_dirs]
if not isinstance(self.crop_size, list):
# make a list so its handled consistently if multiple crops are used
if len(self.data_dirs) > 0: # for test mode, data_dirs is []
self.crop_size = [self.crop_size] * len(self.data_dirs)
else:
self.crop_size = [self.crop_size]
if len(self.data_dirs) > 0 and len(self.crop_size) != len(self.data_dirs):
raise ValueError(
f"If a list of crop sizes or data directories are given,"
f"they must have the same length but got {len(self.crop_size)} "
f"and {len(self.data_dirs)}"
)
# if self.seed is not None:
# np.random.seed(self.seed)
if augmentations and self.mode == "train":
self.augmentations = data_utils.build_augmentations(augmentations)
else:
self.augmentations = None
#
if self.ctc_track_meta is not None:
self.list_df_track_meta = [
pd.read_csv(
gtf,
delimiter=" ",
header=None,
names=["track_id", "start_frame", "end_frame", "parent_id"],
)
for gtf in self.ctc_track_meta
]
else:
self.list_df_track_meta = None
# frame indices for each dataset; list of lists (each list corresponds to a dataset)
self.frame_idx = [torch.arange(len(gt_dataset)) for gt_dataset in self.gt_list]
# Method in BaseDataset. Creates label_idx and chunked_frame_idx to be
# used in call to get_instances()
self.create_chunks_other()
def get_indices(self, idx: int) -> tuple:
"""Retrieve label and frame indices given batch index.
Args:
idx: the index of the batch.
Returns:
the label and frame indices corresponding to a batch,
"""
return self.label_idx[idx], self.chunked_frame_idx[idx]
def get_instances(self, label_idx: list[int], frame_idx: list[int]) -> list[Frame]:
"""Get an element of the dataset.
Args:
label_idx: index of the labels
frame_idx: index of the frames
Returns:
a list of Frame objects containing frame metadata and Instance Objects.
See `dreem.io.data_structures` for more info.
"""
image_paths = self.raw_img_list[label_idx]
gt_paths = self.gt_list[label_idx]
# df_track_meta is currently unused but may be needed for future track metadata processing
# if self.list_df_track_meta is not None:
# df_track_meta = self.list_df_track_meta[label_idx]
# else:
# df_track_meta = None
# get the correct crop size based on the video
video_par_path = Path(image_paths[0]).parent.parent
if len(self.data_dirs) > 0:
crop_size = self.crop_size[0]
for j, data_dir in enumerate(self.data_dirs):
if Path(data_dir) == video_par_path:
crop_size = self.crop_size[j]
break
else:
crop_size = self.crop_size[0]
frames = []
max_crop_h, max_crop_w = 0, 0
for i in frame_idx:
instances, gt_track_ids, centroids, dict_centroids, bboxes, masks = (
[],
[],
[],
{},
[],
[],
)
i = int(i)
img = image_paths[i]
gt_sec = gt_paths[i]
img = np.array(Image.open(img))
gt_sec = np.array(Image.open(gt_sec))
if img.dtype == np.uint16:
img = ((img - img.min()) * (1 / (img.max() - img.min()) * 255)).astype(
np.uint8
)
# if df_track_meta is None:
unique_instances = np.unique(gt_sec)
# else:
# unique_instances = df_track_meta["track_id"].unique()
for instance in unique_instances:
# not all instances are in the frame, and they also label the
# background instance as zero
if instance in gt_sec and instance != 0:
mask = gt_sec == instance
center_of_mass = measurements.center_of_mass(mask)
# scipy returns yx
x, y = center_of_mass[::-1]
if self.use_tight_bbox:
bbox = data_utils.get_tight_bbox_masks(mask)
else:
bbox = data_utils.pad_bbox(
data_utils.get_bbox([int(x), int(y)], crop_size),
padding=self.padding,
)
mask = torch.as_tensor(mask)
gt_track_ids.append(int(instance))
centroids.append([x, y])
dict_centroids[int(instance)] = [x, y]
bboxes.append(bbox)
masks.append(mask)
# albumentations wants (spatial, channels), ensure correct dims
if self.augmentations is not None:
for transform in self.augmentations:
# for occlusion simulation, can remove if we don't want
if isinstance(transform, A.CoarseDropout):
transform.fill_value = random.randint(0, 255)
augmented = self.augmentations(
image=img,
mask=gt_sec, # albumentations ensures geometric transformations are synced between image and mask
keypoints=np.vstack(centroids),
)
img, aug_mask, centroids = (
augmented["image"],
augmented["mask"],
augmented["keypoints"],
)
aug_mask = torch.Tensor(aug_mask).unsqueeze(0)
img = torch.Tensor(img).unsqueeze(0)
for j in range(len(gt_track_ids)):
# just formatting for compatibility with Instance class
instance_centroid = {
"centroid": np.array(dict_centroids[gt_track_ids[j]])
}
pose = {"centroid": dict_centroids[gt_track_ids[j]]} # more formatting
crop_raw = data_utils.crop_bbox(img, bboxes[j])
if self.apply_mask_to_crop:
if self.augmentations is not None:
cropped_mask = data_utils.crop_bbox(aug_mask, bboxes[j])
# filter for the instance of interest
cropped_mask[cropped_mask != gt_track_ids[j]] = 0
else:
# masks[j] is already filtered for the instance of interest
cropped_mask = data_utils.crop_bbox(masks[j], bboxes[j])
cropped_mask[cropped_mask != 0] = 1
# apply mask to crop
crop = crop_raw * cropped_mask
else:
crop = crop_raw
c, h, w = crop.shape
if h > max_crop_h:
max_crop_h = h
if w > max_crop_w:
max_crop_w = w
instances.append(
Instance(
gt_track_id=gt_track_ids[j],
pred_track_id=-1,
centroid=instance_centroid,
skeleton=self.skeleton,
point_scores=np.array([1.0]),
instance_score=np.array([1.0]),
pose=pose,
bbox=bboxes[j],
crop=crop,
mask=masks[j],
)
)
if self.mode == "train":
np.random.shuffle(instances)
frames.append(
Frame(
video_id=label_idx,
frame_id=i,
vid_file=Path(image_paths[0]).parent.name,
img_shape=img.shape,
instances=instances,
)
)
# pad bbox to max size
if self.use_tight_bbox:
# bound the max crop size to the user defined crop size
max_crop_h = crop_size if max_crop_h == 0 else min(max_crop_h, crop_size)
max_crop_w = crop_size if max_crop_w == 0 else min(max_crop_w, crop_size)
# gather all the crops
for frame in frames:
for instance in frame.instances:
data_utils.pad_variable_size_crops(
instance, (max_crop_h, max_crop_w)
)
return frames
__init__(gt_list, raw_img_list, data_dirs=None, padding=5, crop_size=20, chunk=False, clip_length=10, mode='train', augmentations=None, n_chunks=1.0, seed=None, max_batching_gap=15, use_tight_bbox=False, ctc_track_meta=None, apply_mask_to_crop=False, **kwargs)
¶
Initialize CellTrackingDataset.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
gt_list
|
list[list[str]]
|
filepaths of gt label images in a list of lists (each list corresponds to a dataset) |
required |
raw_img_list
|
list[list[str]]
|
filepaths of original tif images in a list of lists (each list corresponds to a dataset) |
required |
data_dirs
|
Optional[list[str]]
|
paths to data directories |
None
|
padding
|
int
|
amount of padding around object crops |
5
|
crop_size
|
int
|
the size of the object crops. Can be either: - An integer specifying a single crop size for all objects - A list of integers specifying different crop sizes for different data directories |
20
|
chunk
|
bool
|
whether or not to chunk the dataset into batches |
False
|
clip_length
|
int
|
the number of frames in each chunk |
10
|
mode
|
str
|
|
'train'
|
augmentations
|
dict | None
|
An optional dict mapping augmentations to parameters. The keys should map directly to augmentation classes in albumentations. Example: augs = { 'Rotate': {'limit': [-90, 90]}, 'GaussianBlur': {'blur_limit': (3, 7), 'sigma_limit': 0}, 'RandomContrast': {'limit': 0.2} } |
None
|
n_chunks
|
int | float
|
Number of chunks to subsample from. Can either a fraction of the dataset (ie (0,1.0]) or number of chunks |
1.0
|
seed
|
int | None
|
set a seed for reproducibility |
None
|
max_batching_gap
|
int
|
the max number of frames that can be unlabelled before starting a new batch |
15
|
use_tight_bbox
|
bool
|
whether to use tight bounding box (around keypoints) instead of the default square bounding box |
False
|
ctc_track_meta
|
list[str] | None
|
filepaths of man_track.txt files in a list of lists (each list corresponds to a dataset) |
None
|
apply_mask_to_crop
|
bool
|
whether to apply the mask to the crop |
False
|
**kwargs
|
Additional keyword arguments (unused but accepted for compatibility) |
{}
|
Source code in dreem/datasets/cell_tracking_dataset.py
def __init__(
self,
gt_list: list[list[str]],
raw_img_list: list[list[str]],
data_dirs: Optional[list[str]] = None,
padding: int = 5,
crop_size: int = 20,
chunk: bool = False,
clip_length: int = 10,
mode: str = "train",
augmentations: dict | None = None,
n_chunks: int | float = 1.0,
seed: int | None = None,
max_batching_gap: int = 15,
use_tight_bbox: bool = False,
ctc_track_meta: list[str] | None = None,
apply_mask_to_crop: bool = False,
**kwargs,
):
"""Initialize CellTrackingDataset.
Args:
gt_list: filepaths of gt label images in a list of lists (each list
corresponds to a dataset)
raw_img_list: filepaths of original tif images in a list of lists
(each list corresponds to a dataset)
data_dirs: paths to data directories
padding: amount of padding around object crops
crop_size: the size of the object crops. Can be either:
- An integer specifying a single crop size for all objects
- A list of integers specifying different crop sizes for
different data directories
chunk: whether or not to chunk the dataset into batches
clip_length: the number of frames in each chunk
mode: `train` or `val`. Determines whether this dataset is used for
training or validation. Currently doesn't affect dataset logic
augmentations: An optional dict mapping augmentations to parameters.
The keys
should map directly to augmentation classes in albumentations. Example:
augs = {
'Rotate': {'limit': [-90, 90]},
'GaussianBlur': {'blur_limit': (3, 7), 'sigma_limit': 0},
'RandomContrast': {'limit': 0.2}
}
n_chunks: Number of chunks to subsample from.
Can either a fraction of the dataset (ie (0,1.0]) or number of chunks
seed: set a seed for reproducibility
max_batching_gap: the max number of frames that can be unlabelled
before starting a new batch
use_tight_bbox: whether to use tight bounding box (around keypoints)
instead of the default square bounding box
ctc_track_meta: filepaths of man_track.txt files in a list of lists
(each list corresponds to a dataset)
apply_mask_to_crop: whether to apply the mask to the crop
**kwargs: Additional keyword arguments (unused but accepted for compatibility)
"""
super().__init__(
gt_list,
raw_img_list,
padding,
crop_size,
chunk,
clip_length,
mode,
augmentations,
n_chunks,
seed,
ctc_track_meta,
)
self.raw_img_list = raw_img_list
self.gt_list = gt_list
self.ctc_track_meta = ctc_track_meta
self.data_dirs = data_dirs
self.chunk = chunk
self.clip_length = clip_length
self.crop_size = crop_size
self.padding = padding
self.mode = mode.lower()
self.n_chunks = n_chunks
self.seed = seed
self.max_batching_gap = max_batching_gap
self.use_tight_bbox = use_tight_bbox
self.skeleton = sio.Skeleton(nodes=["centroid"])
self.apply_mask_to_crop = apply_mask_to_crop
if not isinstance(self.data_dirs, list):
self.data_dirs = [self.data_dirs]
if not isinstance(self.crop_size, list):
# make a list so its handled consistently if multiple crops are used
if len(self.data_dirs) > 0: # for test mode, data_dirs is []
self.crop_size = [self.crop_size] * len(self.data_dirs)
else:
self.crop_size = [self.crop_size]
if len(self.data_dirs) > 0 and len(self.crop_size) != len(self.data_dirs):
raise ValueError(
f"If a list of crop sizes or data directories are given,"
f"they must have the same length but got {len(self.crop_size)} "
f"and {len(self.data_dirs)}"
)
# if self.seed is not None:
# np.random.seed(self.seed)
if augmentations and self.mode == "train":
self.augmentations = data_utils.build_augmentations(augmentations)
else:
self.augmentations = None
#
if self.ctc_track_meta is not None:
self.list_df_track_meta = [
pd.read_csv(
gtf,
delimiter=" ",
header=None,
names=["track_id", "start_frame", "end_frame", "parent_id"],
)
for gtf in self.ctc_track_meta
]
else:
self.list_df_track_meta = None
# frame indices for each dataset; list of lists (each list corresponds to a dataset)
self.frame_idx = [torch.arange(len(gt_dataset)) for gt_dataset in self.gt_list]
# Method in BaseDataset. Creates label_idx and chunked_frame_idx to be
# used in call to get_instances()
self.create_chunks_other()
get_indices(idx)
¶
Retrieve label and frame indices given batch index.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
idx
|
int
|
the index of the batch. |
required |
Returns:
| Type | Description |
|---|---|
tuple
|
the label and frame indices corresponding to a batch, |
Source code in dreem/datasets/cell_tracking_dataset.py
get_instances(label_idx, frame_idx)
¶
Get an element of the dataset.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
label_idx
|
list[int]
|
index of the labels |
required |
frame_idx
|
list[int]
|
index of the frames |
required |
Returns:
| Type | Description |
|---|---|
list[Frame]
|
a list of Frame objects containing frame metadata and Instance Objects.
See |
Source code in dreem/datasets/cell_tracking_dataset.py
def get_instances(self, label_idx: list[int], frame_idx: list[int]) -> list[Frame]:
"""Get an element of the dataset.
Args:
label_idx: index of the labels
frame_idx: index of the frames
Returns:
a list of Frame objects containing frame metadata and Instance Objects.
See `dreem.io.data_structures` for more info.
"""
image_paths = self.raw_img_list[label_idx]
gt_paths = self.gt_list[label_idx]
# df_track_meta is currently unused but may be needed for future track metadata processing
# if self.list_df_track_meta is not None:
# df_track_meta = self.list_df_track_meta[label_idx]
# else:
# df_track_meta = None
# get the correct crop size based on the video
video_par_path = Path(image_paths[0]).parent.parent
if len(self.data_dirs) > 0:
crop_size = self.crop_size[0]
for j, data_dir in enumerate(self.data_dirs):
if Path(data_dir) == video_par_path:
crop_size = self.crop_size[j]
break
else:
crop_size = self.crop_size[0]
frames = []
max_crop_h, max_crop_w = 0, 0
for i in frame_idx:
instances, gt_track_ids, centroids, dict_centroids, bboxes, masks = (
[],
[],
[],
{},
[],
[],
)
i = int(i)
img = image_paths[i]
gt_sec = gt_paths[i]
img = np.array(Image.open(img))
gt_sec = np.array(Image.open(gt_sec))
if img.dtype == np.uint16:
img = ((img - img.min()) * (1 / (img.max() - img.min()) * 255)).astype(
np.uint8
)
# if df_track_meta is None:
unique_instances = np.unique(gt_sec)
# else:
# unique_instances = df_track_meta["track_id"].unique()
for instance in unique_instances:
# not all instances are in the frame, and they also label the
# background instance as zero
if instance in gt_sec and instance != 0:
mask = gt_sec == instance
center_of_mass = measurements.center_of_mass(mask)
# scipy returns yx
x, y = center_of_mass[::-1]
if self.use_tight_bbox:
bbox = data_utils.get_tight_bbox_masks(mask)
else:
bbox = data_utils.pad_bbox(
data_utils.get_bbox([int(x), int(y)], crop_size),
padding=self.padding,
)
mask = torch.as_tensor(mask)
gt_track_ids.append(int(instance))
centroids.append([x, y])
dict_centroids[int(instance)] = [x, y]
bboxes.append(bbox)
masks.append(mask)
# albumentations wants (spatial, channels), ensure correct dims
if self.augmentations is not None:
for transform in self.augmentations:
# for occlusion simulation, can remove if we don't want
if isinstance(transform, A.CoarseDropout):
transform.fill_value = random.randint(0, 255)
augmented = self.augmentations(
image=img,
mask=gt_sec, # albumentations ensures geometric transformations are synced between image and mask
keypoints=np.vstack(centroids),
)
img, aug_mask, centroids = (
augmented["image"],
augmented["mask"],
augmented["keypoints"],
)
aug_mask = torch.Tensor(aug_mask).unsqueeze(0)
img = torch.Tensor(img).unsqueeze(0)
for j in range(len(gt_track_ids)):
# just formatting for compatibility with Instance class
instance_centroid = {
"centroid": np.array(dict_centroids[gt_track_ids[j]])
}
pose = {"centroid": dict_centroids[gt_track_ids[j]]} # more formatting
crop_raw = data_utils.crop_bbox(img, bboxes[j])
if self.apply_mask_to_crop:
if self.augmentations is not None:
cropped_mask = data_utils.crop_bbox(aug_mask, bboxes[j])
# filter for the instance of interest
cropped_mask[cropped_mask != gt_track_ids[j]] = 0
else:
# masks[j] is already filtered for the instance of interest
cropped_mask = data_utils.crop_bbox(masks[j], bboxes[j])
cropped_mask[cropped_mask != 0] = 1
# apply mask to crop
crop = crop_raw * cropped_mask
else:
crop = crop_raw
c, h, w = crop.shape
if h > max_crop_h:
max_crop_h = h
if w > max_crop_w:
max_crop_w = w
instances.append(
Instance(
gt_track_id=gt_track_ids[j],
pred_track_id=-1,
centroid=instance_centroid,
skeleton=self.skeleton,
point_scores=np.array([1.0]),
instance_score=np.array([1.0]),
pose=pose,
bbox=bboxes[j],
crop=crop,
mask=masks[j],
)
)
if self.mode == "train":
np.random.shuffle(instances)
frames.append(
Frame(
video_id=label_idx,
frame_id=i,
vid_file=Path(image_paths[0]).parent.name,
img_shape=img.shape,
instances=instances,
)
)
# pad bbox to max size
if self.use_tight_bbox:
# bound the max crop size to the user defined crop size
max_crop_h = crop_size if max_crop_h == 0 else min(max_crop_h, crop_size)
max_crop_w = crop_size if max_crop_w == 0 else min(max_crop_w, crop_size)
# gather all the crops
for frame in frames:
for instance in frame.instances:
data_utils.pad_variable_size_crops(
instance, (max_crop_h, max_crop_w)
)
return frames
MicroscopyDataset
¶
Bases: BaseDataset
Dataset for loading Microscopy Data.
Methods:
| Name | Description |
|---|---|
__del__ |
Handle file closing before deletion. |
__init__ |
Initialize MicroscopyDataset. |
get_indices |
Retrieve label and frame indices given batch index. |
get_instances |
Get an element of the dataset. |
Source code in dreem/datasets/microscopy_dataset.py
class MicroscopyDataset(BaseDataset):
"""Dataset for loading Microscopy Data."""
def __init__(
self,
videos: list[str],
tracks: list[str],
source: str,
padding: int = 5,
crop_size: int = 20,
chunk: bool = False,
clip_length: int = 10,
mode: str = "Train",
augmentations: dict | None = None,
n_chunks: int | float = 1.0,
seed: int | None = None,
):
"""Initialize MicroscopyDataset.
Args:
videos: paths to raw microscopy videos
tracks: paths to trackmate gt labels (either .xml or .csv)
source: file format of gt labels based on label generator
padding: amount of padding around object crops
crop_size: the size of the object crops
chunk: whether or not to chunk the dataset into batches
clip_length: the number of frames in each chunk
mode: `train` or `val`. Determines whether this dataset is used for
training or validation. Currently doesn't affect dataset logic
augmentations: An optional dict mapping augmentations to parameters. The keys
should map directly to augmentation classes in albumentations. Example:
augs = {
'Rotate': {'limit': [-90, 90]},
'GaussianBlur': {'blur_limit': (3, 7), 'sigma_limit': 0},
'RandomContrast': {'limit': 0.2}
}
n_chunks: Number of chunks to subsample from.
Can either a fraction of the dataset (ie (0,1.0]) or number of chunks
seed: set a seed for reproducibility
"""
super().__init__(
tracks,
videos,
padding,
crop_size,
chunk,
clip_length,
mode,
augmentations,
n_chunks,
seed,
)
self.vid_files = videos
self.tracks = tracks
self.chunk = chunk
self.clip_length = clip_length
self.crop_size = crop_size
self.padding = padding
self.mode = mode.lower()
self.n_chunks = n_chunks
self.seed = seed
# if self.seed is not None:
# np.random.seed(self.seed)
if augmentations and self.mode == "train":
self.augmentations = data_utils.build_augmentations(augmentations)
else:
self.augmentations = None
if source.lower() == "trackmate":
parser = data_utils.parse_trackmate
elif source.lower() in ["icy", "isbi"]:
def parser(x):
return data_utils.parse_synthetic(x, source=source)
else:
raise ValueError(
f"{source} is unsupported! Must be one of [trackmate, icy, isbi]"
)
self.labels = [
parser(self.tracks[video_idx]) for video_idx in range(len(self.tracks))
]
self.videos = []
for vid_file in self.vid_files:
if not isinstance(vid_file, list):
self.videos.append(data_utils.LazyTiffStack(vid_file))
else:
self.videos.append([Image.open(frame_file) for frame_file in vid_file])
self.frame_idx = [
(
torch.arange(Image.open(video).n_frames)
if isinstance(video, str)
else torch.arange(len(video))
)
for video in self.vid_files
]
# Method in BaseDataset. Creates label_idx and chunked_frame_idx to be
# used in call to get_instances()
self.create_chunks_other()
def get_indices(self, idx: int) -> tuple:
"""Retrieve label and frame indices given batch index.
Args:
idx: the index of the batch.
"""
return self.label_idx[idx], self.chunked_frame_idx[idx]
def get_instances(self, label_idx: list[int], frame_idx: list[int]) -> list[Frame]:
"""Get an element of the dataset.
Args:
label_idx: index of the labels
frame_idx: index of the frames
Returns:
A list of Frames containing Instances to be tracked (See `dreem.io.data_structures for more info`)
"""
labels = self.labels[label_idx]
labels = labels.dropna(how="all")
video = self.videos[label_idx]
frames = []
for frame_id in frame_idx:
instances, gt_track_ids, centroids = [], [], []
img = (
video.get_section(frame_id)
if not isinstance(video, list)
else np.array(video[frame_id])
)
lf = labels[labels["FRAME"].astype(int) == frame_id.item()]
for instance in sorted(lf["TRACK_ID"].unique()):
gt_track_ids.append(int(instance))
x = lf[lf["TRACK_ID"] == instance]["POSITION_X"].iloc[0]
y = lf[lf["TRACK_ID"] == instance]["POSITION_Y"].iloc[0]
centroids.append([x, y])
# albumentations wants (spatial, channels), ensure correct dims
if self.augmentations is not None:
for transform in self.augmentations:
# for occlusion simulation, can remove if we don't want
if isinstance(transform, A.CoarseDropout):
transform.fill_value = random.randint(0, 255)
augmented = self.augmentations(
image=img,
keypoints=np.vstack(centroids),
)
img, centroids = augmented["image"], augmented["keypoints"]
img = torch.Tensor(img)
# torch wants (channels, spatial) - ensure correct dims
if len(img.shape) == 2:
img = img.unsqueeze(0)
elif len(img.shape) == 3:
if img.shape[2] == 3:
img = img.T # todo: check for edge cases
for gt_id in range(len(gt_track_ids)):
c = centroids[gt_id]
bbox = data_utils.pad_bbox(
data_utils.get_bbox([int(c[0]), int(c[1])], self.crop_size),
padding=self.padding,
)
crop = data_utils.crop_bbox(img, bbox)
instances.append(
Instance(
gt_track_id=gt_track_ids[gt_id],
pred_track_id=-1,
bbox=bbox,
crop=crop,
)
)
if self.mode == "train":
np.random.shuffle(instances)
frames.append(
Frame(
video_id=label_idx,
frame_id=frame_id,
img_shape=img.shape,
instances=instances,
)
)
return frames
def __del__(self):
"""Handle file closing before deletion."""
for vid_reader in self.videos:
if not isinstance(vid_reader, list):
vid_reader.close()
else:
for frame_reader in vid_reader:
frame_reader.close()
__del__()
¶
__init__(videos, tracks, source, padding=5, crop_size=20, chunk=False, clip_length=10, mode='Train', augmentations=None, n_chunks=1.0, seed=None)
¶
Initialize MicroscopyDataset.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
videos
|
list[str]
|
paths to raw microscopy videos |
required |
tracks
|
list[str]
|
paths to trackmate gt labels (either .xml or .csv) |
required |
source
|
str
|
file format of gt labels based on label generator |
required |
padding
|
int
|
amount of padding around object crops |
5
|
crop_size
|
int
|
the size of the object crops |
20
|
chunk
|
bool
|
whether or not to chunk the dataset into batches |
False
|
clip_length
|
int
|
the number of frames in each chunk |
10
|
mode
|
str
|
|
'Train'
|
augmentations
|
dict | None
|
An optional dict mapping augmentations to parameters. The keys should map directly to augmentation classes in albumentations. Example: augs = { 'Rotate': {'limit': [-90, 90]}, 'GaussianBlur': {'blur_limit': (3, 7), 'sigma_limit': 0}, 'RandomContrast': {'limit': 0.2} } |
None
|
n_chunks
|
int | float
|
Number of chunks to subsample from. Can either a fraction of the dataset (ie (0,1.0]) or number of chunks |
1.0
|
seed
|
int | None
|
set a seed for reproducibility |
None
|
Source code in dreem/datasets/microscopy_dataset.py
def __init__(
self,
videos: list[str],
tracks: list[str],
source: str,
padding: int = 5,
crop_size: int = 20,
chunk: bool = False,
clip_length: int = 10,
mode: str = "Train",
augmentations: dict | None = None,
n_chunks: int | float = 1.0,
seed: int | None = None,
):
"""Initialize MicroscopyDataset.
Args:
videos: paths to raw microscopy videos
tracks: paths to trackmate gt labels (either .xml or .csv)
source: file format of gt labels based on label generator
padding: amount of padding around object crops
crop_size: the size of the object crops
chunk: whether or not to chunk the dataset into batches
clip_length: the number of frames in each chunk
mode: `train` or `val`. Determines whether this dataset is used for
training or validation. Currently doesn't affect dataset logic
augmentations: An optional dict mapping augmentations to parameters. The keys
should map directly to augmentation classes in albumentations. Example:
augs = {
'Rotate': {'limit': [-90, 90]},
'GaussianBlur': {'blur_limit': (3, 7), 'sigma_limit': 0},
'RandomContrast': {'limit': 0.2}
}
n_chunks: Number of chunks to subsample from.
Can either a fraction of the dataset (ie (0,1.0]) or number of chunks
seed: set a seed for reproducibility
"""
super().__init__(
tracks,
videos,
padding,
crop_size,
chunk,
clip_length,
mode,
augmentations,
n_chunks,
seed,
)
self.vid_files = videos
self.tracks = tracks
self.chunk = chunk
self.clip_length = clip_length
self.crop_size = crop_size
self.padding = padding
self.mode = mode.lower()
self.n_chunks = n_chunks
self.seed = seed
# if self.seed is not None:
# np.random.seed(self.seed)
if augmentations and self.mode == "train":
self.augmentations = data_utils.build_augmentations(augmentations)
else:
self.augmentations = None
if source.lower() == "trackmate":
parser = data_utils.parse_trackmate
elif source.lower() in ["icy", "isbi"]:
def parser(x):
return data_utils.parse_synthetic(x, source=source)
else:
raise ValueError(
f"{source} is unsupported! Must be one of [trackmate, icy, isbi]"
)
self.labels = [
parser(self.tracks[video_idx]) for video_idx in range(len(self.tracks))
]
self.videos = []
for vid_file in self.vid_files:
if not isinstance(vid_file, list):
self.videos.append(data_utils.LazyTiffStack(vid_file))
else:
self.videos.append([Image.open(frame_file) for frame_file in vid_file])
self.frame_idx = [
(
torch.arange(Image.open(video).n_frames)
if isinstance(video, str)
else torch.arange(len(video))
)
for video in self.vid_files
]
# Method in BaseDataset. Creates label_idx and chunked_frame_idx to be
# used in call to get_instances()
self.create_chunks_other()
get_indices(idx)
¶
Retrieve label and frame indices given batch index.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
idx
|
int
|
the index of the batch. |
required |
get_instances(label_idx, frame_idx)
¶
Get an element of the dataset.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
label_idx
|
list[int]
|
index of the labels |
required |
frame_idx
|
list[int]
|
index of the frames |
required |
Returns:
| Type | Description |
|---|---|
list[Frame]
|
A list of Frames containing Instances to be tracked (See |
Source code in dreem/datasets/microscopy_dataset.py
def get_instances(self, label_idx: list[int], frame_idx: list[int]) -> list[Frame]:
"""Get an element of the dataset.
Args:
label_idx: index of the labels
frame_idx: index of the frames
Returns:
A list of Frames containing Instances to be tracked (See `dreem.io.data_structures for more info`)
"""
labels = self.labels[label_idx]
labels = labels.dropna(how="all")
video = self.videos[label_idx]
frames = []
for frame_id in frame_idx:
instances, gt_track_ids, centroids = [], [], []
img = (
video.get_section(frame_id)
if not isinstance(video, list)
else np.array(video[frame_id])
)
lf = labels[labels["FRAME"].astype(int) == frame_id.item()]
for instance in sorted(lf["TRACK_ID"].unique()):
gt_track_ids.append(int(instance))
x = lf[lf["TRACK_ID"] == instance]["POSITION_X"].iloc[0]
y = lf[lf["TRACK_ID"] == instance]["POSITION_Y"].iloc[0]
centroids.append([x, y])
# albumentations wants (spatial, channels), ensure correct dims
if self.augmentations is not None:
for transform in self.augmentations:
# for occlusion simulation, can remove if we don't want
if isinstance(transform, A.CoarseDropout):
transform.fill_value = random.randint(0, 255)
augmented = self.augmentations(
image=img,
keypoints=np.vstack(centroids),
)
img, centroids = augmented["image"], augmented["keypoints"]
img = torch.Tensor(img)
# torch wants (channels, spatial) - ensure correct dims
if len(img.shape) == 2:
img = img.unsqueeze(0)
elif len(img.shape) == 3:
if img.shape[2] == 3:
img = img.T # todo: check for edge cases
for gt_id in range(len(gt_track_ids)):
c = centroids[gt_id]
bbox = data_utils.pad_bbox(
data_utils.get_bbox([int(c[0]), int(c[1])], self.crop_size),
padding=self.padding,
)
crop = data_utils.crop_bbox(img, bbox)
instances.append(
Instance(
gt_track_id=gt_track_ids[gt_id],
pred_track_id=-1,
bbox=bbox,
crop=crop,
)
)
if self.mode == "train":
np.random.shuffle(instances)
frames.append(
Frame(
video_id=label_idx,
frame_id=frame_id,
img_shape=img.shape,
instances=instances,
)
)
return frames
SleapDataset
¶
Bases: BaseDataset
Dataset for loading animal behavior data from sleap.
Methods:
| Name | Description |
|---|---|
__del__ |
Handle file closing before garbage collection. |
__init__ |
Initialize SleapDataset. |
get_indices |
Retrieve label and frame indices given batch index. |
get_instances |
Get an element of the dataset. |
Source code in dreem/datasets/sleap_dataset.py
class SleapDataset(BaseDataset):
"""Dataset for loading animal behavior data from sleap."""
def __init__(
self,
slp_files: list[str],
video_files: list[str],
data_dirs: Optional[list[str]] = None,
padding: int = 5,
crop_size: Union[int, list[int]] = 128,
anchors: int | list[str] | str = "",
chunk: bool = True,
clip_length: int = 16,
mode: str = "train",
handle_missing: str = "centroid",
augmentations: dict | None = None,
n_chunks: int | float = 1.0,
seed: int | None = None,
verbose: bool = False,
normalize_image: bool = True,
max_batching_gap: int = 15,
use_tight_bbox: bool = False,
dilation_radius_px: Union[int, list[int]] = 0,
max_detection_overlap: float = 0,
max_tracks: int = inf,
**kwargs,
):
"""Initialize SleapDataset.
Args:
slp_files: a list of .slp files storing tracking annotations
video_files: a list of paths to video files
data_dirs: a path, or a list of paths to data directories. If provided, crop_size should be a list of integers
with the same length as data_dirs.
padding: amount of padding around object crops
crop_size: the size of the object crops. Can be either:
- An integer specifying a single crop size for all objects
- A list of integers specifying different crop sizes for different data directories
anchors: One of:
* a string indicating a single node to center crops around
* a list of skeleton node names to be used as the center of crops
* an int indicating the number of anchors to randomly select
If unavailable then crop around the midpoint between all visible anchors.
chunk: whether or not to chunk the dataset into batches
clip_length: the number of frames in each chunk
mode: `train`, `val`, or `test`. Determines whether this dataset is used for
training, validation/testing/inference.
handle_missing: how to handle missing single nodes. one of `["drop", "ignore", "centroid"]`.
if "drop" then we dont include instances which are missing the `anchor`.
if "ignore" then we use a mask instead of a crop and nan centroids/bboxes.
if "centroid" then we default to the pose centroid as the node to crop around.
augmentations: An optional dict mapping augmentations to parameters. The keys
should map directly to augmentation classes in albumentations. Example:
augmentations = {
'Rotate': {'limit': [-90, 90], 'p': 0.5},
'GaussianBlur': {'blur_limit': (3, 7), 'sigma_limit': 0, 'p': 0.2},
'RandomContrast': {'limit': 0.2, 'p': 0.6}
}
n_chunks: Number of chunks to subsample from.
Can either a fraction of the dataset (ie (0,1.0]) or number of chunks
seed: set a seed for reproducibility
verbose: boolean representing whether to print
normalize_image: whether to normalize the image to [0, 1]
max_batching_gap: the max number of frames that can be unlabelled before starting a new batch
use_tight_bbox: whether to use tight bounding box (around keypoints) instead of the default square bounding box
dilation_radius_px: radius of the keypoints dilation in pixels. 0 means no mask applied
max_detection_overlap: the iom threshold for non-maximum suppression of detections
max_tracks: the maximum number of tracks that can be created while tracking. Remove any detections that exceed this number.
**kwargs: Additional keyword arguments (unused but accepted for compatibility)
"""
super().__init__(
slp_files,
video_files,
padding,
crop_size,
chunk,
clip_length,
mode,
augmentations,
n_chunks,
seed,
)
self.slp_files = slp_files
self.data_dirs = data_dirs
self.video_files = video_files
self.padding = padding
self.crop_size = crop_size
self.chunk = chunk
self.clip_length = clip_length
self.mode = mode.lower()
self.handle_missing = handle_missing.lower()
self.n_chunks = n_chunks
self.seed = seed
self.normalize_image = normalize_image
self.max_batching_gap = max_batching_gap
self.use_tight_bbox = use_tight_bbox
self.dilation_radius_px = dilation_radius_px
self.max_detection_overlap = (
max_detection_overlap if max_detection_overlap is not None else 0
)
self.max_tracks = max_tracks if max_tracks is not None else inf
if isinstance(anchors, int):
self.anchors = anchors
elif isinstance(anchors, str):
self.anchors = [anchors]
else:
self.anchors = anchors
if not isinstance(self.data_dirs, list):
self.data_dirs = [self.data_dirs]
if not isinstance(self.crop_size, list):
# make a list so its handled consistently if multiple crops are used
if len(self.data_dirs) > 0: # for test mode, data_dirs is []
self.crop_size = [self.crop_size] * len(self.data_dirs)
else:
self.crop_size = [self.crop_size]
if not isinstance(self.dilation_radius_px, list):
self.dilation_radius_px = [self.dilation_radius_px] * len(self.data_dirs)
else:
self.dilation_radius_px = [self.dilation_radius_px]
if len(self.data_dirs) > 0 and len(self.crop_size) != len(self.data_dirs):
raise ValueError(
f"If a list of crop sizes or data directories are given,"
f"they must have the same length but got {len(self.crop_size)} "
f"and {len(self.data_dirs)}"
)
if (
isinstance(self.anchors, list) and len(self.anchors) == 0
) or self.anchors == 0:
raise ValueError(f"Must provide at least one anchor but got {self.anchors}")
self.verbose = verbose
# if self.seed is not None:
# np.random.seed(self.seed)
# load_slp is a wrapper around sio.load_slp for frame gap checks
self.labels = []
self.annotated_segments = {}
for slp_file in self.slp_files:
labels, annotated_segments = data_utils.load_slp(slp_file)
self.labels.append(labels)
self.annotated_segments[slp_file] = annotated_segments
self.videos = [imageio.get_reader(vid_file) for vid_file in self.vid_files]
# preprocessors
self.remove_excess_detections = RemoveExcessDetections(max_tracks)
self.non_max_suppression = NonMaxSuppression(max_detection_overlap)
# Method in BaseDataset. Creates label_idx and chunked_frame_idx to be
# used in call to get_instances()
self.create_chunks_slp()
def get_indices(self, idx: int) -> tuple:
"""Retrieve label and frame indices given batch index.
Args:
idx: the index of the batch.
"""
return self.label_idx[idx], self.chunked_frame_idx[idx]
def get_instances(
self, label_idx: list[int], frame_idx: torch.Tensor
) -> list[Frame]:
"""Get an element of the dataset.
Args:
label_idx: index of the labels
frame_idx: indices of the frames to load in to the batch
Returns:
A list of `dreem.io.Frame` objects containing metadata and instance data for the batch/clip.
"""
sleap_labels_obj = self.labels[label_idx]
video_name = self.video_files[label_idx]
# get the correct crop size based on the video
video_par_path = Path(video_name).parent
if len(self.data_dirs) > 0:
crop_size = self.crop_size[0]
dilation_radius_px = self.dilation_radius_px[0]
for j, data_dir in enumerate(self.data_dirs):
if Path(data_dir) == video_par_path:
crop_size = self.crop_size[j]
dilation_radius_px = self.dilation_radius_px[j]
break
else:
crop_size = self.crop_size[0]
dilation_radius_px = self.dilation_radius_px[0]
vid_reader = self.videos[label_idx]
skeleton = sleap_labels_obj.skeletons[-1]
frames = []
max_crop_h, max_crop_w = 0, 0
for i, frame_ind in enumerate(frame_idx):
(
instances,
gt_track_ids,
poses,
shown_poses,
point_scores,
instance_score,
) = ([], [], [], [], [], [])
frame_ind = int(frame_ind)
# sleap-io method for indexing a Labels() object based on the frame's index
lf = sleap_labels_obj[(sleap_labels_obj.video, frame_ind)]
if frame_ind != lf.frame_idx:
logger.warning(f"Frame index mismatch: {frame_ind} != {lf.frame_idx}")
try:
img = vid_reader.get_data(int(frame_ind))
except IndexError as e:
logger.warning(
f"Could not read frame {frame_ind} from {video_name} due to {e}"
)
continue
if len(img.shape) == 2:
img = img.expand_dims(-1)
h, w, c = img.shape
if c == 1:
img = np.concatenate(
[img, img, img], axis=-1
) # convert to grayscale to rgb
if np.issubdtype(img.dtype, np.integer): # convert int to float
img = img.astype(np.float32)
if self.normalize_image:
img = img / 255
n_instances_dropped = 0
gt_instances = []
# don't load instances that have been 'greyed out' i.e. all nans for keypoints
for inst in lf.instances:
pts = np.array([p for p in inst.numpy()])
if np.isnan(pts).all():
continue
else:
gt_instances.append(inst)
dict_instances = {}
no_track_instances = []
for instance in gt_instances:
if instance.track is not None:
gt_track_id = sleap_labels_obj.tracks.index(instance.track)
if gt_track_id not in dict_instances:
dict_instances[gt_track_id] = instance
else:
existing_instance = dict_instances[gt_track_id]
# if existing is PredictedInstance and current is not, then current is a UserInstance and should be used
if isinstance(
existing_instance, sio.PredictedInstance
) and not isinstance(instance, sio.PredictedInstance):
dict_instances[gt_track_id] = instance
else:
no_track_instances.append(instance)
gt_instances = list(dict_instances.values()) + no_track_instances
if self.mode == "train":
np.random.shuffle(gt_instances)
for instance in gt_instances:
if (
np.random.uniform() < self.instance_dropout["p"]
and n_instances_dropped < self.instance_dropout["n"]
):
n_instances_dropped += 1
continue
if instance.track is not None:
gt_track_id = sleap_labels_obj.tracks.index(instance.track)
else:
gt_track_id = -1
gt_track_ids.append(gt_track_id)
poses.append(
dict(
zip(
[n.name for n in instance.skeleton.nodes],
[p for p in instance.numpy()],
)
)
)
shown_poses = [
{
key: val
for key, val in instance.items()
if not np.isnan(val).any()
}
for instance in poses
]
point_scores.append(
np.array(
[
(
1.0 # point scores not reliably available in sleap io PredictedPointsArray
# point.score
# if isinstance(point, sio.PredictedPoint)
# else 1.0
)
for point in instance.numpy()
]
)
)
if isinstance(instance, sio.PredictedInstance):
instance_score.append(instance.score)
else:
instance_score.append(1.0)
# augmentations
if self.augmentations is not None:
for transform in self.augmentations:
if isinstance(transform, A.CoarseDropout):
transform.fill_value = random.randint(0, 255)
if shown_poses:
keypoints = np.vstack([list(s.values()) for s in shown_poses])
else:
keypoints = []
augmented = self.augmentations(image=img, keypoints=keypoints)
img, aug_poses = augmented["image"], augmented["keypoints"]
aug_poses = [
arr
for arr in np.split(
np.array(aug_poses),
np.array([len(s) for s in shown_poses]).cumsum(),
)
if arr.size != 0
]
aug_poses = [
dict(zip(list(pose_dict.keys()), aug_pose_arr.tolist()))
for aug_pose_arr, pose_dict in zip(aug_poses, shown_poses)
]
_ = [
pose.update(aug_pose)
for pose, aug_pose in zip(shown_poses, aug_poses)
]
img = tvf.to_tensor(img)
for j in range(len(gt_track_ids)):
pose = shown_poses[j]
"""Check for anchor"""
crops = []
boxes = []
centroids = {}
if isinstance(self.anchors, int):
anchors_to_choose = list(pose.keys()) + ["midpoint"]
anchors = np.random.choice(anchors_to_choose, self.anchors)
else:
anchors = self.anchors
dropped_anchors = self.node_dropout(anchors)
for anchor in anchors:
if anchor in dropped_anchors:
centroid = np.array([np.nan, np.nan])
elif anchor == "midpoint" or anchor == "centroid":
centroid = np.nanmean(np.array(list(pose.values())), axis=0)
elif anchor in pose:
centroid = np.array(pose[anchor])
if np.isnan(centroid).any():
centroid = np.array([np.nan, np.nan])
elif (
anchor not in pose
and len(anchors) == 1
and self.handle_missing == "centroid"
):
anchor = "midpoint"
centroid = np.nanmean(np.array(list(pose.values())), axis=0)
else:
centroid = np.array([np.nan, np.nan])
arr_pose = np.array(list(pose.values()))
if np.isnan(centroid).all():
bbox = torch.tensor([np.nan, np.nan, np.nan, np.nan])
else:
if self.use_tight_bbox and len(pose) > 1:
# tight bbox, dont allow this for centroid-only poses!
# note bbox will be a different size for each instance; padded at the end of the loop
bbox = data_utils.get_tight_bbox(arr_pose)
else:
bbox = data_utils.pad_bbox(
data_utils.get_bbox(centroid, crop_size),
padding=self.padding,
)
if bbox.isnan().all():
crop = torch.zeros(
c,
crop_size + 2 * self.padding,
crop_size + 2 * self.padding,
dtype=img.dtype,
)
else:
crop = data_utils.crop_bbox(img, bbox)
if dilation_radius_px > 0:
if np.isnan(arr_pose).any():
logger.warning("arr_pose is nan")
mask = data_utils.get_mask_from_keypoints(
arr_pose, crop, dilation_radius_px, bbox
)
crop = crop * mask
# logger.debug(f"Applying mask to crop {frame_ind}_{j}")
crops.append(crop)
# get max h,w for padding for tight bboxes
c, h, w = crop.shape
if h > max_crop_h:
max_crop_h = h
if w > max_crop_w:
max_crop_w = w
centroids[anchor] = centroid
boxes.append(bbox)
if len(crops) > 0:
crops = torch.concat(crops, dim=0)
if len(boxes) > 0:
boxes = torch.stack(boxes, dim=0)
if self.handle_missing == "drop" and boxes.isnan().any():
continue
instance = Instance(
gt_track_id=gt_track_ids[j],
pred_track_id=-1,
crop=crops,
centroid=centroids,
bbox=boxes,
skeleton=skeleton,
pose=poses[j],
point_scores=point_scores[j],
instance_score=instance_score[j],
)
instances.append(instance)
# remove excess detections
if len(instances) > self.max_tracks:
state = self.remove_excess_detections.run(
{
"frame_ind": frame_ind,
"instances": instances,
}
)
instances = state["instances"]
# non-maximum suppression (high overlap bounding boxes)
if self.max_detection_overlap > 0 and len(instances) > 0:
state = self.non_max_suppression.run(
{
"frame_ind": frame_ind,
"instances": instances,
}
)
instances = state["instances"]
frame = Frame(
video_id=label_idx,
frame_id=frame_ind,
vid_file=video_name,
img_shape=img.shape,
instances=instances,
)
frames.append(frame)
return frames
def __del__(self):
"""Handle file closing before garbage collection."""
for reader in self.videos:
reader.close()
__del__()
¶
__init__(slp_files, video_files, data_dirs=None, padding=5, crop_size=128, anchors='', chunk=True, clip_length=16, mode='train', handle_missing='centroid', augmentations=None, n_chunks=1.0, seed=None, verbose=False, normalize_image=True, max_batching_gap=15, use_tight_bbox=False, dilation_radius_px=0, max_detection_overlap=0, max_tracks=inf, **kwargs)
¶
Initialize SleapDataset.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
slp_files
|
list[str]
|
a list of .slp files storing tracking annotations |
required |
video_files
|
list[str]
|
a list of paths to video files |
required |
data_dirs
|
Optional[list[str]]
|
a path, or a list of paths to data directories. If provided, crop_size should be a list of integers with the same length as data_dirs. |
None
|
padding
|
int
|
amount of padding around object crops |
5
|
crop_size
|
Union[int, list[int]]
|
the size of the object crops. Can be either: - An integer specifying a single crop size for all objects - A list of integers specifying different crop sizes for different data directories |
128
|
anchors
|
int | list[str] | str
|
One of: * a string indicating a single node to center crops around * a list of skeleton node names to be used as the center of crops * an int indicating the number of anchors to randomly select If unavailable then crop around the midpoint between all visible anchors. |
''
|
chunk
|
bool
|
whether or not to chunk the dataset into batches |
True
|
clip_length
|
int
|
the number of frames in each chunk |
16
|
mode
|
str
|
|
'train'
|
handle_missing
|
str
|
how to handle missing single nodes. one of |
'centroid'
|
augmentations
|
dict | None
|
An optional dict mapping augmentations to parameters. The keys should map directly to augmentation classes in albumentations. Example: augmentations = { 'Rotate': {'limit': [-90, 90], 'p': 0.5}, 'GaussianBlur': {'blur_limit': (3, 7), 'sigma_limit': 0, 'p': 0.2}, 'RandomContrast': {'limit': 0.2, 'p': 0.6} } |
None
|
n_chunks
|
int | float
|
Number of chunks to subsample from. Can either a fraction of the dataset (ie (0,1.0]) or number of chunks |
1.0
|
seed
|
int | None
|
set a seed for reproducibility |
None
|
verbose
|
bool
|
boolean representing whether to print |
False
|
normalize_image
|
bool
|
whether to normalize the image to [0, 1] |
True
|
max_batching_gap
|
int
|
the max number of frames that can be unlabelled before starting a new batch |
15
|
use_tight_bbox
|
bool
|
whether to use tight bounding box (around keypoints) instead of the default square bounding box |
False
|
dilation_radius_px
|
Union[int, list[int]]
|
radius of the keypoints dilation in pixels. 0 means no mask applied |
0
|
max_detection_overlap
|
float
|
the iom threshold for non-maximum suppression of detections |
0
|
max_tracks
|
int
|
the maximum number of tracks that can be created while tracking. Remove any detections that exceed this number. |
inf
|
**kwargs
|
Additional keyword arguments (unused but accepted for compatibility) |
{}
|
Source code in dreem/datasets/sleap_dataset.py
def __init__(
self,
slp_files: list[str],
video_files: list[str],
data_dirs: Optional[list[str]] = None,
padding: int = 5,
crop_size: Union[int, list[int]] = 128,
anchors: int | list[str] | str = "",
chunk: bool = True,
clip_length: int = 16,
mode: str = "train",
handle_missing: str = "centroid",
augmentations: dict | None = None,
n_chunks: int | float = 1.0,
seed: int | None = None,
verbose: bool = False,
normalize_image: bool = True,
max_batching_gap: int = 15,
use_tight_bbox: bool = False,
dilation_radius_px: Union[int, list[int]] = 0,
max_detection_overlap: float = 0,
max_tracks: int = inf,
**kwargs,
):
"""Initialize SleapDataset.
Args:
slp_files: a list of .slp files storing tracking annotations
video_files: a list of paths to video files
data_dirs: a path, or a list of paths to data directories. If provided, crop_size should be a list of integers
with the same length as data_dirs.
padding: amount of padding around object crops
crop_size: the size of the object crops. Can be either:
- An integer specifying a single crop size for all objects
- A list of integers specifying different crop sizes for different data directories
anchors: One of:
* a string indicating a single node to center crops around
* a list of skeleton node names to be used as the center of crops
* an int indicating the number of anchors to randomly select
If unavailable then crop around the midpoint between all visible anchors.
chunk: whether or not to chunk the dataset into batches
clip_length: the number of frames in each chunk
mode: `train`, `val`, or `test`. Determines whether this dataset is used for
training, validation/testing/inference.
handle_missing: how to handle missing single nodes. one of `["drop", "ignore", "centroid"]`.
if "drop" then we dont include instances which are missing the `anchor`.
if "ignore" then we use a mask instead of a crop and nan centroids/bboxes.
if "centroid" then we default to the pose centroid as the node to crop around.
augmentations: An optional dict mapping augmentations to parameters. The keys
should map directly to augmentation classes in albumentations. Example:
augmentations = {
'Rotate': {'limit': [-90, 90], 'p': 0.5},
'GaussianBlur': {'blur_limit': (3, 7), 'sigma_limit': 0, 'p': 0.2},
'RandomContrast': {'limit': 0.2, 'p': 0.6}
}
n_chunks: Number of chunks to subsample from.
Can either a fraction of the dataset (ie (0,1.0]) or number of chunks
seed: set a seed for reproducibility
verbose: boolean representing whether to print
normalize_image: whether to normalize the image to [0, 1]
max_batching_gap: the max number of frames that can be unlabelled before starting a new batch
use_tight_bbox: whether to use tight bounding box (around keypoints) instead of the default square bounding box
dilation_radius_px: radius of the keypoints dilation in pixels. 0 means no mask applied
max_detection_overlap: the iom threshold for non-maximum suppression of detections
max_tracks: the maximum number of tracks that can be created while tracking. Remove any detections that exceed this number.
**kwargs: Additional keyword arguments (unused but accepted for compatibility)
"""
super().__init__(
slp_files,
video_files,
padding,
crop_size,
chunk,
clip_length,
mode,
augmentations,
n_chunks,
seed,
)
self.slp_files = slp_files
self.data_dirs = data_dirs
self.video_files = video_files
self.padding = padding
self.crop_size = crop_size
self.chunk = chunk
self.clip_length = clip_length
self.mode = mode.lower()
self.handle_missing = handle_missing.lower()
self.n_chunks = n_chunks
self.seed = seed
self.normalize_image = normalize_image
self.max_batching_gap = max_batching_gap
self.use_tight_bbox = use_tight_bbox
self.dilation_radius_px = dilation_radius_px
self.max_detection_overlap = (
max_detection_overlap if max_detection_overlap is not None else 0
)
self.max_tracks = max_tracks if max_tracks is not None else inf
if isinstance(anchors, int):
self.anchors = anchors
elif isinstance(anchors, str):
self.anchors = [anchors]
else:
self.anchors = anchors
if not isinstance(self.data_dirs, list):
self.data_dirs = [self.data_dirs]
if not isinstance(self.crop_size, list):
# make a list so its handled consistently if multiple crops are used
if len(self.data_dirs) > 0: # for test mode, data_dirs is []
self.crop_size = [self.crop_size] * len(self.data_dirs)
else:
self.crop_size = [self.crop_size]
if not isinstance(self.dilation_radius_px, list):
self.dilation_radius_px = [self.dilation_radius_px] * len(self.data_dirs)
else:
self.dilation_radius_px = [self.dilation_radius_px]
if len(self.data_dirs) > 0 and len(self.crop_size) != len(self.data_dirs):
raise ValueError(
f"If a list of crop sizes or data directories are given,"
f"they must have the same length but got {len(self.crop_size)} "
f"and {len(self.data_dirs)}"
)
if (
isinstance(self.anchors, list) and len(self.anchors) == 0
) or self.anchors == 0:
raise ValueError(f"Must provide at least one anchor but got {self.anchors}")
self.verbose = verbose
# if self.seed is not None:
# np.random.seed(self.seed)
# load_slp is a wrapper around sio.load_slp for frame gap checks
self.labels = []
self.annotated_segments = {}
for slp_file in self.slp_files:
labels, annotated_segments = data_utils.load_slp(slp_file)
self.labels.append(labels)
self.annotated_segments[slp_file] = annotated_segments
self.videos = [imageio.get_reader(vid_file) for vid_file in self.vid_files]
# preprocessors
self.remove_excess_detections = RemoveExcessDetections(max_tracks)
self.non_max_suppression = NonMaxSuppression(max_detection_overlap)
# Method in BaseDataset. Creates label_idx and chunked_frame_idx to be
# used in call to get_instances()
self.create_chunks_slp()
get_indices(idx)
¶
Retrieve label and frame indices given batch index.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
idx
|
int
|
the index of the batch. |
required |
get_instances(label_idx, frame_idx)
¶
Get an element of the dataset.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
label_idx
|
list[int]
|
index of the labels |
required |
frame_idx
|
Tensor
|
indices of the frames to load in to the batch |
required |
Returns:
| Type | Description |
|---|---|
list[Frame]
|
A list of |
Source code in dreem/datasets/sleap_dataset.py
def get_instances(
self, label_idx: list[int], frame_idx: torch.Tensor
) -> list[Frame]:
"""Get an element of the dataset.
Args:
label_idx: index of the labels
frame_idx: indices of the frames to load in to the batch
Returns:
A list of `dreem.io.Frame` objects containing metadata and instance data for the batch/clip.
"""
sleap_labels_obj = self.labels[label_idx]
video_name = self.video_files[label_idx]
# get the correct crop size based on the video
video_par_path = Path(video_name).parent
if len(self.data_dirs) > 0:
crop_size = self.crop_size[0]
dilation_radius_px = self.dilation_radius_px[0]
for j, data_dir in enumerate(self.data_dirs):
if Path(data_dir) == video_par_path:
crop_size = self.crop_size[j]
dilation_radius_px = self.dilation_radius_px[j]
break
else:
crop_size = self.crop_size[0]
dilation_radius_px = self.dilation_radius_px[0]
vid_reader = self.videos[label_idx]
skeleton = sleap_labels_obj.skeletons[-1]
frames = []
max_crop_h, max_crop_w = 0, 0
for i, frame_ind in enumerate(frame_idx):
(
instances,
gt_track_ids,
poses,
shown_poses,
point_scores,
instance_score,
) = ([], [], [], [], [], [])
frame_ind = int(frame_ind)
# sleap-io method for indexing a Labels() object based on the frame's index
lf = sleap_labels_obj[(sleap_labels_obj.video, frame_ind)]
if frame_ind != lf.frame_idx:
logger.warning(f"Frame index mismatch: {frame_ind} != {lf.frame_idx}")
try:
img = vid_reader.get_data(int(frame_ind))
except IndexError as e:
logger.warning(
f"Could not read frame {frame_ind} from {video_name} due to {e}"
)
continue
if len(img.shape) == 2:
img = img.expand_dims(-1)
h, w, c = img.shape
if c == 1:
img = np.concatenate(
[img, img, img], axis=-1
) # convert to grayscale to rgb
if np.issubdtype(img.dtype, np.integer): # convert int to float
img = img.astype(np.float32)
if self.normalize_image:
img = img / 255
n_instances_dropped = 0
gt_instances = []
# don't load instances that have been 'greyed out' i.e. all nans for keypoints
for inst in lf.instances:
pts = np.array([p for p in inst.numpy()])
if np.isnan(pts).all():
continue
else:
gt_instances.append(inst)
dict_instances = {}
no_track_instances = []
for instance in gt_instances:
if instance.track is not None:
gt_track_id = sleap_labels_obj.tracks.index(instance.track)
if gt_track_id not in dict_instances:
dict_instances[gt_track_id] = instance
else:
existing_instance = dict_instances[gt_track_id]
# if existing is PredictedInstance and current is not, then current is a UserInstance and should be used
if isinstance(
existing_instance, sio.PredictedInstance
) and not isinstance(instance, sio.PredictedInstance):
dict_instances[gt_track_id] = instance
else:
no_track_instances.append(instance)
gt_instances = list(dict_instances.values()) + no_track_instances
if self.mode == "train":
np.random.shuffle(gt_instances)
for instance in gt_instances:
if (
np.random.uniform() < self.instance_dropout["p"]
and n_instances_dropped < self.instance_dropout["n"]
):
n_instances_dropped += 1
continue
if instance.track is not None:
gt_track_id = sleap_labels_obj.tracks.index(instance.track)
else:
gt_track_id = -1
gt_track_ids.append(gt_track_id)
poses.append(
dict(
zip(
[n.name for n in instance.skeleton.nodes],
[p for p in instance.numpy()],
)
)
)
shown_poses = [
{
key: val
for key, val in instance.items()
if not np.isnan(val).any()
}
for instance in poses
]
point_scores.append(
np.array(
[
(
1.0 # point scores not reliably available in sleap io PredictedPointsArray
# point.score
# if isinstance(point, sio.PredictedPoint)
# else 1.0
)
for point in instance.numpy()
]
)
)
if isinstance(instance, sio.PredictedInstance):
instance_score.append(instance.score)
else:
instance_score.append(1.0)
# augmentations
if self.augmentations is not None:
for transform in self.augmentations:
if isinstance(transform, A.CoarseDropout):
transform.fill_value = random.randint(0, 255)
if shown_poses:
keypoints = np.vstack([list(s.values()) for s in shown_poses])
else:
keypoints = []
augmented = self.augmentations(image=img, keypoints=keypoints)
img, aug_poses = augmented["image"], augmented["keypoints"]
aug_poses = [
arr
for arr in np.split(
np.array(aug_poses),
np.array([len(s) for s in shown_poses]).cumsum(),
)
if arr.size != 0
]
aug_poses = [
dict(zip(list(pose_dict.keys()), aug_pose_arr.tolist()))
for aug_pose_arr, pose_dict in zip(aug_poses, shown_poses)
]
_ = [
pose.update(aug_pose)
for pose, aug_pose in zip(shown_poses, aug_poses)
]
img = tvf.to_tensor(img)
for j in range(len(gt_track_ids)):
pose = shown_poses[j]
"""Check for anchor"""
crops = []
boxes = []
centroids = {}
if isinstance(self.anchors, int):
anchors_to_choose = list(pose.keys()) + ["midpoint"]
anchors = np.random.choice(anchors_to_choose, self.anchors)
else:
anchors = self.anchors
dropped_anchors = self.node_dropout(anchors)
for anchor in anchors:
if anchor in dropped_anchors:
centroid = np.array([np.nan, np.nan])
elif anchor == "midpoint" or anchor == "centroid":
centroid = np.nanmean(np.array(list(pose.values())), axis=0)
elif anchor in pose:
centroid = np.array(pose[anchor])
if np.isnan(centroid).any():
centroid = np.array([np.nan, np.nan])
elif (
anchor not in pose
and len(anchors) == 1
and self.handle_missing == "centroid"
):
anchor = "midpoint"
centroid = np.nanmean(np.array(list(pose.values())), axis=0)
else:
centroid = np.array([np.nan, np.nan])
arr_pose = np.array(list(pose.values()))
if np.isnan(centroid).all():
bbox = torch.tensor([np.nan, np.nan, np.nan, np.nan])
else:
if self.use_tight_bbox and len(pose) > 1:
# tight bbox, dont allow this for centroid-only poses!
# note bbox will be a different size for each instance; padded at the end of the loop
bbox = data_utils.get_tight_bbox(arr_pose)
else:
bbox = data_utils.pad_bbox(
data_utils.get_bbox(centroid, crop_size),
padding=self.padding,
)
if bbox.isnan().all():
crop = torch.zeros(
c,
crop_size + 2 * self.padding,
crop_size + 2 * self.padding,
dtype=img.dtype,
)
else:
crop = data_utils.crop_bbox(img, bbox)
if dilation_radius_px > 0:
if np.isnan(arr_pose).any():
logger.warning("arr_pose is nan")
mask = data_utils.get_mask_from_keypoints(
arr_pose, crop, dilation_radius_px, bbox
)
crop = crop * mask
# logger.debug(f"Applying mask to crop {frame_ind}_{j}")
crops.append(crop)
# get max h,w for padding for tight bboxes
c, h, w = crop.shape
if h > max_crop_h:
max_crop_h = h
if w > max_crop_w:
max_crop_w = w
centroids[anchor] = centroid
boxes.append(bbox)
if len(crops) > 0:
crops = torch.concat(crops, dim=0)
if len(boxes) > 0:
boxes = torch.stack(boxes, dim=0)
if self.handle_missing == "drop" and boxes.isnan().any():
continue
instance = Instance(
gt_track_id=gt_track_ids[j],
pred_track_id=-1,
crop=crops,
centroid=centroids,
bbox=boxes,
skeleton=skeleton,
pose=poses[j],
point_scores=point_scores[j],
instance_score=instance_score[j],
)
instances.append(instance)
# remove excess detections
if len(instances) > self.max_tracks:
state = self.remove_excess_detections.run(
{
"frame_ind": frame_ind,
"instances": instances,
}
)
instances = state["instances"]
# non-maximum suppression (high overlap bounding boxes)
if self.max_detection_overlap > 0 and len(instances) > 0:
state = self.non_max_suppression.run(
{
"frame_ind": frame_ind,
"instances": instances,
}
)
instances = state["instances"]
frame = Frame(
video_id=label_idx,
frame_id=frame_ind,
vid_file=video_name,
img_shape=img.shape,
instances=instances,
)
frames.append(frame)
return frames
TrackingDataset
¶
Bases: LightningDataModule
Lightning dataset used to load dataloaders for train, test and validation.
Nice for wrapping around other data formats.
Methods:
| Name | Description |
|---|---|
__init__ |
Initialize tracking dataset. |
setup |
Set up lightning dataset. |
test_dataloader |
Get. |
train_dataloader |
Get train_dataloader. |
val_dataloader |
Get val dataloader. |
Source code in dreem/datasets/tracking_dataset.py
class TrackingDataset(LightningDataModule):
"""Lightning dataset used to load dataloaders for train, test and validation.
Nice for wrapping around other data formats.
"""
def __init__(
self,
train_ds: SleapDataset | MicroscopyDataset | CellTrackingDataset | None = None,
train_dl: DataLoader | None = None,
val_ds: SleapDataset | MicroscopyDataset | CellTrackingDataset | None = None,
val_dl: DataLoader | None = None,
test_ds: SleapDataset | MicroscopyDataset | CellTrackingDataset | None = None,
test_dl: DataLoader | None = None,
):
"""Initialize tracking dataset.
Args:
train_ds: Sleap or Microscopy training Dataset
train_dl: Training dataloader. Only used for overriding `train_dataloader`.
val_ds: Sleap or Microscopy Validation set
val_dl : Validation dataloader. Only used for overriding `val_dataloader`.
test_ds: Sleap or Microscopy test set
test_dl : Test dataloader. Only used for overriding `test_dataloader`.
"""
super().__init__()
self.train_ds = train_ds
self.train_dl = train_dl
self.val_ds = val_ds
self.val_dl = val_dl
self.test_ds = test_ds
self.test_dl = test_dl
def setup(self, stage=None):
"""Set up lightning dataset.
UNUSED.
"""
pass
def train_dataloader(self) -> DataLoader:
"""Get train_dataloader.
Returns: The Training Dataloader.
"""
if self.train_dl is None and self.train_ds is None:
return None
elif self.train_dl is None:
return DataLoader(
self.train_ds,
batch_size=1,
shuffle=True,
pin_memory=False,
collate_fn=self.train_ds.no_batching_fn,
num_workers=0,
generator=(
torch.Generator(device="cuda")
if torch.cuda.is_available()
else torch.Generator()
),
)
else:
return self.train_dl
def val_dataloader(self) -> DataLoader:
"""Get val dataloader.
Returns: The validation dataloader.
"""
if self.val_dl is None and self.val_ds is None:
return None
elif self.val_dl is None:
return DataLoader(
self.val_ds,
batch_size=1,
shuffle=False,
pin_memory=0,
collate_fn=self.train_ds.no_batching_fn,
num_workers=False,
generator=None,
)
else:
return self.val_dl
def test_dataloader(self) -> DataLoader:
"""Get.
Returns: The test dataloader
"""
if self.test_dl is None and self.test_ds is None:
return None
elif self.test_dl is None:
return DataLoader(
self.test_ds,
batch_size=1,
shuffle=False,
pin_memory=0,
collate_fn=self.train_ds.no_batching_fn,
num_workers=False,
generator=None,
)
else:
return self.test_dl
__init__(train_ds=None, train_dl=None, val_ds=None, val_dl=None, test_ds=None, test_dl=None)
¶
Initialize tracking dataset.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
train_ds
|
SleapDataset | MicroscopyDataset | CellTrackingDataset | None
|
Sleap or Microscopy training Dataset |
None
|
train_dl
|
DataLoader | None
|
Training dataloader. Only used for overriding |
None
|
val_ds
|
SleapDataset | MicroscopyDataset | CellTrackingDataset | None
|
Sleap or Microscopy Validation set |
None
|
val_dl
|
Validation dataloader. Only used for overriding |
required | |
test_ds
|
SleapDataset | MicroscopyDataset | CellTrackingDataset | None
|
Sleap or Microscopy test set |
None
|
test_dl
|
Test dataloader. Only used for overriding |
required |
Source code in dreem/datasets/tracking_dataset.py
def __init__(
self,
train_ds: SleapDataset | MicroscopyDataset | CellTrackingDataset | None = None,
train_dl: DataLoader | None = None,
val_ds: SleapDataset | MicroscopyDataset | CellTrackingDataset | None = None,
val_dl: DataLoader | None = None,
test_ds: SleapDataset | MicroscopyDataset | CellTrackingDataset | None = None,
test_dl: DataLoader | None = None,
):
"""Initialize tracking dataset.
Args:
train_ds: Sleap or Microscopy training Dataset
train_dl: Training dataloader. Only used for overriding `train_dataloader`.
val_ds: Sleap or Microscopy Validation set
val_dl : Validation dataloader. Only used for overriding `val_dataloader`.
test_ds: Sleap or Microscopy test set
test_dl : Test dataloader. Only used for overriding `test_dataloader`.
"""
super().__init__()
self.train_ds = train_ds
self.train_dl = train_dl
self.val_ds = val_ds
self.val_dl = val_dl
self.test_ds = test_ds
self.test_dl = test_dl
setup(stage=None)
¶
test_dataloader()
¶
Get.
Returns: The test dataloader
Source code in dreem/datasets/tracking_dataset.py
def test_dataloader(self) -> DataLoader:
"""Get.
Returns: The test dataloader
"""
if self.test_dl is None and self.test_ds is None:
return None
elif self.test_dl is None:
return DataLoader(
self.test_ds,
batch_size=1,
shuffle=False,
pin_memory=0,
collate_fn=self.train_ds.no_batching_fn,
num_workers=False,
generator=None,
)
else:
return self.test_dl
train_dataloader()
¶
Get train_dataloader.
Returns: The Training Dataloader.
Source code in dreem/datasets/tracking_dataset.py
def train_dataloader(self) -> DataLoader:
"""Get train_dataloader.
Returns: The Training Dataloader.
"""
if self.train_dl is None and self.train_ds is None:
return None
elif self.train_dl is None:
return DataLoader(
self.train_ds,
batch_size=1,
shuffle=True,
pin_memory=False,
collate_fn=self.train_ds.no_batching_fn,
num_workers=0,
generator=(
torch.Generator(device="cuda")
if torch.cuda.is_available()
else torch.Generator()
),
)
else:
return self.train_dl
val_dataloader()
¶
Get val dataloader.
Returns: The validation dataloader.
Source code in dreem/datasets/tracking_dataset.py
def val_dataloader(self) -> DataLoader:
"""Get val dataloader.
Returns: The validation dataloader.
"""
if self.val_dl is None and self.val_ds is None:
return None
elif self.val_dl is None:
return DataLoader(
self.val_ds,
batch_size=1,
shuffle=False,
pin_memory=0,
collate_fn=self.train_ds.no_batching_fn,
num_workers=False,
generator=None,
)
else:
return self.val_dl