io
dreem.io
¶
Module containing input/output data structures for easy storage and manipulation.
Modules:
| Name | Description |
|---|---|
association_matrix |
Module containing class for storing and looking up association scores. |
config |
Data structures for handling config parsing. |
flags |
Module containing flag codes for frames and other data structures. |
frame |
Module containing data classes such as Instances and Frames. |
instance |
Module containing data class for storing detections. |
track |
Module containing data structures for storing instances of the same Track. |
visualize |
Helper functions for visualizing tracking. |
Classes:
| Name | Description |
|---|---|
AssociationMatrix |
Class representing the associations between detections. |
Config |
Class handling loading components based on config params. |
Frame |
Data structure containing metadata for a single frame of a video. |
FrameFlagCode |
Enumeration of flag codes for Frame objects. |
Instance |
Class representing a single instance to be tracked. |
Track |
Object for storing instances of the same track. |
AssociationMatrix
¶
Class representing the associations between detections.
Attributes:
| Name | Type | Description |
|---|---|---|
matrix |
ndarray | Tensor
|
the |
ref_instances |
list[Instance]
|
all instances used to associate against. |
query_instances |
list[Instance]
|
query instances that were associated against ref instances. |
Methods:
| Name | Description |
|---|---|
__getindices__ |
Get the indices of the instance for lookup. |
__getitem__ |
Get elements of the association matrix. |
__repr__ |
Get the string representation of the Association Matrix. |
get_tracks |
Group instances by track. |
numpy |
Convert association matrix to a numpy array. |
reduce |
Aggregate the association matrix by specified dimensions and grouping. |
to |
Move instance to different device or change dtype. (See |
to_dataframe |
Convert the association matrix to a pandas DataFrame. |
Source code in dreem/io/association_matrix.py
@attrs.define
class AssociationMatrix:
"""Class representing the associations between detections.
Attributes:
matrix: the `n_query x n_ref` association matrix`
ref_instances: all instances used to associate against.
query_instances: query instances that were associated against ref instances.
"""
matrix: np.ndarray | torch.Tensor
ref_instances: list[Instance] = attrs.field()
query_instances: list[Instance] = attrs.field()
@ref_instances.validator
def _check_ref_instances(self, attribute, value):
"""Check to ensure that the number of association matrix columns and reference instances match.
Args:
attribute: The ref instances.
value: the list of ref instances.
Raises:
ValueError if the number of columns and reference instances don't match.
"""
if len(value) != self.matrix.shape[-1]:
raise ValueError(
(
"Ref instances must equal number of columns in Association matrix"
f"Found {len(value)} ref instances but {self.matrix.shape[-1]} columns."
)
)
@query_instances.validator
def _check_query_instances(self, attribute, value):
"""Check to ensure that the number of association matrix rows and query instances match.
Args:
attribute: The query instances.
value: the list of query instances.
Raises:
ValueError if the number of rows and query instances don't match.
"""
if len(value) != self.matrix.shape[0]:
raise ValueError(
(
"Query instances must equal number of rows in Association matrix"
f"Found {len(value)} query instances but {self.matrix.shape[0]} rows."
)
)
def __repr__(self) -> str:
"""Get the string representation of the Association Matrix.
Returns:
the string representation of the association matrix.
"""
return (
f"AssociationMatrix({self.matrix},"
f"query_instances={len(self.query_instances)},"
f"ref_instances={len(self.ref_instances)})"
)
def numpy(self) -> np.ndarray:
"""Convert association matrix to a numpy array.
Returns:
The association matrix as a numpy array.
"""
if isinstance(self.matrix, torch.Tensor):
return self.matrix.detach().cpu().numpy()
return self.matrix
def to_dataframe(
self, row_labels: str = "gt", col_labels: str = "gt"
) -> pd.DataFrame:
"""Convert the association matrix to a pandas DataFrame.
Args:
row_labels: How to label the rows(queries).
If list, then must match # of rows/queries
If `"gt"` then label by gt track id.
If `"pred"` then label by pred track id.
Otherwise label by the query_instance indices
col_labels: How to label the columns(references).
If list, then must match # of columns/refs
If `"gt"` then label by gt track id.
If `"pred"` then label by pred track id.
Otherwise label by the ref_instance indices
Returns:
The association matrix as a pandas dataframe.
"""
matrix = self.numpy()
if not isinstance(row_labels, str):
if len(row_labels) == len(self.query_instances):
row_inds = row_labels
else:
raise ValueError(
(
"Mismatched # of rows and labels!",
f"Found {len(row_labels)} with {len(self.query_instances)} rows",
)
)
else:
if row_labels == "gt":
row_inds = [
instance.gt_track_id.item() for instance in self.query_instances
]
elif row_labels == "pred":
row_inds = [
instance.pred_track_id.item() for instance in self.query_instances
]
else:
row_inds = np.arange(len(self.query_instances))
if not isinstance(col_labels, str):
if len(col_labels) == len(self.ref_instances):
col_inds = col_labels
else:
raise ValueError(
(
"Mismatched # of columns and labels!",
f"Found {len(col_labels)} with {len(self.ref_instances)} columns",
)
)
else:
if col_labels == "gt":
col_inds = [
instance.gt_track_id.item() for instance in self.ref_instances
]
elif col_labels == "pred":
col_inds = [
instance.pred_track_id.item() for instance in self.ref_instances
]
else:
col_inds = np.arange(len(self.ref_instances))
asso_df = pd.DataFrame(matrix, index=row_inds, columns=col_inds)
return asso_df
def reduce(
self,
row_dims: str = "instance",
col_dims: str = "track",
row_grouping: str | None = None,
col_grouping: str = "pred",
reduce_method: callable = np.sum,
) -> pd.DataFrame:
"""Aggregate the association matrix by specified dimensions and grouping.
Args:
row_dims: A str indicating how to what dimensions to reduce rows to.
Either "instance" (remains unchanged), or "track" (n_rows=n_traj).
col_dims: A str indicating how to dimensions to reduce rows to.
Either "instance" (remains unchanged), or "track" (n_cols=n_traj)
row_grouping: A str indicating how to group rows when aggregating. Either "pred" or "gt".
col_grouping: A str indicating how to group columns when aggregating. Either "pred" or "gt".
reduce_method: A callable function that operates on numpy matrices and can take an `axis` arg for reducing.
Returns:
The association matrix reduced to an inst/traj x traj/inst association matrix as a dataframe.
"""
n_rows = len(self.query_instances)
n_cols = len(self.ref_instances)
col_tracks = {-1: self.ref_instances}
row_tracks = {-1: self.query_instances}
col_inds = [i for i in range(len(self.ref_instances))]
row_inds = [i for i in range(len(self.query_instances))]
if col_dims == "track":
col_tracks = self.get_tracks(self.ref_instances, col_grouping)
col_inds = list(col_tracks.keys())
n_cols = len(col_inds)
if row_dims == "track":
row_tracks = self.get_tracks(self.query_instances, row_grouping)
row_inds = list(row_tracks.keys())
n_rows = len(row_inds)
reduced_matrix = []
for row_track, row_instances in row_tracks.items():
for col_track, col_instances in col_tracks.items():
asso_matrix = self[row_instances, col_instances]
if col_dims == "track":
asso_matrix = reduce_method(asso_matrix, axis=1)
if row_dims == "track":
asso_matrix = reduce_method(asso_matrix, axis=0)
reduced_matrix.append(asso_matrix)
reduced_matrix = np.array(reduced_matrix).reshape(n_cols, n_rows).T
return pd.DataFrame(reduced_matrix, index=row_inds, columns=col_inds)
def __getitem__(
self, inds: tuple[int | Instance | list[int | Instance]]
) -> np.ndarray:
"""Get elements of the association matrix.
Args:
inds: A tuple of query indices and reference indices.
Indices can be either:
A single instance or integer.
A list of instances or integers.
Returns:
An np.ndarray containing the elements requested.
"""
query_inst, ref_inst = inds
query_ind = self.__getindices__(query_inst, self.query_instances)
ref_ind = self.__getindices__(ref_inst, self.ref_instances)
try:
return self.numpy()[query_ind[:, None], ref_ind].squeeze()
except IndexError as e:
logger.exception(f"Query_insts: {type(query_inst)}")
logger.exception(f"Query_inds: {query_ind}")
logger.exception(f"Ref_insts: {type(ref_inst)}")
logger.exception(f"Ref_ind: {ref_ind}")
logger.exception(e)
raise (e)
def __getindices__(
self,
instance: Instance | int | np.typing.ArrayLike,
instance_lookup: list[Instance],
) -> np.ndarray:
"""Get the indices of the instance for lookup.
Args:
instance: The instance(s) to be retrieved
Can either be a single int/instance or a list of int/instances
instance_lookup: A list of Instances to be used to retrieve indices
Returns:
A np array of indices.
"""
if isinstance(instance, Instance):
ind = np.array([instance_lookup.index(instance)])
elif instance is None:
ind = np.arange(len(instance_lookup))
elif np.isscalar(instance):
ind = np.array([instance])
else:
instances = instance
if not [isinstance(inst, (Instance, int)) for inst in instance]:
raise ValueError(
f"List of indices must be `int` or `Instance`. Found {set([type(inst) for inst in instance])}"
)
ind = np.array(
[
(
instance_lookup.index(instance)
if isinstance(instance, Instance)
else instance
)
for instance in instances
]
)
return ind
def get_tracks(
self, instances: list["Instance"], label: str = "pred"
) -> dict[int, list["Instance"]]:
"""Group instances by track.
Args:
instances: The list of instances to group
label: the track id type to group by. Either `pred` or `gt`.
Returns:
A dictionary of track_id:instances
"""
if label == "pred":
traj_ids = set([instance.pred_track_id.item() for instance in instances])
traj = {
track_id: [
instance
for instance in instances
if instance.pred_track_id.item() == track_id
]
for track_id in traj_ids
}
elif label == "gt":
traj_ids = set(
[instance.gt_track_id.item() for instance in self.ref_instances]
)
traj = {
track_id: [
instance
for instance in self.ref_instances
if instance.gt_track_id.item() == track_id
]
for track_id in traj_ids
}
else:
raise ValueError(f"Unsupported label '{label}'. Expected 'pred' or 'gt'.")
return traj
def to(self, map_location: str | torch.device) -> Self:
"""Move instance to different device or change dtype. (See `torch.to` for more info).
Args:
map_location: Either the device or dtype for the instance to be moved.
Returns:
self: reference to the instance moved to correct device/dtype.
"""
self.matrix = self.matrix.to(map_location)
self.ref_instances = [
instance.to(map_location) for instance in self.ref_instances
]
self.query_instances = [
instance.to(map_location) for instance in self.query_instances
]
return self
__getindices__(instance, instance_lookup)
¶
Get the indices of the instance for lookup.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
instance
|
Instance | int | ArrayLike
|
The instance(s) to be retrieved Can either be a single int/instance or a list of int/instances |
required |
instance_lookup
|
list[Instance]
|
A list of Instances to be used to retrieve indices |
required |
Returns:
| Type | Description |
|---|---|
ndarray
|
A np array of indices. |
Source code in dreem/io/association_matrix.py
def __getindices__(
self,
instance: Instance | int | np.typing.ArrayLike,
instance_lookup: list[Instance],
) -> np.ndarray:
"""Get the indices of the instance for lookup.
Args:
instance: The instance(s) to be retrieved
Can either be a single int/instance or a list of int/instances
instance_lookup: A list of Instances to be used to retrieve indices
Returns:
A np array of indices.
"""
if isinstance(instance, Instance):
ind = np.array([instance_lookup.index(instance)])
elif instance is None:
ind = np.arange(len(instance_lookup))
elif np.isscalar(instance):
ind = np.array([instance])
else:
instances = instance
if not [isinstance(inst, (Instance, int)) for inst in instance]:
raise ValueError(
f"List of indices must be `int` or `Instance`. Found {set([type(inst) for inst in instance])}"
)
ind = np.array(
[
(
instance_lookup.index(instance)
if isinstance(instance, Instance)
else instance
)
for instance in instances
]
)
return ind
__getitem__(inds)
¶
Get elements of the association matrix.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
inds
|
tuple[int | Instance | list[int | Instance]]
|
A tuple of query indices and reference indices. Indices can be either: A single instance or integer. A list of instances or integers. |
required |
Returns:
| Type | Description |
|---|---|
ndarray
|
An np.ndarray containing the elements requested. |
Source code in dreem/io/association_matrix.py
def __getitem__(
self, inds: tuple[int | Instance | list[int | Instance]]
) -> np.ndarray:
"""Get elements of the association matrix.
Args:
inds: A tuple of query indices and reference indices.
Indices can be either:
A single instance or integer.
A list of instances or integers.
Returns:
An np.ndarray containing the elements requested.
"""
query_inst, ref_inst = inds
query_ind = self.__getindices__(query_inst, self.query_instances)
ref_ind = self.__getindices__(ref_inst, self.ref_instances)
try:
return self.numpy()[query_ind[:, None], ref_ind].squeeze()
except IndexError as e:
logger.exception(f"Query_insts: {type(query_inst)}")
logger.exception(f"Query_inds: {query_ind}")
logger.exception(f"Ref_insts: {type(ref_inst)}")
logger.exception(f"Ref_ind: {ref_ind}")
logger.exception(e)
raise (e)
__repr__()
¶
Get the string representation of the Association Matrix.
Returns:
| Type | Description |
|---|---|
str
|
the string representation of the association matrix. |
Source code in dreem/io/association_matrix.py
get_tracks(instances, label='pred')
¶
Group instances by track.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
instances
|
list[Instance]
|
The list of instances to group |
required |
label
|
str
|
the track id type to group by. Either |
'pred'
|
Returns:
| Type | Description |
|---|---|
dict[int, list[Instance]]
|
A dictionary of track_id:instances |
Source code in dreem/io/association_matrix.py
def get_tracks(
self, instances: list["Instance"], label: str = "pred"
) -> dict[int, list["Instance"]]:
"""Group instances by track.
Args:
instances: The list of instances to group
label: the track id type to group by. Either `pred` or `gt`.
Returns:
A dictionary of track_id:instances
"""
if label == "pred":
traj_ids = set([instance.pred_track_id.item() for instance in instances])
traj = {
track_id: [
instance
for instance in instances
if instance.pred_track_id.item() == track_id
]
for track_id in traj_ids
}
elif label == "gt":
traj_ids = set(
[instance.gt_track_id.item() for instance in self.ref_instances]
)
traj = {
track_id: [
instance
for instance in self.ref_instances
if instance.gt_track_id.item() == track_id
]
for track_id in traj_ids
}
else:
raise ValueError(f"Unsupported label '{label}'. Expected 'pred' or 'gt'.")
return traj
numpy()
¶
Convert association matrix to a numpy array.
Returns:
| Type | Description |
|---|---|
ndarray
|
The association matrix as a numpy array. |
reduce(row_dims='instance', col_dims='track', row_grouping=None, col_grouping='pred', reduce_method=np.sum)
¶
Aggregate the association matrix by specified dimensions and grouping.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
row_dims
|
str
|
A str indicating how to what dimensions to reduce rows to. Either "instance" (remains unchanged), or "track" (n_rows=n_traj). |
'instance'
|
col_dims
|
str
|
A str indicating how to dimensions to reduce rows to. Either "instance" (remains unchanged), or "track" (n_cols=n_traj) |
'track'
|
row_grouping
|
str | None
|
A str indicating how to group rows when aggregating. Either "pred" or "gt". |
None
|
col_grouping
|
str
|
A str indicating how to group columns when aggregating. Either "pred" or "gt". |
'pred'
|
reduce_method
|
callable
|
A callable function that operates on numpy matrices and can take an |
sum
|
Returns:
| Type | Description |
|---|---|
DataFrame
|
The association matrix reduced to an inst/traj x traj/inst association matrix as a dataframe. |
Source code in dreem/io/association_matrix.py
def reduce(
self,
row_dims: str = "instance",
col_dims: str = "track",
row_grouping: str | None = None,
col_grouping: str = "pred",
reduce_method: callable = np.sum,
) -> pd.DataFrame:
"""Aggregate the association matrix by specified dimensions and grouping.
Args:
row_dims: A str indicating how to what dimensions to reduce rows to.
Either "instance" (remains unchanged), or "track" (n_rows=n_traj).
col_dims: A str indicating how to dimensions to reduce rows to.
Either "instance" (remains unchanged), or "track" (n_cols=n_traj)
row_grouping: A str indicating how to group rows when aggregating. Either "pred" or "gt".
col_grouping: A str indicating how to group columns when aggregating. Either "pred" or "gt".
reduce_method: A callable function that operates on numpy matrices and can take an `axis` arg for reducing.
Returns:
The association matrix reduced to an inst/traj x traj/inst association matrix as a dataframe.
"""
n_rows = len(self.query_instances)
n_cols = len(self.ref_instances)
col_tracks = {-1: self.ref_instances}
row_tracks = {-1: self.query_instances}
col_inds = [i for i in range(len(self.ref_instances))]
row_inds = [i for i in range(len(self.query_instances))]
if col_dims == "track":
col_tracks = self.get_tracks(self.ref_instances, col_grouping)
col_inds = list(col_tracks.keys())
n_cols = len(col_inds)
if row_dims == "track":
row_tracks = self.get_tracks(self.query_instances, row_grouping)
row_inds = list(row_tracks.keys())
n_rows = len(row_inds)
reduced_matrix = []
for row_track, row_instances in row_tracks.items():
for col_track, col_instances in col_tracks.items():
asso_matrix = self[row_instances, col_instances]
if col_dims == "track":
asso_matrix = reduce_method(asso_matrix, axis=1)
if row_dims == "track":
asso_matrix = reduce_method(asso_matrix, axis=0)
reduced_matrix.append(asso_matrix)
reduced_matrix = np.array(reduced_matrix).reshape(n_cols, n_rows).T
return pd.DataFrame(reduced_matrix, index=row_inds, columns=col_inds)
to(map_location)
¶
Move instance to different device or change dtype. (See torch.to for more info).
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
map_location
|
str | device
|
Either the device or dtype for the instance to be moved. |
required |
Returns:
| Name | Type | Description |
|---|---|---|
self |
Self
|
reference to the instance moved to correct device/dtype. |
Source code in dreem/io/association_matrix.py
def to(self, map_location: str | torch.device) -> Self:
"""Move instance to different device or change dtype. (See `torch.to` for more info).
Args:
map_location: Either the device or dtype for the instance to be moved.
Returns:
self: reference to the instance moved to correct device/dtype.
"""
self.matrix = self.matrix.to(map_location)
self.ref_instances = [
instance.to(map_location) for instance in self.ref_instances
]
self.query_instances = [
instance.to(map_location) for instance in self.query_instances
]
return self
to_dataframe(row_labels='gt', col_labels='gt')
¶
Convert the association matrix to a pandas DataFrame.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
row_labels
|
str
|
How to label the rows(queries).
If list, then must match # of rows/queries
If |
'gt'
|
col_labels
|
str
|
How to label the columns(references).
If list, then must match # of columns/refs
If |
'gt'
|
Returns:
| Type | Description |
|---|---|
DataFrame
|
The association matrix as a pandas dataframe. |
Source code in dreem/io/association_matrix.py
def to_dataframe(
self, row_labels: str = "gt", col_labels: str = "gt"
) -> pd.DataFrame:
"""Convert the association matrix to a pandas DataFrame.
Args:
row_labels: How to label the rows(queries).
If list, then must match # of rows/queries
If `"gt"` then label by gt track id.
If `"pred"` then label by pred track id.
Otherwise label by the query_instance indices
col_labels: How to label the columns(references).
If list, then must match # of columns/refs
If `"gt"` then label by gt track id.
If `"pred"` then label by pred track id.
Otherwise label by the ref_instance indices
Returns:
The association matrix as a pandas dataframe.
"""
matrix = self.numpy()
if not isinstance(row_labels, str):
if len(row_labels) == len(self.query_instances):
row_inds = row_labels
else:
raise ValueError(
(
"Mismatched # of rows and labels!",
f"Found {len(row_labels)} with {len(self.query_instances)} rows",
)
)
else:
if row_labels == "gt":
row_inds = [
instance.gt_track_id.item() for instance in self.query_instances
]
elif row_labels == "pred":
row_inds = [
instance.pred_track_id.item() for instance in self.query_instances
]
else:
row_inds = np.arange(len(self.query_instances))
if not isinstance(col_labels, str):
if len(col_labels) == len(self.ref_instances):
col_inds = col_labels
else:
raise ValueError(
(
"Mismatched # of columns and labels!",
f"Found {len(col_labels)} with {len(self.ref_instances)} columns",
)
)
else:
if col_labels == "gt":
col_inds = [
instance.gt_track_id.item() for instance in self.ref_instances
]
elif col_labels == "pred":
col_inds = [
instance.pred_track_id.item() for instance in self.ref_instances
]
else:
col_inds = np.arange(len(self.ref_instances))
asso_df = pd.DataFrame(matrix, index=row_inds, columns=col_inds)
return asso_df
Config
¶
Class handling loading components based on config params.
Methods:
| Name | Description |
|---|---|
__init__ |
Initialize the class with config from hydra/omega conf. |
__repr__ |
Object representation of config class. |
__str__ |
Return a string representation of config class. |
from_yaml |
Load config directly from yaml. |
get |
Get config item. |
get_checkpointing |
Getter for lightning checkpointing callback. |
get_ctc_paths |
Get file paths from directory. Only for CTC datasets. |
get_data_paths |
Get file paths from directory. Only for SLEAP datasets. |
get_dataloader |
Getter for dataloader. |
get_dataset |
Getter for datasets. |
get_early_stopping |
Getter for lightning early stopping callback. |
get_gtr_runner |
Get lightning module for training, validation, and inference. |
get_logger |
Getter for logging callback. |
get_loss |
Getter for loss functions. |
get_model |
Getter for gtr model. |
get_optimizer |
Getter for optimizer. |
get_scheduler |
Getter for lr scheduler. |
get_tracker_cfg |
Getter for tracker config params. |
get_trainer |
Getter for the lightning trainer. |
set_hparams |
Setter function for overwriting specific hparams. |
Attributes:
| Name | Type | Description |
|---|---|---|
data_paths |
Get data paths. |
Source code in dreem/io/config.py
class Config:
"""Class handling loading components based on config params."""
def __init__(self, cfg: DictConfig, params_cfg: DictConfig | None = None):
"""Initialize the class with config from hydra/omega conf.
First uses `base_param` file then overwrites with specific `params_config`.
Args:
cfg: The `DictConfig` containing all the hyperparameters needed for
training/evaluation.
params_cfg: The `DictConfig` containing subset of hyperparameters to override.
training/evaluation
"""
base_cfg = cfg
logger.info(f"Base Config: {cfg}")
if "params_config" in cfg:
params_cfg = OmegaConf.load(cfg.params_config)
if params_cfg:
logger.info(f"Overwriting base config with {params_cfg}")
with open_dict(base_cfg):
self.cfg = OmegaConf.merge(base_cfg, params_cfg) # merge configs
else:
self.cfg = cfg
OmegaConf.set_struct(self.cfg, False)
self._vid_files = {}
def __repr__(self):
"""Object representation of config class."""
return f"Config({self.cfg})"
def __str__(self):
"""Return a string representation of config class."""
return f"Config({self.cfg})"
@classmethod
def from_yaml(cls, base_cfg_path: str, params_cfg_path: str | None = None) -> None:
"""Load config directly from yaml.
Args:
base_cfg_path: path to base config file.
params_cfg_path: path to override params.
"""
base_cfg = OmegaConf.load(base_cfg_path)
params_cfg = OmegaConf.load(params_cfg_path) if params_cfg_path else None
return cls(base_cfg, params_cfg)
def set_hparams(self, hparams: dict) -> bool:
"""Setter function for overwriting specific hparams.
Useful for changing 1 or 2 hyperparameters such as dataset.
Args:
hparams: A dict containing the hyperparameter to be overwritten and
the value to be changed
Returns:
`True` if config is successfully updated, `False` otherwise
"""
if hparams == {} or hparams is None:
logger.warning("Nothing to update!")
return False
for hparam, val in hparams.items():
try:
OmegaConf.update(self.cfg, hparam, val)
except Exception as e:
logger.exception(f"Failed to update {hparam} to {val} due to {e}")
return False
return True
def get(self, key: str, default=None, cfg: dict = None):
"""Get config item.
Args:
key: key of item to return
default: default value to return if key is missing.
cfg: the config dict from which to retrieve an item
"""
if cfg is None:
cfg = self.cfg
param = cfg.get(key, default)
if isinstance(param, DictConfig):
param = OmegaConf.to_container(param, resolve=True)
return param
def get_model(self) -> GlobalTrackingTransformer:
"""Getter for gtr model.
Returns:
A global tracking transformer with parameters indicated by cfg
"""
from dreem.models import GlobalTrackingTransformer, GTRRunner
model_params = self.get("model", {})
ckpt_path = model_params.pop("ckpt_path", None)
if ckpt_path is not None and len(ckpt_path) > 0:
return GTRRunner.load_from_checkpoint(ckpt_path).model
return GlobalTrackingTransformer(**model_params)
def get_tracker_cfg(self) -> dict:
"""Getter for tracker config params.
Returns:
A dict containing the init params for `Tracker`.
"""
return self.get("tracker", {})
def get_gtr_runner(self, ckpt_path: str | None = None) -> GTRRunner:
"""Get lightning module for training, validation, and inference.
Args:
ckpt_path: path to checkpoint for override
Returns:
a gtr runner model
"""
from dreem.models import GTRRunner
keys = ["tracker", "optimizer", "scheduler", "loss", "runner", "model"]
args = [key + "_cfg" if key != "runner" else key for key in keys]
params = {}
for key, arg in zip(keys, args):
sub_params = self.get(key, {})
# if len(sub_params) == 0:
# logger.warning(
# f"`{key}` not found in config or is empty. Using defaults for {arg}!"
# )
if key == "runner":
runner_params = sub_params
for k, v in runner_params.items():
params[k] = v
else:
params[arg] = sub_params
ckpt_path = params["model_cfg"].pop("ckpt_path", None)
if ckpt_path is not None and ckpt_path != "":
model = GTRRunner.load_from_checkpoint(
ckpt_path, tracker_cfg=params["tracker_cfg"], **runner_params
)
else:
model = GTRRunner(**params)
return model
def get_ctc_paths(
self, list_dir_path: list[str]
) -> tuple[list[str], list[str], list[str]]:
"""Get file paths from directory. Only for CTC datasets.
Args:
list_dir_path: list of directories to search for labels and videos
Returns:
lists of labels file paths and video file paths
"""
gt_list = []
raw_img_list = []
ctc_track_meta = []
# user can specify a list of directories, each of which can contain several subdirectories that come in pairs of (dset_name, dset_name_GT/TRA)
for dir_path in list_dir_path:
for subdir in os.listdir(dir_path):
if subdir.endswith("_GT"):
gt_path = os.path.join(dir_path, subdir, "TRA")
raw_img_path = os.path.join(dir_path, subdir.replace("_GT", ""))
# get filepaths for all tif files in gt_path
gt_list.append(glob.glob(os.path.join(gt_path, "*.tif*")))
# get filepaths for all tif files in raw_img_path
raw_img_list.append(glob.glob(os.path.join(raw_img_path, "*.tif*")))
man_track_file = glob.glob(os.path.join(gt_path, "man_track.txt"))
if len(man_track_file) > 0:
ctc_track_meta.append(man_track_file[0])
else:
logger.debug(
f"No man_track.txt file found in {gt_path}. Continuing..."
)
else:
continue
return gt_list, raw_img_list, ctc_track_meta
def get_data_paths(self, mode: str, data_cfg: dict) -> tuple[list[str], list[str]]:
"""Get file paths from directory. Only for SLEAP datasets.
Args:
mode: [None, "train", "test", "val"]. Indicates whether to use
train, val, or test params for dataset
data_cfg: Config for the dataset containing "dir" key.
Returns:
lists of labels file paths and video file paths respectively
"""
# hack to get around the fact that for test mode, get_data_paths is called before get_dataset.
# also, for train/val mode, data_cfg has had the dir key popped through self.get() called in get_dataset()
if mode == "test":
list_dir_path = data_cfg.get("dir", {}).get("path", None)
if list_dir_path is None:
raise ValueError(
"`dir` is missing from dataset config. Please provide a path to the directory containing the labels and videos."
)
self.labels_suffix = data_cfg.get("dir", {}).get("labels_suffix")
self.vid_suffix = data_cfg.get("dir", {}).get("vid_suffix")
else:
list_dir_path = self.data_dirs
if not isinstance(list_dir_path, list):
list_dir_path = [list_dir_path]
if self.labels_suffix == ".slp":
label_files = []
vid_files = []
for dir_path in list_dir_path:
logger.debug(f"Searching `{dir_path}` directory")
labels_path = f"{dir_path}/*{self.labels_suffix}"
vid_path = f"{dir_path}/*{self.vid_suffix}"
logger.debug(f"Searching for labels matching {labels_path}")
label_files.extend(glob.glob(labels_path))
logger.debug(f"Searching for videos matching {vid_path}")
vid_files.extend(glob.glob(vid_path))
elif self.labels_suffix == ".tif":
label_files, vid_files, ctc_track_meta = self.get_ctc_paths(list_dir_path)
logger.debug(f"Found {len(label_files)} labels and {len(vid_files)} videos")
# backdoor to set label files directly in the configs (i.e. bypass dir.path)
if data_cfg.get("slp_files", None):
logger.debug("Overriding label files with user provided list")
slp_files = data_cfg.get("slp_files")
if len(slp_files) > 0:
label_files = slp_files
if data_cfg.get("video_files", None):
individual_video_files = data_cfg.get("video_files")
if len(individual_video_files) > 0:
vid_files = individual_video_files
return label_files, vid_files
def get_dataset(
self,
mode: str,
label_files: list[str] | None = None,
vid_files: list[str | list[str]] = None,
overrides: dict | None = None,
) -> SleapDataset | CellTrackingDataset:
"""Getter for datasets.
Args:
mode: [None, "train", "test", "val"]. Indicates whether to use
train, val, or test params for dataset
label_files: path to label_files for override
vid_files: path to vid_files for override
overrides: overrides to apply to the dataset config
Returns:
Either a `SleapDataset` or `CellTrackingDataset` with params indicated by cfg
"""
from dreem.datasets import CellTrackingDataset, SleapDataset
dataset_params = self.get("dataset")
if dataset_params is None:
raise KeyError("`dataset` key is missing from cfg!")
if mode.lower() == "train":
dataset_params = self.get("train_dataset", {}, dataset_params)
elif mode.lower() == "val":
dataset_params = self.get("val_dataset", {}, dataset_params)
elif mode.lower() == "test":
dataset_params = self.get("test_dataset", {}, dataset_params)
else:
raise ValueError(
"`mode` must be one of ['train', 'val','test'], not '{mode}'"
)
# input validation
self.data_dirs = dataset_params.get("dir", {}).get("path", None)
self.labels_suffix = dataset_params.get("dir", {}).get("labels_suffix")
self.vid_suffix = dataset_params.get("dir", {}).get("vid_suffix")
if self.data_dirs is None:
raise ValueError(
"`dir` is missing from dataset config. Please provide a path to the directory containing the labels and videos."
)
if self.labels_suffix is None or self.vid_suffix is None:
raise KeyError(
f"Must provide a labels suffix and vid suffix to search for but found {self.labels_suffix} and {self.vid_suffix}"
)
# infer dataset type from the user provided suffix
if self.labels_suffix == ".slp":
# during training, multiple files can be used at once, so label_files is not passed in
# during inference, a single label_files string can be passed in as get_data_paths is
# called before get_dataset, hence the check
if label_files is None or vid_files is None:
label_files, vid_files = self.get_data_paths(mode, dataset_params)
dataset_params["slp_files"] = label_files
dataset_params["video_files"] = vid_files
dataset_params["data_dirs"] = self.data_dirs
self.data_paths = (mode, vid_files)
if overrides:
dataset_params.update(overrides)
return SleapDataset(**dataset_params)
elif self.labels_suffix == ".tif" or self.labels_suffix == ".tiff":
# for CTC datasets, pass in a list of gt and raw image directories, eaech of which contain tifs
ctc_track_meta = None
list_dir_path = self.data_dirs # don't modify self.data_dirs
if not isinstance(list_dir_path, list):
list_dir_path = [list_dir_path]
if label_files is None or vid_files is None:
label_files, vid_files, ctc_track_meta = self.get_ctc_paths(
list_dir_path
)
dataset_params["data_dirs"] = self.data_dirs
# extract filepaths of all raw images and gt images (i.e. labelled masks)
dataset_params["gt_list"] = label_files
dataset_params["raw_img_list"] = vid_files
dataset_params["ctc_track_meta"] = ctc_track_meta
if overrides:
dataset_params.update(overrides)
return CellTrackingDataset(**dataset_params)
else:
raise ValueError(
"Could not resolve dataset type from Config! Only .slp (SLEAP) and .tif (Cell Tracking Challenge) data formats are supported."
)
@property
def data_paths(self):
"""Get data paths."""
return self._vid_files
@data_paths.setter
def data_paths(self, paths: tuple[str, list[str]]):
"""Set data paths.
Args:
paths: A tuple containing (mode, vid_files)
"""
mode, vid_files = paths
self._vid_files[mode] = vid_files
def get_dataloader(
self,
dataset: SleapDataset | MicroscopyDataset | CellTrackingDataset,
mode: str,
) -> torch.utils.data.DataLoader:
"""Getter for dataloader.
Args:
dataset: the Sleap or Microscopy Dataset used to initialize the dataloader
mode: either ["train", "val", or "test"] indicates which dataset
config to use
Returns:
A torch dataloader for `dataset` with parameters configured as specified
"""
dataloader_params = self.get("dataloader", {})
if mode.lower() == "train":
dataloader_params = self.get("train_dataloader", {}, dataloader_params)
elif mode.lower() == "val":
dataloader_params = self.get("val_dataloader", {}, dataloader_params)
elif mode.lower() == "test":
dataloader_params = self.get("test_dataloader", {}, dataloader_params)
else:
raise ValueError(
"`mode` must be one of ['train', 'val','test'], not '{mode}'"
)
if dataloader_params.get("num_workers", 0) > 0:
# prevent too many open files error
pin_memory = True
torch.multiprocessing.set_sharing_strategy("file_system")
else:
pin_memory = False
return torch.utils.data.DataLoader(
dataset=dataset,
batch_size=1,
pin_memory=pin_memory,
collate_fn=dataset.no_batching_fn,
**dataloader_params,
)
def get_optimizer(self, params: Iterable) -> torch.optim.Optimizer:
"""Getter for optimizer.
Args:
params: iterable of model parameters to optimize or dicts defining
parameter groups
Returns:
A torch Optimizer with specified params
"""
from dreem.models.model_utils import init_optimizer
optimizer_params = self.get("optimizer")
return init_optimizer(params, optimizer_params)
def get_scheduler(
self, optimizer: torch.optim.Optimizer
) -> torch.optim.lr_scheduler.LRScheduler | None:
"""Getter for lr scheduler.
Args:
optimizer: The optimizer to wrap the scheduler around
Returns:
A torch learning rate scheduler with specified params
"""
from dreem.models.model_utils import init_scheduler
lr_scheduler_params = self.get("scheduler")
if lr_scheduler_params is None:
logger.warning(
"`scheduler` key not found in cfg or is empty. No scheduler will be returned!"
)
return None
return init_scheduler(optimizer, lr_scheduler_params)
def get_loss(self) -> AssoLoss:
"""Getter for loss functions.
Returns:
An AssoLoss with specified params
"""
from dreem.training.losses import AssoLoss
loss_params = self.get("loss", {})
if len(loss_params) == 0:
logger.warning(
"`loss` key not found in cfg. Using default params for `AssoLoss`"
)
return AssoLoss(**loss_params)
def get_logger(self) -> pl.loggers.Logger:
"""Getter for logging callback.
Returns:
A Logger with specified params
"""
from dreem.models.model_utils import init_logger
logger_params = self.get("logging", {})
if len(logger_params) == 0:
logger.warning(
"`logging` key not found in cfg. No logger will be configured!"
)
return init_logger(
logger_params, OmegaConf.to_container(self.cfg, resolve=True)
)
def get_early_stopping(self) -> pl.callbacks.EarlyStopping:
"""Getter for lightning early stopping callback.
Returns:
A lightning early stopping callback with specified params
"""
early_stopping_params = self.get("early_stopping", None)
if early_stopping_params is None:
logger.warning(
"`early_stopping` was not found in cfg or was `null`. Early stopping will not be used!"
)
return None
elif len(early_stopping_params) == 0:
logger.warning("`early_stopping` cfg is empty! Using defaults")
return pl.callbacks.EarlyStopping(**early_stopping_params)
def get_checkpointing(self) -> pl.callbacks.ModelCheckpoint:
"""Getter for lightning checkpointing callback.
Returns:
A lightning checkpointing callback with specified params
"""
# convert to dict to enable extracting/removing params
checkpoint_params = self.get("checkpointing", {})
logging_params = self.get("logging", {})
dirpath = checkpoint_params.pop("dirpath", None)
if dirpath is None:
dirpath = f"./models/{self.get('group', '', logging_params)}/{self.get('name', '', logging_params)}"
dirpath = Path(dirpath).resolve()
if not Path(dirpath).exists():
try:
Path(dirpath).mkdir(parents=True, exist_ok=True)
except OSError as e:
logger.exception(
f"Cannot create a new folder!. Check the permissions to {dirpath}. \n {e}"
)
_ = checkpoint_params.pop("dirpath", None)
monitor = checkpoint_params.pop("monitor", ["val_loss"])
checkpointers = []
logger.info(
f"Saving checkpoints to `{dirpath}` based on the following metrics: {monitor}"
)
if len(checkpoint_params) == 0:
logger.warning(
"""`checkpointing` key was not found in cfg or was empty!
Configuring checkpointing to use default params!"""
)
for metric in monitor:
checkpointer = pl.callbacks.ModelCheckpoint(
monitor=metric,
dirpath=dirpath,
filename=f"{{epoch}}-{{{metric}}}",
**checkpoint_params,
)
checkpointer.CHECKPOINT_NAME_LAST = f"{{epoch}}-final-{{{metric}}}"
checkpointers.append(checkpointer)
return checkpointers
def get_trainer(
self,
callbacks: list[pl.callbacks.Callback] | None = None,
logger: pl.loggers.WandbLogger | None = None,
devices: int = 1,
accelerator: str = "auto",
) -> pl.Trainer:
"""Getter for the lightning trainer.
Args:
callbacks: a list of lightning callbacks preconfigured to be used
for training
logger: the Wandb logger used for logging during training
devices: The number of gpus to be used. 0 means cpu
accelerator: either "gpu" or "cpu" specifies which device to use
Returns:
A lightning Trainer with specified params
"""
trainer_params = self.get("trainer", {})
profiler = trainer_params.pop("profiler", None)
# if len(trainer_params) == 0:
# print(
# "`trainer` key was not found in cfg or was empty. Using defaults for `pl.Trainer`!"
# )
if "accelerator" not in trainer_params:
trainer_params["accelerator"] = accelerator
if "devices" not in trainer_params:
trainer_params["devices"] = devices
map_profiler = {
"advanced": pl.profilers.AdvancedProfiler,
"simple": pl.profilers.SimpleProfiler,
"pytorch": pl.profilers.PyTorchProfiler,
"passthrough": pl.profilers.PassThroughProfiler,
"xla": pl.profilers.XLAProfiler,
}
if profiler:
if profiler in map_profiler:
profiler = map_profiler[profiler](filename="profile")
else:
raise ValueError(
f"Profiler {profiler} not supported! Please use one of {list(map_profiler.keys())}"
)
return pl.Trainer(
callbacks=callbacks,
logger=logger,
profiler=profiler,
**trainer_params,
)
data_paths
property
writable
¶
Get data paths.
__init__(cfg, params_cfg=None)
¶
Initialize the class with config from hydra/omega conf.
First uses base_param file then overwrites with specific params_config.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
cfg
|
DictConfig
|
The |
required |
params_cfg
|
DictConfig | None
|
The |
None
|
Source code in dreem/io/config.py
def __init__(self, cfg: DictConfig, params_cfg: DictConfig | None = None):
"""Initialize the class with config from hydra/omega conf.
First uses `base_param` file then overwrites with specific `params_config`.
Args:
cfg: The `DictConfig` containing all the hyperparameters needed for
training/evaluation.
params_cfg: The `DictConfig` containing subset of hyperparameters to override.
training/evaluation
"""
base_cfg = cfg
logger.info(f"Base Config: {cfg}")
if "params_config" in cfg:
params_cfg = OmegaConf.load(cfg.params_config)
if params_cfg:
logger.info(f"Overwriting base config with {params_cfg}")
with open_dict(base_cfg):
self.cfg = OmegaConf.merge(base_cfg, params_cfg) # merge configs
else:
self.cfg = cfg
OmegaConf.set_struct(self.cfg, False)
self._vid_files = {}
__repr__()
¶
__str__()
¶
from_yaml(base_cfg_path, params_cfg_path=None)
classmethod
¶
Load config directly from yaml.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
base_cfg_path
|
str
|
path to base config file. |
required |
params_cfg_path
|
str | None
|
path to override params. |
None
|
Source code in dreem/io/config.py
@classmethod
def from_yaml(cls, base_cfg_path: str, params_cfg_path: str | None = None) -> None:
"""Load config directly from yaml.
Args:
base_cfg_path: path to base config file.
params_cfg_path: path to override params.
"""
base_cfg = OmegaConf.load(base_cfg_path)
params_cfg = OmegaConf.load(params_cfg_path) if params_cfg_path else None
return cls(base_cfg, params_cfg)
get(key, default=None, cfg=None)
¶
Get config item.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
key
|
str
|
key of item to return |
required |
default
|
default value to return if key is missing. |
None
|
|
cfg
|
dict
|
the config dict from which to retrieve an item |
None
|
Source code in dreem/io/config.py
def get(self, key: str, default=None, cfg: dict = None):
"""Get config item.
Args:
key: key of item to return
default: default value to return if key is missing.
cfg: the config dict from which to retrieve an item
"""
if cfg is None:
cfg = self.cfg
param = cfg.get(key, default)
if isinstance(param, DictConfig):
param = OmegaConf.to_container(param, resolve=True)
return param
get_checkpointing()
¶
Getter for lightning checkpointing callback.
Returns:
| Type | Description |
|---|---|
ModelCheckpoint
|
A lightning checkpointing callback with specified params |
Source code in dreem/io/config.py
def get_checkpointing(self) -> pl.callbacks.ModelCheckpoint:
"""Getter for lightning checkpointing callback.
Returns:
A lightning checkpointing callback with specified params
"""
# convert to dict to enable extracting/removing params
checkpoint_params = self.get("checkpointing", {})
logging_params = self.get("logging", {})
dirpath = checkpoint_params.pop("dirpath", None)
if dirpath is None:
dirpath = f"./models/{self.get('group', '', logging_params)}/{self.get('name', '', logging_params)}"
dirpath = Path(dirpath).resolve()
if not Path(dirpath).exists():
try:
Path(dirpath).mkdir(parents=True, exist_ok=True)
except OSError as e:
logger.exception(
f"Cannot create a new folder!. Check the permissions to {dirpath}. \n {e}"
)
_ = checkpoint_params.pop("dirpath", None)
monitor = checkpoint_params.pop("monitor", ["val_loss"])
checkpointers = []
logger.info(
f"Saving checkpoints to `{dirpath}` based on the following metrics: {monitor}"
)
if len(checkpoint_params) == 0:
logger.warning(
"""`checkpointing` key was not found in cfg or was empty!
Configuring checkpointing to use default params!"""
)
for metric in monitor:
checkpointer = pl.callbacks.ModelCheckpoint(
monitor=metric,
dirpath=dirpath,
filename=f"{{epoch}}-{{{metric}}}",
**checkpoint_params,
)
checkpointer.CHECKPOINT_NAME_LAST = f"{{epoch}}-final-{{{metric}}}"
checkpointers.append(checkpointer)
return checkpointers
get_ctc_paths(list_dir_path)
¶
Get file paths from directory. Only for CTC datasets.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
list_dir_path
|
list[str]
|
list of directories to search for labels and videos |
required |
Returns:
| Type | Description |
|---|---|
tuple[list[str], list[str], list[str]]
|
lists of labels file paths and video file paths |
Source code in dreem/io/config.py
def get_ctc_paths(
self, list_dir_path: list[str]
) -> tuple[list[str], list[str], list[str]]:
"""Get file paths from directory. Only for CTC datasets.
Args:
list_dir_path: list of directories to search for labels and videos
Returns:
lists of labels file paths and video file paths
"""
gt_list = []
raw_img_list = []
ctc_track_meta = []
# user can specify a list of directories, each of which can contain several subdirectories that come in pairs of (dset_name, dset_name_GT/TRA)
for dir_path in list_dir_path:
for subdir in os.listdir(dir_path):
if subdir.endswith("_GT"):
gt_path = os.path.join(dir_path, subdir, "TRA")
raw_img_path = os.path.join(dir_path, subdir.replace("_GT", ""))
# get filepaths for all tif files in gt_path
gt_list.append(glob.glob(os.path.join(gt_path, "*.tif*")))
# get filepaths for all tif files in raw_img_path
raw_img_list.append(glob.glob(os.path.join(raw_img_path, "*.tif*")))
man_track_file = glob.glob(os.path.join(gt_path, "man_track.txt"))
if len(man_track_file) > 0:
ctc_track_meta.append(man_track_file[0])
else:
logger.debug(
f"No man_track.txt file found in {gt_path}. Continuing..."
)
else:
continue
return gt_list, raw_img_list, ctc_track_meta
get_data_paths(mode, data_cfg)
¶
Get file paths from directory. Only for SLEAP datasets.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
mode
|
str
|
[None, "train", "test", "val"]. Indicates whether to use train, val, or test params for dataset |
required |
data_cfg
|
dict
|
Config for the dataset containing "dir" key. |
required |
Returns:
| Type | Description |
|---|---|
tuple[list[str], list[str]]
|
lists of labels file paths and video file paths respectively |
Source code in dreem/io/config.py
def get_data_paths(self, mode: str, data_cfg: dict) -> tuple[list[str], list[str]]:
"""Get file paths from directory. Only for SLEAP datasets.
Args:
mode: [None, "train", "test", "val"]. Indicates whether to use
train, val, or test params for dataset
data_cfg: Config for the dataset containing "dir" key.
Returns:
lists of labels file paths and video file paths respectively
"""
# hack to get around the fact that for test mode, get_data_paths is called before get_dataset.
# also, for train/val mode, data_cfg has had the dir key popped through self.get() called in get_dataset()
if mode == "test":
list_dir_path = data_cfg.get("dir", {}).get("path", None)
if list_dir_path is None:
raise ValueError(
"`dir` is missing from dataset config. Please provide a path to the directory containing the labels and videos."
)
self.labels_suffix = data_cfg.get("dir", {}).get("labels_suffix")
self.vid_suffix = data_cfg.get("dir", {}).get("vid_suffix")
else:
list_dir_path = self.data_dirs
if not isinstance(list_dir_path, list):
list_dir_path = [list_dir_path]
if self.labels_suffix == ".slp":
label_files = []
vid_files = []
for dir_path in list_dir_path:
logger.debug(f"Searching `{dir_path}` directory")
labels_path = f"{dir_path}/*{self.labels_suffix}"
vid_path = f"{dir_path}/*{self.vid_suffix}"
logger.debug(f"Searching for labels matching {labels_path}")
label_files.extend(glob.glob(labels_path))
logger.debug(f"Searching for videos matching {vid_path}")
vid_files.extend(glob.glob(vid_path))
elif self.labels_suffix == ".tif":
label_files, vid_files, ctc_track_meta = self.get_ctc_paths(list_dir_path)
logger.debug(f"Found {len(label_files)} labels and {len(vid_files)} videos")
# backdoor to set label files directly in the configs (i.e. bypass dir.path)
if data_cfg.get("slp_files", None):
logger.debug("Overriding label files with user provided list")
slp_files = data_cfg.get("slp_files")
if len(slp_files) > 0:
label_files = slp_files
if data_cfg.get("video_files", None):
individual_video_files = data_cfg.get("video_files")
if len(individual_video_files) > 0:
vid_files = individual_video_files
return label_files, vid_files
get_dataloader(dataset, mode)
¶
Getter for dataloader.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
dataset
|
SleapDataset | MicroscopyDataset | CellTrackingDataset
|
the Sleap or Microscopy Dataset used to initialize the dataloader |
required |
mode
|
str
|
either ["train", "val", or "test"] indicates which dataset config to use |
required |
Returns:
| Type | Description |
|---|---|
DataLoader
|
A torch dataloader for |
Source code in dreem/io/config.py
def get_dataloader(
self,
dataset: SleapDataset | MicroscopyDataset | CellTrackingDataset,
mode: str,
) -> torch.utils.data.DataLoader:
"""Getter for dataloader.
Args:
dataset: the Sleap or Microscopy Dataset used to initialize the dataloader
mode: either ["train", "val", or "test"] indicates which dataset
config to use
Returns:
A torch dataloader for `dataset` with parameters configured as specified
"""
dataloader_params = self.get("dataloader", {})
if mode.lower() == "train":
dataloader_params = self.get("train_dataloader", {}, dataloader_params)
elif mode.lower() == "val":
dataloader_params = self.get("val_dataloader", {}, dataloader_params)
elif mode.lower() == "test":
dataloader_params = self.get("test_dataloader", {}, dataloader_params)
else:
raise ValueError(
"`mode` must be one of ['train', 'val','test'], not '{mode}'"
)
if dataloader_params.get("num_workers", 0) > 0:
# prevent too many open files error
pin_memory = True
torch.multiprocessing.set_sharing_strategy("file_system")
else:
pin_memory = False
return torch.utils.data.DataLoader(
dataset=dataset,
batch_size=1,
pin_memory=pin_memory,
collate_fn=dataset.no_batching_fn,
**dataloader_params,
)
get_dataset(mode, label_files=None, vid_files=None, overrides=None)
¶
Getter for datasets.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
mode
|
str
|
[None, "train", "test", "val"]. Indicates whether to use train, val, or test params for dataset |
required |
label_files
|
list[str] | None
|
path to label_files for override |
None
|
vid_files
|
list[str | list[str]]
|
path to vid_files for override |
None
|
overrides
|
dict | None
|
overrides to apply to the dataset config |
None
|
Returns:
Either a SleapDataset or CellTrackingDataset with params indicated by cfg
Source code in dreem/io/config.py
def get_dataset(
self,
mode: str,
label_files: list[str] | None = None,
vid_files: list[str | list[str]] = None,
overrides: dict | None = None,
) -> SleapDataset | CellTrackingDataset:
"""Getter for datasets.
Args:
mode: [None, "train", "test", "val"]. Indicates whether to use
train, val, or test params for dataset
label_files: path to label_files for override
vid_files: path to vid_files for override
overrides: overrides to apply to the dataset config
Returns:
Either a `SleapDataset` or `CellTrackingDataset` with params indicated by cfg
"""
from dreem.datasets import CellTrackingDataset, SleapDataset
dataset_params = self.get("dataset")
if dataset_params is None:
raise KeyError("`dataset` key is missing from cfg!")
if mode.lower() == "train":
dataset_params = self.get("train_dataset", {}, dataset_params)
elif mode.lower() == "val":
dataset_params = self.get("val_dataset", {}, dataset_params)
elif mode.lower() == "test":
dataset_params = self.get("test_dataset", {}, dataset_params)
else:
raise ValueError(
"`mode` must be one of ['train', 'val','test'], not '{mode}'"
)
# input validation
self.data_dirs = dataset_params.get("dir", {}).get("path", None)
self.labels_suffix = dataset_params.get("dir", {}).get("labels_suffix")
self.vid_suffix = dataset_params.get("dir", {}).get("vid_suffix")
if self.data_dirs is None:
raise ValueError(
"`dir` is missing from dataset config. Please provide a path to the directory containing the labels and videos."
)
if self.labels_suffix is None or self.vid_suffix is None:
raise KeyError(
f"Must provide a labels suffix and vid suffix to search for but found {self.labels_suffix} and {self.vid_suffix}"
)
# infer dataset type from the user provided suffix
if self.labels_suffix == ".slp":
# during training, multiple files can be used at once, so label_files is not passed in
# during inference, a single label_files string can be passed in as get_data_paths is
# called before get_dataset, hence the check
if label_files is None or vid_files is None:
label_files, vid_files = self.get_data_paths(mode, dataset_params)
dataset_params["slp_files"] = label_files
dataset_params["video_files"] = vid_files
dataset_params["data_dirs"] = self.data_dirs
self.data_paths = (mode, vid_files)
if overrides:
dataset_params.update(overrides)
return SleapDataset(**dataset_params)
elif self.labels_suffix == ".tif" or self.labels_suffix == ".tiff":
# for CTC datasets, pass in a list of gt and raw image directories, eaech of which contain tifs
ctc_track_meta = None
list_dir_path = self.data_dirs # don't modify self.data_dirs
if not isinstance(list_dir_path, list):
list_dir_path = [list_dir_path]
if label_files is None or vid_files is None:
label_files, vid_files, ctc_track_meta = self.get_ctc_paths(
list_dir_path
)
dataset_params["data_dirs"] = self.data_dirs
# extract filepaths of all raw images and gt images (i.e. labelled masks)
dataset_params["gt_list"] = label_files
dataset_params["raw_img_list"] = vid_files
dataset_params["ctc_track_meta"] = ctc_track_meta
if overrides:
dataset_params.update(overrides)
return CellTrackingDataset(**dataset_params)
else:
raise ValueError(
"Could not resolve dataset type from Config! Only .slp (SLEAP) and .tif (Cell Tracking Challenge) data formats are supported."
)
get_early_stopping()
¶
Getter for lightning early stopping callback.
Returns:
| Type | Description |
|---|---|
EarlyStopping
|
A lightning early stopping callback with specified params |
Source code in dreem/io/config.py
def get_early_stopping(self) -> pl.callbacks.EarlyStopping:
"""Getter for lightning early stopping callback.
Returns:
A lightning early stopping callback with specified params
"""
early_stopping_params = self.get("early_stopping", None)
if early_stopping_params is None:
logger.warning(
"`early_stopping` was not found in cfg or was `null`. Early stopping will not be used!"
)
return None
elif len(early_stopping_params) == 0:
logger.warning("`early_stopping` cfg is empty! Using defaults")
return pl.callbacks.EarlyStopping(**early_stopping_params)
get_gtr_runner(ckpt_path=None)
¶
Get lightning module for training, validation, and inference.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
ckpt_path
|
str | None
|
path to checkpoint for override |
None
|
Returns:
| Type | Description |
|---|---|
GTRRunner
|
a gtr runner model |
Source code in dreem/io/config.py
def get_gtr_runner(self, ckpt_path: str | None = None) -> GTRRunner:
"""Get lightning module for training, validation, and inference.
Args:
ckpt_path: path to checkpoint for override
Returns:
a gtr runner model
"""
from dreem.models import GTRRunner
keys = ["tracker", "optimizer", "scheduler", "loss", "runner", "model"]
args = [key + "_cfg" if key != "runner" else key for key in keys]
params = {}
for key, arg in zip(keys, args):
sub_params = self.get(key, {})
# if len(sub_params) == 0:
# logger.warning(
# f"`{key}` not found in config or is empty. Using defaults for {arg}!"
# )
if key == "runner":
runner_params = sub_params
for k, v in runner_params.items():
params[k] = v
else:
params[arg] = sub_params
ckpt_path = params["model_cfg"].pop("ckpt_path", None)
if ckpt_path is not None and ckpt_path != "":
model = GTRRunner.load_from_checkpoint(
ckpt_path, tracker_cfg=params["tracker_cfg"], **runner_params
)
else:
model = GTRRunner(**params)
return model
get_logger()
¶
Getter for logging callback.
Returns:
| Type | Description |
|---|---|
Logger
|
A Logger with specified params |
Source code in dreem/io/config.py
def get_logger(self) -> pl.loggers.Logger:
"""Getter for logging callback.
Returns:
A Logger with specified params
"""
from dreem.models.model_utils import init_logger
logger_params = self.get("logging", {})
if len(logger_params) == 0:
logger.warning(
"`logging` key not found in cfg. No logger will be configured!"
)
return init_logger(
logger_params, OmegaConf.to_container(self.cfg, resolve=True)
)
get_loss()
¶
Getter for loss functions.
Returns:
| Type | Description |
|---|---|
AssoLoss
|
An AssoLoss with specified params |
Source code in dreem/io/config.py
def get_loss(self) -> AssoLoss:
"""Getter for loss functions.
Returns:
An AssoLoss with specified params
"""
from dreem.training.losses import AssoLoss
loss_params = self.get("loss", {})
if len(loss_params) == 0:
logger.warning(
"`loss` key not found in cfg. Using default params for `AssoLoss`"
)
return AssoLoss(**loss_params)
get_model()
¶
Getter for gtr model.
Returns:
| Type | Description |
|---|---|
GlobalTrackingTransformer
|
A global tracking transformer with parameters indicated by cfg |
Source code in dreem/io/config.py
def get_model(self) -> GlobalTrackingTransformer:
"""Getter for gtr model.
Returns:
A global tracking transformer with parameters indicated by cfg
"""
from dreem.models import GlobalTrackingTransformer, GTRRunner
model_params = self.get("model", {})
ckpt_path = model_params.pop("ckpt_path", None)
if ckpt_path is not None and len(ckpt_path) > 0:
return GTRRunner.load_from_checkpoint(ckpt_path).model
return GlobalTrackingTransformer(**model_params)
get_optimizer(params)
¶
Getter for optimizer.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
params
|
Iterable
|
iterable of model parameters to optimize or dicts defining parameter groups |
required |
Returns:
| Type | Description |
|---|---|
Optimizer
|
A torch Optimizer with specified params |
Source code in dreem/io/config.py
def get_optimizer(self, params: Iterable) -> torch.optim.Optimizer:
"""Getter for optimizer.
Args:
params: iterable of model parameters to optimize or dicts defining
parameter groups
Returns:
A torch Optimizer with specified params
"""
from dreem.models.model_utils import init_optimizer
optimizer_params = self.get("optimizer")
return init_optimizer(params, optimizer_params)
get_scheduler(optimizer)
¶
Getter for lr scheduler.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
optimizer
|
Optimizer
|
The optimizer to wrap the scheduler around |
required |
Returns:
| Type | Description |
|---|---|
LRScheduler | None
|
A torch learning rate scheduler with specified params |
Source code in dreem/io/config.py
def get_scheduler(
self, optimizer: torch.optim.Optimizer
) -> torch.optim.lr_scheduler.LRScheduler | None:
"""Getter for lr scheduler.
Args:
optimizer: The optimizer to wrap the scheduler around
Returns:
A torch learning rate scheduler with specified params
"""
from dreem.models.model_utils import init_scheduler
lr_scheduler_params = self.get("scheduler")
if lr_scheduler_params is None:
logger.warning(
"`scheduler` key not found in cfg or is empty. No scheduler will be returned!"
)
return None
return init_scheduler(optimizer, lr_scheduler_params)
get_tracker_cfg()
¶
Getter for tracker config params.
Returns:
| Type | Description |
|---|---|
dict
|
A dict containing the init params for |
get_trainer(callbacks=None, logger=None, devices=1, accelerator='auto')
¶
Getter for the lightning trainer.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
callbacks
|
list[Callback] | None
|
a list of lightning callbacks preconfigured to be used for training |
None
|
logger
|
WandbLogger | None
|
the Wandb logger used for logging during training |
None
|
devices
|
int
|
The number of gpus to be used. 0 means cpu |
1
|
accelerator
|
str
|
either "gpu" or "cpu" specifies which device to use |
'auto'
|
Returns:
| Type | Description |
|---|---|
Trainer
|
A lightning Trainer with specified params |
Source code in dreem/io/config.py
def get_trainer(
self,
callbacks: list[pl.callbacks.Callback] | None = None,
logger: pl.loggers.WandbLogger | None = None,
devices: int = 1,
accelerator: str = "auto",
) -> pl.Trainer:
"""Getter for the lightning trainer.
Args:
callbacks: a list of lightning callbacks preconfigured to be used
for training
logger: the Wandb logger used for logging during training
devices: The number of gpus to be used. 0 means cpu
accelerator: either "gpu" or "cpu" specifies which device to use
Returns:
A lightning Trainer with specified params
"""
trainer_params = self.get("trainer", {})
profiler = trainer_params.pop("profiler", None)
# if len(trainer_params) == 0:
# print(
# "`trainer` key was not found in cfg or was empty. Using defaults for `pl.Trainer`!"
# )
if "accelerator" not in trainer_params:
trainer_params["accelerator"] = accelerator
if "devices" not in trainer_params:
trainer_params["devices"] = devices
map_profiler = {
"advanced": pl.profilers.AdvancedProfiler,
"simple": pl.profilers.SimpleProfiler,
"pytorch": pl.profilers.PyTorchProfiler,
"passthrough": pl.profilers.PassThroughProfiler,
"xla": pl.profilers.XLAProfiler,
}
if profiler:
if profiler in map_profiler:
profiler = map_profiler[profiler](filename="profile")
else:
raise ValueError(
f"Profiler {profiler} not supported! Please use one of {list(map_profiler.keys())}"
)
return pl.Trainer(
callbacks=callbacks,
logger=logger,
profiler=profiler,
**trainer_params,
)
set_hparams(hparams)
¶
Setter function for overwriting specific hparams.
Useful for changing 1 or 2 hyperparameters such as dataset.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
hparams
|
dict
|
A dict containing the hyperparameter to be overwritten and the value to be changed |
required |
Returns:
| Type | Description |
|---|---|
bool
|
|
Source code in dreem/io/config.py
def set_hparams(self, hparams: dict) -> bool:
"""Setter function for overwriting specific hparams.
Useful for changing 1 or 2 hyperparameters such as dataset.
Args:
hparams: A dict containing the hyperparameter to be overwritten and
the value to be changed
Returns:
`True` if config is successfully updated, `False` otherwise
"""
if hparams == {} or hparams is None:
logger.warning("Nothing to update!")
return False
for hparam, val in hparams.items():
try:
OmegaConf.update(self.cfg, hparam, val)
except Exception as e:
logger.exception(f"Failed to update {hparam} to {val} due to {e}")
return False
return True
Frame
¶
Data structure containing metadata for a single frame of a video.
Attributes:
| Name | Type | Description |
|---|---|---|
video_id |
Tensor
|
The video index in the dataset. |
frame_id |
Tensor
|
The index of the frame in a video. |
vid_file |
Tensor
|
The path to the video the frame is from. |
img_shape |
Size
|
The shape of the original frame (not the crop). |
instances |
list['Instance']
|
A list of Instance objects that appear in the frame. |
asso_output |
AssociationMatrix
|
The association matrix between instances output directly from the transformer. |
matches |
tuple
|
matches from LSA algorithm between the instances and available trajectories during tracking. |
traj_score |
tuple
|
Either a dict containing the association matrix between instances and trajectories along postprocessing pipeline or a single association matrix. |
device |
str
|
The device the frame should be moved to. |
is_flagged |
bool
|
Whether the frame has been flagged for any reason. |
flag_reasons |
set[FrameFlagCode]
|
Set of FrameFlagCode values indicating why the frame was flagged. |
Methods:
| Name | Description |
|---|---|
__attrs_post_init__ |
Handle more intricate default initializations and moving to device. |
__repr__ |
Return String representation of the Frame. |
add_flag |
Add a flag reason to the frame. |
add_traj_score |
Add trajectory score to dictionary. |
clear_flags |
Clear all flags from the frame. |
from_slp |
Convert |
get_anchors |
Get the anchor names of instances in the frame. |
get_bboxes |
Get the bounding boxes of all instances in the frame. |
get_centroids |
Get the centroids around which each instance's crop was formed. |
get_crops |
Get the crops of all instances in the frame. |
get_features |
Get the reid feature vectors of all instances in the frame. |
get_gt_track_ids |
Get the gt track ids of all instances in the frame. |
get_pred_track_ids |
Get the pred track ids of all instances in the frame. |
get_traj_score |
Get dictionary containing association matrix between instances and trajectories along postprocessing pipeline. |
has_asso_output |
Determine whether the frame has an association matrix computed. |
has_bboxes |
Check if any of frames instances has a bounding box. |
has_crops |
Check if any of frames instances has a crop. |
has_features |
Check if any of frames instances has reid features already computed. |
has_flag |
Check if the frame has a specific flag reason. |
has_gt_track_ids |
Check if any of frames instances has a gt track id. |
has_instances |
Determine whether there are instances in the frame. |
has_matches |
Check whether or not matches have been computed for frame. |
has_pred_track_ids |
Check if any of frames instances has a pred track id. |
has_traj_score |
Check if any trajectory association matrix has been saved. |
remove_flag |
Remove a specific flag reason from the frame. |
to |
Move frame to different device or dtype (See |
to_h5 |
Convert frame to h5py group. |
to_slp |
Convert Frame to sleap_io.LabeledFrame object. |
Source code in dreem/io/frame.py
@attrs.define(eq=False)
class Frame:
"""Data structure containing metadata for a single frame of a video.
Attributes:
video_id: The video index in the dataset.
frame_id: The index of the frame in a video.
vid_file: The path to the video the frame is from.
img_shape: The shape of the original frame (not the crop).
instances: A list of Instance objects that appear in the frame.
asso_output: The association matrix between instances
output directly from the transformer.
matches: matches from LSA algorithm between the instances and
available trajectories during tracking.
traj_score: Either a dict containing the association matrix
between instances and trajectories along postprocessing pipeline
or a single association matrix.
device: The device the frame should be moved to.
is_flagged: Whether the frame has been flagged for any reason.
flag_reasons: Set of FrameFlagCode values indicating why the frame was flagged.
"""
_video_id: int = attrs.field(alias="video_id", converter=_to_tensor)
_frame_id: int = attrs.field(alias="frame_id", converter=_to_tensor)
_video: str = attrs.field(alias="vid_file", default="")
_img_shape: torch.Size = attrs.field(
alias="img_shape", converter=_to_size, factory=lambda: torch.Size([])
)
_instances: list["Instance"] = attrs.field(alias="instances", factory=list)
_asso_output: AssociationMatrix | None = attrs.field(
alias="asso_output", default=None
)
_matches: tuple = attrs.field(alias="matches", factory=tuple)
_traj_score: dict = attrs.field(alias="traj_score", factory=dict)
_device: str | torch.device | None = attrs.field(alias="device", default=None)
_is_flagged: bool = attrs.field(alias="is_flagged", default=False)
_flag_reasons: set[FrameFlagCode] = attrs.field(alias="flag_reasons", factory=set)
def __attrs_post_init__(self) -> None:
"""Handle more intricate default initializations and moving to device."""
if len(self.img_shape) == 0:
self.img_shape = torch.Size([0, 0, 0])
for instance in self.instances:
instance.frame = self
self.to(self.device)
def __repr__(self) -> str:
"""Return String representation of the Frame.
Returns:
The string representation of the frame.
"""
return (
"Frame("
f"video={self._video.filename if isinstance(self._video, sio.Video) else self._video}, "
f"video_id={self._video_id.item()}, "
f"frame_id={self._frame_id.item()}, "
f"img_shape={self._img_shape}, "
f"num_detected={self.num_detected}, "
f"asso_output={self._asso_output}, "
f"traj_score={self._traj_score}, "
f"matches={self._matches}, "
f"instances={self._instances}, "
f"device={self._device}, "
f"is_flagged={self._is_flagged}, "
f"flag_reasons={self._flag_reasons}"
")"
)
def to(self, map_location: str | torch.device) -> Self:
"""Move frame to different device or dtype (See `torch.to` for more info).
Args:
map_location: A string representing the device to move to.
Returns:
The frame moved to a different device/dtype.
"""
self._video_id = self._video_id.to(map_location)
self._frame_id = self._frame_id.to(map_location)
# torch.Size is immutable and doesn't need device movement
if isinstance(self._asso_output, torch.Tensor):
self._asso_output = self._asso_output.to(map_location)
if isinstance(self._matches, torch.Tensor):
self._matches = self._matches.to(map_location)
for key, val in self._traj_score.items():
if isinstance(val, torch.Tensor):
self._traj_score[key] = val.to(map_location)
for instance in self.instances:
instance = instance.to(map_location)
if isinstance(map_location, (str, torch.device)):
self._device = map_location
return self
@classmethod
def from_slp(
cls,
lf: sio.LabeledFrame,
video_id: int = 0,
device: str | None = None,
**kwargs,
) -> Self:
"""Convert `sio.LabeledFrame` to `dreem.io.Frame`.
Args:
lf: A sio.LabeledFrame object
video_id: The ID of the video containing this frame.
device: The device to use for tensor operations.
**kwargs: Additional keyword arguments passed to Instance creation.
Returns:
A dreem.io.Frame object
"""
from dreem.io.instance import Instance
img_shape = lf.image.shape
if len(img_shape) == 2:
img_shape = (1, *img_shape)
elif len(img_shape) > 2 and img_shape[-1] <= 3:
img_shape = (lf.image.shape[-1], lf.image.shape[0], lf.image.shape[1])
return cls(
video_id=video_id,
frame_id=(
lf.frame_idx.astype(np.int32)
if isinstance(lf.frame_idx, np.number)
else lf.frame_idx
),
vid_file=lf.video.filename,
img_shape=torch.Size(img_shape),
instances=[Instance.from_slp(instance, **kwargs) for instance in lf],
device=device,
)
def to_slp(
self,
track_lookup: dict[int, sio.Track] | None = None,
video: sio.Video | None = None,
) -> tuple[sio.LabeledFrame, dict[int, sio.Track]]:
"""Convert Frame to sleap_io.LabeledFrame object.
Args:
track_lookup: A lookup dictionary containing the track_id and sio.Track for persistence
video: An sio.Video object used for overriding.
Returns: A tuple containing a LabeledFrame object with necessary metadata and
a lookup dictionary containing the track_id and sio.Track for persistence
"""
if track_lookup is None:
track_lookup = {}
slp_instances = []
for instance in self.instances:
slp_instance, track_lookup = instance.to_slp(track_lookup=track_lookup)
slp_instances.append(slp_instance)
if video is None:
video = (
self.video
if isinstance(self.video, sio.Video)
else sio.load_video(self.video)
)
return (
sio.LabeledFrame(
video=video,
frame_idx=self.frame_id.item(),
instances=slp_instances,
),
track_lookup,
)
def to_h5(
self,
clip_group: h5py.Group,
instance_labels: list | None = None,
save: dict[str, bool] | None = None,
) -> h5py.Group:
"""Convert frame to h5py group.
Args:
clip_group: the h5py group representing the clip (e.g batch/video) the frame belongs to
instance_labels: the labels used to create instance group names
save: whether to save crops, features and embeddings for the instance
Returns:
An h5py group containing the frame
"""
if save is None:
save = {"crop": False, "features": False, "embeddings": False}
frame_group = clip_group.require_group(f"frame_{self.frame_id.item()}")
frame_group.attrs.create("frame_id", self.frame_id.item())
frame_group.attrs.create("vid_id", self.video_id.item())
frame_group.attrs.create("vid_name", self.vid_name)
frame_group.create_dataset(
"asso_matrix",
data=self.asso_output.numpy() if self.asso_output is not None else [],
)
asso_group = frame_group.require_group("traj_scores")
for key, value in self.get_traj_score().items():
asso_group.create_dataset(
key, data=value.to_numpy() if value is not None else []
)
if instance_labels is None:
instance_labels = self.get_gt_track_ids.cpu().numpy()
for instance_label, instance in zip(instance_labels, self.instances):
kwargs = {}
if save.get("crop", False):
kwargs["crop"] = instance.crop.cpu().numpy()
if save.get("features", False):
kwargs["features"] = instance.features.cpu().numpy()
if save.get("embeddings", False):
for key, val in instance.get_embedding().items():
kwargs[f"{key}_emb"] = val.cpu().numpy()
_ = instance.to_h5(frame_group, f"instance_{instance_label}", **kwargs)
return frame_group
@property
def device(self) -> str:
"""The device the frame is on.
Returns:
The string representation of the device the frame is on.
"""
return self._device
@device.setter
def device(self, device: str) -> None:
"""Set the device.
Note: Do not set `frame.device = device` normally. Use `frame.to(device)` instead.
Args:
device: the device the function should be on.
"""
self._device = device
@property
def video_id(self) -> torch.Tensor:
"""The index of the video the frame comes from.
Returns:
A tensor containing the video index.
"""
return self._video_id
@video_id.setter
def video_id(self, video_id: int) -> None:
"""Set the video index.
Note: Generally the video_id should be immutable after initialization.
Args:
video_id: an int representing the index of the video that the frame came from.
"""
self._video_id = torch.tensor([video_id])
@property
def frame_id(self) -> torch.Tensor:
"""The index of the frame in a full video.
Returns:
A torch tensor containing the index of the frame in the video.
"""
return self._frame_id
@frame_id.setter
def frame_id(self, frame_id: int) -> None:
"""Set the frame index of the frame.
Note: The frame_id should generally be immutable after initialization.
Args:
frame_id: The int index of the frame in the full video.
"""
self._frame_id = torch.tensor([frame_id])
@property
def video(self) -> sio.Video | str:
"""Get the video associated with the frame.
Returns: An sio.Video object representing the video or a placeholder string
if it is not possible to create the sio.Video
"""
return self._video
@video.setter
def video(self, video: sio.Video | str) -> None:
"""Set the video associated with the frame.
Note: we try to store the video in an sio.Video object.
However, if this is not possible (e.g. incompatible format or missing filepath)
then we simply store the string.
Args:
video: sio.Video containing the vid reader or string path to video_file
"""
if isinstance(video, sio.Video):
self._video = video
else:
try:
self._video = sio.load_video(video)
except ValueError:
self._video = video
@property
def vid_name(self) -> str:
"""Get the path to the video corresponding to this frame.
Returns: A str file path corresponding to the frame.
"""
if isinstance(self.video, str):
return self.video
else:
return self.video.name
@property
def img_shape(self) -> torch.Size:
"""The shape of the pre-cropped frame.
Returns:
A torch.Size object containing the shape of the frame. Should generally be (c, h, w)
"""
return self._img_shape
@img_shape.setter
def img_shape(self, img_shape: ArrayLike | torch.Size) -> None:
"""Set the shape of the frame image.
Note: the img_shape should generally be immutable after initialization.
Args:
img_shape: an ArrayLike object or torch.Size containing the shape of the frame image.
"""
self._img_shape = _to_size(img_shape)
@property
def instances(self) -> list["Instance"]:
"""A list of instances in the frame.
Returns:
The list of instances that appear in the frame.
"""
return self._instances
@instances.setter
def instances(self, instances: list["Instance"]) -> None:
"""Set the frame's instance.
Args:
instances: A list of Instances that appear in the frame.
"""
for instance in instances:
instance.frame = self
self._instances = instances
def has_instances(self) -> bool:
"""Determine whether there are instances in the frame.
Returns:
True if there are instances in the frame, otherwise False.
"""
if self.num_detected == 0:
return False
return True
@property
def num_detected(self) -> int:
"""The number of instances in the frame.
Returns:
the number of instances in the frame.
"""
return len(self.instances)
@property
def asso_output(self) -> AssociationMatrix:
"""The association matrix between instances outputted directly by transformer.
Returns:
An arraylike (n_query, n_nonquery) association matrix between instances.
"""
return self._asso_output
def has_asso_output(self) -> bool:
"""Determine whether the frame has an association matrix computed.
Returns:
True if the frame has an association matrix otherwise, False.
"""
if self._asso_output is None or len(self._asso_output.matrix) == 0:
return False
return True
@asso_output.setter
def asso_output(self, asso_output: AssociationMatrix) -> None:
"""Set the association matrix of a frame.
Args:
asso_output: An arraylike (n_query, n_nonquery) association matrix between instances.
"""
self._asso_output = asso_output
@property
def matches(self) -> tuple:
"""Matches between frame instances and available trajectories.
Returns:
A tuple containing the instance idx and trajectory idx for the matched instance.
"""
return self._matches
@matches.setter
def matches(self, matches: tuple) -> None:
"""Set the frame matches.
Args:
matches: A tuple containing the instance idx and trajectory idx for the matched instance.
"""
self._matches = matches
def has_matches(self) -> bool:
"""Check whether or not matches have been computed for frame.
Returns:
True if frame contains matches otherwise False.
"""
if self._matches is not None and len(self._matches) > 0:
return True
return False
def get_traj_score(self, key: str | None = None) -> dict | ArrayLike | None:
"""Get dictionary containing association matrix between instances and trajectories along postprocessing pipeline.
Args:
key: The key of the trajectory score to be accessed.
Can be one of {None, 'initial', 'decay_time', 'max_center_dist', 'iou', 'final'}
Returns:
- dictionary containing all trajectory scores if key is None
- trajectory score associated with key
- None if the key is not found
"""
if key is None:
return self._traj_score
else:
try:
return self._traj_score[key]
except KeyError as e:
logger.exception(f"Could not access {key} traj_score due to {e}")
return None
def add_traj_score(self, key: str, traj_score: ArrayLike) -> None:
"""Add trajectory score to dictionary.
Args:
key: key associated with traj score to be used in dictionary
traj_score: association matrix between instances and trajectories
"""
self._traj_score[key] = traj_score
def has_traj_score(self) -> bool:
"""Check if any trajectory association matrix has been saved.
Returns:
True there is at least one association matrix otherwise, false.
"""
if len(self._traj_score) == 0:
return False
return True
def has_gt_track_ids(self) -> bool:
"""Check if any of frames instances has a gt track id.
Returns:
True if at least 1 instance has a gt track id otherwise False.
"""
if self.has_instances():
return any([instance.has_gt_track_id() for instance in self.instances])
return False
def get_gt_track_ids(self) -> torch.Tensor:
"""Get the gt track ids of all instances in the frame.
Returns:
an (N,) shaped tensor with the gt track ids of each instance in the frame.
"""
if not self.has_instances():
return torch.tensor([])
return torch.cat([instance.gt_track_id for instance in self.instances])
def has_pred_track_ids(self) -> bool:
"""Check if any of frames instances has a pred track id.
Returns:
True if at least 1 instance has a pred track id otherwise False.
"""
if self.has_instances():
return any([instance.has_pred_track_id() for instance in self.instances])
return False
def get_pred_track_ids(self) -> torch.Tensor:
"""Get the pred track ids of all instances in the frame.
Returns:
an (N,) shaped tensor with the pred track ids of each instance in the frame.
"""
if not self.has_instances():
return torch.tensor([])
return torch.cat([instance.pred_track_id for instance in self.instances])
def has_bboxes(self) -> bool:
"""Check if any of frames instances has a bounding box.
Returns:
True if at least 1 instance has a bounding box otherwise False.
"""
if self.has_instances():
return any([instance.has_bboxes() for instance in self.instances])
return False
def get_bboxes(self) -> torch.Tensor:
"""Get the bounding boxes of all instances in the frame.
Returns:
an (N,4) shaped tensor with bounding boxes of each instance in the frame.
"""
if not self.has_instances():
return torch.empty(0, 4)
return torch.cat([instance.bbox for instance in self.instances], dim=0)
def has_crops(self) -> bool:
"""Check if any of frames instances has a crop.
Returns:
True if at least 1 instance has a crop otherwise False.
"""
if self.has_instances():
return any([instance.has_crop() for instance in self.instances])
return False
def get_crops(self) -> torch.Tensor:
"""Get the crops of all instances in the frame.
Returns:
an (N, C, H, W) shaped tensor with crops of each instance in the frame.
"""
if not self.has_instances():
return torch.tensor([])
return torch.cat([instance.crop for instance in self.instances], dim=0)
def has_features(self) -> bool:
"""Check if any of frames instances has reid features already computed.
Returns:
True if at least 1 instance have reid features otherwise False.
"""
if self.has_instances():
return any([instance.has_features() for instance in self.instances])
return False
def get_features(self) -> torch.Tensor:
"""Get the reid feature vectors of all instances in the frame.
Returns:
an (N, D) shaped tensor with reid feature vectors of each instance in the frame.
"""
if not self.has_instances():
return torch.tensor([])
return torch.cat([instance.features for instance in self.instances], dim=0)
def get_anchors(self) -> list[str]:
"""Get the anchor names of instances in the frame.
Returns:
A list of anchor names used by the instances to get the crop.
"""
return [instance.anchor for instance in self.instances]
def get_centroids(self) -> tuple[list[str], ArrayLike]:
"""Get the centroids around which each instance's crop was formed.
Returns:
anchors: the node names for the corresponding point
points: an n_instances x 2 array containing the centroids
"""
anchors = [
anchor for instance in self.instances for anchor in instance.centroid.keys()
]
points = np.array(
[
point
for instance in self.instances
for point in instance.centroid.values()
]
)
return (anchors, points)
@property
def is_flagged(self) -> bool:
"""Whether the frame has been flagged for any reason.
Returns:
True if the frame has been flagged, otherwise False.
"""
return self._is_flagged
@is_flagged.setter
def is_flagged(self, value: bool) -> None:
"""Set the flagged status of the frame.
Args:
value: True to flag the frame, False to unflag it.
"""
self._is_flagged = value
if not value:
self._flag_reasons.clear()
@property
def flag_reasons(self) -> set[FrameFlagCode]:
"""List of flag codes indicating why the frame was flagged.
Returns:
A set of FrameFlagCode values indicating the reasons for flagging.
"""
return self._flag_reasons.copy()
def add_flag(self, flag_code: FrameFlagCode) -> None:
"""Add a flag reason to the frame.
Args:
flag_code: The FrameFlagCode indicating why the frame is being flagged.
"""
if flag_code not in self._flag_reasons:
self._flag_reasons.add(flag_code)
self._is_flagged = True
def remove_flag(self, flag_code: FrameFlagCode) -> None:
"""Remove a specific flag reason from the frame.
Args:
flag_code: The FrameFlagCode to remove.
"""
if flag_code in self._flag_reasons:
self._flag_reasons.discard(flag_code)
if not self._flag_reasons:
self._is_flagged = False
def has_flag(self, flag_code: FrameFlagCode) -> bool:
"""Check if the frame has a specific flag reason.
Args:
flag_code: The FrameFlagCode to check for.
Returns:
True if the frame has the specified flag code, otherwise False.
"""
return flag_code in self._flag_reasons
def clear_flags(self) -> None:
"""Clear all flags from the frame."""
self._flag_reasons.clear()
self._is_flagged = False
asso_output
property
writable
¶
The association matrix between instances outputted directly by transformer.
Returns:
| Type | Description |
|---|---|
AssociationMatrix
|
An arraylike (n_query, n_nonquery) association matrix between instances. |
device
property
writable
¶
The device the frame is on.
Returns:
| Type | Description |
|---|---|
str
|
The string representation of the device the frame is on. |
flag_reasons
property
¶
List of flag codes indicating why the frame was flagged.
Returns:
| Type | Description |
|---|---|
set[FrameFlagCode]
|
A set of FrameFlagCode values indicating the reasons for flagging. |
frame_id
property
writable
¶
The index of the frame in a full video.
Returns:
| Type | Description |
|---|---|
Tensor
|
A torch tensor containing the index of the frame in the video. |
img_shape
property
writable
¶
The shape of the pre-cropped frame.
Returns:
| Type | Description |
|---|---|
Size
|
A torch.Size object containing the shape of the frame. Should generally be (c, h, w) |
instances
property
writable
¶
A list of instances in the frame.
Returns:
| Type | Description |
|---|---|
list['Instance']
|
The list of instances that appear in the frame. |
is_flagged
property
writable
¶
Whether the frame has been flagged for any reason.
Returns:
| Type | Description |
|---|---|
bool
|
True if the frame has been flagged, otherwise False. |
matches
property
writable
¶
Matches between frame instances and available trajectories.
Returns:
| Type | Description |
|---|---|
tuple
|
A tuple containing the instance idx and trajectory idx for the matched instance. |
num_detected
property
¶
The number of instances in the frame.
Returns:
| Type | Description |
|---|---|
int
|
the number of instances in the frame. |
vid_name
property
¶
Get the path to the video corresponding to this frame.
Returns: A str file path corresponding to the frame.
video
property
writable
¶
Get the video associated with the frame.
Returns: An sio.Video object representing the video or a placeholder string if it is not possible to create the sio.Video
video_id
property
writable
¶
The index of the video the frame comes from.
Returns:
| Type | Description |
|---|---|
Tensor
|
A tensor containing the video index. |
__attrs_post_init__()
¶
Handle more intricate default initializations and moving to device.
__repr__()
¶
Return String representation of the Frame.
Returns:
| Type | Description |
|---|---|
str
|
The string representation of the frame. |
Source code in dreem/io/frame.py
def __repr__(self) -> str:
"""Return String representation of the Frame.
Returns:
The string representation of the frame.
"""
return (
"Frame("
f"video={self._video.filename if isinstance(self._video, sio.Video) else self._video}, "
f"video_id={self._video_id.item()}, "
f"frame_id={self._frame_id.item()}, "
f"img_shape={self._img_shape}, "
f"num_detected={self.num_detected}, "
f"asso_output={self._asso_output}, "
f"traj_score={self._traj_score}, "
f"matches={self._matches}, "
f"instances={self._instances}, "
f"device={self._device}, "
f"is_flagged={self._is_flagged}, "
f"flag_reasons={self._flag_reasons}"
")"
)
add_flag(flag_code)
¶
Add a flag reason to the frame.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
flag_code
|
FrameFlagCode
|
The FrameFlagCode indicating why the frame is being flagged. |
required |
Source code in dreem/io/frame.py
add_traj_score(key, traj_score)
¶
Add trajectory score to dictionary.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
key
|
str
|
key associated with traj score to be used in dictionary |
required |
traj_score
|
ArrayLike
|
association matrix between instances and trajectories |
required |
Source code in dreem/io/frame.py
clear_flags()
¶
from_slp(lf, video_id=0, device=None, **kwargs)
classmethod
¶
Convert sio.LabeledFrame to dreem.io.Frame.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
lf
|
LabeledFrame
|
A sio.LabeledFrame object |
required |
video_id
|
int
|
The ID of the video containing this frame. |
0
|
device
|
str | None
|
The device to use for tensor operations. |
None
|
**kwargs
|
Additional keyword arguments passed to Instance creation. |
{}
|
Returns:
| Type | Description |
|---|---|
Self
|
A dreem.io.Frame object |
Source code in dreem/io/frame.py
@classmethod
def from_slp(
cls,
lf: sio.LabeledFrame,
video_id: int = 0,
device: str | None = None,
**kwargs,
) -> Self:
"""Convert `sio.LabeledFrame` to `dreem.io.Frame`.
Args:
lf: A sio.LabeledFrame object
video_id: The ID of the video containing this frame.
device: The device to use for tensor operations.
**kwargs: Additional keyword arguments passed to Instance creation.
Returns:
A dreem.io.Frame object
"""
from dreem.io.instance import Instance
img_shape = lf.image.shape
if len(img_shape) == 2:
img_shape = (1, *img_shape)
elif len(img_shape) > 2 and img_shape[-1] <= 3:
img_shape = (lf.image.shape[-1], lf.image.shape[0], lf.image.shape[1])
return cls(
video_id=video_id,
frame_id=(
lf.frame_idx.astype(np.int32)
if isinstance(lf.frame_idx, np.number)
else lf.frame_idx
),
vid_file=lf.video.filename,
img_shape=torch.Size(img_shape),
instances=[Instance.from_slp(instance, **kwargs) for instance in lf],
device=device,
)
get_anchors()
¶
Get the anchor names of instances in the frame.
Returns:
| Type | Description |
|---|---|
list[str]
|
A list of anchor names used by the instances to get the crop. |
get_bboxes()
¶
Get the bounding boxes of all instances in the frame.
Returns:
| Type | Description |
|---|---|
Tensor
|
an (N,4) shaped tensor with bounding boxes of each instance in the frame. |
Source code in dreem/io/frame.py
def get_bboxes(self) -> torch.Tensor:
"""Get the bounding boxes of all instances in the frame.
Returns:
an (N,4) shaped tensor with bounding boxes of each instance in the frame.
"""
if not self.has_instances():
return torch.empty(0, 4)
return torch.cat([instance.bbox for instance in self.instances], dim=0)
get_centroids()
¶
Get the centroids around which each instance's crop was formed.
Returns:
| Name | Type | Description |
|---|---|---|
anchors |
tuple[list[str], ArrayLike]
|
the node names for the corresponding point points: an n_instances x 2 array containing the centroids |
Source code in dreem/io/frame.py
def get_centroids(self) -> tuple[list[str], ArrayLike]:
"""Get the centroids around which each instance's crop was formed.
Returns:
anchors: the node names for the corresponding point
points: an n_instances x 2 array containing the centroids
"""
anchors = [
anchor for instance in self.instances for anchor in instance.centroid.keys()
]
points = np.array(
[
point
for instance in self.instances
for point in instance.centroid.values()
]
)
return (anchors, points)
get_crops()
¶
Get the crops of all instances in the frame.
Returns:
| Type | Description |
|---|---|
Tensor
|
an (N, C, H, W) shaped tensor with crops of each instance in the frame. |
Source code in dreem/io/frame.py
get_features()
¶
Get the reid feature vectors of all instances in the frame.
Returns:
| Type | Description |
|---|---|
Tensor
|
an (N, D) shaped tensor with reid feature vectors of each instance in the frame. |
Source code in dreem/io/frame.py
def get_features(self) -> torch.Tensor:
"""Get the reid feature vectors of all instances in the frame.
Returns:
an (N, D) shaped tensor with reid feature vectors of each instance in the frame.
"""
if not self.has_instances():
return torch.tensor([])
return torch.cat([instance.features for instance in self.instances], dim=0)
get_gt_track_ids()
¶
Get the gt track ids of all instances in the frame.
Returns:
| Type | Description |
|---|---|
Tensor
|
an (N,) shaped tensor with the gt track ids of each instance in the frame. |
Source code in dreem/io/frame.py
def get_gt_track_ids(self) -> torch.Tensor:
"""Get the gt track ids of all instances in the frame.
Returns:
an (N,) shaped tensor with the gt track ids of each instance in the frame.
"""
if not self.has_instances():
return torch.tensor([])
return torch.cat([instance.gt_track_id for instance in self.instances])
get_pred_track_ids()
¶
Get the pred track ids of all instances in the frame.
Returns:
| Type | Description |
|---|---|
Tensor
|
an (N,) shaped tensor with the pred track ids of each instance in the frame. |
Source code in dreem/io/frame.py
def get_pred_track_ids(self) -> torch.Tensor:
"""Get the pred track ids of all instances in the frame.
Returns:
an (N,) shaped tensor with the pred track ids of each instance in the frame.
"""
if not self.has_instances():
return torch.tensor([])
return torch.cat([instance.pred_track_id for instance in self.instances])
get_traj_score(key=None)
¶
Get dictionary containing association matrix between instances and trajectories along postprocessing pipeline.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
key
|
str | None
|
The key of the trajectory score to be accessed. Can be one of {None, 'initial', 'decay_time', 'max_center_dist', 'iou', 'final'} |
None
|
Returns:
| Type | Description |
|---|---|
dict | ArrayLike | None
|
|
Source code in dreem/io/frame.py
def get_traj_score(self, key: str | None = None) -> dict | ArrayLike | None:
"""Get dictionary containing association matrix between instances and trajectories along postprocessing pipeline.
Args:
key: The key of the trajectory score to be accessed.
Can be one of {None, 'initial', 'decay_time', 'max_center_dist', 'iou', 'final'}
Returns:
- dictionary containing all trajectory scores if key is None
- trajectory score associated with key
- None if the key is not found
"""
if key is None:
return self._traj_score
else:
try:
return self._traj_score[key]
except KeyError as e:
logger.exception(f"Could not access {key} traj_score due to {e}")
return None
has_asso_output()
¶
Determine whether the frame has an association matrix computed.
Returns:
| Type | Description |
|---|---|
bool
|
True if the frame has an association matrix otherwise, False. |
Source code in dreem/io/frame.py
has_bboxes()
¶
Check if any of frames instances has a bounding box.
Returns:
| Type | Description |
|---|---|
bool
|
True if at least 1 instance has a bounding box otherwise False. |
Source code in dreem/io/frame.py
has_crops()
¶
Check if any of frames instances has a crop.
Returns:
| Type | Description |
|---|---|
bool
|
True if at least 1 instance has a crop otherwise False. |
has_features()
¶
Check if any of frames instances has reid features already computed.
Returns:
| Type | Description |
|---|---|
bool
|
True if at least 1 instance have reid features otherwise False. |
Source code in dreem/io/frame.py
has_flag(flag_code)
¶
Check if the frame has a specific flag reason.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
flag_code
|
FrameFlagCode
|
The FrameFlagCode to check for. |
required |
Returns:
| Type | Description |
|---|---|
bool
|
True if the frame has the specified flag code, otherwise False. |
Source code in dreem/io/frame.py
has_gt_track_ids()
¶
Check if any of frames instances has a gt track id.
Returns:
| Type | Description |
|---|---|
bool
|
True if at least 1 instance has a gt track id otherwise False. |
Source code in dreem/io/frame.py
has_instances()
¶
Determine whether there are instances in the frame.
Returns:
| Type | Description |
|---|---|
bool
|
True if there are instances in the frame, otherwise False. |
has_matches()
¶
Check whether or not matches have been computed for frame.
Returns:
| Type | Description |
|---|---|
bool
|
True if frame contains matches otherwise False. |
has_pred_track_ids()
¶
Check if any of frames instances has a pred track id.
Returns:
| Type | Description |
|---|---|
bool
|
True if at least 1 instance has a pred track id otherwise False. |
Source code in dreem/io/frame.py
has_traj_score()
¶
Check if any trajectory association matrix has been saved.
Returns:
| Type | Description |
|---|---|
bool
|
True there is at least one association matrix otherwise, false. |
remove_flag(flag_code)
¶
Remove a specific flag reason from the frame.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
flag_code
|
FrameFlagCode
|
The FrameFlagCode to remove. |
required |
Source code in dreem/io/frame.py
to(map_location)
¶
Move frame to different device or dtype (See torch.to for more info).
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
map_location
|
str | device
|
A string representing the device to move to. |
required |
Returns:
| Type | Description |
|---|---|
Self
|
The frame moved to a different device/dtype. |
Source code in dreem/io/frame.py
def to(self, map_location: str | torch.device) -> Self:
"""Move frame to different device or dtype (See `torch.to` for more info).
Args:
map_location: A string representing the device to move to.
Returns:
The frame moved to a different device/dtype.
"""
self._video_id = self._video_id.to(map_location)
self._frame_id = self._frame_id.to(map_location)
# torch.Size is immutable and doesn't need device movement
if isinstance(self._asso_output, torch.Tensor):
self._asso_output = self._asso_output.to(map_location)
if isinstance(self._matches, torch.Tensor):
self._matches = self._matches.to(map_location)
for key, val in self._traj_score.items():
if isinstance(val, torch.Tensor):
self._traj_score[key] = val.to(map_location)
for instance in self.instances:
instance = instance.to(map_location)
if isinstance(map_location, (str, torch.device)):
self._device = map_location
return self
to_h5(clip_group, instance_labels=None, save=None)
¶
Convert frame to h5py group.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
clip_group
|
Group
|
the h5py group representing the clip (e.g batch/video) the frame belongs to |
required |
instance_labels
|
list | None
|
the labels used to create instance group names |
None
|
save
|
dict[str, bool] | None
|
whether to save crops, features and embeddings for the instance |
None
|
Returns: An h5py group containing the frame
Source code in dreem/io/frame.py
def to_h5(
self,
clip_group: h5py.Group,
instance_labels: list | None = None,
save: dict[str, bool] | None = None,
) -> h5py.Group:
"""Convert frame to h5py group.
Args:
clip_group: the h5py group representing the clip (e.g batch/video) the frame belongs to
instance_labels: the labels used to create instance group names
save: whether to save crops, features and embeddings for the instance
Returns:
An h5py group containing the frame
"""
if save is None:
save = {"crop": False, "features": False, "embeddings": False}
frame_group = clip_group.require_group(f"frame_{self.frame_id.item()}")
frame_group.attrs.create("frame_id", self.frame_id.item())
frame_group.attrs.create("vid_id", self.video_id.item())
frame_group.attrs.create("vid_name", self.vid_name)
frame_group.create_dataset(
"asso_matrix",
data=self.asso_output.numpy() if self.asso_output is not None else [],
)
asso_group = frame_group.require_group("traj_scores")
for key, value in self.get_traj_score().items():
asso_group.create_dataset(
key, data=value.to_numpy() if value is not None else []
)
if instance_labels is None:
instance_labels = self.get_gt_track_ids.cpu().numpy()
for instance_label, instance in zip(instance_labels, self.instances):
kwargs = {}
if save.get("crop", False):
kwargs["crop"] = instance.crop.cpu().numpy()
if save.get("features", False):
kwargs["features"] = instance.features.cpu().numpy()
if save.get("embeddings", False):
for key, val in instance.get_embedding().items():
kwargs[f"{key}_emb"] = val.cpu().numpy()
_ = instance.to_h5(frame_group, f"instance_{instance_label}", **kwargs)
return frame_group
to_slp(track_lookup=None, video=None)
¶
Convert Frame to sleap_io.LabeledFrame object.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
track_lookup
|
dict[int, Track] | None
|
A lookup dictionary containing the track_id and sio.Track for persistence |
None
|
video
|
Video | None
|
An sio.Video object used for overriding. |
None
|
Returns: A tuple containing a LabeledFrame object with necessary metadata and a lookup dictionary containing the track_id and sio.Track for persistence
Source code in dreem/io/frame.py
def to_slp(
self,
track_lookup: dict[int, sio.Track] | None = None,
video: sio.Video | None = None,
) -> tuple[sio.LabeledFrame, dict[int, sio.Track]]:
"""Convert Frame to sleap_io.LabeledFrame object.
Args:
track_lookup: A lookup dictionary containing the track_id and sio.Track for persistence
video: An sio.Video object used for overriding.
Returns: A tuple containing a LabeledFrame object with necessary metadata and
a lookup dictionary containing the track_id and sio.Track for persistence
"""
if track_lookup is None:
track_lookup = {}
slp_instances = []
for instance in self.instances:
slp_instance, track_lookup = instance.to_slp(track_lookup=track_lookup)
slp_instances.append(slp_instance)
if video is None:
video = (
self.video
if isinstance(self.video, sio.Video)
else sio.load_video(self.video)
)
return (
sio.LabeledFrame(
video=video,
frame_idx=self.frame_id.item(),
instances=slp_instances,
),
track_lookup,
)
FrameFlagCode
¶
Bases: Enum
Enumeration of flag codes for Frame objects.
Each flag code represents a specific reason why a frame might be flagged. This enum can be extended with additional flag codes as needed.
Attributes:
| Name | Type | Description |
|---|---|---|
LOW_CONFIDENCE |
Frame contains instances with low confidence scores (below confidence threshold). |
|
HIGH_ENTROPY |
Frame contains instances with high entropy in association scores, indicating uncertain tracking assignments. |
|
MISSING_DETECTIONS |
Frame has no detected instances when some were expected. |
|
TRACKING_FAILURE |
Frame failed to be assigned to any trajectory. |
Methods:
| Name | Description |
|---|---|
__repr__ |
Return the representation of the flag code. |
__str__ |
Return the string value of the flag code. |
Source code in dreem/io/flags.py
class FrameFlagCode(Enum):
"""Enumeration of flag codes for Frame objects.
Each flag code represents a specific reason why a frame might be flagged.
This enum can be extended with additional flag codes as needed.
Attributes:
LOW_CONFIDENCE: Frame contains instances with low confidence scores
(below confidence threshold).
HIGH_ENTROPY: Frame contains instances with high entropy in association
scores, indicating uncertain tracking assignments.
MISSING_DETECTIONS: Frame has no detected instances when some were expected.
TRACKING_FAILURE: Frame failed to be assigned to any trajectory.
"""
LOW_CONFIDENCE = "low_confidence"
def __str__(self) -> str:
"""Return the string value of the flag code."""
return self.value
def __repr__(self) -> str:
"""Return the representation of the flag code."""
return f"FrameFlagCode.{self.name}"
Instance
¶
Class representing a single instance to be tracked.
Attributes:
| Name | Type | Description |
|---|---|---|
gt_track_id |
Tensor
|
Ground truth track id - only used for train/eval. |
pred_track_id |
Tensor
|
Predicted track id. Untracked instance is represented by -1. |
bbox |
Tensor
|
The bounding box coordinate of the instance. Defaults to an empty tensor. |
crop |
Tensor
|
The crop of the instance. |
centroid |
dict[str, ArrayLike]
|
the centroid around which the bbox was cropped. |
features |
Tensor
|
The reid features extracted from the CNN backbone used in the transformer. |
track_score |
float
|
The track score output from the association matrix. |
point_scores |
ArrayLike
|
The point scores from sleap. |
instance_score |
float
|
The instance scores from sleap. |
skeleton |
Skeleton
|
The sleap skeleton used for the instance. |
pose |
dict[str, ArrayLike]
|
A dictionary containing the node name and corresponding point. |
device |
str
|
String representation of the device the instance should be on. |
Methods:
| Name | Description |
|---|---|
__attrs_post_init__ |
Handle dimensionality and more intricate default initializations post-init. |
__repr__ |
Return string representation of the Instance. |
add_embedding |
Save embedding to instance embedding dictionary. |
from_slp |
Convert a slp instance to a dreem instance. |
get_embedding |
Retrieve instance's spatial/temporal embedding. |
has_bbox |
Determine if the instance has a bbox. |
has_crop |
Determine if the instance has a crop. |
has_embedding |
Determine if the instance has embedding type requested. |
has_features |
Determine if the instance has computed reid features. |
has_gt_track_id |
Determine if instance has a gt track assignment. |
has_pose |
Check if the instance has a pose. |
has_pred_track_id |
Determine whether instance has predicted track id. |
to |
Move instance to different device or change dtype. (See |
to_h5 |
Convert instance to an h5 group". |
to_slp |
Convert instance to sleap_io.PredictedInstance object. |
Source code in dreem/io/instance.py
@attrs.define(eq=False)
class Instance:
"""Class representing a single instance to be tracked.
Attributes:
gt_track_id: Ground truth track id - only used for train/eval.
pred_track_id: Predicted track id. Untracked instance is represented by -1.
bbox: The bounding box coordinate of the instance. Defaults to an empty tensor.
crop: The crop of the instance.
centroid: the centroid around which the bbox was cropped.
features: The reid features extracted from the CNN backbone used in the transformer.
track_score: The track score output from the association matrix.
point_scores: The point scores from sleap.
instance_score: The instance scores from sleap.
skeleton: The sleap skeleton used for the instance.
pose: A dictionary containing the node name and corresponding point.
device: String representation of the device the instance should be on.
"""
_gt_track_id: int = attrs.field(
alias="gt_track_id", default=-1, converter=_to_tensor
)
_pred_track_id: int = attrs.field(
alias="pred_track_id", default=-1, converter=_to_tensor
)
_bbox: ArrayLike = attrs.field(alias="bbox", factory=list, converter=_to_tensor)
_crop: ArrayLike = attrs.field(alias="crop", factory=list, converter=_to_tensor)
_centroid: dict[str, ArrayLike] = attrs.field(alias="centroid", factory=dict)
_features: ArrayLike = attrs.field(
alias="features", factory=list, converter=_to_tensor
)
_embeddings: dict = attrs.field(alias="embeddings", factory=dict)
_track_score: float = attrs.field(alias="track_score", default=-1.0)
_instance_score: float = attrs.field(alias="instance_score", default=-1.0)
_point_scores: ArrayLike | None = attrs.field(alias="point_scores", default=None)
_skeleton: sio.Skeleton | None = attrs.field(alias="skeleton", default=None)
_mask: ArrayLike | None = attrs.field(
alias="mask", converter=_to_tensor, default=None
)
_pose: dict[str, ArrayLike] = attrs.field(alias="pose", factory=dict)
_device: str | torch.device | None = attrs.field(alias="device", default=None)
_frame: Optional["Frame"] = None
def __attrs_post_init__(self) -> None:
"""Handle dimensionality and more intricate default initializations post-init."""
self.bbox = _expand_to_rank(self.bbox, 3)
self.crop = _expand_to_rank(self.crop, 4)
self.features = _expand_to_rank(self.features, 2)
if self.skeleton is None:
self.skeleton = sio.Skeleton(["centroid"])
if self.bbox.shape[-1] == 0:
self.bbox = torch.empty([1, 0, 4])
if self.crop.shape[-1] == 0 and self.bbox.shape[1] != 0:
y1, x1, y2, x2 = self.bbox.squeeze(dim=0).nanmean(dim=0)
self.centroid = {"centroid": np.array([(x1 + x2) / 2, (y1 + y2) / 2])}
if len(self.pose) == 0 and self.bbox.shape[1]:
y1, x1, y2, x2 = self.bbox.squeeze(dim=0).mean(dim=0)
self._pose = {"centroid": np.array([(x1 + x2) / 2, (y1 + y2) / 2])}
if self.point_scores is None and len(self.pose) != 0:
self._point_scores = np.zeros((len(self.pose), 2))
self.to(self.device)
def __repr__(self) -> str:
"""Return string representation of the Instance."""
return (
"Instance("
f"gt_track_id={self._gt_track_id.item()}, "
f"pred_track_id={self._pred_track_id.item()}, "
f"bbox={self._bbox}, "
f"centroid={self._centroid}, "
f"crop={self._crop.shape}, "
f"features={self._features.shape}, "
f"device={self._device}"
")"
)
def to(self, map_location: str | torch.device) -> Self:
"""Move instance to different device or change dtype. (See `torch.to` for more info).
Args:
map_location: Either the device or dtype for the instance to be moved.
Returns:
self: reference to the instance moved to correct device/dtype.
"""
if map_location is not None and map_location != "":
self._gt_track_id = self._gt_track_id.to(map_location)
self._pred_track_id = self._pred_track_id.to(map_location)
self._bbox = self._bbox.to(map_location)
self._crop = self._crop.to(map_location)
self._features = self._features.to(map_location)
if isinstance(map_location, (str, torch.device)):
self.device = map_location
return self
@classmethod
def from_slp(
cls,
slp_instance: sio.PredictedInstance | sio.Instance,
bbox_size: int | tuple[int, int] = 64,
crop: ArrayLike | None = None,
device: str | None = None,
) -> Self:
"""Convert a slp instance to a dreem instance.
Args:
slp_instance: A `sleap_io.Instance` object representing a detection
bbox_size: size of the pose-centered bbox to form.
crop: The corresponding crop of the bbox
device: which device to keep the instance on
Returns:
A dreem.Instance object with a pose-centered bbox and no crop.
"""
try:
track_id = int(slp_instance.track.name)
except ValueError:
track_id = int(
"".join([str(ord(c)) for c in slp_instance.track.name])
) # better way to handle this?
if isinstance(bbox_size, int):
bbox_size = (bbox_size, bbox_size)
track_score = -1.0
point_scores = np.full(len(slp_instance.points), -1)
instance_score = -1
if isinstance(slp_instance, sio.PredictedInstance):
track_score = slp_instance.tracking_score
point_scores = slp_instance.numpy()[:, -1]
instance_score = slp_instance.score
centroid = np.nanmean(slp_instance.numpy(), axis=1)
bbox = [
centroid[1] - bbox_size[1],
centroid[0] - bbox_size[0],
centroid[1] + bbox_size[1],
centroid[0] + bbox_size[0],
]
return cls(
gt_track_id=track_id,
bbox=bbox,
crop=crop,
centroid={"centroid": centroid},
track_score=track_score,
point_scores=point_scores,
instance_score=instance_score,
skeleton=slp_instance.skeleton,
pose={
node.name: point.numpy() for node, point in slp_instance.points.items()
},
device=device,
)
def to_slp(
self, track_lookup: dict[int, sio.Track] = {}
) -> tuple[sio.PredictedInstance, dict[int, sio.Track]]:
"""Convert instance to sleap_io.PredictedInstance object.
Args:
track_lookup: A track look up dictionary containing track_id:sio.Track.
Returns: A sleap_io.PredictedInstance with necessary metadata
and a track_lookup dictionary to persist tracks.
"""
try:
track_id = self.pred_track_id.item()
if track_id not in track_lookup:
track_lookup[track_id] = sio.Track(name=self.pred_track_id.item())
track = track_lookup[track_id]
return (
sio.PredictedInstance.from_numpy(
points_data=np.array(list(self.pose.values())),
skeleton=self.skeleton,
point_scores=self.point_scores,
score=self.instance_score,
tracking_score=self.track_score,
track=track,
),
track_lookup,
)
except Exception as e:
logger.exception(
f"Pose: {np.array(list(self.pose.values())).shape}, Pose score shape {self.point_scores.shape}"
)
raise RuntimeError(f"Failed to convert to sio.PredictedInstance: {e}")
def to_h5(
self, frame_group: h5py.Group, label: Any = None, **kwargs: dict
) -> h5py.Group:
"""Convert instance to an h5 group".
By default we always save:
- the gt/pred track id
- bbox
- centroid
- pose
- instance/traj/points score
Larger arrays (crops/features/embeddings) can be saved by passing as kwargs
Args:
frame_group: the h5py group representing the frame the instance appears on
label: the name of the instance group that will be created
**kwargs: additional key:value pairs to be saved as datasets.
Returns:
The h5 group representing this instance.
"""
if label is None:
if self.pred_track_id != -1:
label = f"instance_{self.pred_track_id.item()}"
else:
label = f"instance_{self.gt_track_id.item()}"
instance_group = frame_group.create_group(label)
instance_group.attrs.create("gt_track_id", self.gt_track_id.item())
instance_group.attrs.create("pred_track_id", self.pred_track_id.item())
instance_group.attrs.create("track_score", self.track_score)
instance_group.attrs.create("instance_score", self.instance_score)
instance_group.create_dataset("bbox", data=self.bbox.cpu().numpy())
pose_group = instance_group.create_group("pose")
pose_group.create_dataset("points", data=np.array(list(self.pose.values())))
pose_group.attrs.create("nodes", list(self.pose.keys()))
pose_group.create_dataset("scores", data=self.point_scores)
for key, value in kwargs.items():
if "emb" in key:
emb_group = instance_group.require_group("emb")
emb_group.create_dataset(key, data=value)
else:
instance_group.create_dataset(key, data=value)
return instance_group
@property
def device(self) -> str:
"""The device the instance is on.
Returns:
The str representation of the device the gpu is on.
"""
return self._device
@device.setter
def device(self, device) -> None:
"""Set for the device property.
Args:
device: The str representation of the device.
"""
self._device = device
@property
def gt_track_id(self) -> torch.Tensor:
"""The ground truth track id of the instance.
Returns:
A tensor containing the ground truth track id
"""
return self._gt_track_id
@gt_track_id.setter
def gt_track_id(self, track: int):
"""Set the instance ground-truth track id.
Args:
track: An int representing the ground-truth track id.
"""
if track is not None:
self._gt_track_id = torch.tensor([track])
else:
self._gt_track_id = torch.tensor([])
def has_gt_track_id(self) -> bool:
"""Determine if instance has a gt track assignment.
Returns:
True if the gt track id is set, otherwise False.
"""
if self._gt_track_id.shape[0] == 0:
return False
else:
return True
@property
def pred_track_id(self) -> torch.Tensor:
"""The track id predicted by the tracker using asso_output from model.
Returns:
A tensor containing the predicted track id.
"""
return self._pred_track_id
@pred_track_id.setter
def pred_track_id(self, track: int) -> None:
"""Set predicted track id.
Args:
track: an int representing the predicted track id.
"""
if track is not None:
self._pred_track_id = torch.tensor([track])
else:
self._pred_track_id = torch.tensor([])
def has_pred_track_id(self) -> bool:
"""Determine whether instance has predicted track id.
Returns:
True if instance has a pred track id, False otherwise.
"""
if self._pred_track_id.item() == -1 or self._pred_track_id.shape[0] == 0:
return False
else:
return True
@property
def bbox(self) -> torch.Tensor:
"""The bounding box coordinates of the instance in the original frame.
Returns:
A (1,4) tensor containing the bounding box coordinates.
"""
return self._bbox
@bbox.setter
def bbox(self, bbox: ArrayLike) -> None:
"""Set the instance bounding box.
Args:
bbox: an arraylike object containing the bounding box coordinates.
"""
if bbox is None or len(bbox) == 0:
self._bbox = torch.empty((0, 4))
else:
if not isinstance(bbox, torch.Tensor):
self._bbox = torch.tensor(bbox)
else:
self._bbox = bbox
if self._bbox.shape[0] and len(self._bbox.shape) == 1:
self._bbox = self._bbox.unsqueeze(0)
if self._bbox.shape[1] and len(self._bbox.shape) == 2:
self._bbox = self._bbox.unsqueeze(0)
def has_bbox(self) -> bool:
"""Determine if the instance has a bbox.
Returns:
True if the instance has a bounding box, false otherwise.
"""
if self._bbox.shape[1] == 0:
return False
else:
return True
@property
def centroid(self) -> dict[str, ArrayLike]:
"""The centroid around which the crop was formed.
Returns:
A dict containing the anchor name and the x, y bbox midpoint.
"""
return self._centroid
@centroid.setter
def centroid(self, centroid: dict[str, ArrayLike]) -> None:
"""Set the centroid of the instance.
Args:
centroid: A dict containing the anchor name and points.
"""
self._centroid = centroid
@property
def anchor(self) -> list[str]:
"""The anchor node name around which the crop was formed.
Returns:
the list of anchors around which each crop was formed
the list of anchors around which each crop was formed
"""
if self.centroid:
return list(self.centroid.keys())
return ""
@property
def mask(self) -> torch.Tensor:
"""The mask of the instance.
Returns:
A (h, w) tensor containing the mask of the instance.
"""
return self._mask
@mask.setter
def mask(self, mask: ArrayLike) -> None:
"""Set the mask of the instance.
Args:
mask: an arraylike object containing the mask of the instance.
"""
if mask is None or len(mask) == 0:
self._mask = torch.tensor([])
else:
if not isinstance(mask, torch.Tensor):
self._mask = torch.tensor(mask)
else:
self._mask = mask
@property
def crop(self) -> torch.Tensor:
"""The crop of the instance.
Returns:
A (1, c, h , w) tensor containing the cropped image centered around the instance.
"""
return self._crop
@crop.setter
def crop(self, crop: ArrayLike) -> None:
"""Set the crop of the instance.
Args:
crop: an arraylike object containing the cropped image of the centered instance.
"""
if crop is None or len(crop) == 0:
self._crop = torch.tensor([])
else:
if not isinstance(crop, torch.Tensor):
self._crop = torch.tensor(crop)
else:
self._crop = crop
if len(self._crop.shape) == 2:
self._crop = self._crop.unsqueeze(0)
if len(self._crop.shape) == 3:
self._crop = self._crop.unsqueeze(0)
def has_crop(self) -> bool:
"""Determine if the instance has a crop.
Returns:
True if the instance has an image otherwise False.
"""
if self._crop.shape[-1] == 0:
return False
else:
return True
@property
def features(self) -> torch.Tensor:
"""Re-ID feature vector from backbone model to be used as input to transformer.
Returns:
a (1, d) tensor containing the reid feature vector.
"""
return self._features
@features.setter
def features(self, features: ArrayLike) -> None:
"""Set the reid feature vector of the instance.
Args:
features: a (1,d) array like object containing the reid features for the instance.
"""
if features is None or len(features) == 0:
self._features = torch.tensor([])
elif not isinstance(features, torch.Tensor):
self._features = torch.tensor(features)
else:
self._features = features
if self._features.shape[0] and len(self._features.shape) == 1:
self._features = self._features.unsqueeze(0)
def has_features(self) -> bool:
"""Determine if the instance has computed reid features.
Returns:
True if the instance has reid features, False otherwise.
"""
if self._features.shape[-1] == 0:
return False
else:
return True
def has_embedding(self, emb_type: str | None = None) -> bool:
"""Determine if the instance has embedding type requested.
Args:
emb_type: The key to check in the embedding dictionary.
Returns:
True if `emb_type` in embedding_dict else false
"""
return emb_type in self._embeddings
def get_embedding(
self, emb_type: str = "all"
) -> dict[str, torch.Tensor] | torch.Tensor | None:
"""Retrieve instance's spatial/temporal embedding.
Args:
emb_type: The string key of the embedding to retrieve. Should be "pos", "temp"
Returns:
* A torch tensor representing the spatial/temporal location of the instance.
* None if the embedding is not stored
"""
if emb_type.lower() == "all":
return self._embeddings
else:
try:
return self._embeddings[emb_type]
except KeyError:
logger.exception(
f"{emb_type} not saved! Only {list(self._embeddings.keys())} are available"
)
return None
def add_embedding(self, emb_type: str, embedding: torch.Tensor) -> None:
"""Save embedding to instance embedding dictionary.
Args:
emb_type: Key/embedding type to be saved to dictionary
embedding: The actual torch tensor embedding.
"""
embedding = _expand_to_rank(embedding, 2)
self._embeddings[emb_type] = embedding
@property
def frame(self) -> "Frame":
"""Get the frame the instance belongs to.
Returns:
The back reference to the `Frame` that this `Instance` belongs to.
"""
return self._frame
@frame.setter
def frame(self, frame: "Frame") -> None:
"""Set the back reference to the `Frame` that this `Instance` belongs to.
This field is set when instances are added to `Frame` object.
Args:
frame: A `Frame` object containing the metadata for the frame that the instance belongs to
"""
self._frame = frame
@property
def pose(self) -> dict[str, ArrayLike]:
"""Get the pose of the instance.
Returns:
A dictionary containing the node and corresponding x,y points
"""
return self._pose
@pose.setter
def pose(self, pose: dict[str, ArrayLike]) -> None:
"""Set the pose of the instance.
Args:
pose: A nodes x 2 array containing the pose coordinates.
"""
if pose is not None:
self._pose = pose
elif self.bbox.shape[0]:
y1, x1, y2, x2 = self.bbox.squeeze()
self._pose = {"centroid": np.array([(x1 + x2) / 2, (y1 + y2) / 2])}
else:
self._pose = {}
def has_pose(self) -> bool:
"""Check if the instance has a pose.
Returns True if the instance has a pose.
"""
if len(self.pose):
return True
return False
@property
def shown_pose(self) -> dict[str, ArrayLike]:
"""Get the pose with shown nodes only.
Returns: A dictionary filtered by nodes that are shown (points are not nan).
"""
pose = self.pose
return {node: point for node, point in pose.items() if not np.isna(point).any()}
@property
def skeleton(self) -> sio.Skeleton:
"""Get the skeleton associated with the instance.
Returns: The sio.Skeleton associated with the instance.
"""
return self._skeleton
@skeleton.setter
def skeleton(self, skeleton: sio.Skeleton) -> None:
"""Set the skeleton associated with the instance.
Args:
skeleton: The sio.Skeleton associated with the instance.
"""
self._skeleton = skeleton
@property
def point_scores(self) -> ArrayLike:
"""Get the point scores associated with the pose prediction.
Returns: a vector of shape n containing the point scores outputted from sleap associated with pose predictions.
"""
return self._point_scores
@point_scores.setter
def point_scores(self, point_scores: ArrayLike) -> None:
"""Set the point scores associated with the pose prediction.
Args:
point_scores: a vector of shape n containing the point scores
outputted from sleap associated with pose predictions.
"""
self._point_scores = point_scores
@property
def instance_score(self) -> float:
"""Get the pose prediction score associated with the instance.
Returns: a float from 0-1 representing an instance_score.
"""
return self._instance_score
@instance_score.setter
def instance_score(self, instance_score: float) -> None:
"""Set the pose prediction score associated with the instance.
Args:
instance_score: a float from 0-1 representing an instance_score.
"""
self._instance_score = instance_score
@property
def track_score(self) -> float:
"""Get the track_score of the instance.
Returns: A float from 0-1 representing the output used in the tracker for assignment.
"""
return self._track_score
@track_score.setter
def track_score(self, track_score: float) -> None:
"""Set the track_score of the instance.
Args:
track_score: A float from 0-1 representing the output used in the tracker for assignment.
"""
self._track_score = track_score
anchor
property
¶
The anchor node name around which the crop was formed.
Returns:
| Type | Description |
|---|---|
list[str]
|
the list of anchors around which each crop was formed the list of anchors around which each crop was formed |
bbox
property
writable
¶
The bounding box coordinates of the instance in the original frame.
Returns:
| Type | Description |
|---|---|
Tensor
|
A (1,4) tensor containing the bounding box coordinates. |
centroid
property
writable
¶
The centroid around which the crop was formed.
Returns:
| Type | Description |
|---|---|
dict[str, ArrayLike]
|
A dict containing the anchor name and the x, y bbox midpoint. |
crop
property
writable
¶
The crop of the instance.
Returns:
| Type | Description |
|---|---|
Tensor
|
A (1, c, h , w) tensor containing the cropped image centered around the instance. |
device
property
writable
¶
The device the instance is on.
Returns:
| Type | Description |
|---|---|
str
|
The str representation of the device the gpu is on. |
features
property
writable
¶
Re-ID feature vector from backbone model to be used as input to transformer.
Returns:
| Type | Description |
|---|---|
Tensor
|
a (1, d) tensor containing the reid feature vector. |
frame
property
writable
¶
Get the frame the instance belongs to.
Returns:
| Type | Description |
|---|---|
Frame
|
The back reference to the |
gt_track_id
property
writable
¶
The ground truth track id of the instance.
Returns:
| Type | Description |
|---|---|
Tensor
|
A tensor containing the ground truth track id |
instance_score
property
writable
¶
Get the pose prediction score associated with the instance.
Returns: a float from 0-1 representing an instance_score.
mask
property
writable
¶
The mask of the instance.
Returns:
| Type | Description |
|---|---|
Tensor
|
A (h, w) tensor containing the mask of the instance. |
point_scores
property
writable
¶
Get the point scores associated with the pose prediction.
Returns: a vector of shape n containing the point scores outputted from sleap associated with pose predictions.
pose
property
writable
¶
Get the pose of the instance.
Returns:
| Type | Description |
|---|---|
dict[str, ArrayLike]
|
A dictionary containing the node and corresponding x,y points |
pred_track_id
property
writable
¶
The track id predicted by the tracker using asso_output from model.
Returns:
| Type | Description |
|---|---|
Tensor
|
A tensor containing the predicted track id. |
shown_pose
property
¶
Get the pose with shown nodes only.
Returns: A dictionary filtered by nodes that are shown (points are not nan).
skeleton
property
writable
¶
Get the skeleton associated with the instance.
Returns: The sio.Skeleton associated with the instance.
track_score
property
writable
¶
Get the track_score of the instance.
Returns: A float from 0-1 representing the output used in the tracker for assignment.
__attrs_post_init__()
¶
Handle dimensionality and more intricate default initializations post-init.
Source code in dreem/io/instance.py
def __attrs_post_init__(self) -> None:
"""Handle dimensionality and more intricate default initializations post-init."""
self.bbox = _expand_to_rank(self.bbox, 3)
self.crop = _expand_to_rank(self.crop, 4)
self.features = _expand_to_rank(self.features, 2)
if self.skeleton is None:
self.skeleton = sio.Skeleton(["centroid"])
if self.bbox.shape[-1] == 0:
self.bbox = torch.empty([1, 0, 4])
if self.crop.shape[-1] == 0 and self.bbox.shape[1] != 0:
y1, x1, y2, x2 = self.bbox.squeeze(dim=0).nanmean(dim=0)
self.centroid = {"centroid": np.array([(x1 + x2) / 2, (y1 + y2) / 2])}
if len(self.pose) == 0 and self.bbox.shape[1]:
y1, x1, y2, x2 = self.bbox.squeeze(dim=0).mean(dim=0)
self._pose = {"centroid": np.array([(x1 + x2) / 2, (y1 + y2) / 2])}
if self.point_scores is None and len(self.pose) != 0:
self._point_scores = np.zeros((len(self.pose), 2))
self.to(self.device)
__repr__()
¶
Return string representation of the Instance.
Source code in dreem/io/instance.py
def __repr__(self) -> str:
"""Return string representation of the Instance."""
return (
"Instance("
f"gt_track_id={self._gt_track_id.item()}, "
f"pred_track_id={self._pred_track_id.item()}, "
f"bbox={self._bbox}, "
f"centroid={self._centroid}, "
f"crop={self._crop.shape}, "
f"features={self._features.shape}, "
f"device={self._device}"
")"
)
add_embedding(emb_type, embedding)
¶
Save embedding to instance embedding dictionary.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
emb_type
|
str
|
Key/embedding type to be saved to dictionary |
required |
embedding
|
Tensor
|
The actual torch tensor embedding. |
required |
Source code in dreem/io/instance.py
def add_embedding(self, emb_type: str, embedding: torch.Tensor) -> None:
"""Save embedding to instance embedding dictionary.
Args:
emb_type: Key/embedding type to be saved to dictionary
embedding: The actual torch tensor embedding.
"""
embedding = _expand_to_rank(embedding, 2)
self._embeddings[emb_type] = embedding
from_slp(slp_instance, bbox_size=64, crop=None, device=None)
classmethod
¶
Convert a slp instance to a dreem instance.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
slp_instance
|
PredictedInstance | Instance
|
A |
required |
bbox_size
|
int | tuple[int, int]
|
size of the pose-centered bbox to form. |
64
|
crop
|
ArrayLike | None
|
The corresponding crop of the bbox |
None
|
device
|
str | None
|
which device to keep the instance on |
None
|
Returns: A dreem.Instance object with a pose-centered bbox and no crop.
Source code in dreem/io/instance.py
@classmethod
def from_slp(
cls,
slp_instance: sio.PredictedInstance | sio.Instance,
bbox_size: int | tuple[int, int] = 64,
crop: ArrayLike | None = None,
device: str | None = None,
) -> Self:
"""Convert a slp instance to a dreem instance.
Args:
slp_instance: A `sleap_io.Instance` object representing a detection
bbox_size: size of the pose-centered bbox to form.
crop: The corresponding crop of the bbox
device: which device to keep the instance on
Returns:
A dreem.Instance object with a pose-centered bbox and no crop.
"""
try:
track_id = int(slp_instance.track.name)
except ValueError:
track_id = int(
"".join([str(ord(c)) for c in slp_instance.track.name])
) # better way to handle this?
if isinstance(bbox_size, int):
bbox_size = (bbox_size, bbox_size)
track_score = -1.0
point_scores = np.full(len(slp_instance.points), -1)
instance_score = -1
if isinstance(slp_instance, sio.PredictedInstance):
track_score = slp_instance.tracking_score
point_scores = slp_instance.numpy()[:, -1]
instance_score = slp_instance.score
centroid = np.nanmean(slp_instance.numpy(), axis=1)
bbox = [
centroid[1] - bbox_size[1],
centroid[0] - bbox_size[0],
centroid[1] + bbox_size[1],
centroid[0] + bbox_size[0],
]
return cls(
gt_track_id=track_id,
bbox=bbox,
crop=crop,
centroid={"centroid": centroid},
track_score=track_score,
point_scores=point_scores,
instance_score=instance_score,
skeleton=slp_instance.skeleton,
pose={
node.name: point.numpy() for node, point in slp_instance.points.items()
},
device=device,
)
get_embedding(emb_type='all')
¶
Retrieve instance's spatial/temporal embedding.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
emb_type
|
str
|
The string key of the embedding to retrieve. Should be "pos", "temp" |
'all'
|
Returns:
| Type | Description |
|---|---|
dict[str, Tensor] | Tensor | None
|
|
Source code in dreem/io/instance.py
def get_embedding(
self, emb_type: str = "all"
) -> dict[str, torch.Tensor] | torch.Tensor | None:
"""Retrieve instance's spatial/temporal embedding.
Args:
emb_type: The string key of the embedding to retrieve. Should be "pos", "temp"
Returns:
* A torch tensor representing the spatial/temporal location of the instance.
* None if the embedding is not stored
"""
if emb_type.lower() == "all":
return self._embeddings
else:
try:
return self._embeddings[emb_type]
except KeyError:
logger.exception(
f"{emb_type} not saved! Only {list(self._embeddings.keys())} are available"
)
return None
has_bbox()
¶
Determine if the instance has a bbox.
Returns:
| Type | Description |
|---|---|
bool
|
True if the instance has a bounding box, false otherwise. |
has_crop()
¶
Determine if the instance has a crop.
Returns:
| Type | Description |
|---|---|
bool
|
True if the instance has an image otherwise False. |
has_embedding(emb_type=None)
¶
Determine if the instance has embedding type requested.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
emb_type
|
str | None
|
The key to check in the embedding dictionary. |
None
|
Returns:
| Type | Description |
|---|---|
bool
|
True if |
Source code in dreem/io/instance.py
has_features()
¶
Determine if the instance has computed reid features.
Returns:
| Type | Description |
|---|---|
bool
|
True if the instance has reid features, False otherwise. |
has_gt_track_id()
¶
Determine if instance has a gt track assignment.
Returns:
| Type | Description |
|---|---|
bool
|
True if the gt track id is set, otherwise False. |
has_pose()
¶
has_pred_track_id()
¶
Determine whether instance has predicted track id.
Returns:
| Type | Description |
|---|---|
bool
|
True if instance has a pred track id, False otherwise. |
to(map_location)
¶
Move instance to different device or change dtype. (See torch.to for more info).
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
map_location
|
str | device
|
Either the device or dtype for the instance to be moved. |
required |
Returns:
| Name | Type | Description |
|---|---|---|
self |
Self
|
reference to the instance moved to correct device/dtype. |
Source code in dreem/io/instance.py
def to(self, map_location: str | torch.device) -> Self:
"""Move instance to different device or change dtype. (See `torch.to` for more info).
Args:
map_location: Either the device or dtype for the instance to be moved.
Returns:
self: reference to the instance moved to correct device/dtype.
"""
if map_location is not None and map_location != "":
self._gt_track_id = self._gt_track_id.to(map_location)
self._pred_track_id = self._pred_track_id.to(map_location)
self._bbox = self._bbox.to(map_location)
self._crop = self._crop.to(map_location)
self._features = self._features.to(map_location)
if isinstance(map_location, (str, torch.device)):
self.device = map_location
return self
to_h5(frame_group, label=None, **kwargs)
¶
Convert instance to an h5 group".
By default we always save
- the gt/pred track id
- bbox
- centroid
- pose
- instance/traj/points score
Larger arrays (crops/features/embeddings) can be saved by passing as kwargs
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
frame_group
|
Group
|
the h5py group representing the frame the instance appears on |
required |
label
|
Any
|
the name of the instance group that will be created |
None
|
**kwargs
|
dict
|
additional key:value pairs to be saved as datasets. |
{}
|
Returns:
| Type | Description |
|---|---|
Group
|
The h5 group representing this instance. |
Source code in dreem/io/instance.py
def to_h5(
self, frame_group: h5py.Group, label: Any = None, **kwargs: dict
) -> h5py.Group:
"""Convert instance to an h5 group".
By default we always save:
- the gt/pred track id
- bbox
- centroid
- pose
- instance/traj/points score
Larger arrays (crops/features/embeddings) can be saved by passing as kwargs
Args:
frame_group: the h5py group representing the frame the instance appears on
label: the name of the instance group that will be created
**kwargs: additional key:value pairs to be saved as datasets.
Returns:
The h5 group representing this instance.
"""
if label is None:
if self.pred_track_id != -1:
label = f"instance_{self.pred_track_id.item()}"
else:
label = f"instance_{self.gt_track_id.item()}"
instance_group = frame_group.create_group(label)
instance_group.attrs.create("gt_track_id", self.gt_track_id.item())
instance_group.attrs.create("pred_track_id", self.pred_track_id.item())
instance_group.attrs.create("track_score", self.track_score)
instance_group.attrs.create("instance_score", self.instance_score)
instance_group.create_dataset("bbox", data=self.bbox.cpu().numpy())
pose_group = instance_group.create_group("pose")
pose_group.create_dataset("points", data=np.array(list(self.pose.values())))
pose_group.attrs.create("nodes", list(self.pose.keys()))
pose_group.create_dataset("scores", data=self.point_scores)
for key, value in kwargs.items():
if "emb" in key:
emb_group = instance_group.require_group("emb")
emb_group.create_dataset(key, data=value)
else:
instance_group.create_dataset(key, data=value)
return instance_group
to_slp(track_lookup={})
¶
Convert instance to sleap_io.PredictedInstance object.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
track_lookup
|
dict[int, Track]
|
A track look up dictionary containing track_id:sio.Track. |
{}
|
Returns: A sleap_io.PredictedInstance with necessary metadata and a track_lookup dictionary to persist tracks.
Source code in dreem/io/instance.py
def to_slp(
self, track_lookup: dict[int, sio.Track] = {}
) -> tuple[sio.PredictedInstance, dict[int, sio.Track]]:
"""Convert instance to sleap_io.PredictedInstance object.
Args:
track_lookup: A track look up dictionary containing track_id:sio.Track.
Returns: A sleap_io.PredictedInstance with necessary metadata
and a track_lookup dictionary to persist tracks.
"""
try:
track_id = self.pred_track_id.item()
if track_id not in track_lookup:
track_lookup[track_id] = sio.Track(name=self.pred_track_id.item())
track = track_lookup[track_id]
return (
sio.PredictedInstance.from_numpy(
points_data=np.array(list(self.pose.values())),
skeleton=self.skeleton,
point_scores=self.point_scores,
score=self.instance_score,
tracking_score=self.track_score,
track=track,
),
track_lookup,
)
except Exception as e:
logger.exception(
f"Pose: {np.array(list(self.pose.values())).shape}, Pose score shape {self.point_scores.shape}"
)
raise RuntimeError(f"Failed to convert to sio.PredictedInstance: {e}")
Track
¶
Object for storing instances of the same track.
Attributes:
| Name | Type | Description |
|---|---|---|
id |
the track label. |
|
instances |
list['Instance']
|
A list of instances belonging to the track. |
Methods:
| Name | Description |
|---|---|
__getitem__ |
Get an instance from the track. |
__len__ |
Get the length of the track. |
__repr__ |
Get the string representation of the track. |
Source code in dreem/io/track.py
@attrs.define(eq=False)
class Track:
"""Object for storing instances of the same track.
Attributes:
id: the track label.
instances: A list of instances belonging to the track.
"""
_id: int = attrs.field(alias="id")
_instances: list["Instance"] = attrs.field(alias="instances", factory=list)
def __repr__(self) -> str:
"""Get the string representation of the track.
Returns:
the string representation of the Track.
"""
return f"Track(id={self.id}, len={len(self)})"
@property
def track_id(self) -> int:
"""Get the id of the track.
Returns:
The integer id of the track.
"""
return self._id
@track_id.setter
def track_id(self, track_id: int) -> None:
"""Set the id of the track.
Args:
track_id: the int id of the track.
"""
self._id = track_id
@property
def instances(self) -> list["Instance"]:
"""Get the instances belonging to this track.
Returns:
A list of instances with this track id.
"""
return self._instances
@instances.setter
def instances(self, instances) -> None:
"""Set the instances belonging to this track.
Args:
instances: A list of instances that belong to the same track.
"""
self._instances = instances
@property
def frames(self) -> set[Frame]:
"""Get the frames where this track appears.
Returns:
A set of `Frame` objects where this track appears.
"""
return set([instance.frame for instance in self.instances])
def __len__(self) -> int:
"""Get the length of the track.
Returns:
The number of instances/frames in the track.
"""
return len(self.instances)
def __getitem__(self, ind: int | list[int]) -> "Instance" | list["Instance"]:
"""Get an instance from the track.
Args:
ind: Either a single int or list of int indices.
Returns:
the instance at that index of the track.instances.
"""
if isinstance(ind, int):
return self.instances[ind]
elif isinstance(ind, list):
return [self.instances[i] for i in ind]
else:
raise ValueError(f"Ind must be an int or list of ints, found {type(ind)}")
frames
property
¶
Get the frames where this track appears.
Returns:
| Type | Description |
|---|---|
set[Frame]
|
A set of |
instances
property
writable
¶
Get the instances belonging to this track.
Returns:
| Type | Description |
|---|---|
list['Instance']
|
A list of instances with this track id. |
track_id
property
writable
¶
Get the id of the track.
Returns:
| Type | Description |
|---|---|
int
|
The integer id of the track. |
__getitem__(ind)
¶
Get an instance from the track.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
ind
|
int | list[int]
|
Either a single int or list of int indices. |
required |
Returns:
| Type | Description |
|---|---|
'Instance' | list['Instance']
|
the instance at that index of the track.instances. |
Source code in dreem/io/track.py
def __getitem__(self, ind: int | list[int]) -> "Instance" | list["Instance"]:
"""Get an instance from the track.
Args:
ind: Either a single int or list of int indices.
Returns:
the instance at that index of the track.instances.
"""
if isinstance(ind, int):
return self.instances[ind]
elif isinstance(ind, list):
return [self.instances[i] for i in ind]
else:
raise ValueError(f"Ind must be an int or list of ints, found {type(ind)}")
__len__()
¶
__repr__()
¶
Get the string representation of the track.
Returns:
| Type | Description |
|---|---|
str
|
the string representation of the Track. |