dreem
dreem
¶
Top-level package for dreem.
Modules:
| Name | Description |
|---|---|
cli |
DREEM CLI - Global Tracking Transformer for biological multi-object tracking. |
datasets |
Data loading and preprocessing. |
inference |
Tracking Inference using GTR Model. |
io |
Module containing input/output data structures for easy storage and manipulation. |
models |
Model architectures and layers. |
training |
Initialize training module. |
utils |
Utility modules for DREEM processing pipelines. |
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. |
GTRRunner |
A lightning wrapper around GTR model. |
GlobalTrackingTransformer |
Modular GTR model composed of visual encoder + transformer used for tracking. |
Instance |
Class representing a single instance to be tracked. |
Tracker |
Tracker class used for assignment based on sliding inference from GTR. |
Transformer |
Transformer class. |
VisualEncoder |
Class wrapping around a visual feature extractor backbone. |
Functions:
| Name | Description |
|---|---|
annotate_video |
Annotate video frames with labels. |
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,
)
GTRRunner
¶
Bases: LightningModule
A lightning wrapper around GTR model.
Used for training, validation and inference.
Methods:
| Name | Description |
|---|---|
__init__ |
Initialize a lightning module for GTR. |
configure_optimizers |
Get optimizers and schedulers for training. |
forward |
Execute forward pass of the lightning module. |
log_metrics |
Log metrics computed during evaluation. |
on_test_end |
Run inference and metrics pipeline to compute metrics for test set. |
on_validation_epoch_end |
Execute hook for validation end. |
predict_step |
Run inference for model. |
setup_tracking |
Setup the tracker for tracking. |
test_step |
Execute single test step for model. Performs eval in addition to tracking. |
training_step |
Execute single training step for model. |
validation_step |
Execute single val step for model. |
Source code in dreem/models/gtr_runner.py
class GTRRunner(LightningModule):
"""A lightning wrapper around GTR model.
Used for training, validation and inference.
"""
DEFAULT_METRICS = {
"train": [],
"val": [],
"test": ["motmetrics"],
}
DEFAULT_TRACKING = {
"train": False,
"val": False,
"test": True,
}
DEFAULT_SAVE = {"train": False, "val": False, "test": False}
def __init__(
self,
model_cfg: dict | None = None,
tracker_cfg: dict | None = None,
loss_cfg: dict | None = None,
optimizer_cfg: dict | None = None,
scheduler_cfg: dict | None = None,
metrics: dict[str, list[str]] | None = None,
test_save_path: str = "./test_results.h5",
save_frame_meta: bool = False,
):
"""Initialize a lightning module for GTR.
Args:
model_cfg: hyperparameters for GlobalTrackingTransformer
tracker_cfg: The parameters used for the tracker post-processing
loss_cfg: hyperparameters for AssoLoss
optimizer_cfg: hyper parameters used for optimizer.
Only used to overwrite `configure_optimizer`
scheduler_cfg: hyperparameters for lr_scheduler used to overwrite `configure_optimizer
metrics: a dict containing the metrics to be computed during train, val, and test.
test_save_path: path to a directory to save the eval and tracking results to
save_frame_meta: Whether to save frame metadata to the h5 file
"""
super().__init__()
self.save_hyperparameters()
self.model_cfg = model_cfg if model_cfg else {}
self.loss_cfg = loss_cfg if loss_cfg else {}
self.tracker_cfg = tracker_cfg if tracker_cfg else {}
self.model = GlobalTrackingTransformer(**self.model_cfg)
self.loss = AssoLoss(**self.loss_cfg)
from dreem.inference.tracker import Tracker
self.tracker = Tracker(**self.tracker_cfg)
self.optimizer_cfg = optimizer_cfg
self.scheduler_cfg = scheduler_cfg
self.metrics = metrics if metrics is not None else self.DEFAULT_METRICS
self.test_results = {
"preds": [],
"save_path": test_save_path,
"save_frame_meta": save_frame_meta,
}
def forward(
self,
ref_instances: list["Instance"],
query_instances: list["Instance"] | None = None,
) -> list["AssociationMatrix"]:
"""Execute forward pass of the lightning module.
Args:
ref_instances: a list of `Instance` objects containing crops and other data needed for transformer model
query_instances: a list of `Instance` objects used as queries in the decoder. Mostly used for inference.
Returns:
An association matrix between objects
"""
asso_preds = self.model(ref_instances, query_instances)
return asso_preds
def training_step(
self, train_batch: list[list["Frame"]], batch_idx: int
) -> dict[str, float]:
"""Execute single training step for model.
Args:
train_batch: A single batch from the dataset which is a list of `Frame` objects
with length `clip_length` containing Instances and other metadata.
batch_idx: the batch number used by lightning
Returns:
A dict containing the train loss plus any other metrics specified
"""
result = self._shared_eval_step(train_batch[0], mode="train")
self.log_metrics(result, len(train_batch[0]), "train")
return result
def validation_step(
self, val_batch: list[list["Frame"]], batch_idx: int
) -> dict[str, float]:
"""Execute single val step for model.
Args:
val_batch: A single batch from the dataset which is a list of `Frame` objects
with length `clip_length` containing Instances and other metadata.
batch_idx: the batch number used by lightning
Returns:
A dict containing the val loss plus any other metrics specified
"""
result = self._shared_eval_step(val_batch[0], mode="val")
self.log_metrics(result, len(val_batch[0]), "val")
return result
def test_step(
self, test_batch: list[list["Frame"]], batch_idx: int
) -> dict[str, float]:
"""Execute single test step for model. Performs eval in addition to tracking.
Args:
test_batch: A single batch from the dataset which is a list of `Frame` objects
with length `clip_length` containing Instances and other metadata.
batch_idx: the batch number used by lightning
Returns:
None
"""
frames_pred = self.tracker(self.model, test_batch[0])
self.test_results["preds"].extend([frame.to("cpu") for frame in frames_pred])
self.log_metrics(None, len(test_batch[0]), "test")
return None
def predict_step(self, batch: list[list["Frame"]], batch_idx: int) -> list["Frame"]:
"""Run inference for model.
Computes association + assignment.
Args:
batch: A single batch from the dataset which is a list of `Frame` objects
with length `clip_length` containing Instances and other metadata.
batch_idx: the batch number used by lightning
Returns:
A list of dicts where each dict is a frame containing the predicted track ids
"""
frames_pred = self.tracker(self.model, batch[0])
for frame in frames_pred:
frame = frame.to("cpu")
return frames_pred
def _shared_eval_step(self, frames: list["Frame"], mode: str) -> dict[str, float]:
"""Run evaluation used by train, test, and val steps.
Args:
frames: A list of dicts where each dict is a frame containing gt data
mode: which metrics to compute and whether to use persistent tracking or not
Returns:
a dict containing the loss and any other metrics specified by `eval_metrics`
"""
try:
instances = [instance for frame in frames for instance in frame.instances]
if len(instances) == 0:
return None
logits = self(instances)
logits = [asso.matrix for asso in logits]
loss = self.loss(logits, frames)
return_metrics = {"loss": loss}
return_metrics["batch_size"] = len(frames)
except Exception as e:
logger.exception(
f"Failed on frame {frames[0].frame_id} of video {frames[0].video_id}"
)
logger.exception(e)
raise (e)
return return_metrics
def configure_optimizers(self) -> dict:
"""Get optimizers and schedulers for training.
Is overridden by config but defaults to Adam + ReduceLROnPlateau.
Returns:
an optimizer config dict containing the optimizer, scheduler, and scheduler params
"""
# todo: init from config
if self.optimizer_cfg is None:
optimizer = torch.optim.Adam(self.parameters(), lr=1e-4, betas=(0.9, 0.999))
else:
optimizer = init_optimizer(self.parameters(), self.optimizer_cfg)
if self.scheduler_cfg is None:
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
optimizer, "min", 0.5, 10
)
else:
scheduler = init_scheduler(optimizer, self.scheduler_cfg)
return {
"optimizer": optimizer,
"lr_scheduler": {
"scheduler": scheduler,
"monitor": "val_loss",
"interval": "epoch",
"frequency": 1,
},
}
def log_metrics(self, result: dict, batch_size: int, mode: str) -> None:
"""Log metrics computed during evaluation.
Args:
result: A dict containing metrics to be logged.
batch_size: the size of the batch used to compute the metrics
mode: One of {'train', 'test' or 'val'}. Used as prefix while logging.
"""
if result:
batch_size = result.pop("batch_size")
for metric, val in result.items():
if isinstance(val, torch.Tensor):
val = val.item()
self.log(f"{mode}_{metric}", val, batch_size=batch_size)
def on_validation_epoch_end(self):
"""Execute hook for validation end.
Currently, we simply clear the gpu cache and do garbage collection.
"""
gc.collect()
torch.cuda.empty_cache()
def setup_tracking(self, tracker_cfg: DictConfig, mode: str = "inference"):
"""Setup the tracker for tracking.
Args:
tracker_cfg: The configuration for the tracker.
mode: The mode to run the tracker in.
Returns:
A dictionary of overrides for the tracker.
"""
from dreem.inference.tracker import Tracker
save_frame_meta = tracker_cfg.cfg.get("save_frame_meta", False)
self.tracker_cfg = tracker_cfg.get_tracker_cfg()
self.tracker_cfg["enable_crop_saving"] = (
save_frame_meta # to save frame metadata, need to disable crop=None in GTR (memory saving)
)
self.tracker = Tracker(**self.tracker_cfg)
logger.info("Using the following tracker:")
logger.info(self.tracker)
if mode == "eval":
self.metrics["test"] = tracker_cfg.get("metrics", {}).get("test", "all")
logger.info("Computing the following metrics:")
logger.info(self.metrics["test"])
self.test_results["save_frame_meta"] = save_frame_meta
self.test_results["save_path"] = tracker_cfg.get("outdir", ".")
os.makedirs(self.test_results["save_path"], exist_ok=True)
overrides_dict = {
"max_tracks": self.tracker_cfg.get("max_tracks", inf),
"save_frame_meta": save_frame_meta,
}
return overrides_dict
def on_test_end(self):
"""Run inference and metrics pipeline to compute metrics for test set.
Args:
test_results: dict containing predictions and metrics to be filled out in metrics.evaluate
metrics: list of metrics to compute
"""
# input validation
metrics_to_compute = self.metrics[
"test"
] # list of metrics to compute, or "all"
if metrics_to_compute == "all":
metrics_to_compute = ["motmetrics"]
if isinstance(metrics_to_compute, str):
metrics_to_compute = [metrics_to_compute]
for metric in metrics_to_compute:
if metric not in ["motmetrics"]:
raise ValueError(
f"Metric {metric} not supported. Please select from 'motmetrics'"
)
preds = self.test_results["preds"]
# results is a dict with key being the metric name, and value being the metric value computed
results = metrics.evaluate(preds, metrics_to_compute)
# save metrics and frame metadata to hdf5
# Get the video name from the first frame
vid_name = Path(preds[0].vid_name).stem
# save the results to an hdf5 file
fname = os.path.join(
self.test_results["save_path"], f"{vid_name}.dreem_metrics.h5"
)
logger.info(f"Saving metrics to {fname}")
# Check if the h5 file exists and add a suffix to prevent name collision
suffix_counter = 0
original_fname = fname
while os.path.exists(fname):
suffix_counter += 1
fname = original_fname.replace(
".dreem_metrics.h5", f"_{suffix_counter}.dreem_metrics.h5"
)
if suffix_counter > 0:
logger.info(f"File already exists. Saving to {fname} instead")
with h5py.File(fname, "a") as results_file:
# Create a group for this video
vid_group = results_file.require_group(vid_name)
# Save each metric
for metric_name, value in results.items():
if metric_name == "motmetrics":
# For num_switches, save mot_summary and mot_events separately
mot_summary = value[0]
frame_switch_map = value[1]
motevents = value[2]
mot_summary.to_csv(
os.path.join(
self.test_results["save_path"], f"{vid_name}.motmetrics.csv"
),
index=True,
)
motevents.to_csv(
os.path.join(
self.test_results["save_path"], f"{vid_name}.motevents.csv"
),
index=False,
)
mot_summary_group = vid_group.require_group("mot_summary")
# Loop through each row in mot_summary and save as attributes
for _, row in mot_summary.iterrows():
mot_summary_group.attrs[row.name] = row["acc"]
if self.test_results["save_frame_meta"]:
# save frame metadata for every frame, specifically assoc matrices
frame_meta_group = vid_group.require_group("frame_meta")
switch_group = frame_meta_group.require_group("switches")
for frame in preds:
frame = frame.to("cpu")
_ = frame.to_h5(
frame_meta_group, frame.get_gt_track_ids().cpu().numpy()
)
if frame.frame_id.item() in frame_switch_map:
if frame_switch_map[frame.frame_id.item()]:
switch_group.attrs[
"frame_" + str(frame.frame_id.item())
] = True
else:
switch_group.attrs[
"frame_" + str(frame.frame_id.item())
] = False
elif metric_name == "global_tracking_accuracy":
gta_by_gt_track = value
gta_group = vid_group.require_group("global_tracking_accuracy")
# save as a key value pair with gt track id: gta
for gt_track_id, gta in gta_by_gt_track.items():
gta_group.attrs[f"track_{gt_track_id}"] = gta
# save the tracking results to a slp/labelled masks file
if isinstance(self.trainer.test_dataloaders.dataset, CellTrackingDataset):
outpath = os.path.join(
self.test_results["save_path"],
f"{vid_name}.dreem_inference.{datetime.now().strftime('%m-%d-%Y-%H-%M-%S')}.tif",
)
pred_imgs = []
for frame in preds:
frame_masks = []
for instance in frame.instances:
# centroid = instance.centroid["centroid"] # Currently unused but available if needed
mask = instance.mask.cpu().numpy()
track_id = instance.pred_track_id.cpu().numpy().item()
mask = mask.astype(np.uint8)
mask[mask != 0] = track_id # label the mask with the track id
frame_masks.append(mask)
frame_mask = np.max(frame_masks, axis=0)
pred_imgs.append(frame_mask)
pred_imgs = np.stack(pred_imgs)
tifffile.imwrite(outpath, pred_imgs.astype(np.uint16))
else:
outpath = os.path.join(
self.test_results["save_path"],
f"{vid_name}.dreem_inference.{datetime.now().strftime('%m-%d-%Y-%H-%M-%S')}.slp",
)
pred_slp = []
suggestions = []
logger.info(f"Saving inference results to {outpath}")
# save the tracking results to a slp file
tracks = {}
for frame in preds:
if frame.frame_id.item() == 0:
video = (
sio.Video(frame.video)
if isinstance(frame.video, str)
else sio.Video
)
if frame.has_flag(FrameFlagCode.LOW_CONFIDENCE):
suggestion = SuggestionFrame(
video=video, frame_idx=frame.frame_id.item()
)
suggestions.append(suggestion)
lf, tracks = frame.to_slp(tracks, video=video)
pred_slp.append(lf)
pred_slp = sio.Labels(pred_slp, suggestions=suggestions)
pred_slp.save(outpath)
# clear the preds
self.test_results["preds"] = []
__init__(model_cfg=None, tracker_cfg=None, loss_cfg=None, optimizer_cfg=None, scheduler_cfg=None, metrics=None, test_save_path='./test_results.h5', save_frame_meta=False)
¶
Initialize a lightning module for GTR.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
model_cfg
|
dict | None
|
hyperparameters for GlobalTrackingTransformer |
None
|
tracker_cfg
|
dict | None
|
The parameters used for the tracker post-processing |
None
|
loss_cfg
|
dict | None
|
hyperparameters for AssoLoss |
None
|
optimizer_cfg
|
dict | None
|
hyper parameters used for optimizer.
Only used to overwrite |
None
|
scheduler_cfg
|
dict | None
|
hyperparameters for lr_scheduler used to overwrite `configure_optimizer |
None
|
metrics
|
dict[str, list[str]] | None
|
a dict containing the metrics to be computed during train, val, and test. |
None
|
test_save_path
|
str
|
path to a directory to save the eval and tracking results to |
'./test_results.h5'
|
save_frame_meta
|
bool
|
Whether to save frame metadata to the h5 file |
False
|
Source code in dreem/models/gtr_runner.py
def __init__(
self,
model_cfg: dict | None = None,
tracker_cfg: dict | None = None,
loss_cfg: dict | None = None,
optimizer_cfg: dict | None = None,
scheduler_cfg: dict | None = None,
metrics: dict[str, list[str]] | None = None,
test_save_path: str = "./test_results.h5",
save_frame_meta: bool = False,
):
"""Initialize a lightning module for GTR.
Args:
model_cfg: hyperparameters for GlobalTrackingTransformer
tracker_cfg: The parameters used for the tracker post-processing
loss_cfg: hyperparameters for AssoLoss
optimizer_cfg: hyper parameters used for optimizer.
Only used to overwrite `configure_optimizer`
scheduler_cfg: hyperparameters for lr_scheduler used to overwrite `configure_optimizer
metrics: a dict containing the metrics to be computed during train, val, and test.
test_save_path: path to a directory to save the eval and tracking results to
save_frame_meta: Whether to save frame metadata to the h5 file
"""
super().__init__()
self.save_hyperparameters()
self.model_cfg = model_cfg if model_cfg else {}
self.loss_cfg = loss_cfg if loss_cfg else {}
self.tracker_cfg = tracker_cfg if tracker_cfg else {}
self.model = GlobalTrackingTransformer(**self.model_cfg)
self.loss = AssoLoss(**self.loss_cfg)
from dreem.inference.tracker import Tracker
self.tracker = Tracker(**self.tracker_cfg)
self.optimizer_cfg = optimizer_cfg
self.scheduler_cfg = scheduler_cfg
self.metrics = metrics if metrics is not None else self.DEFAULT_METRICS
self.test_results = {
"preds": [],
"save_path": test_save_path,
"save_frame_meta": save_frame_meta,
}
configure_optimizers()
¶
Get optimizers and schedulers for training.
Is overridden by config but defaults to Adam + ReduceLROnPlateau.
Returns:
| Type | Description |
|---|---|
dict
|
an optimizer config dict containing the optimizer, scheduler, and scheduler params |
Source code in dreem/models/gtr_runner.py
def configure_optimizers(self) -> dict:
"""Get optimizers and schedulers for training.
Is overridden by config but defaults to Adam + ReduceLROnPlateau.
Returns:
an optimizer config dict containing the optimizer, scheduler, and scheduler params
"""
# todo: init from config
if self.optimizer_cfg is None:
optimizer = torch.optim.Adam(self.parameters(), lr=1e-4, betas=(0.9, 0.999))
else:
optimizer = init_optimizer(self.parameters(), self.optimizer_cfg)
if self.scheduler_cfg is None:
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
optimizer, "min", 0.5, 10
)
else:
scheduler = init_scheduler(optimizer, self.scheduler_cfg)
return {
"optimizer": optimizer,
"lr_scheduler": {
"scheduler": scheduler,
"monitor": "val_loss",
"interval": "epoch",
"frequency": 1,
},
}
forward(ref_instances, query_instances=None)
¶
Execute forward pass of the lightning module.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
ref_instances
|
list[Instance]
|
a list of |
required |
query_instances
|
list[Instance] | None
|
a list of |
None
|
Returns:
| Type | Description |
|---|---|
list[AssociationMatrix]
|
An association matrix between objects |
Source code in dreem/models/gtr_runner.py
def forward(
self,
ref_instances: list["Instance"],
query_instances: list["Instance"] | None = None,
) -> list["AssociationMatrix"]:
"""Execute forward pass of the lightning module.
Args:
ref_instances: a list of `Instance` objects containing crops and other data needed for transformer model
query_instances: a list of `Instance` objects used as queries in the decoder. Mostly used for inference.
Returns:
An association matrix between objects
"""
asso_preds = self.model(ref_instances, query_instances)
return asso_preds
log_metrics(result, batch_size, mode)
¶
Log metrics computed during evaluation.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
result
|
dict
|
A dict containing metrics to be logged. |
required |
batch_size
|
int
|
the size of the batch used to compute the metrics |
required |
mode
|
str
|
One of {'train', 'test' or 'val'}. Used as prefix while logging. |
required |
Source code in dreem/models/gtr_runner.py
def log_metrics(self, result: dict, batch_size: int, mode: str) -> None:
"""Log metrics computed during evaluation.
Args:
result: A dict containing metrics to be logged.
batch_size: the size of the batch used to compute the metrics
mode: One of {'train', 'test' or 'val'}. Used as prefix while logging.
"""
if result:
batch_size = result.pop("batch_size")
for metric, val in result.items():
if isinstance(val, torch.Tensor):
val = val.item()
self.log(f"{mode}_{metric}", val, batch_size=batch_size)
on_test_end()
¶
Run inference and metrics pipeline to compute metrics for test set.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
test_results
|
dict containing predictions and metrics to be filled out in metrics.evaluate |
required | |
metrics
|
list of metrics to compute |
required |
Source code in dreem/models/gtr_runner.py
def on_test_end(self):
"""Run inference and metrics pipeline to compute metrics for test set.
Args:
test_results: dict containing predictions and metrics to be filled out in metrics.evaluate
metrics: list of metrics to compute
"""
# input validation
metrics_to_compute = self.metrics[
"test"
] # list of metrics to compute, or "all"
if metrics_to_compute == "all":
metrics_to_compute = ["motmetrics"]
if isinstance(metrics_to_compute, str):
metrics_to_compute = [metrics_to_compute]
for metric in metrics_to_compute:
if metric not in ["motmetrics"]:
raise ValueError(
f"Metric {metric} not supported. Please select from 'motmetrics'"
)
preds = self.test_results["preds"]
# results is a dict with key being the metric name, and value being the metric value computed
results = metrics.evaluate(preds, metrics_to_compute)
# save metrics and frame metadata to hdf5
# Get the video name from the first frame
vid_name = Path(preds[0].vid_name).stem
# save the results to an hdf5 file
fname = os.path.join(
self.test_results["save_path"], f"{vid_name}.dreem_metrics.h5"
)
logger.info(f"Saving metrics to {fname}")
# Check if the h5 file exists and add a suffix to prevent name collision
suffix_counter = 0
original_fname = fname
while os.path.exists(fname):
suffix_counter += 1
fname = original_fname.replace(
".dreem_metrics.h5", f"_{suffix_counter}.dreem_metrics.h5"
)
if suffix_counter > 0:
logger.info(f"File already exists. Saving to {fname} instead")
with h5py.File(fname, "a") as results_file:
# Create a group for this video
vid_group = results_file.require_group(vid_name)
# Save each metric
for metric_name, value in results.items():
if metric_name == "motmetrics":
# For num_switches, save mot_summary and mot_events separately
mot_summary = value[0]
frame_switch_map = value[1]
motevents = value[2]
mot_summary.to_csv(
os.path.join(
self.test_results["save_path"], f"{vid_name}.motmetrics.csv"
),
index=True,
)
motevents.to_csv(
os.path.join(
self.test_results["save_path"], f"{vid_name}.motevents.csv"
),
index=False,
)
mot_summary_group = vid_group.require_group("mot_summary")
# Loop through each row in mot_summary and save as attributes
for _, row in mot_summary.iterrows():
mot_summary_group.attrs[row.name] = row["acc"]
if self.test_results["save_frame_meta"]:
# save frame metadata for every frame, specifically assoc matrices
frame_meta_group = vid_group.require_group("frame_meta")
switch_group = frame_meta_group.require_group("switches")
for frame in preds:
frame = frame.to("cpu")
_ = frame.to_h5(
frame_meta_group, frame.get_gt_track_ids().cpu().numpy()
)
if frame.frame_id.item() in frame_switch_map:
if frame_switch_map[frame.frame_id.item()]:
switch_group.attrs[
"frame_" + str(frame.frame_id.item())
] = True
else:
switch_group.attrs[
"frame_" + str(frame.frame_id.item())
] = False
elif metric_name == "global_tracking_accuracy":
gta_by_gt_track = value
gta_group = vid_group.require_group("global_tracking_accuracy")
# save as a key value pair with gt track id: gta
for gt_track_id, gta in gta_by_gt_track.items():
gta_group.attrs[f"track_{gt_track_id}"] = gta
# save the tracking results to a slp/labelled masks file
if isinstance(self.trainer.test_dataloaders.dataset, CellTrackingDataset):
outpath = os.path.join(
self.test_results["save_path"],
f"{vid_name}.dreem_inference.{datetime.now().strftime('%m-%d-%Y-%H-%M-%S')}.tif",
)
pred_imgs = []
for frame in preds:
frame_masks = []
for instance in frame.instances:
# centroid = instance.centroid["centroid"] # Currently unused but available if needed
mask = instance.mask.cpu().numpy()
track_id = instance.pred_track_id.cpu().numpy().item()
mask = mask.astype(np.uint8)
mask[mask != 0] = track_id # label the mask with the track id
frame_masks.append(mask)
frame_mask = np.max(frame_masks, axis=0)
pred_imgs.append(frame_mask)
pred_imgs = np.stack(pred_imgs)
tifffile.imwrite(outpath, pred_imgs.astype(np.uint16))
else:
outpath = os.path.join(
self.test_results["save_path"],
f"{vid_name}.dreem_inference.{datetime.now().strftime('%m-%d-%Y-%H-%M-%S')}.slp",
)
pred_slp = []
suggestions = []
logger.info(f"Saving inference results to {outpath}")
# save the tracking results to a slp file
tracks = {}
for frame in preds:
if frame.frame_id.item() == 0:
video = (
sio.Video(frame.video)
if isinstance(frame.video, str)
else sio.Video
)
if frame.has_flag(FrameFlagCode.LOW_CONFIDENCE):
suggestion = SuggestionFrame(
video=video, frame_idx=frame.frame_id.item()
)
suggestions.append(suggestion)
lf, tracks = frame.to_slp(tracks, video=video)
pred_slp.append(lf)
pred_slp = sio.Labels(pred_slp, suggestions=suggestions)
pred_slp.save(outpath)
# clear the preds
self.test_results["preds"] = []
on_validation_epoch_end()
¶
Execute hook for validation end.
Currently, we simply clear the gpu cache and do garbage collection.
predict_step(batch, batch_idx)
¶
Run inference for model.
Computes association + assignment.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
batch
|
list[list[Frame]]
|
A single batch from the dataset which is a list of |
required |
batch_idx
|
int
|
the batch number used by lightning |
required |
Returns:
| Type | Description |
|---|---|
list[Frame]
|
A list of dicts where each dict is a frame containing the predicted track ids |
Source code in dreem/models/gtr_runner.py
def predict_step(self, batch: list[list["Frame"]], batch_idx: int) -> list["Frame"]:
"""Run inference for model.
Computes association + assignment.
Args:
batch: A single batch from the dataset which is a list of `Frame` objects
with length `clip_length` containing Instances and other metadata.
batch_idx: the batch number used by lightning
Returns:
A list of dicts where each dict is a frame containing the predicted track ids
"""
frames_pred = self.tracker(self.model, batch[0])
for frame in frames_pred:
frame = frame.to("cpu")
return frames_pred
setup_tracking(tracker_cfg, mode='inference')
¶
Setup the tracker for tracking.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
tracker_cfg
|
DictConfig
|
The configuration for the tracker. |
required |
mode
|
str
|
The mode to run the tracker in. |
'inference'
|
Returns:
| Type | Description |
|---|---|
|
A dictionary of overrides for the tracker. |
Source code in dreem/models/gtr_runner.py
def setup_tracking(self, tracker_cfg: DictConfig, mode: str = "inference"):
"""Setup the tracker for tracking.
Args:
tracker_cfg: The configuration for the tracker.
mode: The mode to run the tracker in.
Returns:
A dictionary of overrides for the tracker.
"""
from dreem.inference.tracker import Tracker
save_frame_meta = tracker_cfg.cfg.get("save_frame_meta", False)
self.tracker_cfg = tracker_cfg.get_tracker_cfg()
self.tracker_cfg["enable_crop_saving"] = (
save_frame_meta # to save frame metadata, need to disable crop=None in GTR (memory saving)
)
self.tracker = Tracker(**self.tracker_cfg)
logger.info("Using the following tracker:")
logger.info(self.tracker)
if mode == "eval":
self.metrics["test"] = tracker_cfg.get("metrics", {}).get("test", "all")
logger.info("Computing the following metrics:")
logger.info(self.metrics["test"])
self.test_results["save_frame_meta"] = save_frame_meta
self.test_results["save_path"] = tracker_cfg.get("outdir", ".")
os.makedirs(self.test_results["save_path"], exist_ok=True)
overrides_dict = {
"max_tracks": self.tracker_cfg.get("max_tracks", inf),
"save_frame_meta": save_frame_meta,
}
return overrides_dict
test_step(test_batch, batch_idx)
¶
Execute single test step for model. Performs eval in addition to tracking.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
test_batch
|
list[list[Frame]]
|
A single batch from the dataset which is a list of |
required |
batch_idx
|
int
|
the batch number used by lightning |
required |
Returns:
| Type | Description |
|---|---|
dict[str, float]
|
None |
Source code in dreem/models/gtr_runner.py
def test_step(
self, test_batch: list[list["Frame"]], batch_idx: int
) -> dict[str, float]:
"""Execute single test step for model. Performs eval in addition to tracking.
Args:
test_batch: A single batch from the dataset which is a list of `Frame` objects
with length `clip_length` containing Instances and other metadata.
batch_idx: the batch number used by lightning
Returns:
None
"""
frames_pred = self.tracker(self.model, test_batch[0])
self.test_results["preds"].extend([frame.to("cpu") for frame in frames_pred])
self.log_metrics(None, len(test_batch[0]), "test")
return None
training_step(train_batch, batch_idx)
¶
Execute single training step for model.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
train_batch
|
list[list[Frame]]
|
A single batch from the dataset which is a list of |
required |
batch_idx
|
int
|
the batch number used by lightning |
required |
Returns:
| Type | Description |
|---|---|
dict[str, float]
|
A dict containing the train loss plus any other metrics specified |
Source code in dreem/models/gtr_runner.py
def training_step(
self, train_batch: list[list["Frame"]], batch_idx: int
) -> dict[str, float]:
"""Execute single training step for model.
Args:
train_batch: A single batch from the dataset which is a list of `Frame` objects
with length `clip_length` containing Instances and other metadata.
batch_idx: the batch number used by lightning
Returns:
A dict containing the train loss plus any other metrics specified
"""
result = self._shared_eval_step(train_batch[0], mode="train")
self.log_metrics(result, len(train_batch[0]), "train")
return result
validation_step(val_batch, batch_idx)
¶
Execute single val step for model.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
val_batch
|
list[list[Frame]]
|
A single batch from the dataset which is a list of |
required |
batch_idx
|
int
|
the batch number used by lightning |
required |
Returns:
| Type | Description |
|---|---|
dict[str, float]
|
A dict containing the val loss plus any other metrics specified |
Source code in dreem/models/gtr_runner.py
def validation_step(
self, val_batch: list[list["Frame"]], batch_idx: int
) -> dict[str, float]:
"""Execute single val step for model.
Args:
val_batch: A single batch from the dataset which is a list of `Frame` objects
with length `clip_length` containing Instances and other metadata.
batch_idx: the batch number used by lightning
Returns:
A dict containing the val loss plus any other metrics specified
"""
result = self._shared_eval_step(val_batch[0], mode="val")
self.log_metrics(result, len(val_batch[0]), "val")
return result
GlobalTrackingTransformer
¶
Bases: Module
Modular GTR model composed of visual encoder + transformer used for tracking.
Methods:
| Name | Description |
|---|---|
__init__ |
Initialize GTR. |
extract_features |
Extract features from instances using visual encoder backbone. |
forward |
Execute forward pass of GTR Model to get asso matrix. |
Source code in dreem/models/global_tracking_transformer.py
class GlobalTrackingTransformer(torch.nn.Module):
"""Modular GTR model composed of visual encoder + transformer used for tracking."""
def __init__(
self,
encoder_cfg: dict | None = None,
d_model: int = 1024,
nhead: int = 8,
num_encoder_layers: int = 6,
num_decoder_layers: int = 6,
dropout: int = 0.1,
activation: str = "relu",
return_intermediate_dec: bool = False,
norm: bool = False,
num_layers_attn_head: int = 2,
dropout_attn_head: int = 0.1,
embedding_meta: dict | None = None,
return_embedding: bool = False,
decoder_self_attn: bool = False,
):
"""Initialize GTR.
Args:
encoder_cfg: Dictionary of arguments to pass to the CNN constructor,
e.g: `cfg = {"model_name": "resnet18", "pretrained": False, "in_chans": 3}`
d_model: The number of features in the encoder/decoder inputs.
nhead: The number of heads in the transformer encoder/decoder.
num_encoder_layers: The number of encoder-layers in the encoder.
num_decoder_layers: The number of decoder-layers in the decoder.
dropout: Dropout value applied to the output of transformer layers.
activation: Activation function to use.
return_intermediate_dec: Return intermediate layers from decoder.
norm: If True, normalize output of encoder and decoder.
num_layers_attn_head: The number of layers in the attention head.
dropout_attn_head: Dropout value for the attention_head.
embedding_meta: Metadata for positional embeddings. See below.
return_embedding: Whether to return the positional embeddings
decoder_self_attn: If True, use decoder self attention.
More details on `embedding_meta`:
By default this will be an empty dict and indicate
that no positional embeddings should be used. To use the positional embeddings
pass in a dictionary containing a "pos" and "temp" key with subdictionaries for correct parameters ie:
`{"pos": {'mode': 'learned', 'emb_num': 16, 'over_boxes: True},
"temp": {'mode': 'learned', 'emb_num': 16}}`. (see `dreem.models.embeddings.Embedding.EMB_TYPES`
and `dreem.models.embeddings.Embedding.EMB_MODES` for embedding parameters).
"""
super().__init__()
if not encoder_cfg:
encoder_cfg = {}
self.visual_encoder = create_visual_encoder(d_model=d_model, **encoder_cfg)
self.transformer = Transformer(
d_model=d_model,
nhead=nhead,
num_encoder_layers=num_encoder_layers,
num_decoder_layers=num_decoder_layers,
dropout=dropout,
activation=activation,
return_intermediate_dec=return_intermediate_dec,
norm=norm,
num_layers_attn_head=num_layers_attn_head,
dropout_attn_head=dropout_attn_head,
embedding_meta=embedding_meta,
return_embedding=return_embedding,
decoder_self_attn=decoder_self_attn,
encoder_cfg=encoder_cfg,
)
def forward(
self,
ref_instances: list["Instance"],
query_instances: list["Instance"] = None,
retain_crops: bool = False,
) -> list["AssociationMatrix"]:
"""Execute forward pass of GTR Model to get asso matrix.
Args:
ref_instances: List of instances from chunk containing crops of objects + gt label info
query_instances: list of instances used as query in decoder.
Returns:
An N_T x N association matrix
"""
# Extract feature representations with pre-trained encoder.
self.extract_features(ref_instances, retain_crops=retain_crops)
if query_instances:
self.extract_features(query_instances, retain_crops=retain_crops)
asso_preds = self.transformer(ref_instances, query_instances)
return asso_preds
def extract_features(
self,
instances: list["Instance"],
force_recompute: bool = False,
retain_crops: bool = False,
) -> None:
"""Extract features from instances using visual encoder backbone.
Args:
instances: A list of instances to compute features for
force_recompute: indicate whether to compute features for all instances regardless of if they have instances
"""
if not force_recompute:
instances_to_compute = [
instance
for instance in instances
if instance.has_crop() and not instance.has_features()
]
else:
instances_to_compute = instances
if len(instances_to_compute) == 0:
return
elif len(instances_to_compute) == 1: # handle batch norm error when B=1
instances_to_compute = instances
crops = torch.concatenate([instance.crop for instance in instances_to_compute])
features = self.visual_encoder(crops)
for i, z_i in enumerate(features):
instances_to_compute[i].features = z_i
if not retain_crops: # crops are an attribute of instance by default
instances_to_compute[i].crop = None
__init__(encoder_cfg=None, d_model=1024, nhead=8, num_encoder_layers=6, num_decoder_layers=6, dropout=0.1, activation='relu', return_intermediate_dec=False, norm=False, num_layers_attn_head=2, dropout_attn_head=0.1, embedding_meta=None, return_embedding=False, decoder_self_attn=False)
¶
Initialize GTR.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
encoder_cfg
|
dict | None
|
Dictionary of arguments to pass to the CNN constructor,
e.g: |
None
|
d_model
|
int
|
The number of features in the encoder/decoder inputs. |
1024
|
nhead
|
int
|
The number of heads in the transformer encoder/decoder. |
8
|
num_encoder_layers
|
int
|
The number of encoder-layers in the encoder. |
6
|
num_decoder_layers
|
int
|
The number of decoder-layers in the decoder. |
6
|
dropout
|
int
|
Dropout value applied to the output of transformer layers. |
0.1
|
activation
|
str
|
Activation function to use. |
'relu'
|
return_intermediate_dec
|
bool
|
Return intermediate layers from decoder. |
False
|
norm
|
bool
|
If True, normalize output of encoder and decoder. |
False
|
num_layers_attn_head
|
int
|
The number of layers in the attention head. |
2
|
dropout_attn_head
|
int
|
Dropout value for the attention_head. |
0.1
|
embedding_meta
|
dict | None
|
Metadata for positional embeddings. See below. |
None
|
return_embedding
|
bool
|
Whether to return the positional embeddings |
False
|
decoder_self_attn
|
bool
|
If True, use decoder self attention. More details on |
False
|
Source code in dreem/models/global_tracking_transformer.py
def __init__(
self,
encoder_cfg: dict | None = None,
d_model: int = 1024,
nhead: int = 8,
num_encoder_layers: int = 6,
num_decoder_layers: int = 6,
dropout: int = 0.1,
activation: str = "relu",
return_intermediate_dec: bool = False,
norm: bool = False,
num_layers_attn_head: int = 2,
dropout_attn_head: int = 0.1,
embedding_meta: dict | None = None,
return_embedding: bool = False,
decoder_self_attn: bool = False,
):
"""Initialize GTR.
Args:
encoder_cfg: Dictionary of arguments to pass to the CNN constructor,
e.g: `cfg = {"model_name": "resnet18", "pretrained": False, "in_chans": 3}`
d_model: The number of features in the encoder/decoder inputs.
nhead: The number of heads in the transformer encoder/decoder.
num_encoder_layers: The number of encoder-layers in the encoder.
num_decoder_layers: The number of decoder-layers in the decoder.
dropout: Dropout value applied to the output of transformer layers.
activation: Activation function to use.
return_intermediate_dec: Return intermediate layers from decoder.
norm: If True, normalize output of encoder and decoder.
num_layers_attn_head: The number of layers in the attention head.
dropout_attn_head: Dropout value for the attention_head.
embedding_meta: Metadata for positional embeddings. See below.
return_embedding: Whether to return the positional embeddings
decoder_self_attn: If True, use decoder self attention.
More details on `embedding_meta`:
By default this will be an empty dict and indicate
that no positional embeddings should be used. To use the positional embeddings
pass in a dictionary containing a "pos" and "temp" key with subdictionaries for correct parameters ie:
`{"pos": {'mode': 'learned', 'emb_num': 16, 'over_boxes: True},
"temp": {'mode': 'learned', 'emb_num': 16}}`. (see `dreem.models.embeddings.Embedding.EMB_TYPES`
and `dreem.models.embeddings.Embedding.EMB_MODES` for embedding parameters).
"""
super().__init__()
if not encoder_cfg:
encoder_cfg = {}
self.visual_encoder = create_visual_encoder(d_model=d_model, **encoder_cfg)
self.transformer = Transformer(
d_model=d_model,
nhead=nhead,
num_encoder_layers=num_encoder_layers,
num_decoder_layers=num_decoder_layers,
dropout=dropout,
activation=activation,
return_intermediate_dec=return_intermediate_dec,
norm=norm,
num_layers_attn_head=num_layers_attn_head,
dropout_attn_head=dropout_attn_head,
embedding_meta=embedding_meta,
return_embedding=return_embedding,
decoder_self_attn=decoder_self_attn,
encoder_cfg=encoder_cfg,
)
extract_features(instances, force_recompute=False, retain_crops=False)
¶
Extract features from instances using visual encoder backbone.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
instances
|
list[Instance]
|
A list of instances to compute features for |
required |
force_recompute
|
bool
|
indicate whether to compute features for all instances regardless of if they have instances |
False
|
Source code in dreem/models/global_tracking_transformer.py
def extract_features(
self,
instances: list["Instance"],
force_recompute: bool = False,
retain_crops: bool = False,
) -> None:
"""Extract features from instances using visual encoder backbone.
Args:
instances: A list of instances to compute features for
force_recompute: indicate whether to compute features for all instances regardless of if they have instances
"""
if not force_recompute:
instances_to_compute = [
instance
for instance in instances
if instance.has_crop() and not instance.has_features()
]
else:
instances_to_compute = instances
if len(instances_to_compute) == 0:
return
elif len(instances_to_compute) == 1: # handle batch norm error when B=1
instances_to_compute = instances
crops = torch.concatenate([instance.crop for instance in instances_to_compute])
features = self.visual_encoder(crops)
for i, z_i in enumerate(features):
instances_to_compute[i].features = z_i
if not retain_crops: # crops are an attribute of instance by default
instances_to_compute[i].crop = None
forward(ref_instances, query_instances=None, retain_crops=False)
¶
Execute forward pass of GTR Model to get asso matrix.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
ref_instances
|
list[Instance]
|
List of instances from chunk containing crops of objects + gt label info |
required |
query_instances
|
list[Instance]
|
list of instances used as query in decoder. |
None
|
Returns:
| Type | Description |
|---|---|
list[AssociationMatrix]
|
An N_T x N association matrix |
Source code in dreem/models/global_tracking_transformer.py
def forward(
self,
ref_instances: list["Instance"],
query_instances: list["Instance"] = None,
retain_crops: bool = False,
) -> list["AssociationMatrix"]:
"""Execute forward pass of GTR Model to get asso matrix.
Args:
ref_instances: List of instances from chunk containing crops of objects + gt label info
query_instances: list of instances used as query in decoder.
Returns:
An N_T x N association matrix
"""
# Extract feature representations with pre-trained encoder.
self.extract_features(ref_instances, retain_crops=retain_crops)
if query_instances:
self.extract_features(query_instances, retain_crops=retain_crops)
asso_preds = self.transformer(ref_instances, query_instances)
return asso_preds
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}")
Tracker
¶
Tracker class used for assignment based on sliding inference from GTR.
Methods:
| Name | Description |
|---|---|
__call__ |
Wrap around |
__init__ |
Initialize a tracker to run inference. |
__repr__ |
Get string representation of tracker. |
sliding_inference |
Perform sliding inference on the input video (instances) with a given window size. |
track |
Run tracker and get predicted trajectories. |
Source code in dreem/inference/tracker.py
class Tracker:
"""Tracker class used for assignment based on sliding inference from GTR."""
def __init__(
self,
window_size: int = 8,
use_vis_feats: bool = True,
overlap_thresh: float = 0.01,
mult_thresh: bool = True,
decay_time: float | None = None,
iou: str | None = None,
max_center_dist: float | None = None,
distance_penalty_multiplier: float = 1.0,
angle_diff_penalty_multiplier: float = 1.0,
max_gap: int = inf,
max_tracks: int = inf,
verbose: bool = False,
confidence_threshold: float = 0.0,
temperature: float = 0.1,
max_angle_diff: float = 0.0,
front_nodes: list[str] | None = None,
back_nodes: list[str] | None = None,
enable_crop_saving: bool = False,
**kwargs,
):
"""Initialize a tracker to run inference.
Args:
window_size: the size of the window used during sliding inference.
use_vis_feats: Whether or not to use visual feature extractor.
overlap_thresh: the trajectory overlap threshold to be used for assignment.
mult_thresh: Whether or not to use weight threshold.
decay_time: weight for `decay_time` postprocessing.
iou: Either [None, '', "mult" or "max"]
Whether to use multiplicative or max iou reweighting.
max_center_dist: distance threshold for filtering trajectory score matrix.
max_gap: the max number of frames a trajectory can be missing before termination.
max_tracks: the maximum number of tracks that can be created while tracking.
We force the tracker to assign instances to a track instead of creating a new track if max_tracks has been reached.
verbose: Whether or not to turn on debug printing after each operation.
confidence_threshold: threshold for filtering out instances with high confidence. Set to 0 to disable confidence thresholding.
temperature: temperature for softmax.
max_angle_diff: maximum angle difference between pose principal axes when considering association between two instances. Set to 0 to disable angle difference filtering.
front_nodes: list of skeleton node names to be used to determine the orientation of the object. If None, computes using all available nodes.
back_nodes: list of skeleton node names to be used to determine the orientation of the object. If None, computes using all available nodes.
enable_crop_saving: Whether to save crops to frame metadata.
**kwargs: Additional keyword arguments (unused but accepted for compatibility).
"""
self.track_queue = TrackQueue(
window_size=window_size, max_gap=max_gap, verbose=verbose
)
self.use_vis_feats = use_vis_feats
self.overlap_thresh = overlap_thresh
self.mult_thresh = mult_thresh
self.decay_time = decay_time
self.iou = iou
self.max_center_dist = max_center_dist if max_center_dist is not None else inf
self.verbose = verbose
self.max_tracks = max_tracks
self.confidence_threshold = confidence_threshold
self.temperature = temperature
self.front_nodes = front_nodes
self.back_nodes = back_nodes
self.enable_crop_saving = enable_crop_saving
self.max_angle_diff = (
deg2rad(max_angle_diff) if max_angle_diff is not None else inf
)
self.orientation_weighting = OrientationWeighting(
self.max_angle_diff, angle_diff_penalty_multiplier
)
self.distance_weighting = DistanceWeighting(
self.max_center_dist, distance_penalty_multiplier
)
self.iou_weighting = IOUWeighting(iou)
self.confidence_flagging = ConfidenceFlagging(confidence_threshold)
def __call__(
self, model: GlobalTrackingTransformer, frames: list[Frame]
) -> list[Frame]:
"""Wrap around `track` to enable `tracker()` instead of `tracker.track()`.
Args:
model: the pretrained GlobalTrackingTransformer to be used for inference
frames: list of Frames to run inference on
Returns:
List of frames containing association matrix scores and instances populated with pred track ids.
"""
return self.track(model, frames)
def __repr__(self) -> str:
"""Get string representation of tracker.
Returns: the string representation of the tracker
"""
return (
"Tracker("
f"max_tracks={self.max_tracks}, "
f"use_vis_feats={self.use_vis_feats}, "
f"overlap_thresh={self.overlap_thresh}, "
f"mult_thresh={self.mult_thresh}, "
f"decay_time={self.decay_time}, "
f"max_center_dist={self.max_center_dist}, "
f"verbose={self.verbose}, "
f"queue={self.track_queue}, "
f"temperature={self.temperature}"
f"queue={self.track_queue}, "
f"temperature={self.temperature}"
)
def track(
self, model: GlobalTrackingTransformer, frames: list[dict]
) -> list[Frame]:
"""Run tracker and get predicted trajectories.
Args:
model: the pretrained GlobalTrackingTransformer to be used for inference
frames: data dict to run inference on
Returns:
List of Frames populated with pred track ids and association matrix scores
"""
_ = model.eval()
instances_pred = self.sliding_inference(model, frames)
return instances_pred
def sliding_inference(
self, model: GlobalTrackingTransformer, frames: list[Frame]
) -> list[Frame]:
"""Perform sliding inference on the input video (instances) with a given window size.
Args:
model: the pretrained GlobalTrackingTransformer to be used for inference
frames: A list of Frames (See `dreem.io.Frame` for more info).
Returns:
frames: A list of Frames populated with pred_track_ids and asso_matrices
"""
# B: batch size.
# D: embedding dimension.
# nc: number of channels.
# H: height.
# W: width.
for batch_idx, frame_to_track in enumerate(frames):
tracked_frames = self.track_queue.collate_tracks(
device=frame_to_track.device
)
logger.debug(f"Current number of tracks is {self.track_queue.n_tracks}")
if frame_to_track.frame_id == 0: # clear queue on new video
logger.debug("New Video! Resetting Track Queue.")
self.track_queue.end_tracks()
# Initialize tracks on first frame where detections appear
if len(self.track_queue) == 0:
if frame_to_track.has_instances():
logger.debug(
f"Initializing track on clip ind {batch_idx} frame {frame_to_track.frame_id.item()}"
)
curr_track_id = 0
for i, instance in enumerate(frames[batch_idx].instances):
instance.pred_track_id = instance.gt_track_id
curr_track_id = max(curr_track_id, instance.pred_track_id)
for i, instance in enumerate(frames[batch_idx].instances):
if instance.pred_track_id == -1:
instance.pred_track_id = curr_track_id
curr_track_id += 1
else:
if frame_to_track.has_instances(): # Check if there are detections. If there are skip and increment gap count
frames_to_track = tracked_frames + [
frame_to_track
] # better var name?
query_ind = len(frames_to_track) - 1
frame_to_track = self._run_global_tracker(
model,
frames_to_track,
query_ind=query_ind,
)
del frames_to_track
if frame_to_track.has_instances():
_ = self.track_queue.add_frame(frame_to_track)
else:
self.track_queue.increment_gaps([])
frames[batch_idx] = frame_to_track
del frame_to_track, tracked_frames
torch.cuda.empty_cache()
return frames
def _run_global_tracker(
self, model: GlobalTrackingTransformer, frames: list[Frame], query_ind: int
) -> Frame:
"""Run global tracker performs the actual tracking.
Uses Hungarian algorithm to do track assigning.
Args:
model: the pretrained GlobalTrackingTransformer to be used for inference
frames: A list of Frames containing reid features. See `dreem.io.data_structures` for more info.
query_ind: An integer for the query frame within the window of instances.
Returns:
query_frame: The query frame now populated with the pred_track_ids.
"""
# *: each item in frames is a frame in the window. So it follows
# that each frame in the window has * detected instances.
# D: embedding dimension.
# total_instances: number of instances in the window.
# N_i: number of detected instances in i-th frame of window.
# instances_per_frame: a list of number of instances in each frame of the window.
# n_query: number of instances in current/query frame (rightmost frame of the window).
# n_nonquery: number of instances in the window excluding the current/query frame.
# window_size: length of window.
# L: number of decoder blocks.
# n_traj: number of existing tracks within the window so far.
# Number of instances in each frame of the window.
# E.g.: instances_per_frame: [4, 5, 6, 7]; window of length 4 with 4 detected instances in the first frame of the window.
_ = model.eval()
query_frame = frames[query_ind]
_, h, w = query_frame.img_shape
query_instances = query_frame.instances
all_instances = [instance for frame in frames for instance in frame.instances]
logger.debug(f"Frame {query_frame.frame_id.item()}")
instances_per_frame = [frame.num_detected for frame in frames]
total_instances = sum(instances_per_frame) # Number of instances in window
logger.debug(f"total_instances: {total_instances}")
overlap_thresh = self.overlap_thresh
mult_thresh = self.mult_thresh
n_traj = self.track_queue.n_tracks
curr_tracks = self.track_queue.curr_track
query_poses = [
(instance.pose, instance.pred_track_id.item(), instance.crop.clone().cpu())
for instance in query_frame.instances
]
nonquery_poses = [
(instance.pose, instance.pred_track_id.item(), instance.crop.clone().cpu())
for nonquery_frame in frames
if nonquery_frame.frame_id != query_frame.frame_id.item()
for instance in nonquery_frame.instances
]
enable_crop_saving = self.enable_crop_saving or is_pose_centroid_only(
query_poses[0][0]
)
if (
is_pose_centroid_only(query_poses[0][0])
and query_frame.frame_id.item() == 1
and self.max_angle_diff > 0
):
logger.warning(
"Crop saving is enabled due to max_angle_diff > 0 and a centroid-only skeleton. This can increase memory usage. If you experience memory issues, consider setting max_angle_diff = 0. This will skip orientation based post processing."
)
with torch.no_grad():
asso_matrix = model(
all_instances, query_instances, retain_crops=enable_crop_saving
)
asso_output = asso_matrix[-1].matrix.split(
instances_per_frame, dim=1
) # (window_size, n_query, N_i)
asso_output = model_utils.softmax_asso(
asso_output
) # (window_size, n_query, N_i)
asso_output = torch.cat(asso_output, dim=1).cpu() # (n_query, total_instances)
asso_output_df = pd.DataFrame(
asso_output.clone().numpy(),
columns=[f"Instance {i}" for i in range(asso_output.shape[-1])],
)
asso_output_df.index.name = "Instances"
asso_output_df.columns.name = "Instances"
query_frame.add_traj_score("asso_output", asso_output_df)
query_frame.asso_output = asso_matrix[-1]
n_query = (
query_frame.num_detected
) # Number of instances in the current/query frame.
n_nonquery = (
total_instances - n_query
) # Number of instances in the window not including the current/query frame.
logger.debug(f"n_nonquery: {n_nonquery}")
logger.debug(f"n_query: {n_query}")
instance_ids = (
torch.cat(
[
x.get_pred_track_ids()
for batch_idx, x in enumerate(frames)
if batch_idx != query_ind
],
dim=0,
)
.view(n_nonquery)
.cpu()
) # (n_nonquery,)
query_inds = [
x
for x in range(
sum(instances_per_frame[:query_ind]),
sum(instances_per_frame[: query_ind + 1]),
)
]
nonquery_inds = [i for i in range(total_instances) if i not in query_inds]
# instead should we do model(nonquery_instances, query_instances)?
asso_nonquery = asso_output[:, nonquery_inds] # (n_query, n_nonquery)
asso_nonquery_df = pd.DataFrame(
asso_nonquery.clone().numpy(), columns=nonquery_inds
)
asso_nonquery_df.index.name = "Current Frame Instances"
asso_nonquery_df.columns.name = "Nonquery Instances"
query_frame.add_traj_score("asso_nonquery", asso_nonquery_df)
# get raw bbox coords of prev frame instances from frame.instances_per_frame
query_boxes_px = torch.cat(
[instance.bbox for instance in query_frame.instances], dim=0
).cpu()
pred_boxes = model_utils.get_boxes(all_instances)
query_boxes = pred_boxes[query_inds] # n_k x 4
nonquery_boxes = pred_boxes[nonquery_inds] # n_nonquery x 4
unique_ids = torch.unique(instance_ids).cpu() # (n_nonquery,)
logger.debug(f"Instance IDs: {instance_ids}")
logger.debug(f"unique ids: {unique_ids}")
id_inds = (
unique_ids[None, :] == instance_ids[:, None]
).float() # (n_nonquery, n_traj)
################################################################################
# (n_query x n_nonquery) x (n_nonquery x n_traj) --> n_query x n_traj
traj_score = torch.mm(asso_nonquery, id_inds.cpu()) # (n_query, n_traj)
assoc_col_pred_id_map = {
i: unique_ids[i].item() for i in range(unique_ids.shape[0])
}
traj_score_df = pd.DataFrame(
traj_score.clone().numpy(), columns=unique_ids.cpu().numpy()
)
traj_score_df.index.name = "Current Frame Instances"
traj_score_df.columns.name = "Unique IDs"
query_frame.add_traj_score("traj_score", traj_score_df)
################################################################################
# with iou -> combining with location in tracker, they set to True
# todo -> should also work without pos_embed
if id_inds.numel() > 0:
# this throws error, think we need to slice?
# last_inds = (id_inds * torch.arange(
# n_nonquery, device=id_inds.device)[:, None]).max(dim=0)[1] # n_traj
last_inds = (id_inds * torch.arange(n_nonquery)[:, None]).max(dim=0)[1] # M
last_boxes = nonquery_boxes[last_inds] # n_traj x 4
last_ious = _pairwise_iou(Boxes(query_boxes), Boxes(last_boxes)) # n_k x M
else:
last_ious = traj_score.new_zeros(traj_score.shape)
state = self.iou_weighting.run(
{
"traj_score": traj_score,
"last_ious": last_ious.cpu(),
}
)
traj_score = state["traj_score"]
if self.iou is not None and self.iou != "":
iou_traj_score = pd.DataFrame(
traj_score.clone().numpy(), columns=unique_ids.cpu().numpy()
)
iou_traj_score.index.name = "Current Frame Instances"
iou_traj_score.columns.name = "Unique IDs"
query_frame.add_traj_score("weight_iou", iou_traj_score)
################################################################################
last_poses = []
for ind in last_inds.cpu(): # its important to index nonquery_poses by last_inds as this maintains ordering that matches the association matrix ordering
last_poses.append(nonquery_poses[ind])
last_pred_ids = [pred_track_id for _, pred_track_id, _ in last_poses]
if self.max_center_dist is not None and self.max_center_dist > 0:
last_boxes_px = last_boxes.cpu()
last_boxes_px[:, :, [0, 2]] *= w
last_boxes_px[:, :, [1, 3]] *= h
state = self.distance_weighting.run(
{
"traj_score": traj_score,
"query_boxes_px": query_boxes_px,
"last_boxes_px": last_boxes_px,
"h": h,
"w": w,
}
)
traj_score = state["traj_score"]
# update metadata
max_center_dist_traj_score = pd.DataFrame(
traj_score.clone().numpy(), columns=unique_ids.cpu().numpy()
)
max_center_dist_traj_score.index.name = "Current Frame Instances"
max_center_dist_traj_score.columns.name = "Unique IDs"
query_frame.add_traj_score("max_center_dist", max_center_dist_traj_score)
################################################################################
if self.max_angle_diff > 0:
query_principal_axes = []
last_principal_axes = []
successes = []
for instance, pred_track_id, crop in query_poses:
instance_principal_axes, success = get_principal_axis_with_fallback(
instance,
self.front_nodes,
self.back_nodes,
crop,
logger,
query_frame.frame_id.item(),
)
query_principal_axes.append(instance_principal_axes)
successes.append(success)
query_principal_axes = torch.stack(query_principal_axes) # (n_query, 2)
for instance, pred_track_id, crop in last_poses:
instance_principal_axes, success = get_principal_axis_with_fallback(
instance,
self.front_nodes,
self.back_nodes,
crop,
logger,
query_frame.frame_id.item(),
)
last_principal_axes.append(instance_principal_axes)
successes.append(success)
last_principal_axes = torch.stack(last_principal_axes) # (n_traj, 2)
if all(successes):
state = self.orientation_weighting.run(
{
"traj_score": traj_score,
"query_principal_axes": query_principal_axes,
"last_principal_axes": last_principal_axes,
}
)
traj_score = state["traj_score"]
# update metadata
angle_diff_weight_traj_score = pd.DataFrame(
traj_score.clone().numpy(), columns=unique_ids.cpu().numpy()
)
angle_diff_weight_traj_score.index.name = "Current Frame Instances"
angle_diff_weight_traj_score.columns.name = "Unique IDs"
query_frame.add_traj_score(
"angle_diff_weight", angle_diff_weight_traj_score
)
else:
logger.warning(
f"Some instances had no principal axis in frame {query_frame.frame_id.item()}. Skipping instance angle based post processing."
)
################################################################################
scaled_traj_score = torch.nn.functional.log_softmax(
traj_score / self.temperature, dim=1
)
scaled_traj_score_df = pd.DataFrame(
scaled_traj_score.numpy(), columns=unique_ids.cpu().numpy()
)
scaled_traj_score_df.index.name = "Current Frame Instances"
scaled_traj_score_df.columns.name = "Unique IDs"
query_frame.add_traj_score("scaled", scaled_traj_score_df)
################################################################################
# Flag frames with low confidence
if self.confidence_threshold > 0:
state = self.confidence_flagging.run(
{
"scaled_traj_score": scaled_traj_score,
"n_query": n_query,
"query_frame": query_frame,
}
)
scaled_traj_score = state["scaled_traj_score"]
match_i, match_j = linear_sum_assignment((-traj_score))
track_ids = instance_ids.new_full((n_query,), -1)
for i, j in zip(match_i, match_j):
# The overlap threshold is multiplied by the number of times the unique track j is matched to an
# instance out of all instances in the window excluding the current frame.
#
# So if this is correct, the threshold is higher for matching an instance from the current frame
# to an existing track if that track has already been matched several times.
# So if an existing track in the window has been matched a lot, it gets harder to match to that track.
thresh = (
overlap_thresh * id_inds[:, j].sum() if mult_thresh else overlap_thresh
)
# if we've already reached max tracks, just assign it to one of the existing tracks
if n_traj >= self.max_tracks or traj_score[i, j] > thresh:
logger.debug(
f"Assigning instance {i} to track {j} with id {unique_ids[j]}"
)
track_ids[i] = unique_ids[j]
query_frame.instances[i].track_score = scaled_traj_score[i, j].item()
logger.debug(f"track_ids: {track_ids}")
for i in range(n_query):
# True if association score was below the threshold, and we haven't reached max tracks
if track_ids[i] < 0 and n_traj < self.max_tracks:
max_track_id = max(curr_tracks)
logger.debug(f"Creating new track {max_track_id + 1}")
curr_tracks.add(max_track_id + 1)
track_ids[i] = max_track_id + 1
query_frame.matches = (match_i, match_j)
for instance, track_id in zip(query_frame.instances, track_ids):
instance.pred_track_id = track_id
final_traj_score = pd.DataFrame(
traj_score.clone().numpy(), columns=unique_ids.cpu().numpy()
)
final_traj_score.index.name = "Current Frame Instances"
final_traj_score.columns.name = "Unique IDs"
query_frame.add_traj_score("final", final_traj_score)
return query_frame
__call__(model, frames)
¶
Wrap around track to enable tracker() instead of tracker.track().
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
model
|
GlobalTrackingTransformer
|
the pretrained GlobalTrackingTransformer to be used for inference |
required |
frames
|
list[Frame]
|
list of Frames to run inference on |
required |
Returns:
| Type | Description |
|---|---|
list[Frame]
|
List of frames containing association matrix scores and instances populated with pred track ids. |
Source code in dreem/inference/tracker.py
def __call__(
self, model: GlobalTrackingTransformer, frames: list[Frame]
) -> list[Frame]:
"""Wrap around `track` to enable `tracker()` instead of `tracker.track()`.
Args:
model: the pretrained GlobalTrackingTransformer to be used for inference
frames: list of Frames to run inference on
Returns:
List of frames containing association matrix scores and instances populated with pred track ids.
"""
return self.track(model, frames)
__init__(window_size=8, use_vis_feats=True, overlap_thresh=0.01, mult_thresh=True, decay_time=None, iou=None, max_center_dist=None, distance_penalty_multiplier=1.0, angle_diff_penalty_multiplier=1.0, max_gap=inf, max_tracks=inf, verbose=False, confidence_threshold=0.0, temperature=0.1, max_angle_diff=0.0, front_nodes=None, back_nodes=None, enable_crop_saving=False, **kwargs)
¶
Initialize a tracker to run inference.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
window_size
|
int
|
the size of the window used during sliding inference. |
8
|
use_vis_feats
|
bool
|
Whether or not to use visual feature extractor. |
True
|
overlap_thresh
|
float
|
the trajectory overlap threshold to be used for assignment. |
0.01
|
mult_thresh
|
bool
|
Whether or not to use weight threshold. |
True
|
decay_time
|
float | None
|
weight for |
None
|
iou
|
str | None
|
Either [None, '', "mult" or "max"] Whether to use multiplicative or max iou reweighting. |
None
|
max_center_dist
|
float | None
|
distance threshold for filtering trajectory score matrix. |
None
|
max_gap
|
int
|
the max number of frames a trajectory can be missing before termination. |
inf
|
max_tracks
|
int
|
the maximum number of tracks that can be created while tracking. We force the tracker to assign instances to a track instead of creating a new track if max_tracks has been reached. |
inf
|
verbose
|
bool
|
Whether or not to turn on debug printing after each operation. |
False
|
confidence_threshold
|
float
|
threshold for filtering out instances with high confidence. Set to 0 to disable confidence thresholding. |
0.0
|
temperature
|
float
|
temperature for softmax. |
0.1
|
max_angle_diff
|
float
|
maximum angle difference between pose principal axes when considering association between two instances. Set to 0 to disable angle difference filtering. |
0.0
|
front_nodes
|
list[str] | None
|
list of skeleton node names to be used to determine the orientation of the object. If None, computes using all available nodes. |
None
|
back_nodes
|
list[str] | None
|
list of skeleton node names to be used to determine the orientation of the object. If None, computes using all available nodes. |
None
|
enable_crop_saving
|
bool
|
Whether to save crops to frame metadata. |
False
|
**kwargs
|
Additional keyword arguments (unused but accepted for compatibility). |
{}
|
Source code in dreem/inference/tracker.py
def __init__(
self,
window_size: int = 8,
use_vis_feats: bool = True,
overlap_thresh: float = 0.01,
mult_thresh: bool = True,
decay_time: float | None = None,
iou: str | None = None,
max_center_dist: float | None = None,
distance_penalty_multiplier: float = 1.0,
angle_diff_penalty_multiplier: float = 1.0,
max_gap: int = inf,
max_tracks: int = inf,
verbose: bool = False,
confidence_threshold: float = 0.0,
temperature: float = 0.1,
max_angle_diff: float = 0.0,
front_nodes: list[str] | None = None,
back_nodes: list[str] | None = None,
enable_crop_saving: bool = False,
**kwargs,
):
"""Initialize a tracker to run inference.
Args:
window_size: the size of the window used during sliding inference.
use_vis_feats: Whether or not to use visual feature extractor.
overlap_thresh: the trajectory overlap threshold to be used for assignment.
mult_thresh: Whether or not to use weight threshold.
decay_time: weight for `decay_time` postprocessing.
iou: Either [None, '', "mult" or "max"]
Whether to use multiplicative or max iou reweighting.
max_center_dist: distance threshold for filtering trajectory score matrix.
max_gap: the max number of frames a trajectory can be missing before termination.
max_tracks: the maximum number of tracks that can be created while tracking.
We force the tracker to assign instances to a track instead of creating a new track if max_tracks has been reached.
verbose: Whether or not to turn on debug printing after each operation.
confidence_threshold: threshold for filtering out instances with high confidence. Set to 0 to disable confidence thresholding.
temperature: temperature for softmax.
max_angle_diff: maximum angle difference between pose principal axes when considering association between two instances. Set to 0 to disable angle difference filtering.
front_nodes: list of skeleton node names to be used to determine the orientation of the object. If None, computes using all available nodes.
back_nodes: list of skeleton node names to be used to determine the orientation of the object. If None, computes using all available nodes.
enable_crop_saving: Whether to save crops to frame metadata.
**kwargs: Additional keyword arguments (unused but accepted for compatibility).
"""
self.track_queue = TrackQueue(
window_size=window_size, max_gap=max_gap, verbose=verbose
)
self.use_vis_feats = use_vis_feats
self.overlap_thresh = overlap_thresh
self.mult_thresh = mult_thresh
self.decay_time = decay_time
self.iou = iou
self.max_center_dist = max_center_dist if max_center_dist is not None else inf
self.verbose = verbose
self.max_tracks = max_tracks
self.confidence_threshold = confidence_threshold
self.temperature = temperature
self.front_nodes = front_nodes
self.back_nodes = back_nodes
self.enable_crop_saving = enable_crop_saving
self.max_angle_diff = (
deg2rad(max_angle_diff) if max_angle_diff is not None else inf
)
self.orientation_weighting = OrientationWeighting(
self.max_angle_diff, angle_diff_penalty_multiplier
)
self.distance_weighting = DistanceWeighting(
self.max_center_dist, distance_penalty_multiplier
)
self.iou_weighting = IOUWeighting(iou)
self.confidence_flagging = ConfidenceFlagging(confidence_threshold)
__repr__()
¶
Get string representation of tracker.
Returns: the string representation of the tracker
Source code in dreem/inference/tracker.py
def __repr__(self) -> str:
"""Get string representation of tracker.
Returns: the string representation of the tracker
"""
return (
"Tracker("
f"max_tracks={self.max_tracks}, "
f"use_vis_feats={self.use_vis_feats}, "
f"overlap_thresh={self.overlap_thresh}, "
f"mult_thresh={self.mult_thresh}, "
f"decay_time={self.decay_time}, "
f"max_center_dist={self.max_center_dist}, "
f"verbose={self.verbose}, "
f"queue={self.track_queue}, "
f"temperature={self.temperature}"
f"queue={self.track_queue}, "
f"temperature={self.temperature}"
)
sliding_inference(model, frames)
¶
Perform sliding inference on the input video (instances) with a given window size.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
model
|
GlobalTrackingTransformer
|
the pretrained GlobalTrackingTransformer to be used for inference |
required |
frames
|
list[Frame]
|
A list of Frames (See |
required |
Returns:
| Name | Type | Description |
|---|---|---|
frames |
list[Frame]
|
A list of Frames populated with pred_track_ids and asso_matrices |
Source code in dreem/inference/tracker.py
def sliding_inference(
self, model: GlobalTrackingTransformer, frames: list[Frame]
) -> list[Frame]:
"""Perform sliding inference on the input video (instances) with a given window size.
Args:
model: the pretrained GlobalTrackingTransformer to be used for inference
frames: A list of Frames (See `dreem.io.Frame` for more info).
Returns:
frames: A list of Frames populated with pred_track_ids and asso_matrices
"""
# B: batch size.
# D: embedding dimension.
# nc: number of channels.
# H: height.
# W: width.
for batch_idx, frame_to_track in enumerate(frames):
tracked_frames = self.track_queue.collate_tracks(
device=frame_to_track.device
)
logger.debug(f"Current number of tracks is {self.track_queue.n_tracks}")
if frame_to_track.frame_id == 0: # clear queue on new video
logger.debug("New Video! Resetting Track Queue.")
self.track_queue.end_tracks()
# Initialize tracks on first frame where detections appear
if len(self.track_queue) == 0:
if frame_to_track.has_instances():
logger.debug(
f"Initializing track on clip ind {batch_idx} frame {frame_to_track.frame_id.item()}"
)
curr_track_id = 0
for i, instance in enumerate(frames[batch_idx].instances):
instance.pred_track_id = instance.gt_track_id
curr_track_id = max(curr_track_id, instance.pred_track_id)
for i, instance in enumerate(frames[batch_idx].instances):
if instance.pred_track_id == -1:
instance.pred_track_id = curr_track_id
curr_track_id += 1
else:
if frame_to_track.has_instances(): # Check if there are detections. If there are skip and increment gap count
frames_to_track = tracked_frames + [
frame_to_track
] # better var name?
query_ind = len(frames_to_track) - 1
frame_to_track = self._run_global_tracker(
model,
frames_to_track,
query_ind=query_ind,
)
del frames_to_track
if frame_to_track.has_instances():
_ = self.track_queue.add_frame(frame_to_track)
else:
self.track_queue.increment_gaps([])
frames[batch_idx] = frame_to_track
del frame_to_track, tracked_frames
torch.cuda.empty_cache()
return frames
track(model, frames)
¶
Run tracker and get predicted trajectories.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
model
|
GlobalTrackingTransformer
|
the pretrained GlobalTrackingTransformer to be used for inference |
required |
frames
|
list[dict]
|
data dict to run inference on |
required |
Returns:
| Type | Description |
|---|---|
list[Frame]
|
List of Frames populated with pred track ids and association matrix scores |
Source code in dreem/inference/tracker.py
def track(
self, model: GlobalTrackingTransformer, frames: list[dict]
) -> list[Frame]:
"""Run tracker and get predicted trajectories.
Args:
model: the pretrained GlobalTrackingTransformer to be used for inference
frames: data dict to run inference on
Returns:
List of Frames populated with pred track ids and association matrix scores
"""
_ = model.eval()
instances_pred = self.sliding_inference(model, frames)
return instances_pred
Transformer
¶
Bases: Module
Transformer class.
Methods:
| Name | Description |
|---|---|
__init__ |
Initialize Transformer. |
forward |
Execute a forward pass through the transformer and attention head. |
Source code in dreem/models/transformer.py
class Transformer(torch.nn.Module):
"""Transformer class."""
def __init__(
self,
d_model: int = 1024,
nhead: int = 8,
num_encoder_layers: int = 6,
num_decoder_layers: int = 6,
dropout: float = 0.1,
activation: str = "relu",
return_intermediate_dec: bool = False,
norm: bool = False,
num_layers_attn_head: int = 2,
dropout_attn_head: float = 0.1,
embedding_meta: dict | None = None,
return_embedding: bool = False,
decoder_self_attn: bool = False,
encoder_cfg: dict | None = None,
) -> None:
"""Initialize Transformer.
Args:
d_model: The number of features in the encoder/decoder inputs.
nhead: The number of heads in the transformer encoder/decoder.
num_encoder_layers: The number of encoder-layers in the encoder.
num_decoder_layers: The number of decoder-layers in the decoder.
dropout: Dropout value applied to the output of transformer layers.
activation: Activation function to use.
return_intermediate_dec: Return intermediate layers from decoder.
norm: If True, normalize output of encoder and decoder.
num_layers_attn_head: The number of layers in the attention head.
dropout_attn_head: Dropout value for the attention_head.
embedding_meta: Metadata for positional embeddings. See below.
return_embedding: Whether to return the positional embeddings
decoder_self_attn: If True, use decoder self attention.
encoder_cfg: Encoder configuration.
More details on `embedding_meta`:
By default this will be an empty dict and indicate
that no positional embeddings should be used. To use the positional embeddings
pass in a dictionary containing a "pos" and "temp" key with subdictionaries for correct parameters ie:
{"pos": {'mode': 'learned', 'emb_num': 16, 'over_boxes: 'True'},
"temp": {'mode': 'learned', 'emb_num': 16}}. (see `dreem.models.embeddings.Embedding.EMB_TYPES`
and `dreem.models.embeddings.Embedding.EMB_MODES` for embedding parameters).
"""
super().__init__()
self.d_model = dim_feedforward = feature_dim_attn_head = d_model
self.embedding_meta = embedding_meta
self.return_embedding = return_embedding
self.encoder_cfg = encoder_cfg
self.pos_emb = Embedding(emb_type="off", mode="off", features=self.d_model)
self.temp_emb = Embedding(emb_type="off", mode="off", features=self.d_model)
if self.embedding_meta:
if "pos" in self.embedding_meta:
pos_emb_cfg = self.embedding_meta["pos"]
if pos_emb_cfg:
self.pos_emb = Embedding(
emb_type="pos", features=self.d_model, **pos_emb_cfg
)
if "temp" in self.embedding_meta:
temp_emb_cfg = self.embedding_meta["temp"]
if temp_emb_cfg:
self.temp_emb = Embedding(
emb_type="temp", features=self.d_model, **temp_emb_cfg
)
self.fourier_embeddings = FourierPositionalEmbeddings(
n_components=8, d_model=d_model
)
# Transformer Encoder
encoder_layer = TransformerEncoderLayer(
d_model, nhead, dim_feedforward, dropout, activation, norm
)
encoder_norm = nn.LayerNorm(d_model) if (norm) else None
# only used if using descriptor visual encoder; default resnet encoder uses d_model directly
if self.encoder_cfg and "encoder_type" in self.encoder_cfg:
self.visual_feat_dim = (
self.encoder_cfg["ndim"] if "ndim" in self.encoder_cfg else 5
) # 5 is default for descriptor
self.fourier_proj = nn.Linear(self.d_model + self.visual_feat_dim, d_model)
self.fourier_norm = nn.LayerNorm(self.d_model)
self.encoder = TransformerEncoder(
encoder_layer, num_encoder_layers, encoder_norm
)
# Transformer Decoder
decoder_layer = TransformerDecoderLayer(
d_model,
nhead,
dim_feedforward,
dropout,
activation,
norm,
decoder_self_attn,
)
decoder_norm = nn.LayerNorm(d_model) if (norm) else None
self.decoder = TransformerDecoder(
decoder_layer, num_decoder_layers, return_intermediate_dec, decoder_norm
)
# Transformer attention head
self.attn_head = ATTWeightHead(
feature_dim=feature_dim_attn_head,
num_layers=num_layers_attn_head,
dropout=dropout_attn_head,
)
self._reset_parameters()
def _reset_parameters(self):
"""Initialize model weights from xavier distribution."""
for p in self.parameters():
if not torch.nn.parameter.is_lazy(p) and p.dim() > 1:
try:
nn.init.xavier_uniform_(p)
except ValueError as e:
print(f"Failed Trying to initialize {p}")
raise (e)
def forward(
self,
ref_instances: list[Instance],
query_instances: list[Instance] | None = None,
) -> list[AssociationMatrix]:
"""Execute a forward pass through the transformer and attention head.
Args:
ref_instances: A list of instance objects (See `dreem.io.Instance` for more info.)
query_instances: An set of instances to be used as decoder queries.
Returns:
asso_output: A list of torch.Tensors of shape (L, n_query, total_instances) where:
L: number of decoder blocks
n_query: number of instances in current query/frame
total_instances: number of instances in window
"""
ref_features = torch.cat(
[instance.features for instance in ref_instances], dim=0
).unsqueeze(0)
# window_length = len(frames)
# instances_per_frame = [frame.num_detected for frame in frames]
total_instances = len(ref_instances)
embed_dim = self.d_model
# print(f'T: {window_length}; N: {total_instances}; N_t: {instances_per_frame} n_reid: {reid_features.shape}')
ref_boxes = get_boxes(ref_instances) # total_instances, 4
ref_boxes = torch.nan_to_num(ref_boxes, -1.0)
ref_times, query_times = get_times(ref_instances, query_instances)
# window_length = len(ref_times.unique()) # Currently unused but may be useful for debugging
ref_temp_emb = self.temp_emb(ref_times)
ref_pos_emb = self.pos_emb(ref_boxes)
if self.return_embedding:
for i, instance in enumerate(ref_instances):
instance.add_embedding("pos", ref_pos_emb[i])
instance.add_embedding("temp", ref_temp_emb[i])
ref_emb = (ref_pos_emb + ref_temp_emb) / 2.0
ref_emb = ref_emb.view(1, total_instances, embed_dim)
ref_emb = ref_emb.permute(1, 0, 2) # (total_instances, batch_size, embed_dim)
batch_size, total_instances = ref_features.shape[:-1]
ref_features = ref_features.permute(
1, 0, 2
) # (total_instances, batch_size, embed_dim)
encoder_queries = ref_features
# apply fourier embeddings if using fourier rope, OR if using descriptor (compact) visual encoder
if (
self.embedding_meta
and "use_fourier" in self.embedding_meta
and self.embedding_meta["use_fourier"]
) or (
self.encoder_cfg
and "encoder_type" in self.encoder_cfg
and self.encoder_cfg["encoder_type"] == "descriptor"
):
encoder_queries = apply_fourier_embeddings(
encoder_queries,
ref_times,
self.d_model,
self.fourier_embeddings,
self.fourier_proj,
self.fourier_norm,
)
encoder_features = self.encoder(
encoder_queries, pos_emb=ref_emb
) # (total_instances, batch_size, embed_dim)
n_query = total_instances
query_features = ref_features
query_pos_emb = ref_pos_emb
query_temp_emb = ref_temp_emb
query_emb = ref_emb
if query_instances is not None:
n_query = len(query_instances)
query_features = torch.cat(
[instance.features for instance in query_instances], dim=0
).unsqueeze(0)
query_features = query_features.permute(
1, 0, 2
) # (n_query, batch_size, embed_dim)
query_boxes = get_boxes(query_instances)
query_boxes = torch.nan_to_num(query_boxes, -1.0)
query_temp_emb = self.temp_emb(query_times)
query_pos_emb = self.pos_emb(query_boxes)
query_emb = (query_pos_emb + query_temp_emb) / 2.0
query_emb = query_emb.view(1, n_query, embed_dim)
query_emb = query_emb.permute(1, 0, 2) # (n_query, batch_size, embed_dim)
else:
query_instances = ref_instances
query_times = ref_times
if self.return_embedding:
for i, instance in enumerate(query_instances):
instance.add_embedding("pos", query_pos_emb[i])
instance.add_embedding("temp", query_temp_emb[i])
# apply fourier embeddings if using fourier rope, OR if using descriptor (compact) visual encoder
if (
self.embedding_meta
and "use_fourier" in self.embedding_meta
and self.embedding_meta["use_fourier"]
) or (
self.encoder_cfg
and "encoder_type" in self.encoder_cfg
and self.encoder_cfg["encoder_type"] == "descriptor"
):
query_features = apply_fourier_embeddings(
query_features,
query_times,
self.d_model,
self.fourier_embeddings,
self.fourier_proj,
self.fourier_norm,
)
decoder_features = self.decoder(
query_features,
encoder_features,
ref_pos_emb=ref_emb,
query_pos_emb=query_emb,
) # (L, n_query, batch_size, embed_dim)
decoder_features = decoder_features.transpose(
1, 2
) # # (L, batch_size, n_query, embed_dim)
encoder_features = encoder_features.permute(1, 0, 2).view(
batch_size, total_instances, embed_dim
) # (batch_size, total_instances, embed_dim)
asso_output = []
for frame_features in decoder_features:
asso_matrix = self.attn_head(frame_features, encoder_features).view(
n_query, total_instances
)
asso_matrix = AssociationMatrix(asso_matrix, ref_instances, query_instances)
asso_output.append(asso_matrix)
# (L=1, n_query, total_instances)
return asso_output
__init__(d_model=1024, nhead=8, num_encoder_layers=6, num_decoder_layers=6, dropout=0.1, activation='relu', return_intermediate_dec=False, norm=False, num_layers_attn_head=2, dropout_attn_head=0.1, embedding_meta=None, return_embedding=False, decoder_self_attn=False, encoder_cfg=None)
¶
Initialize Transformer.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
d_model
|
int
|
The number of features in the encoder/decoder inputs. |
1024
|
nhead
|
int
|
The number of heads in the transformer encoder/decoder. |
8
|
num_encoder_layers
|
int
|
The number of encoder-layers in the encoder. |
6
|
num_decoder_layers
|
int
|
The number of decoder-layers in the decoder. |
6
|
dropout
|
float
|
Dropout value applied to the output of transformer layers. |
0.1
|
activation
|
str
|
Activation function to use. |
'relu'
|
return_intermediate_dec
|
bool
|
Return intermediate layers from decoder. |
False
|
norm
|
bool
|
If True, normalize output of encoder and decoder. |
False
|
num_layers_attn_head
|
int
|
The number of layers in the attention head. |
2
|
dropout_attn_head
|
float
|
Dropout value for the attention_head. |
0.1
|
embedding_meta
|
dict | None
|
Metadata for positional embeddings. See below. |
None
|
return_embedding
|
bool
|
Whether to return the positional embeddings |
False
|
decoder_self_attn
|
bool
|
If True, use decoder self attention. |
False
|
encoder_cfg
|
dict | None
|
Encoder configuration. More details on |
None
|
Source code in dreem/models/transformer.py
def __init__(
self,
d_model: int = 1024,
nhead: int = 8,
num_encoder_layers: int = 6,
num_decoder_layers: int = 6,
dropout: float = 0.1,
activation: str = "relu",
return_intermediate_dec: bool = False,
norm: bool = False,
num_layers_attn_head: int = 2,
dropout_attn_head: float = 0.1,
embedding_meta: dict | None = None,
return_embedding: bool = False,
decoder_self_attn: bool = False,
encoder_cfg: dict | None = None,
) -> None:
"""Initialize Transformer.
Args:
d_model: The number of features in the encoder/decoder inputs.
nhead: The number of heads in the transformer encoder/decoder.
num_encoder_layers: The number of encoder-layers in the encoder.
num_decoder_layers: The number of decoder-layers in the decoder.
dropout: Dropout value applied to the output of transformer layers.
activation: Activation function to use.
return_intermediate_dec: Return intermediate layers from decoder.
norm: If True, normalize output of encoder and decoder.
num_layers_attn_head: The number of layers in the attention head.
dropout_attn_head: Dropout value for the attention_head.
embedding_meta: Metadata for positional embeddings. See below.
return_embedding: Whether to return the positional embeddings
decoder_self_attn: If True, use decoder self attention.
encoder_cfg: Encoder configuration.
More details on `embedding_meta`:
By default this will be an empty dict and indicate
that no positional embeddings should be used. To use the positional embeddings
pass in a dictionary containing a "pos" and "temp" key with subdictionaries for correct parameters ie:
{"pos": {'mode': 'learned', 'emb_num': 16, 'over_boxes: 'True'},
"temp": {'mode': 'learned', 'emb_num': 16}}. (see `dreem.models.embeddings.Embedding.EMB_TYPES`
and `dreem.models.embeddings.Embedding.EMB_MODES` for embedding parameters).
"""
super().__init__()
self.d_model = dim_feedforward = feature_dim_attn_head = d_model
self.embedding_meta = embedding_meta
self.return_embedding = return_embedding
self.encoder_cfg = encoder_cfg
self.pos_emb = Embedding(emb_type="off", mode="off", features=self.d_model)
self.temp_emb = Embedding(emb_type="off", mode="off", features=self.d_model)
if self.embedding_meta:
if "pos" in self.embedding_meta:
pos_emb_cfg = self.embedding_meta["pos"]
if pos_emb_cfg:
self.pos_emb = Embedding(
emb_type="pos", features=self.d_model, **pos_emb_cfg
)
if "temp" in self.embedding_meta:
temp_emb_cfg = self.embedding_meta["temp"]
if temp_emb_cfg:
self.temp_emb = Embedding(
emb_type="temp", features=self.d_model, **temp_emb_cfg
)
self.fourier_embeddings = FourierPositionalEmbeddings(
n_components=8, d_model=d_model
)
# Transformer Encoder
encoder_layer = TransformerEncoderLayer(
d_model, nhead, dim_feedforward, dropout, activation, norm
)
encoder_norm = nn.LayerNorm(d_model) if (norm) else None
# only used if using descriptor visual encoder; default resnet encoder uses d_model directly
if self.encoder_cfg and "encoder_type" in self.encoder_cfg:
self.visual_feat_dim = (
self.encoder_cfg["ndim"] if "ndim" in self.encoder_cfg else 5
) # 5 is default for descriptor
self.fourier_proj = nn.Linear(self.d_model + self.visual_feat_dim, d_model)
self.fourier_norm = nn.LayerNorm(self.d_model)
self.encoder = TransformerEncoder(
encoder_layer, num_encoder_layers, encoder_norm
)
# Transformer Decoder
decoder_layer = TransformerDecoderLayer(
d_model,
nhead,
dim_feedforward,
dropout,
activation,
norm,
decoder_self_attn,
)
decoder_norm = nn.LayerNorm(d_model) if (norm) else None
self.decoder = TransformerDecoder(
decoder_layer, num_decoder_layers, return_intermediate_dec, decoder_norm
)
# Transformer attention head
self.attn_head = ATTWeightHead(
feature_dim=feature_dim_attn_head,
num_layers=num_layers_attn_head,
dropout=dropout_attn_head,
)
self._reset_parameters()
forward(ref_instances, query_instances=None)
¶
Execute a forward pass through the transformer and attention head.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
ref_instances
|
list[Instance]
|
A list of instance objects (See |
required |
query_instances
|
list[Instance] | None
|
An set of instances to be used as decoder queries. |
None
|
Returns:
| Name | Type | Description |
|---|---|---|
asso_output |
list[AssociationMatrix]
|
A list of torch.Tensors of shape (L, n_query, total_instances) where: L: number of decoder blocks n_query: number of instances in current query/frame total_instances: number of instances in window |
Source code in dreem/models/transformer.py
def forward(
self,
ref_instances: list[Instance],
query_instances: list[Instance] | None = None,
) -> list[AssociationMatrix]:
"""Execute a forward pass through the transformer and attention head.
Args:
ref_instances: A list of instance objects (See `dreem.io.Instance` for more info.)
query_instances: An set of instances to be used as decoder queries.
Returns:
asso_output: A list of torch.Tensors of shape (L, n_query, total_instances) where:
L: number of decoder blocks
n_query: number of instances in current query/frame
total_instances: number of instances in window
"""
ref_features = torch.cat(
[instance.features for instance in ref_instances], dim=0
).unsqueeze(0)
# window_length = len(frames)
# instances_per_frame = [frame.num_detected for frame in frames]
total_instances = len(ref_instances)
embed_dim = self.d_model
# print(f'T: {window_length}; N: {total_instances}; N_t: {instances_per_frame} n_reid: {reid_features.shape}')
ref_boxes = get_boxes(ref_instances) # total_instances, 4
ref_boxes = torch.nan_to_num(ref_boxes, -1.0)
ref_times, query_times = get_times(ref_instances, query_instances)
# window_length = len(ref_times.unique()) # Currently unused but may be useful for debugging
ref_temp_emb = self.temp_emb(ref_times)
ref_pos_emb = self.pos_emb(ref_boxes)
if self.return_embedding:
for i, instance in enumerate(ref_instances):
instance.add_embedding("pos", ref_pos_emb[i])
instance.add_embedding("temp", ref_temp_emb[i])
ref_emb = (ref_pos_emb + ref_temp_emb) / 2.0
ref_emb = ref_emb.view(1, total_instances, embed_dim)
ref_emb = ref_emb.permute(1, 0, 2) # (total_instances, batch_size, embed_dim)
batch_size, total_instances = ref_features.shape[:-1]
ref_features = ref_features.permute(
1, 0, 2
) # (total_instances, batch_size, embed_dim)
encoder_queries = ref_features
# apply fourier embeddings if using fourier rope, OR if using descriptor (compact) visual encoder
if (
self.embedding_meta
and "use_fourier" in self.embedding_meta
and self.embedding_meta["use_fourier"]
) or (
self.encoder_cfg
and "encoder_type" in self.encoder_cfg
and self.encoder_cfg["encoder_type"] == "descriptor"
):
encoder_queries = apply_fourier_embeddings(
encoder_queries,
ref_times,
self.d_model,
self.fourier_embeddings,
self.fourier_proj,
self.fourier_norm,
)
encoder_features = self.encoder(
encoder_queries, pos_emb=ref_emb
) # (total_instances, batch_size, embed_dim)
n_query = total_instances
query_features = ref_features
query_pos_emb = ref_pos_emb
query_temp_emb = ref_temp_emb
query_emb = ref_emb
if query_instances is not None:
n_query = len(query_instances)
query_features = torch.cat(
[instance.features for instance in query_instances], dim=0
).unsqueeze(0)
query_features = query_features.permute(
1, 0, 2
) # (n_query, batch_size, embed_dim)
query_boxes = get_boxes(query_instances)
query_boxes = torch.nan_to_num(query_boxes, -1.0)
query_temp_emb = self.temp_emb(query_times)
query_pos_emb = self.pos_emb(query_boxes)
query_emb = (query_pos_emb + query_temp_emb) / 2.0
query_emb = query_emb.view(1, n_query, embed_dim)
query_emb = query_emb.permute(1, 0, 2) # (n_query, batch_size, embed_dim)
else:
query_instances = ref_instances
query_times = ref_times
if self.return_embedding:
for i, instance in enumerate(query_instances):
instance.add_embedding("pos", query_pos_emb[i])
instance.add_embedding("temp", query_temp_emb[i])
# apply fourier embeddings if using fourier rope, OR if using descriptor (compact) visual encoder
if (
self.embedding_meta
and "use_fourier" in self.embedding_meta
and self.embedding_meta["use_fourier"]
) or (
self.encoder_cfg
and "encoder_type" in self.encoder_cfg
and self.encoder_cfg["encoder_type"] == "descriptor"
):
query_features = apply_fourier_embeddings(
query_features,
query_times,
self.d_model,
self.fourier_embeddings,
self.fourier_proj,
self.fourier_norm,
)
decoder_features = self.decoder(
query_features,
encoder_features,
ref_pos_emb=ref_emb,
query_pos_emb=query_emb,
) # (L, n_query, batch_size, embed_dim)
decoder_features = decoder_features.transpose(
1, 2
) # # (L, batch_size, n_query, embed_dim)
encoder_features = encoder_features.permute(1, 0, 2).view(
batch_size, total_instances, embed_dim
) # (batch_size, total_instances, embed_dim)
asso_output = []
for frame_features in decoder_features:
asso_matrix = self.attn_head(frame_features, encoder_features).view(
n_query, total_instances
)
asso_matrix = AssociationMatrix(asso_matrix, ref_instances, query_instances)
asso_output.append(asso_matrix)
# (L=1, n_query, total_instances)
return asso_output
VisualEncoder
¶
Bases: Module
Class wrapping around a visual feature extractor backbone.
Currently CNN only.
Methods:
| Name | Description |
|---|---|
__init__ |
Initialize Visual Encoder. |
encoder_dim |
Compute dummy forward pass of encoder model and get embedding dimension. |
forward |
Forward pass of feature extractor to get feature vector. |
select_feature_extractor |
Select the appropriate feature extractor based on config. |
Source code in dreem/models/visual_encoder.py
class VisualEncoder(torch.nn.Module):
"""Class wrapping around a visual feature extractor backbone.
Currently CNN only.
"""
def __init__(
self,
model_name: str = "resnet18",
d_model: int = 512,
in_chans: int = 3,
backend: int = "timm",
**kwargs: Any | None,
):
"""Initialize Visual Encoder.
Args:
model_name (str): Name of the CNN architecture to use (e.g. "resnet18", "resnet50").
d_model (int): Output embedding dimension.
in_chans: the number of input channels of the image.
backend: Which model backend to use. One of {"timm", "torchvision"}
kwargs: see `timm.create_model` and `torchvision.models.resnetX` for kwargs.
"""
super().__init__()
self.model_name = model_name.lower()
self.d_model = d_model
self.backend = backend
if in_chans == 1:
self.in_chans = 3
else:
self.in_chans = in_chans
self.feature_extractor = self.select_feature_extractor(
model_name=self.model_name,
in_chans=self.in_chans,
backend=self.backend,
**kwargs,
)
self.out_layer = torch.nn.Linear(
self.encoder_dim(self.feature_extractor), self.d_model
)
def select_feature_extractor(
self, model_name: str, in_chans: int, backend: str, **kwargs: Any
) -> torch.nn.Module:
"""Select the appropriate feature extractor based on config.
Args:
model_name (str): Name of the CNN architecture to use (e.g. "resnet18", "resnet50").
in_chans: the number of input channels of the image.
backend: Which model backend to use. One of {"timm", "torchvision"}
kwargs: see `timm.create_model` and `torchvision.models.resnetX` for kwargs.
Returns:
a CNN encoder based on the config and backend selected.
"""
if "timm" in backend.lower():
feature_extractor = timm.create_model(
model_name=self.model_name,
in_chans=self.in_chans,
num_classes=0,
**kwargs,
)
elif "torch" in backend.lower():
if model_name.lower() == "resnet18":
feature_extractor = torchvision.models.resnet18(**kwargs)
elif model_name.lower() == "resnet50":
feature_extractor = torchvision.models.resnet50(**kwargs)
else:
raise ValueError(
f"Only `[resnet18, resnet50]` are available when backend is {backend}. Found {model_name}"
)
feature_extractor = torch.nn.Sequential(
*list(feature_extractor.children())[:-1]
)
input_layer = feature_extractor[0]
if in_chans != 3:
feature_extractor[0] = torch.nn.Conv2d(
in_channels=in_chans,
out_channels=input_layer.out_channels,
kernel_size=input_layer.kernel_size,
stride=input_layer.stride,
padding=input_layer.padding,
dilation=input_layer.dilation,
groups=input_layer.groups,
bias=input_layer.bias,
padding_mode=input_layer.padding_mode,
)
else:
raise ValueError(
f"Only ['timm', 'torch'] backends are available! Found {backend}."
)
return feature_extractor
def encoder_dim(self, model: torch.nn.Module) -> int:
"""Compute dummy forward pass of encoder model and get embedding dimension.
Args:
model: a vision encoder model.
Returns:
The embedding dimension size.
"""
_ = model.eval()
dummy_output = model(torch.randn(1, self.in_chans, 224, 224)).squeeze()
_ = model.train() # to be safe
return dummy_output.shape[-1]
def forward(self, img: torch.Tensor) -> torch.Tensor:
"""Forward pass of feature extractor to get feature vector.
Args:
img: Input image tensor of shape (B, C, H, W).
Returns:
feats: Normalized output tensor of shape (B, d_model).
"""
# If grayscale, tile the image to 3 channels.
if img.shape[1] == 1:
img = img.repeat([1, 3, 1, 1]) # (B, nc=3, H, W)
b, c, h, w = img.shape
if c != self.in_chans:
raise ValueError(
f"""Found {c} channels in image but model was configured for {self.in_chans} channels! \n
Hint: have you set the number of anchors in your dataset > 1? \n
If so, make sure to set `in_chans=3 * n_anchors`"""
)
feats = self.feature_extractor(
img
) # (B, out_dim, 1, 1) if using resnet18 backbone.
# Reshape feature vectors
feats = feats.reshape([img.shape[0], -1]) # (B, out_dim)
# Map feature vectors to output dimension using linear layer.
feats = self.out_layer(feats) # (B, d_model)
# Normalize output feature vectors.
feats = F.normalize(feats) # (B, d_model)
return feats
__init__(model_name='resnet18', d_model=512, in_chans=3, backend='timm', **kwargs)
¶
Initialize Visual Encoder.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
model_name
|
str
|
Name of the CNN architecture to use (e.g. "resnet18", "resnet50"). |
'resnet18'
|
d_model
|
int
|
Output embedding dimension. |
512
|
in_chans
|
int
|
the number of input channels of the image. |
3
|
backend
|
int
|
Which model backend to use. One of {"timm", "torchvision"} |
'timm'
|
kwargs
|
Any | None
|
see |
{}
|
Source code in dreem/models/visual_encoder.py
def __init__(
self,
model_name: str = "resnet18",
d_model: int = 512,
in_chans: int = 3,
backend: int = "timm",
**kwargs: Any | None,
):
"""Initialize Visual Encoder.
Args:
model_name (str): Name of the CNN architecture to use (e.g. "resnet18", "resnet50").
d_model (int): Output embedding dimension.
in_chans: the number of input channels of the image.
backend: Which model backend to use. One of {"timm", "torchvision"}
kwargs: see `timm.create_model` and `torchvision.models.resnetX` for kwargs.
"""
super().__init__()
self.model_name = model_name.lower()
self.d_model = d_model
self.backend = backend
if in_chans == 1:
self.in_chans = 3
else:
self.in_chans = in_chans
self.feature_extractor = self.select_feature_extractor(
model_name=self.model_name,
in_chans=self.in_chans,
backend=self.backend,
**kwargs,
)
self.out_layer = torch.nn.Linear(
self.encoder_dim(self.feature_extractor), self.d_model
)
encoder_dim(model)
¶
Compute dummy forward pass of encoder model and get embedding dimension.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
model
|
Module
|
a vision encoder model. |
required |
Returns:
| Type | Description |
|---|---|
int
|
The embedding dimension size. |
Source code in dreem/models/visual_encoder.py
def encoder_dim(self, model: torch.nn.Module) -> int:
"""Compute dummy forward pass of encoder model and get embedding dimension.
Args:
model: a vision encoder model.
Returns:
The embedding dimension size.
"""
_ = model.eval()
dummy_output = model(torch.randn(1, self.in_chans, 224, 224)).squeeze()
_ = model.train() # to be safe
return dummy_output.shape[-1]
forward(img)
¶
Forward pass of feature extractor to get feature vector.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
img
|
Tensor
|
Input image tensor of shape (B, C, H, W). |
required |
Returns:
| Name | Type | Description |
|---|---|---|
feats |
Tensor
|
Normalized output tensor of shape (B, d_model). |
Source code in dreem/models/visual_encoder.py
def forward(self, img: torch.Tensor) -> torch.Tensor:
"""Forward pass of feature extractor to get feature vector.
Args:
img: Input image tensor of shape (B, C, H, W).
Returns:
feats: Normalized output tensor of shape (B, d_model).
"""
# If grayscale, tile the image to 3 channels.
if img.shape[1] == 1:
img = img.repeat([1, 3, 1, 1]) # (B, nc=3, H, W)
b, c, h, w = img.shape
if c != self.in_chans:
raise ValueError(
f"""Found {c} channels in image but model was configured for {self.in_chans} channels! \n
Hint: have you set the number of anchors in your dataset > 1? \n
If so, make sure to set `in_chans=3 * n_anchors`"""
)
feats = self.feature_extractor(
img
) # (B, out_dim, 1, 1) if using resnet18 backbone.
# Reshape feature vectors
feats = feats.reshape([img.shape[0], -1]) # (B, out_dim)
# Map feature vectors to output dimension using linear layer.
feats = self.out_layer(feats) # (B, d_model)
# Normalize output feature vectors.
feats = F.normalize(feats) # (B, d_model)
return feats
select_feature_extractor(model_name, in_chans, backend, **kwargs)
¶
Select the appropriate feature extractor based on config.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
model_name
|
str
|
Name of the CNN architecture to use (e.g. "resnet18", "resnet50"). |
required |
in_chans
|
int
|
the number of input channels of the image. |
required |
backend
|
str
|
Which model backend to use. One of {"timm", "torchvision"} |
required |
kwargs
|
Any
|
see |
{}
|
Returns:
| Type | Description |
|---|---|
Module
|
a CNN encoder based on the config and backend selected. |
Source code in dreem/models/visual_encoder.py
def select_feature_extractor(
self, model_name: str, in_chans: int, backend: str, **kwargs: Any
) -> torch.nn.Module:
"""Select the appropriate feature extractor based on config.
Args:
model_name (str): Name of the CNN architecture to use (e.g. "resnet18", "resnet50").
in_chans: the number of input channels of the image.
backend: Which model backend to use. One of {"timm", "torchvision"}
kwargs: see `timm.create_model` and `torchvision.models.resnetX` for kwargs.
Returns:
a CNN encoder based on the config and backend selected.
"""
if "timm" in backend.lower():
feature_extractor = timm.create_model(
model_name=self.model_name,
in_chans=self.in_chans,
num_classes=0,
**kwargs,
)
elif "torch" in backend.lower():
if model_name.lower() == "resnet18":
feature_extractor = torchvision.models.resnet18(**kwargs)
elif model_name.lower() == "resnet50":
feature_extractor = torchvision.models.resnet50(**kwargs)
else:
raise ValueError(
f"Only `[resnet18, resnet50]` are available when backend is {backend}. Found {model_name}"
)
feature_extractor = torch.nn.Sequential(
*list(feature_extractor.children())[:-1]
)
input_layer = feature_extractor[0]
if in_chans != 3:
feature_extractor[0] = torch.nn.Conv2d(
in_channels=in_chans,
out_channels=input_layer.out_channels,
kernel_size=input_layer.kernel_size,
stride=input_layer.stride,
padding=input_layer.padding,
dilation=input_layer.dilation,
groups=input_layer.groups,
bias=input_layer.bias,
padding_mode=input_layer.padding_mode,
)
else:
raise ValueError(
f"Only ['timm', 'torch'] backends are available! Found {backend}."
)
return feature_extractor
annotate_video(video, labels, key, color_palette=palette, trails=2, boxes=(64, 64), names=True, track_scores=0.5, centroids=4, poses=False, save_path='debug_animal.mp4', fps=30, alpha=0.2)
¶
Annotate video frames with labels.
Labels video with bboxes, centroids, trajectory trails, and/or poses.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
video
|
Reader
|
The video to be annotated in an ndarray |
required |
labels
|
DataFrame
|
The pandas dataframe containing the centroid and/or pose locations of the instances |
required |
key
|
str
|
The key where labels are stored in the dataframe - mostly used for choosing whether to annotate based on pred or gt labels |
required |
color_palette
|
list | str
|
The matplotlib colorpalette to use for annotating the video. Defaults to |
palette
|
trails
|
int
|
The size of the trajectory trail. If trails size <= 0 or None then it is not added |
2
|
boxes
|
int
|
The size of the bbox. If bbox size <= 0 or None then it is not added |
(64, 64)
|
names
|
bool
|
Whether or not to annotate with name |
True
|
centroids
|
int
|
The size of the centroid. If centroid size <= 0 or None then it is not added |
4
|
poses
|
bool
|
Whether or not to annotate with poses |
False
|
save_path
|
str
|
The path to save the annotated video. |
'debug_animal.mp4'
|
fps
|
int
|
The frame rate of the generated video |
30
|
track_scores
|
Minimum track score threshold for displaying tracks |
0.5
|
|
alpha
|
float
|
The opacity of the annotations. |
0.2
|
Returns:
| Type | Description |
|---|---|
list
|
A list of annotated video frames |
Source code in dreem/io/visualize.py
def annotate_video(
video: "imageio.core.format.Reader",
labels: pd.DataFrame,
key: str,
color_palette: list | str = palette,
trails: int = 2,
boxes: int = (64, 64),
names: bool = True,
track_scores=0.5,
centroids: int = 4,
poses: bool = False,
save_path: str = "debug_animal.mp4",
fps: int = 30,
alpha: float = 0.2,
) -> list:
"""Annotate video frames with labels.
Labels video with bboxes, centroids, trajectory trails, and/or poses.
Args:
video: The video to be annotated in an ndarray
labels: The pandas dataframe containing the centroid and/or pose locations of the instances
key: The key where labels are stored in the dataframe - mostly used for choosing whether to annotate based on pred or gt labels
color_palette: The matplotlib colorpalette to use for annotating the video. Defaults to `tab10`
trails: The size of the trajectory trail. If trails size <= 0 or None then it is not added
boxes: The size of the bbox. If bbox size <= 0 or None then it is not added
names: Whether or not to annotate with name
centroids: The size of the centroid. If centroid size <= 0 or None then it is not added
poses: Whether or not to annotate with poses
save_path: The path to save the annotated video.
fps: The frame rate of the generated video
track_scores: Minimum track score threshold for displaying tracks
alpha: The opacity of the annotations.
Returns:
A list of annotated video frames
"""
try:
import cv2
except ImportError as e:
raise ImportError(
"opencv-python is required for video annotation. "
"Install with: pip install opencv-python"
) from e
writer = imageio.get_writer(save_path, fps=fps)
color_palette = (
sns.color_palette(color_palette)
if isinstance(color_palette, str)
else deepcopy(color_palette)
)
if trails:
track_trails = {}
try:
for i in tqdm(sorted(labels["Frame"].unique()), desc="Frame", unit="Frame"):
frame = video.get_data(i)
if frame.shape[0] == 1 or frame.shape[-1] == 1:
frame = cv2.cvtColor(frame, cv2.COLOR_GRAY2RGB)
# else:
# frame = frame.copy()
lf = labels[labels["Frame"] == i]
for idx, instance in lf.iterrows():
if not trails:
track_trails = {}
if poses:
# TODO figure out best way to store poses (maybe pass a slp labels file too?)
trails = False
centroids = False
for idx, (pose, edge) in enumerate(
zip(instance["poses"], instance["edges"])
):
pose = fill_missing(pose.numpy())
pred_track_id = instance[key][idx].numpy().tolist()
# Add midpt to track trail.
if pred_track_id not in list(track_trails.keys()):
track_trails[pred_track_id] = []
# Select a color based on track_id.
track_color_idx = pred_track_id % len(color_palette)
track_color = (
(np.array(color_palette[track_color_idx]) * 255)
.astype(np.uint8)
.tolist()[::-1]
)
for p in pose:
# try:
# p = tuple([int(i) for i in p.numpy()][::-1])
# except:
# continue
p = tuple(int(i) for i in p)[::-1]
track_trails[pred_track_id].append(p)
frame = cv2.circle(
frame, p, radius=2, color=track_color, thickness=-1
)
for e in edge:
source = tuple(int(i) for i in pose[int(e[0])])[::-1]
target = tuple(int(i) for i in pose[int(e[1])])[::-1]
frame = cv2.line(frame, source, target, track_color, 1)
if (boxes) or centroids:
# Get coordinates for detected objects in the current frame.
if isinstance(boxes, int):
boxes = (boxes, boxes)
box_w, box_h = boxes
x = instance["X"]
y = instance["Y"]
min_x, min_y, max_x, max_y = (
int(x - box_w / 2),
int(y - box_h / 2),
int(x + box_w / 2),
int(y + box_h / 2),
)
midpt = (int(x), int(y))
pred_track_id = instance[key]
if "Track_score" in instance.index:
track_score = instance["Track_score"]
else:
track_scores = 0
# Add midpt to track trail.
if pred_track_id not in list(track_trails.keys()):
track_trails[pred_track_id] = []
track_trails[pred_track_id].append(midpt)
# Select a color based on track_id.
track_color_idx = int(pred_track_id) % len(color_palette)
track_color = (
(np.array(color_palette[track_color_idx]) * 255)
.astype(np.uint8)
.tolist()[::-1]
)
# Bbox.
if boxes is not None:
frame = cv2.rectangle(
frame,
(min_x, min_y),
(max_x, max_y),
color=track_color,
thickness=2,
)
# Track trail.
if centroids:
frame = cv2.circle(
frame,
midpt,
radius=centroids,
color=track_color,
thickness=-1,
)
for i in range(0, len(track_trails[pred_track_id]) - 1):
frame = cv2.addWeighted(
cv2.circle(
frame, # .copy(),
track_trails[pred_track_id][i],
radius=4,
color=track_color,
thickness=-1,
),
alpha,
frame,
1 - alpha,
0,
)
if trails:
frame = cv2.line(
frame,
track_trails[pred_track_id][i],
track_trails[pred_track_id][i + 1],
color=track_color,
thickness=trails,
)
# Track name.
name_str = ""
if names:
name_str += f"track_{pred_track_id}"
if names and track_scores:
name_str += " | "
if track_scores:
name_str += f"score: {track_score:0.3f}"
if len(name_str) > 0:
frame = cv2.putText(
frame,
# f"idx:{idx} | track_{pred_track_id}",
name_str,
org=(int(min_x), max(0, int(min_y) - 10)),
fontFace=cv2.FONT_HERSHEY_SIMPLEX,
fontScale=0.9,
color=track_color,
thickness=2,
)
writer.append_data(frame)
# if i % fps == 0:
# gc.collect()
except Exception as e:
writer.close()
logger.exception(e)
return False
writer.close()
return True
setup_logging()
¶
Setup logging based on logging.yaml.
Source code in dreem/__init__.py
def setup_logging():
"""Setup logging based on `logging.yaml`."""
import logging
import logging.config
import os
import yaml
package_directory = os.path.dirname(os.path.abspath(__file__))
with open(os.path.join(package_directory, "..", "logging.yaml"), "r") as stream:
logging_cfg = yaml.load(stream, Loader=yaml.FullLoader)
logging.config.dictConfig(logging_cfg)