import abc
import dataclasses
from typing import Dict, Optional
from typing import Any, Dict, List, Optional, Tuple, Union
import pluggy
import torch
from torch import nn
from torchtyping import TensorType

PluggyHookRelay = pluggy._hooks._HookRelay  # Type alias for more readable function signatures

ConditioningOutput = TensorType["batch_size", "n_objects", "object_dim"]  # noqa: F821


class Conditioning(nn.Module, metaclass=abc.ABCMeta):
    """Base class for conditioning perceptual grouping."""

    @abc.abstractmethod
    def forward(self, *args) -> ConditioningOutput:
        pass


@dataclasses.dataclass
class FrameFeatures:
    """Features associated with a single frame."""

    features: TensorType["batch_size", "n_spatial_features", "feature_dim"]  # noqa: F821
    positions: TensorType["n_spatial_features", "spatial_dims"]  # noqa: F821


@dataclasses.dataclass
class FeatureExtractorOutput:
    """Output of feature extractor."""

    features: TensorType["batch_size", "frames", "n_spatial_features", "feature_dim"]  # noqa: F821
    positions: TensorType["n_spatial_features", "spatial_dims"]  # noqa: F821
    aux_features: Optional[Dict[str, torch.Tensor]] = None

    def __iter__(self):
        """Iterate over features and positions per frame."""
        for frame_features in torch.split(self.features, 1, dim=1):
            yield FrameFeatures(frame_features.squeeze(1), self.positions)


class FeatureExtractor(nn.Module, metaclass=abc.ABCMeta):
    """Abstract base class for Feature Extractors.

    We expect that the forward method returns a flattened representation of the features, to make
    outputs consistent and not dependent on equal spacing or the dimensionality of the spatial
    information.
    """

    @property
    @abc.abstractmethod
    def feature_dim(self):
        """Get dimensionality of the features.

        Returns:
            int: The dimensionality of the features.
        """

    @abc.abstractmethod
    def forward(self, inputs: torch.Tensor) -> FeatureExtractorOutput:
        pass


@dataclasses.dataclass
class PerceptualGroupingOutput:
    """Output of a perceptual grouping algorithm."""

    objects: TensorType["batch_size", "n_objects", "object_dim"]  # noqa: F821
    is_empty: Optional[TensorType["batch_size", "n_objects"]] = None  # noqa: F821
    feature_attributions: Optional[
        TensorType["batch_size", "n_objects", "n_spatial_features"]  # noqa: F821
    ] = None
    value: Optional[
        TensorType["batch_size", "n_objects", "extra_dim", "n_spatial_features"]  # noqa: F821
    ] = None


class PerceptualGrouping(nn.Module, metaclass=abc.ABCMeta):
    """Abstract base class of a perceptual grouping algorithm."""

    @abc.abstractmethod
    def forward(self, extracted_features: FeatureExtractorOutput) -> PerceptualGroupingOutput:
        pass

    @property
    @abc.abstractmethod
    def object_dim(self):
        pass

class Instances:
    """Modified from Detectron2 (https://github.com/facebookresearch/detectron2).

    This class represents a list of instances in an image.
    It stores the attributes of instances (e.g., boxes, masks, labels, scores) as "fields".
    All fields must have the same ``__len__`` which is the number of instances.

    All other (non-field) attributes of this class are considered private:
    they must start with '_' and are not modifiable by a user.

    Some basic usage:

    1. Set/get/check a field:

       .. code-block:: python

          instances.gt_boxes = Boxes(...)
          print(instances.pred_masks)  # a tensor of shape (N, H, W)
          print('gt_masks' in instances)

    2. ``len(instances)`` returns the number of instances
    3. Indexing: ``instances[indices]`` will apply the indexing on all the fields
       and returns a new :class:`Instances`.
       Typically, ``indices`` is a integer vector of indices,
       or a binary mask of length ``num_instances``

       .. code-block:: python

          category_3_detections = instances[instances.pred_classes == 3]
          confident_detections = instances[instances.scores > 0.9]
    """

    def __init__(self, image_size: Tuple[int, int], **kwargs: Any):
        """Init function.

        Args:
            image_size (height, width): the spatial size of the image.
            kwargs: fields to add to this `Instances`.
        """
        self._image_size = image_size
        self._fields: Dict[str, Any] = {}
        for k, v in kwargs.items():
            self.set(k, v)

    @property
    def image_size(self) -> Tuple[int, int]:
        return self._image_size

    def __setattr__(self, name: str, val: Any) -> None:
        if name.startswith("_"):
            super().__setattr__(name, val)
        else:
            self.set(name, val)

    def __getattr__(self, name: str) -> Any:
        if name == "_fields" or name not in self._fields:
            raise AttributeError("Cannot find field '{}' in the given Instances!".format(name))
        return self._fields[name]

    def set(self, name: str, value: Any) -> None:
        """Set the field named `name` to `value`.

        The length of `value` must be the number of instances,
        and must agree with other existing fields in this object.
        """
        data_len = len(value)
        if len(self._fields):
            assert (
                len(self) == data_len
            ), "Adding a field of length {} to a Instances of length {}".format(data_len, len(self))
        self._fields[name] = value

    def has(self, name: str) -> bool:
        """Returns whether the field called `name` exists."""
        return name in self._fields

    def remove(self, name: str) -> None:
        """Remove the field called `name`."""
        del self._fields[name]

    def get(self, name: str) -> Any:
        """Returns the field called `name`."""
        return self._fields[name]

    def get_fields(self) -> Dict[str, Any]:
        """Get field.

        Returns:
            dict: a dict which maps names (str) to data of the fields

        Modifying the returned dict will modify this instance.
        """
        return self._fields

    # Tensor-like methods
    def to(self, *args: Any, **kwargs: Any) -> "Instances":
        """To device.

        Returns:
            Instances: all fields are called with a `to(device)`, if the field has this method.
        """
        ret = Instances(self._image_size)
        for k, v in self._fields.items():
            if hasattr(v, "to"):
                v = v.to(*args, **kwargs)
            ret.set(k, v)
        return ret

    def numpy(self):
        ret = Instances(self._image_size)
        for k, v in self._fields.items():
            if hasattr(v, "numpy"):
                v = v.numpy()
            ret.set(k, v)
        return ret

    def __getitem__(self, item: Union[int, slice, torch.BoolTensor]) -> "Instances":
        """Get entry.

        Args:
            item: an index-like object and will be used to index all the fields.

        Returns:
            If `item` is a string, return the data in the corresponding field.
            Otherwise, returns an `Instances` where all fields are indexed by `item`.
        """
        if type(item) == int:
            if item >= len(self) or item < -len(self):
                raise IndexError("Instances index out of range!")
            else:
                item = slice(item, None, len(self))

        ret = Instances(self._image_size)
        for k, v in self._fields.items():
            ret.set(k, v[item])
        return ret

    def __len__(self) -> int:
        for v in self._fields.values():
            # use __len__ because len() has to be int and is not friendly to tracing
            return v.__len__()
        raise NotImplementedError("Empty Instances does not support __len__!")

    def __iter__(self):
        raise NotImplementedError("`Instances` object is not iterable!")

    @staticmethod
    def cat(instance_lists: List["Instances"]) -> "Instances":
        """Concatenate instances.

        Args:
            instance_lists (list[Instances])

        Returns:
            Instances
        """
        assert all(isinstance(i, Instances) for i in instance_lists)
        assert len(instance_lists) > 0
        if len(instance_lists) == 1:
            return instance_lists[0]

        image_size = instance_lists[0].image_size
        for i in instance_lists[1:]:
            assert i.image_size == image_size
        ret = Instances(image_size)
        for k in instance_lists[0]._fields.keys():
            values = [i.get(k) for i in instance_lists]
            v0 = values[0]
            if isinstance(v0, torch.Tensor):
                values = torch.cat(values, dim=0)
            elif isinstance(v0, list):
                values = list(itertools.chain(*values))
            elif hasattr(type(v0), "cat"):
                values = type(v0).cat(values)
            else:
                raise ValueError("Unsupported type {} for concatenation".format(type(v0)))
            ret.set(k, values)
        return ret

    def __str__(self) -> str:
        s = self.__class__.__name__ + "("
        s += "num_instances={}, ".format(len(self))
        s += "image_height={}, ".format(self._image_size[0])
        s += "image_width={}, ".format(self._image_size[1])
        s += "fields=[{}])".format(", ".join((f"{k}: {v}" for k, v in self._fields.items())))
        return s

    __repr__ = __str__
ocl/base.py
amazon-science/object-centric-multiple-object-tracking
import glob
import os

from hydra.core.hydra_config import HydraConfig


def get_commandline_config_path():
    """Get the path of a config path specified on the command line."""
    hydra_cfg = HydraConfig.get()
    config_sources = hydra_cfg.runtime.config_sources
    config_path = None
    for source in config_sources:
        if source.schema == "file" and source.provider == "command-line":
            config_path = source.path
            break
    return config_path


def find_checkpoint(path):
    """Find checkpoint in output path of previous run."""
    checkpoints = glob.glob(
        os.path.join(path, "lightning_logs", "version_*", "checkpoints", "*.ckpt")
    )
    checkpoints.sort()
    # Return the last checkpoint.
    # TODO (hornmax): If more than one checkpoint is stored this might not lead to the most recent
    # checkpoint being loaded. Generally, I think this is ok as we still allow people to set the
    # checkpoint manually.
    return checkpoints[-1]

ocl/cli/cli_utils.py
amazon-science/object-centric-multiple-object-tracking
"""Script to compute the size of a dataset.

This is useful when subsampling data using transformations in order to determine the final dataset
size.  The size of the dataset is typically need when running distributed training in order to
ensure that all nodes and gpu training processes are presented with the same number of batches.
"""
import dataclasses
import logging
import os
from typing import Dict

import hydra
import hydra_zen
import tqdm
from pluggy import PluginManager

import ocl.hooks
from ocl.config.datasets import DataModuleConfig
from ocl.config.plugins import PluginConfig


@dataclasses.dataclass
class ComputeSizeConfig:
    """Configuration of a training run."""

    dataset: DataModuleConfig
    plugins: Dict[str, PluginConfig] = dataclasses.field(default_factory=dict)


hydra.core.config_store.ConfigStore.instance().store(
    name="compute_size_config",
    node=ComputeSizeConfig,
)


@hydra.main(config_name="compute_size_config", config_path="../../configs", version_base="1.1")
def compute_size(config: ComputeSizeConfig):
    pm = PluginManager("ocl")
    pm.add_hookspecs(ocl.hooks)

    datamodule = hydra_zen.instantiate(config.dataset, hooks=pm.hook)
    pm.register(datamodule)

    plugins = hydra_zen.instantiate(config.plugins)
    for plugin in plugins.values():
        pm.register(plugin)

    # Compute dataset sizes
    # TODO(hornmax): This is needed for webdataset shuffling, is there a way to make this more
    # elegant and less specific?
    os.environ["WDS_EPOCH"] = str(0)
    train_size = sum(
        1
        for _ in tqdm.tqdm(
            datamodule.train_data_iterator(), desc="Reading train split", unit="samples"
        )
    )
    logging.info("Train split size: %d", train_size)
    val_size = sum(
        1
        for _ in tqdm.tqdm(
            datamodule.val_data_iterator(), desc="Reading validation split", unit="samples"
        )
    )
    logging.info("Validation split size: %d", val_size)
    test_size = sum(
        1
        for _ in tqdm.tqdm(
            datamodule.test_data_iterator(), desc="Reading test split", unit="samples"
        )
    )
    logging.info("Test split size: %d", test_size)


if __name__ == "__main__":
    compute_size()

ocl/cli/compute_dataset_size.py
amazon-science/object-centric-multiple-object-tracking
"""Train a slot attention type model."""
import dataclasses
from typing import Any, Dict, Optional
from pytorch_lightning.utilities.cloud_io import load as pl_load
import hydra
import hydra_zen
import pytorch_lightning as pl
from pytorch_lightning.callbacks import ModelCheckpoint
from pluggy import PluginManager
import torch
import ocl.hooks
from ocl import base
from ocl.combined_model import CombinedModel
from ocl.config.datasets import DataModuleConfig
from ocl.config.metrics import MetricConfig
from ocl.config.plugins import PluginConfig
from ocl.plugins import Plugin
from ocl.utils import Combined, Recurrent, RoutableMixin
from ocl.cli import cli_utils, eval_utils, train
from typing import Union

TrainerConf = hydra_zen.builds(
    pl.Trainer, max_epochs=100, zen_partial=False, populate_full_signature=True
)


@dataclasses.dataclass
class TrainingConfig:
    """Configuration of a training run."""

    dataset: DataModuleConfig
    models: Any  # When provided with dict wrap in `utils.Combined`, otherwise interpret as model.
    losses: Dict[str, Any]
    visualizations: Dict[str, Any] = dataclasses.field(default_factory=dict)
    plugins: Dict[str, PluginConfig] = dataclasses.field(default_factory=dict)
    trainer: TrainerConf = TrainerConf
    training_vis_frequency: Optional[int] = None
    training_metrics: Optional[Dict[str, MetricConfig]] = None
    evaluation_metrics: Optional[Dict[str, MetricConfig]] = None
    load_checkpoint: Optional[str] = None
    seed: Optional[int] = None
    experiment: Optional[Any] = None
    root_output_folder: Optional[str] = None


hydra.core.config_store.ConfigStore.instance().store(
    name="training_config",
    node=TrainingConfig,
)


def create_plugin_manager() -> PluginManager:
    pm = PluginManager("ocl")
    pm.add_hookspecs(ocl.hooks)
    return pm


def build_and_register_datamodule_from_config(
        config: TrainingConfig,
        hooks: base.PluggyHookRelay,
        plugin_manager: Optional[PluginManager] = None,
        **datamodule_kwargs,
) -> pl.LightningDataModule:
    datamodule = hydra_zen.instantiate(
        config.dataset, hooks=hooks, _convert_="all", **datamodule_kwargs
    )

    if plugin_manager:
        plugin_manager.register(datamodule)

    return datamodule


def build_and_register_plugins_from_config(
        config: TrainingConfig, plugin_manager: Optional[PluginManager] = None
) -> Dict[str, Plugin]:
    plugins = hydra_zen.instantiate(config.plugins)
    # Use lexicographical sorting to allow to influence registration order. This is necessary in
    # some cases as certain plugins might need to be called before others. Pluggy calls hooks
    # according to FILO (first in last out) and this is slightly unintuitive. We thus register
    # plugins in reverse order to their sorting position, leading to a FIFO (first in first out)
    # behavior with regard to the sorted position.
    if plugin_manager:
        for plugin_name in sorted(plugins.keys())[::-1]:
            plugin_manager.register(plugins[plugin_name])

    return plugins


def build_model_from_config(
        config: TrainingConfig,
        hooks: base.PluggyHookRelay,
        checkpoint_path: Optional[str] = None,
) -> pl.LightningModule:
    models = hydra_zen.instantiate(config.models, _convert_="all")
    losses = hydra_zen.instantiate(config.losses, _convert_="all")
    visualizations = hydra_zen.instantiate(config.visualizations, _convert_="all")

    training_metrics = hydra_zen.instantiate(config.training_metrics)
    evaluation_metrics = hydra_zen.instantiate(config.evaluation_metrics)

    train_vis_freq = config.training_vis_frequency if config.training_vis_frequency else 100

    if checkpoint_path is None:
        model = CombinedModel(
            models=models,
            losses=losses,
            visualizations=visualizations,
            hooks=hooks,
            training_metrics=training_metrics,
            evaluation_metrics=evaluation_metrics,
            vis_log_frequency=train_vis_freq,
        )
    else:
        model = CombinedModel.load_from_checkpoint(
            checkpoint_path,
            strict=False,
            models=models,
            losses=losses,
            visualizations=visualizations,
            hooks=hooks,
            training_metrics=training_metrics,
            evaluation_metrics=evaluation_metrics,
            vis_log_frequency=train_vis_freq,
        )

    return model


@hydra.main(config_name="training_config", config_path="../../configs/", version_base="1.1")
def train(config: TrainingConfig):
    # Set all relevant random seeds. If `config.seed is None`, the function samples a random value.
    # The function takes care of correctly distributing the seed across nodes in multi-node training,
    # and assigns each dataloader worker a different random seed.
    # IMPORTANTLY, we need to take care not to set a custom `worker_init_fn` function on the
    # dataloaders (or take care of worker seeding ourselves).
    pl.seed_everything(config.seed, workers=True)

    pm = create_plugin_manager()

    # datamodule = build_and_register_datamodule_from_config(config, pm.hook, pm)
    checkpoint_path = hydra.utils.to_absolute_path(config.load_checkpoint)
    datamodule, model, pm = eval_utils.build_from_train_config(
        config, checkpoint_path
    )

    trainer: pl.Trainer = hydra_zen.instantiate(
        config.trainer,
        _convert_="all",
        enable_progress_bar=True,
        gpus=[0],
    )

    print("******start validate model******")
    trainer.validate(model, datamodule.val_dataloader())


if __name__ == "__main__":
    train()


ocl/cli/eval.py
amazon-science/object-centric-multiple-object-tracking
#!/usr/bin/env python
"""Evaluate a trained model for object discovery by clustering object representations.

Given a set of images, each with a set of ground truth masks and a set of object masks and
representations, we perform the following steps:
    1) Assign each object a cluster id by clustering the corresponding representations over all
    images.
    2) Merge object masks with the same cluster id on the same image to form a semantic mask.
    3) Compute IoU between masks of predicted clusters and ground truth classes over all images.
    4) Assign clusters to classes based on the IoU and a matching strategy.
"""
import dataclasses
import enum
import json
import logging
import os
from collections import OrderedDict, defaultdict
from functools import partial
from typing import Any, Dict, List, Optional

import hydra
import hydra_zen
import numpy as np
import omegaconf
import pytorch_lightning as pl
import torch
import tqdm

from ocl import feature_extractors, metrics, path_defaults, tree_utils
from ocl.cli import cli_utils, eval_utils, train
from ocl.config.datasets import DataModuleConfig
from ocl.config.metrics import SklearnClusteringConfig
from ocl.config.plugins import PluginConfig

logger = logging.getLogger("eval_cluster_metrics")


class RepresentationType(enum.Enum):
    NONE = enum.auto()
    SLOTS = enum.auto()
    MASK_WEIGHTED_FEATURES = enum.auto()
    CLS_ON_MASKED_INPUT = enum.auto()


@dataclasses.dataclass
class EvaluationClusteringConfig:
    """Configuration for evaluation."""

    # Path to training configuration file or configuration dir. If dir, train_config_name
    # needs to be set as well.
    train_config_path: str

    # Number of classes. Note that on COCO, this should be one larger than the maximum class ID that
    # can appear, which does not correspond to the real number of classes.
    n_classes: int
    # Clustering methods to get cluster ID per object by clustering representations
    clusterings: Optional[Dict[str, SklearnClusteringConfig]] = None
    # Paths for model outputs to get cluster ID per object
    model_clusterings: Optional[Dict[str, str]] = None

    train_config_overrides: Optional[List[str]] = None
    train_config_name: Optional[str] = None
    checkpoint_path: Optional[str] = None
    output_dir: Optional[str] = None
    report_filename: str = "clustering_metrics.json"

    batch_size: int = 25
    class_discovery_threshold: float = 0.02
    use_mask_threshold: bool = False
    mask_threshold: float = 0.5
    ignore_background: bool = False
    use_unmatched_as_background: bool = False
    n_min_mask_pixels: int = 1  # Minimum number of pixels a mask must occupy to be considered valid
    n_min_max_mask_values: float = 1e-4  # Mask must have at least one value above threshold

    # Type of representation to use for clustering.
    representation_type: RepresentationType = RepresentationType.SLOTS

    # Setting this allows to add modules to the model that are executed during evaluation
    modules: Optional[Dict[str, Any]] = None
    # Setting this allows to evaluate on a different dataset than the model was trained on
    dataset: Optional[DataModuleConfig] = None
    # Setting this allows to change the plugins (e.g. to change the preprocessing pipeline)
    plugins: Optional[Dict[str, PluginConfig]] = None

    # Path to slot representations
    slots_path: str = path_defaults.OBJECTS
    # Path to feature representations
    features_path: str = f"{path_defaults.FEATURES}.features"
    # Path to slot masks, image shaped
    masks_path: str = f"{path_defaults.OBJECT_DECODER}.masks_as_image"
    # Path to slot masks, but flattened to match the size of features
    masks_flat_path: str = f"{path_defaults.OBJECT_DECODER}.masks"
    # Path to reference masks
    target_masks_path: str = f"{path_defaults.INPUT}.segmentation_mask"
    # Path under which representation to cluster is stored
    cluster_representation_path: str = "representation"
    # Path under which empty slot mask is stored
    empty_slots_path: str = "empty_slots"

    class_name_by_category_id: Optional[Dict[int, str]] = None


@dataclasses.dataclass
class Results:
    iou_per_class: np.ndarray
    accuracy: np.ndarray
    empty_classes: np.ndarray
    n_classes: int
    n_clusters: int
    has_background: bool
    matching: str
    clustering: Optional[SklearnClusteringConfig] = None
    model_clustering: Optional[str] = None

    def mean_iou(self) -> float:
        return np.mean(self.iou_per_class[~self.empty_classes])

    def num_discovered(self, threshold) -> int:
        iou_non_empty = self.iou_per_class[~self.empty_classes]
        return len(iou_non_empty[iou_non_empty > threshold])

    def mean_iou_discovered(self, threshold) -> float:
        iou_non_empty = self.iou_per_class[~self.empty_classes]
        iou_discovered = iou_non_empty[iou_non_empty > threshold]
        if len(iou_discovered) > 0:
            return np.mean(iou_discovered)
        else:
            return 0.0

    def mean_iou_without_bg(self) -> float:
        if self.has_background:
            return np.mean(self.iou_per_class[1:][~self.empty_classes[1:]])
        else:
            return self.mean_iou()


hydra.core.config_store.ConfigStore.instance().store(
    name="evaluation_clustering_config",
    node=EvaluationClusteringConfig,
)


def get_transform_for_representation_type(
    representation_type: RepresentationType, config: EvaluationClusteringConfig, model
):
    def cast_masks(mask: torch.Tensor):
        # Transform float masks to uint8 masks to save memory
        if config.use_mask_threshold:
            mask = mask > config.mask_threshold
        else:
            maximum, indices = torch.max(mask, dim=1)
            mask = torch.nn.functional.one_hot(indices, num_classes=mask.shape[1])
            mask[:, :, :, 0][maximum == 0.0] = 0
            mask = mask.transpose(-1, -2).transpose(-2, -3)

        return mask.to(torch.uint8)

    def filter_empty_masks(masks: torch.Tensor):
        # Count how many pixels each mask occupies
        max_per_class = masks.max(dim=2).values
        max_per_pixel = masks.max(dim=1, keepdim=True).values
        n_maxs = torch.sum((max_per_pixel == masks) & (max_per_pixel > 0), dim=-1)
        return (n_maxs < config.n_min_mask_pixels) | (max_per_class < config.n_min_max_mask_values)

    if representation_type == RepresentationType.MASK_WEIGHTED_FEATURES:
        required_paths = [config.masks_path, config.features_path, config.masks_flat_path]

        def transform(data):
            data[config.masks_path] = cast_masks(data[config.masks_path])

            # Compute a weighted mean of the features using the slot masks as weights.
            features = data[config.features_path]  # batch, [frames,] n_tokens, n_features
            masks = data[config.masks_flat_path]  # batch, [frames,] n_slots, n_tokens

            if features.ndim == 4:
                features = features.flatten(0, 1)
                masks = masks.flatten(0, 1)

            objects = torch.einsum("btd, bst -> bsd", features, masks)
            data[config.cluster_representation_path] = objects
            data[config.empty_slots_path] = filter_empty_masks(masks)
            return data

    elif representation_type == RepresentationType.SLOTS:
        required_paths = [config.masks_path, config.slots_path, config.masks_flat_path]

        def transform(data):
            data[config.masks_path] = cast_masks(data[config.masks_path])

            objects = data[config.slots_path]
            masks = data[config.masks_flat_path]
            if objects.ndim == 4:
                objects = objects.flatten(0, 1)
                masks = masks.flatten(0, 1)

            data[config.cluster_representation_path] = objects
            data[config.empty_slots_path] = filter_empty_masks(masks)
            return data

    elif representation_type == RepresentationType.CLS_ON_MASKED_INPUT:
        required_paths = [
            "input.image",
            config.masks_path,
            config.masks_flat_path,
            "object_decoder.masks_as_image",
        ]
        # TODO: Make this more flexible.
        feature_extractor = None
        for module in model.modules():
            if isinstance(module, feature_extractors.TimmFeatureExtractor):
                feature_extractor = module

        if feature_extractor is None:
            raise RuntimeError("Unable to find TimmFeatureExtractor in model.")

        def transform(data):
            data[config.masks_path] = cast_masks(data[config.masks_path])

            # Mask out part of the input image.
            input_image = data["input.image"]  # batch, C, H, W
            masks = cast_masks(data["object_decoder.masks_as_image"])

            masked_input = masks.unsqueeze(2).to(dtype=input_image.dtype) * input_image.unsqueeze(1)
            bs, n_objects = masked_input.shape[:2]
            masked_input = masked_input.flatten(0, 1)  # Merge object dim into batch dim.

            feature_extractor.eval()
            with torch.no_grad():
                aux_features = feature_extractor.forward_images(masked_input)[2]
            cls_tokens = aux_features[config.features_path].unflatten(0, (bs, n_objects))

            data[config.cluster_representation_path] = cls_tokens
            data[config.empty_slots_path] = filter_empty_masks(masks)
            return data

    elif representation_type == RepresentationType.NONE:
        required_paths = [config.masks_path, config.masks_flat_path]

        def transform(data):
            data[config.masks_path] = cast_masks(data[config.masks_path])
            masks = data[config.masks_flat_path]
            if masks.ndim == 4:
                masks = masks.flatten(0, 1)

            data[config.empty_slots_path] = filter_empty_masks(masks)

    else:
        raise ValueError(f"Unknown representation type {representation_type}.")

    return transform, required_paths


class ExtractDataFromPredictions(pl.callbacks.Callback):
    def __init__(self, paths, output_paths=None, transform=None):
        self.paths = paths
        self.output_paths = output_paths if output_paths is not None else paths
        self.transform = transform
        self.outputs = defaultdict(list)

    def on_predict_start(self, trainer, pl_module):
        self.outputs = defaultdict(list)

    def on_predict_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx):
        data = {path: tree_utils.get_tree_element(outputs, path.split(".")) for path in self.paths}

        if self.transform:
            data = self.transform(data)

        for path in self.output_paths:
            self.outputs[path].append(data[path].cpu())


def get_data_from_model(model, datamodule, input_paths, output_paths, transform=None):
    data_extractor = ExtractDataFromPredictions(input_paths, output_paths, transform)
    trainer = pl.Trainer(
        accelerator="auto",
        devices=1,
        auto_select_gpus=True,
        callbacks=[data_extractor],
        logger=False,
    )
    trainer.predict(model, datamodule.val_dataloader(), return_predictions=False)
    outputs = data_extractor.outputs

    return outputs


def get_cluster_ids_from_clustering(outputs, clustering, cluster_repr_path, empty_slots_path):
    objects = torch.cat(outputs[cluster_repr_path])
    objects_flat = objects.flatten(0, 1)

    if outputs[empty_slots_path] is not None:
        nonempty_objects = ~torch.cat(outputs[empty_slots_path]).flatten(0, 1).to(torch.bool)
    else:
        nonempty_objects = torch.ones_like(objects_flat[:, 0]).to(torch.bool)

    objects_selected = objects_flat[nonempty_objects]
    if len(objects_selected) > 0:
        object_ids_selected = clustering.fit_predict(objects_selected)
        object_ids = torch.zeros_like(objects_flat[:, 0], dtype=torch.int64)
        object_ids[nonempty_objects] = object_ids_selected + 1
    else:
        object_ids = torch.zeros_like(objects_flat[:, 0], dtype=torch.int64)

    object_ids = object_ids.unflatten(0, objects.shape[:2])
    object_ids = object_ids.split([len(obj) for obj in outputs[cluster_repr_path]])
    return object_ids


def get_cluster_ids_from_outputs(outputs, ids_path):
    return outputs[ids_path]


def report_from_results(results: Dict[str, Results], config: EvaluationClusteringConfig):
    output = OrderedDict()
    output["train_config_path"] = config.train_config_path
    output["checkpoint_path"] = config.checkpoint_path
    output["settings"] = OrderedDict(
        use_mask_threshold=config.use_mask_threshold,
        mask_threshold=config.mask_threshold,
        representation_type=config.representation_type.name,
        ignore_background=config.ignore_background,
        use_unmatched_as_background=config.use_unmatched_as_background,
    )

    for method in sorted(results):
        result = results[method]
        method_outputs = OrderedDict()
        settings = OrderedDict(
            n_classes=result.n_classes,
            n_clusters=result.n_clusters,
            matching=result.matching,
        )
        if result.clustering is not None:
            clustering = result.clustering
            settings["clustering"] = (clustering.method,)
            if clustering.clustering_kwargs is not None:
                clustering_kwargs = clustering.clustering_kwargs
                if omegaconf.OmegaConf.is_config(clustering_kwargs):
                    clustering_kwargs = omegaconf.OmegaConf.to_container(clustering_kwargs)
                settings["clustering_kwargs"] = clustering_kwargs
            if clustering.use_l2_normalization:
                settings["clustering_use_l2_normalization"] = True
            if clustering.use_pca:
                settings["clustering_use_pca"] = True
                settings["clustering_pca_dims"] = clustering.pca_dimensions
                if clustering.pca_kwargs is not None:
                    pca_kwargs = clustering.pca_kwargs
                    if omegaconf.OmegaConf.is_config(pca_kwargs):
                        pca_kwargs = omegaconf.OmegaConf.to_container(pca_kwargs)
                    settings["clustering_pca_kwargs"] = pca_kwargs
        elif result.model_clustering is not None:
            settings["model_clustering"] = result.model_clustering

        method_outputs["settings"] = settings
        method_outputs["iou_mean"] = result.mean_iou()
        method_outputs["iou_without_background"] = result.mean_iou_without_bg()
        method_outputs["accuracy"] = float(result.accuracy)
        method_outputs["num_discovered"] = result.num_discovered(config.class_discovery_threshold)
        method_outputs["iou_discovered"] = result.mean_iou_discovered(
            config.class_discovery_threshold
        )

        start_id = 1 if config.ignore_background else 0
        per_class_iou = result.iou_per_class
        class_name_by_category_id = config.class_name_by_category_id
        if class_name_by_category_id is None:
            class_name_by_category_id = {i: f"{i}" for i in range(start_id, config.n_classes + 1)}

        named_per_class_iou = OrderedDict()
        for category_id, class_name in class_name_by_category_id.items():
            if (config.ignore_background and category_id == 0) or category_id > config.n_classes:
                continue
            named_per_class_iou[class_name] = per_class_iou[category_id - start_id]

        method_outputs["iou_per_class"] = named_per_class_iou
        output[method] = method_outputs

    return output


def generate_and_save_report(
    config: EvaluationClusteringConfig, results_by_method: Dict[str, Results]
):
    report = report_from_results(results_by_method, config)
    metrics_file = os.path.join(config.output_dir, config.report_filename)
    with open(metrics_file, "w") as f:
        json.dump(report, f, indent=2)

    return metrics_file


@hydra.main(
    config_name="evaluation_clustering_config", config_path="../../configs", version_base="1.1"
)
def evaluate(config: EvaluationClusteringConfig):
    config.train_config_path = hydra.utils.to_absolute_path(config.train_config_path)
    if config.train_config_path.endswith(".yaml"):
        config_dir, config_name = os.path.split(config.train_config_path)
    else:
        config_dir, config_name = config.train_config_path, config.train_config_name

    if not os.path.exists(config_dir):
        raise ValueError(f"Inferred config dir at {config_dir} does not exist.")

    if config.checkpoint_path is None:
        try:
            run_dir = os.path.dirname(config_dir)
            checkpoint_path = cli_utils.find_checkpoint(run_dir)
            config.checkpoint_path = checkpoint_path
            logger.info(f"Automatically derived checkpoint path: {checkpoint_path}")
        except (TypeError, IndexError):
            raise ValueError(
                "Unable to automatically derive checkpoint from command line provided config file "
                "path. You can manually specify a checkpoint using the `checkpoint_path` argument."
            )
    else:
        config.checkpoint_path = hydra.utils.to_absolute_path(config.checkpoint_path)
        if not os.path.exists(config.checkpoint_path):
            raise ValueError(f"Checkpoint at {config.checkpoint_path} does not exist.")

    if config.output_dir is None:
        config.output_dir = run_dir
        if not os.path.exists(config.output_dir):
            os.mkdir(config.output_dir)
        logger.info(f"Using {config.output_dir} as output directory.")

    hydra.core.global_hydra.GlobalHydra.instance().clear()
    with hydra.initialize_config_dir(config_dir=config_dir):
        overrides = config.train_config_overrides if config.train_config_overrides else []
        train_config = hydra.compose(os.path.splitext(config_name)[0], overrides=overrides)
        train_config.dataset.batch_size = config.batch_size

        datamodule, model, _ = eval_utils.build_from_train_config(
            train_config, config.checkpoint_path
        )

    if config.modules is not None:
        modules = hydra_zen.instantiate(config.modules, _convert_="all")
        for key, module in modules.items():
            model.models[key] = module

    if config.dataset is not None or config.plugins is not None:
        if config.plugins is not None:
            pm = train.create_plugin_manager()
            train.build_and_register_plugins_from_config(config, pm)
            model.hooks = pm.hook
            datamodule.hooks = pm.hook
        if config.dataset is not None:
            datamodule = train.build_and_register_datamodule_from_config(
                config,
                pm.hook,
                pm,
                batch_size=train_config.dataset.batch_size,
                eval_batch_size=config.batch_size,
            )

    config.clusterings = config.clusterings if config.clusterings else {}
    config.model_clusterings = config.model_clusterings if config.model_clusterings else {}

    clusterings = hydra_zen.instantiate(config.clusterings)

    transform, transform_paths = get_transform_for_representation_type(
        config.representation_type, config, model
    )

    model_clustering_paths = list(config.model_clusterings.values())
    pretransform_paths = (
        model_clustering_paths + transform_paths + [config.masks_path, config.target_masks_path]
    )
    output_paths = model_clustering_paths + [
        config.masks_path,
        config.target_masks_path,
        config.empty_slots_path,
    ]
    if config.representation_type != RepresentationType.NONE:
        output_paths.append(config.cluster_representation_path)

    outputs = get_data_from_model(model, datamodule, pretransform_paths, output_paths, transform)

    cluster_fns = {
        name: partial(
            get_cluster_ids_from_clustering,
            clustering=clustering,
            cluster_repr_path=config.cluster_representation_path,
            empty_slots_path=config.empty_slots_path,
        )
        for name, clustering in clusterings.items()
    }
    cluster_fns.update(
        {
            name: partial(get_cluster_ids_from_outputs, ids_path=path)
            for name, path in config.model_clusterings.items()
        }
    )

    n_clusters_per_clustering = {
        name: clustering.n_clusters for name, clustering in clusterings.items()
    }
    n_clusters_per_clustering.update({name: None for name in config.model_clusterings})

    results_by_method = {}
    for clustering_name, cluster_fn in cluster_fns.items():
        logging.info(f"Fitting clustering {clustering_name}.")
        object_ids = cluster_fn(outputs)

        if n_clusters_per_clustering[clustering_name] is not None:
            n_clusters = n_clusters_per_clustering[clustering_name]
        else:
            # Get number of clusters from data: use maximum cluster ID
            n_clusters = max(int(ids.max()) for ids in object_ids)

        logging.info("Computing IoU over dataset.")
        metric_hung = metrics.DatasetSemanticMaskIoUMetric(
            n_predicted_classes=n_clusters,
            n_classes=config.n_classes,
            matching="hungarian",
            use_threshold=config.use_mask_threshold,
            threshold=config.mask_threshold,
            ignore_background=config.ignore_background,
            use_unmatched_as_background=config.use_unmatched_as_background,
        )
        for predictions, targets, ids in tqdm.tqdm(
            zip(outputs[config.masks_path], outputs[config.target_masks_path], object_ids),
            total=len(object_ids),
        ):
            metric_hung.update(predictions, targets, ids)

        per_class_iou, accuracy, empty_classes = metric_hung.compute()
        results = Results(
            per_class_iou.cpu().numpy(),
            accuracy.cpu().numpy(),
            empty_classes.cpu().numpy(),
            n_classes=config.n_classes,
            n_clusters=n_clusters,
            has_background=not config.ignore_background,
            clustering=config.clusterings.get(clustering_name),
            model_clustering=config.model_clusterings.get(clustering_name),
            matching="hungarian",
        )
        results_by_method[f"{clustering_name}_hungarian"] = results
        logging.info(
            f"{clustering_name}_hungarian: mIoU {results.mean_iou():.4f}, Acc {results.accuracy:.4f}"
        )

        metric_maj = metrics.DatasetSemanticMaskIoUMetric(
            n_predicted_classes=n_clusters,
            n_classes=config.n_classes,
            matching="majority",
            use_threshold=config.use_mask_threshold,
            threshold=config.mask_threshold,
            ignore_background=config.ignore_background,
        )
        metric_maj.load_state_dict(metric_hung.state_dict())  # Load state to avoid recomputing
        metric_maj._update_called = True  # Silence warning about calling compute before update

        per_class_iou, accuracy, empty_classes = metric_maj.compute()
        results = Results(
            per_class_iou.cpu().numpy(),
            accuracy.cpu().numpy(),
            empty_classes.cpu().numpy(),
            n_classes=config.n_classes,
            n_clusters=n_clusters,
            has_background=not config.ignore_background,
            clustering=config.clusterings.get(clustering_name),
            model_clustering=config.model_clusterings.get(clustering_name),
            matching="majority",
        )
        results_by_method[f"{clustering_name}_majority"] = results
        logging.info(
            f"{clustering_name}_majority: mIoU {results.mean_iou():.4f}, Acc {results.accuracy:.4f}"
        )

        # Already save intermediate reports.
        report_path = generate_and_save_report(config, results_by_method)

    logging.info(f"Report saved to {report_path}.")


if __name__ == "__main__":
    evaluate()

ocl/cli/eval_cluster_metrics.py
amazon-science/object-centric-multiple-object-tracking
from typing import Optional

import pytorch_lightning

from ocl.cli import train


def build_from_train_config(
    config: train.TrainingConfig, checkpoint_path: Optional[str], seed: bool = True
):
    if seed:
        pytorch_lightning.seed_everything(config.seed, workers=True)

    pm = train.create_plugin_manager()
    datamodule = train.build_and_register_datamodule_from_config(config, pm.hook, pm)
    train.build_and_register_plugins_from_config(config, pm)
    model = train.build_model_from_config(config, pm.hook, checkpoint_path)

    return datamodule, model, pm

ocl/cli/eval_utils.py
amazon-science/object-centric-multiple-object-tracking
"""Train a slot attention type model."""
import dataclasses
from typing import Any, Dict, Optional
from pytorch_lightning.utilities.cloud_io import load as pl_load
import hydra
import hydra_zen
import pytorch_lightning as pl
from pytorch_lightning.callbacks import ModelCheckpoint
from pluggy import PluginManager
import torch
import ocl.hooks
from ocl import base
from ocl.combined_model import CombinedModel
from ocl.config.datasets import DataModuleConfig
from ocl.config.metrics import MetricConfig
from ocl.config.plugins import PluginConfig
from ocl.plugins import Plugin
from ocl.utils import Combined, Recurrent, RoutableMixin
from typing import Union
TrainerConf = hydra_zen.builds(
    pl.Trainer, max_epochs=100, zen_partial=False, populate_full_signature=True
)


@dataclasses.dataclass
class TrainingConfig:
    """Configuration of a training run."""

    dataset: DataModuleConfig
    models: Any # When provided with dict wrap in `utils.Combined`, otherwise interpret as model.
    losses: Dict[str, Any]
    visualizations: Dict[str, Any] = dataclasses.field(default_factory=dict)
    plugins: Dict[str, PluginConfig] = dataclasses.field(default_factory=dict)
    trainer: TrainerConf = TrainerConf
    training_vis_frequency: Optional[int] = None
    training_metrics: Optional[Dict[str, MetricConfig]] = None
    evaluation_metrics: Optional[Dict[str, MetricConfig]] = None
    load_checkpoint: Optional[str] = None
    seed: Optional[int] = None
    experiment: Optional[Any] = None
    root_output_folder: Optional[str] = None


hydra.core.config_store.ConfigStore.instance().store(
    name="training_config",
    node=TrainingConfig,
)


def create_plugin_manager() -> PluginManager:
    pm = PluginManager("ocl")
    pm.add_hookspecs(ocl.hooks)
    return pm


def build_and_register_datamodule_from_config(
    config: TrainingConfig,
    hooks: base.PluggyHookRelay,
    plugin_manager: Optional[PluginManager] = None,
    **datamodule_kwargs,
) -> pl.LightningDataModule:
    datamodule = hydra_zen.instantiate(
        config.dataset, hooks=hooks, _convert_="all", **datamodule_kwargs
    )

    if plugin_manager:
        plugin_manager.register(datamodule)

    return datamodule


def build_and_register_plugins_from_config(
    config: TrainingConfig, plugin_manager: Optional[PluginManager] = None
) -> Dict[str, Plugin]:
    plugins = hydra_zen.instantiate(config.plugins)
    # Use lexicographical sorting to allow to influence registration order. This is necessary in
    # some cases as certain plugins might need to be called before others. Pluggy calls hooks
    # according to FILO (first in last out) and this is slightly unintuitive. We thus register
    # plugins in reverse order to their sorting position, leading to a FIFO (first in first out)
    # behavior with regard to the sorted position.
    if plugin_manager:
        for plugin_name in sorted(plugins.keys())[::-1]:
            plugin_manager.register(plugins[plugin_name])

    return plugins


def build_model_from_config(
    config: TrainingConfig,
    hooks: base.PluggyHookRelay,
    checkpoint_path: Optional[str] = None,
) -> pl.LightningModule:
    models = hydra_zen.instantiate(config.models, _convert_="all")
    losses = hydra_zen.instantiate(config.losses, _convert_="all")
    visualizations = hydra_zen.instantiate(config.visualizations, _convert_="all")

    training_metrics = hydra_zen.instantiate(config.training_metrics)
    evaluation_metrics = hydra_zen.instantiate(config.evaluation_metrics)

    # train_vis_freq = config.training_vis_frequency if config.training_vis_frequency else 10
    train_vis_freq = 10

    if checkpoint_path is None:
        model = CombinedModel(
            models=models,
            losses=losses,
            visualizations=visualizations,
            hooks=hooks,
            training_metrics=training_metrics,
            evaluation_metrics=evaluation_metrics,
            vis_log_frequency=train_vis_freq,
        )
    else:
        model = CombinedModel.load_from_checkpoint(
            checkpoint_path,
            strict=False,
            models=models,
            losses=losses,
            visualizations=visualizations,
            hooks=hooks,
            training_metrics=training_metrics,
            evaluation_metrics=evaluation_metrics,
            vis_log_frequency=train_vis_freq,
        )
        # checkpoint = pl_load(checkpoint_path)
        # model.load_state_dict(checkpoint['state_dict'], strict=False)
    return model


@hydra.main(config_name="training_config", config_path="../../configs/", version_base="1.1")
def train(config: TrainingConfig):
    # Set all relevant random seeds. If `config.seed is None`, the function samples a random value.
    # The function takes care of correctly distributing the seed across nodes in multi-node training,
    # and assigns each dataloader worker a different random seed.
    # IMPORTANTLY, we need to take care not to set a custom `worker_init_fn` function on the
    # dataloaders (or take care of worker seeding ourselves).
    pl.seed_everything(config.seed, workers=True)

    pm = create_plugin_manager()

    datamodule = build_and_register_datamodule_from_config(config, pm.hook, pm)

    build_and_register_plugins_from_config(config, pm)

    if config.load_checkpoint:
        checkpoint_path = hydra.utils.to_absolute_path(config.load_checkpoint)
    else:
        checkpoint_path = None

    model = build_model_from_config(config, pm.hook, checkpoint_path)

    callbacks = []
    if config.trainer.logger is not False:
        lr_monitor = pl.callbacks.LearningRateMonitor(logging_interval="step")
        callbacks.append(lr_monitor)

    checkpoint_path_ = 'checkpoints'

    checkpoint_callback  = ModelCheckpoint(dirpath = checkpoint_path_, every_n_epochs = 10)
    callbacks.append(checkpoint_callback)

    trainer: pl.Trainer = hydra_zen.instantiate(config.trainer, callbacks=callbacks, _convert_="all", gpus=[4,5,6,7]) #gpus=[5,6]

    trainer.fit(model, datamodule=datamodule, ckpt_path=None)


if __name__ == "__main__":
    train()

ocl/cli/train.py
amazon-science/object-centric-multiple-object-tracking
"""Implementation of combined model."""
from __future__ import annotations

from functools import partial
from typing import TYPE_CHECKING, Any, Dict

import pytorch_lightning as pl
import torch
from torch.profiler import profile, record_function, ProfilerActivity
from ocl import base, path_defaults, tree_utils, utils
from ocl.visualization_types import Visualization

# from slot_attention.tasks import Task

if TYPE_CHECKING:
    import torchmetrics


class CombinedModel(pl.LightningModule):
    def __init__(
        self,
        models: Dict[str, Any],
        losses: Dict[str, Any],
        visualizations: Dict[str, Any],
        hooks: base.PluggyHookRelay,
        training_metrics: Dict[str, torchmetrics.Metric] = None,
        evaluation_metrics: Dict[str, torchmetrics.Metric] = None,
        vis_log_frequency: int = 10,  #100
    ):
        super().__init__()
        if isinstance(models, Dict):
            models = utils.Combined(models)
        self.models = models
        self.losses = losses
        self.visualizations = visualizations
        self.hooks = hooks
        self.vis_log_frequency = vis_log_frequency

        if training_metrics is None:
            training_metrics = {}
        self.training_metrics = torch.nn.ModuleDict(training_metrics)

        if evaluation_metrics is None:
            evaluation_metrics = {}
        self.evaluation_metrics = torch.nn.ModuleDict(evaluation_metrics)

    def configure_optimizers(self):
        return self.hooks.configure_optimizers(model=self)

    def __getattribute__(self, name):
        """Forward pytorch lightning module hooks to the plugin manager.

        We need to implement `__getattribute__` as the model hooks are defined in a superclass of
        `pl.LightningModule` and thus `__getattr__` would never get called for them. This makes the
        call a bit more clumsy.
        """
        if not name.startswith("__") and hasattr(pl.core.hooks.ModelHooks, name):
            # A pytorch lighting hook is being called.
            try:
                hook_caller = getattr(self.hooks, name)
                return partial(hook_caller, model=self)
            except AttributeError:
                pass
        return super().__getattribute__(name)

    def forward(self, input_data: dict):
        # Maybe we should use something like a read only dict to prevent existing keys from being
        # overwritten.
        data: Dict[str, Any]
        data = {path_defaults.INPUT: input_data, path_defaults.GLOBAL_STEP: self.global_step}
        phase = input_data['phase']
        return self.models.forward(inputs=data, phase=phase)

    def _compute_losses(self, inputs, phase="train"):
        quantities_to_log = {}
        for name, loss in self.losses.items():
            quantities_to_log[f"{phase}/{name}"] = loss(inputs=inputs)

        losses = []
        for loss in quantities_to_log.values():
            losses.append(loss)
        total_loss = torch.stack(losses).sum()

        # Log total loss only if there is more than one task
        if len(losses) > 1:
            quantities_to_log[f"{phase}/loss_total"] = total_loss

        return total_loss, quantities_to_log

    def training_step(self, batch, batch_idx):
        batch_size = batch["batch_size"]
        batch['phase'] = 'train'
        # with profile(activities=[
        #     ProfilerActivity.CPU, ProfilerActivity.CUDA], record_shapes=True) as prof:
        #     with record_function("model_train"):
        #         outputs = self(batch)
        #         total_loss, quantities_to_log = self._compute_losses(outputs)
        # print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=15))
        outputs = self(batch)
        total_loss, quantities_to_log = self._compute_losses(outputs)
        quantities_to_log.update(self._compute_metrics(outputs, self.training_metrics))
        self.log_dict(quantities_to_log, on_step=True, on_epoch=False, batch_size=batch_size)

        if self.trainer.global_step % self.vis_log_frequency == 0:
            self._log_visualizations(outputs)

        return total_loss

    def validation_step(self, batch, batch_idx):
        batch_size = batch["batch_size"]
        batch['phase'] = 'val'
        outputs = self(batch)
        total_loss, quantities_to_log = self._compute_losses(outputs, phase="val")

        quantities_to_log.update(
            self._compute_metrics(outputs, self.evaluation_metrics, phase="val")
        )
        self.log_dict(
            quantities_to_log, on_step=False, on_epoch=True, prog_bar=True, batch_size=batch_size
        )

        if batch_idx == 0:
            self._log_visualizations(outputs, phase="val")

    def _compute_metrics(self, outputs, metric_fns, phase="train"):
        metrics = {}
        if len(metric_fns) > 0:
            for metric_name, metric in metric_fns.items():
                if phase == "val":
                    # Call update instead of forward to avoid unnecessary metric compute on batch.
                    metric.update(**outputs)
                else:
                    metric(**outputs)
                metrics[f"{phase}/{metric_name}"] = metric

        return metrics

    def _log_visualizations(self, outputs, phase="train"):
        if self.logger is None:
            return
        logger_experiment = self.logger.experiment
        visualizations = {}
        for name, vis in self.visualizations.items():
            visualizations[name] = vis(inputs=outputs)

        visualization_iterator = tree_utils.walk_tree_with_paths(
            visualizations, path=None, instance_check=lambda a: isinstance(a, Visualization)
        )
        for path, vis in visualization_iterator:
            str_path = ".".join(path)
            vis.add_to_experiment(
                experiment=logger_experiment,
                tag=f"{phase}/{str_path}",
                global_step=self.trainer.global_step,
            )

ocl/combined_model.py
amazon-science/object-centric-multiple-object-tracking
"""Implementation of conditioning approaches for slots."""
from typing import Callable, Optional, Tuple

import numpy as np
import torch
from torch import nn

from ocl import base, path_defaults
from ocl.utils import RoutableMixin
from ocl.predictor import MLP

class RandomConditioning(base.Conditioning, RoutableMixin):
    """Random conditioning with potentially learnt mean and stddev."""

    def __init__(
        self,
        object_dim: int,
        n_slots: int,
        learn_mean: bool = True,
        learn_std: bool = True,
        mean_init: Optional[Callable[[torch.Tensor], None]] = None,
        logsigma_init: Optional[Callable[[torch.Tensor], None]] = None,
        batch_size_path: Optional[str] = path_defaults.BATCH_SIZE,
    ):
        base.Conditioning.__init__(self)
        RoutableMixin.__init__(self, {"batch_size": batch_size_path})
        self.n_slots = n_slots
        self.object_dim = object_dim

        if learn_mean:
            self.slots_mu = nn.Parameter(torch.zeros(1, 1, object_dim))
        else:
            self.register_buffer("slots_mu", torch.zeros(1, 1, object_dim))

        if learn_std:
            self.slots_logsigma = nn.Parameter(torch.zeros(1, 1, object_dim))
        else:
            self.register_buffer("slots_logsigma", torch.zeros(1, 1, object_dim))

        if mean_init is None:
            mean_init = nn.init.xavier_uniform_
        if logsigma_init is None:
            logsigma_init = nn.init.xavier_uniform_

        with torch.no_grad():
            mean_init(self.slots_mu)
            logsigma_init(self.slots_logsigma)

    @RoutableMixin.route
    def forward(self, batch_size: int) -> base.ConditioningOutput:
        mu = self.slots_mu.expand(batch_size, self.n_slots, -1)
        sigma = self.slots_logsigma.exp().expand(batch_size, self.n_slots, -1)
        return mu + sigma * torch.randn_like(mu)


class LearntConditioning(base.Conditioning, RoutableMixin):
    """Conditioning with a learnt set of slot initializations, similar to DETR."""

    def __init__(
        self,
        object_dim: int,
        n_slots: int,
        slot_init: Optional[Callable[[torch.Tensor], None]] = None,
        batch_size_path: Optional[str] = path_defaults.BATCH_SIZE,
    ):
        base.Conditioning.__init__(self)
        RoutableMixin.__init__(self, {"batch_size": batch_size_path})
        self.n_slots = n_slots
        self.object_dim = object_dim

        self.slots = nn.Parameter(torch.zeros(1, n_slots, object_dim))

        if slot_init is None:
            slot_init = nn.init.normal_

        with torch.no_grad():
            slot_init(self.slots)

    @RoutableMixin.route
    def forward(self, batch_size: int) -> base.ConditioningOutput:
        return self.slots.expand(batch_size, -1, -1)


class RandomConditioningWithQMCSampling(RandomConditioning):
    """Random conditioning with learnt mean and stddev using Quasi-Monte Carlo (QMC) samples."""

    def __init__(
        self,
        object_dim: int,
        n_slots: int,
        learn_mean: bool = True,
        learn_std: bool = True,
        mean_init: Optional[Callable[[torch.Tensor], None]] = None,
        logsigma_init: Optional[Callable[[torch.Tensor], None]] = None,
        batch_size_path: Optional[str] = path_defaults.BATCH_SIZE,
    ):
        super().__init__(
            object_dim,
            n_slots,
            learn_mean,
            learn_std,
            mean_init,
            logsigma_init,
            batch_size_path=batch_size_path,
        )

        import scipy.stats  # Import lazily because scipy takes some time to import

        self.randn_rng = scipy.stats.qmc.MultivariateNormalQMC(mean=np.zeros(object_dim))

    def _randn(self, *args: Tuple[int]) -> torch.Tensor:
        n_elements = np.prod(args)
        # QMC sampler needs to sample powers of 2 numbers at a time
        n_elements_rounded2 = 2 ** int(np.ceil(np.log2(n_elements)))
        z = self.randn_rng.random(n_elements_rounded2)[:n_elements]

        return torch.from_numpy(z).view(*args, -1)

    @RoutableMixin.route
    def forward(self, batch_size: int) -> base.ConditioningOutput:
        mu = self.slots_mu.expand(batch_size, self.n_slots, -1)
        sigma = self.slots_logsigma.exp().expand(batch_size, self.n_slots, -1)

        z = self._randn(batch_size, self.n_slots).to(mu, non_blocking=True)
        return mu + sigma * z


class SlotwiseLearntConditioning(base.Conditioning, RoutableMixin):
    """Random conditioning with learnt mean and stddev for each slot.

    Removes permutation equivariance compared to the original slot attention conditioning.
    """

    def __init__(
        self,
        object_dim: int,
        n_slots: int,
        mean_init: Optional[Callable[[torch.Tensor], None]] = None,
        logsigma_init: Optional[Callable[[torch.Tensor], None]] = None,
        batch_size_path: Optional[str] = path_defaults.BATCH_SIZE,
    ):
        base.Conditioning.__init__(self)
        RoutableMixin.__init__(self, {"batch_size": batch_size_path})
        self.n_slots = n_slots
        self.object_dim = object_dim

        self.slots_mu = nn.Parameter(torch.zeros(1, n_slots, object_dim))
        self.slots_logsigma = nn.Parameter(torch.zeros(1, n_slots, object_dim))

        if mean_init is None:
            mean_init = nn.init.normal_
        if logsigma_init is None:
            logsigma_init = nn.init.xavier_uniform_

        with torch.no_grad():
            mean_init(self.slots_mu)
            logsigma_init(self.slots_logsigma)

    @RoutableMixin.route
    def forward(self, batch_size: int) -> base.ConditioningOutput:
        mu = self.slots_mu.expand(batch_size, -1, -1)
        sigma = self.slots_logsigma.exp().expand(batch_size, -1, -1)
        return mu + sigma * torch.randn_like(mu)


class CoordinateEncoderStateInit(base.Conditioning, RoutableMixin):
    """State init that encodes bounding box corrdinates as conditional input.
    Attributes:
        embedding_transform: A nn.Module that is applied on inputs (bounding boxes).
        prepend_background: Boolean flag' whether to prepend a special, zero-valued
            background bounding box to the input. Default: False.
        center_of_mass: Boolean flag; whether to convert bounding boxes to center
            of mass coordinates. Default: False.
        background_value: Default value to fill in the background.
    """
    def __init__(self,
                 object_dim: int,
                 prepend_background: bool = True,
                 center_of_mass: bool = False,
                 background_value: float = 0. ,
                 batch_size_path: Optional[str] = path_defaults.BATCH_SIZE,
                 ):
        base.Conditioning.__init__(self)
        RoutableMixin.__init__(self, {"batch_size": batch_size_path})

        self.embedding_transform = MLP(input_size=4, hidden_size=256,output_size=128, layernorm=None)
        self.prepend_background = prepend_background
        self.center_of_mass = center_of_mass
        self.background_value = background_value
        self.object_dim = object_dim

    @RoutableMixin.route
    def forward(self, target_bbox: torch.Tensor, batch_size: int) -> base.ConditioningOutput:
        del batch_size  # Unused.

        # inputs.shape = (batch_size, seq_len, bboxes, 4)
        inputs = target_bbox[:, 0]  # Only condition on first time step.
        # inputs.shape = (batch_size, bboxes, 4)
        if self.prepend_background:
            # Adds a fake background box [0, 0, 0, 0] at the beginning.
            batch_size = inputs.shape[0]

            # Encode the background as specified by the background_value.
            background = torch.full(
                (batch_size, 1, 4), self.background_value, dtype=inputs.dtype,
                device=inputs.get_device())

            inputs = torch.cat([background, inputs], dim=1)
            # inputs = torch.cat([inputs, background], dim=1)

        if self.center_of_mass:
            y_pos = (inputs[:, :, 0] + inputs[:, :, 2]) / 2
            x_pos = (inputs[:, :, 1] + inputs[:, :, 3]) / 2
            inputs = torch.stack([y_pos, x_pos], dim=-1)

        slots = self.embedding_transform(inputs)
        # duplicated_slots = torch.cat([slots, slots], dim=1)

        return slots

class CoordinateEncoderState_learned_Init(base.Conditioning, RoutableMixin):
    """State init that encodes bounding box corrdinates as conditional input.
    Attributes:
        embedding_transform: A nn.Module that is applied on inputs (bounding boxes).
        prepend_background: Boolean flag' whether to prepend a special, zero-valued
            background bounding box to the input. Default: False.
        center_of_mass: Boolean flag; whether to convert bounding boxes to center
            of mass coordinates. Default: False.
        background_value: Default value to fill in the background.
    """
    def __init__(self,
                 object_dim: int,
                 n_slots: int,
                 prepend_background: bool = True,
                 center_of_mass: bool = False,
                 background_value: float = 0. ,
                 slot_init: Optional[Callable[[torch.Tensor], None]] = None,
                 batch_size_path: Optional[str] = path_defaults.BATCH_SIZE,
                 ):
        base.Conditioning.__init__(self)
        RoutableMixin.__init__(self, {"batch_size": batch_size_path})

        self.embedding_transform = MLP(input_size=4, hidden_size=256,output_size=128, layernorm=None)
        self.prepend_background = prepend_background
        self.center_of_mass = center_of_mass
        self.background_value = background_value
        self.object_dim = object_dim
        self.slots = nn.Parameter(torch.zeros(1, n_slots, object_dim))
        if slot_init is None:
            slot_init = nn.init.normal_

        with torch.no_grad():
            slot_init(self.slots)


    @RoutableMixin.route
    def forward(self, target_bbox: torch.Tensor, batch_size: int) -> base.ConditioningOutput:
        del batch_size  # Unused.
        inputs = target_bbox[:, 0]  # Only condition on first time step.
        inputs_box = inputs
        batch_size, object_num, _ = inputs.shape
        if self.prepend_background:
            # Adds a fake background box [0, 0, 0, 0] at the beginning.
            batch_size = inputs.shape[0]

            # Encode the background as specified by the background_value.
            background = torch.full(
                (batch_size, 1, 4), self.background_value, dtype=inputs.dtype,
                device=inputs.get_device())

            inputs = torch.cat([background, inputs], dim=1)

        if self.center_of_mass:
            y_pos = (inputs[:, :, 0] + inputs[:, :, 2]) / 2
            x_pos = (inputs[:, :, 1] + inputs[:, :, 3]) / 2
            inputs = torch.stack([y_pos, x_pos], dim=-1)

        slots = self.embedding_transform(inputs)
        slots_learned = self.slots.expand(batch_size, -1, -1)

        # slots_output = torch.zeros(slots.shape).to(target_bbox.device)
        #
        # for i in range(batch_size):
        #     for j in range(object_num):
        #         if j == 0:
        #             slots_output[i][0] = slots[i][0]
        #         else:
        #             if torch.sum(inputs_box[i][j]) == 0:
        #                 slots_output[i][j] = slots_learned[i][j]
        #             else:
        #                 slots_output[i][j] = slots[i][j]

        slots_output = torch.cat([slots, slots_learned], dim = 1)
        return slots_output

class CoordinateEncoderState_learned_m_2(base.Conditioning, RoutableMixin):
    """State init that encodes bounding box corrdinates as conditional input.
    Attributes:
        embedding_transform: A nn.Module that is applied on inputs (bounding boxes).
        prepend_background: Boolean flag' whether to prepend a special, zero-valued
            background bounding box to the input. Default: False.
        center_of_mass: Boolean flag; whether to convert bounding boxes to center
            of mass coordinates. Default: False.
        background_value: Default value to fill in the background.
    """
    def __init__(self,
                 object_dim: int,
                 n_slots: int,
                 prepend_background: bool = True,
                 center_of_mass: bool = False,
                 background_value: float = 0. ,
                 slot_init: Optional[Callable[[torch.Tensor], None]] = None,
                 batch_size_path: Optional[str] = path_defaults.BATCH_SIZE,
                 ):
        base.Conditioning.__init__(self)
        RoutableMixin.__init__(self, {"batch_size": batch_size_path})

        self.embedding_transform = MLP(input_size=4, hidden_size=256,output_size=128, layernorm=None)
        self.prepend_background = prepend_background
        self.center_of_mass = center_of_mass
        self.background_value = background_value
        self.object_dim = object_dim
        self.slots = nn.Parameter(torch.zeros(1, 1, n_slots, object_dim))
        if slot_init is None:
            slot_init = nn.init.normal_

        with torch.no_grad():
            slot_init(self.slots)


    @RoutableMixin.route
    def forward(self, target_bbox: torch.Tensor, batch_size: int) -> base.ConditioningOutput:
        del batch_size  # Unused.
        inputs = target_bbox[:, 0]  # Only condition on first time step.
        inputs_box = inputs
        batch_size, object_num, _ = inputs.shape
        if self.prepend_background:
            # Adds a fake background box [0, 0, 0, 0] at the beginning.
            batch_size = inputs.shape[0]

            # Encode the background as specified by the background_value.
            background = torch.full(
                (batch_size, 1, 4), self.background_value, dtype=inputs.dtype,
                device=inputs.get_device())

            inputs = torch.cat([background, inputs], dim=1)

        if self.center_of_mass:
            y_pos = (inputs[:, :, 0] + inputs[:, :, 2]) / 2
            x_pos = (inputs[:, :, 1] + inputs[:, :, 3]) / 2
            inputs = torch.stack([y_pos, x_pos], dim=-1)

        slots = self.embedding_transform(inputs).unsqueeze(1)
        slots_learned = self.slots.expand(batch_size, -1, -1, -1)
        slots_output = torch.cat([slots, slots_learned], dim=1)
        return slots_output


ocl/conditioning.py
amazon-science/object-centric-multiple-object-tracking
from hydra.core.config_store import ConfigStore
from omegaconf import OmegaConf

from ocl.config import (
    conditioning,
    datasets,
    feature_extractors,
    metrics,
    neural_networks,
    optimizers,
    perceptual_groupings,
    plugins,
    predictor,
    memory,
    utils,
)

config_store = ConfigStore.instance()

conditioning.register_configs(config_store)

datasets.register_configs(config_store)
datasets.register_resolvers(OmegaConf)

feature_extractors.register_configs(config_store)

metrics.register_configs(config_store)

neural_networks.register_configs(config_store)

optimizers.register_configs(config_store)

perceptual_groupings.register_configs(config_store)
predictor.register_configs(config_store)
memory.register_configs(config_store)

plugins.register_configs(config_store)
plugins.register_resolvers(OmegaConf)

utils.register_configs(config_store)
utils.register_resolvers(OmegaConf)

ocl/config/__init__.py
amazon-science/object-centric-multiple-object-tracking
"""Configuration of slot conditionings."""
import dataclasses

from hydra_zen import builds
from omegaconf import SI

from ocl import conditioning


@dataclasses.dataclass
class ConditioningConfig:
    """Base class for conditioning module configuration."""


# Unfortunately, we cannot define object_dim as part of the base config class as this prevents using
# required positional arguments in all subclasses. We thus instead pass them here.
LearntConditioningConfig = builds(
    conditioning.LearntConditioning,
    object_dim=SI("${perceptual_grouping.object_dim}"),
    builds_bases=(ConditioningConfig,),
    populate_full_signature=True,
)

RandomConditioningConfig = builds(
    conditioning.RandomConditioning,
    object_dim=SI("${perceptual_grouping.object_dim}"),
    builds_bases=(ConditioningConfig,),
    populate_full_signature=True,
)

RandomConditioningWithQMCSamplingConfig = builds(
    conditioning.RandomConditioningWithQMCSampling,
    object_dim=SI("${perceptual_grouping.object_dim}"),
    builds_bases=(ConditioningConfig,),
    populate_full_signature=True,
)

SlotwiseLearntConditioningConfig = builds(
    conditioning.SlotwiseLearntConditioning,
    object_dim=SI("${perceptual_grouping.object_dim}"),
    builds_bases=(ConditioningConfig,),
    populate_full_signature=True,
)
CoordinateEncoderStateInitConfig = builds(
    conditioning.CoordinateEncoderStateInit,
    object_dim=SI("${perceptual_grouping.object_dim}"),
    builds_bases=(ConditioningConfig,),
    populate_full_signature=True,
)

def register_configs(config_store):
    config_store.store(group="schemas", name="conditioning", node=ConditioningConfig)

    config_store.store(group="conditioning", name="learnt", node=LearntConditioningConfig)
    config_store.store(group="conditioning", name="random", node=RandomConditioningConfig)
    config_store.store(
        group="conditioning",
        name="random_with_qmc_sampling",
        node=RandomConditioningWithQMCSamplingConfig,
    )
    config_store.store(
        group="conditioning", name="slotwise_learnt_random", node=SlotwiseLearntConditioningConfig
    )
    config_store.store(group="conditioning", name="boxhint", node=CoordinateEncoderStateInitConfig)
ocl/config/conditioning.py
amazon-science/object-centric-multiple-object-tracking
"""Register all dataset related configs."""
import dataclasses
import os

from hydra_zen import builds

from ocl import datasets


def get_region():
    """Determine the region this EC2 instance is running in.

    Returns None if not running on an EC2 instance.
    """
    import requests

    try:
        r = requests.get(
            "http://169.254.169.254/latest/dynamic/instance-identity/document", timeout=0.5
        )
        response_json = r.json()
        return response_json.get("region")
    except Exception:
        # Not running on an ec2 instance.
        return None


# Detemine region name and select bucket accordingly.
AWS_REGION = get_region()
if AWS_REGION in ["us-east-2", "us-west-2", "eu-west-1"]:
    # Select bucket in same region.
    DEFAULT_S3_PATH = f"s3://object-centric-datasets-{AWS_REGION}"
else:
    # Use MRAP to find closest bucket.
    DEFAULT_S3_PATH = "s3://arn:aws:s3::436622332146:accesspoint/m6p4hmmybeu97.mrap"


@dataclasses.dataclass
class DataModuleConfig:
    """Base class for PyTorch Lightning DataModules.

    This class does not actually do anything but ensures that datasets behave like pytorch lightning
    datamodules.
    """


def dataset_prefix(path):
    # prefix = os.environ.get("DATASET_PREFIX")
    prefix = '/home/ubuntu/data'
    if prefix:
        return f"{prefix}/{path}"
    # Use the path to the multi-region bucket if no override is specified.
    return f"pipe:aws s3 cp --quiet {DEFAULT_S3_PATH}/{path} -"


WebdatasetDataModuleConfig = builds(
    datasets.WebdatasetDataModule, populate_full_signature=True, builds_bases=(DataModuleConfig,)
)
DummyDataModuleConfig = builds(
    datasets.DummyDataModule, populate_full_signature=True, builds_bases=(DataModuleConfig,)
)


def register_configs(config_store):
    config_store.store(group="schemas", name="dataset", node=DataModuleConfig)
    config_store.store(group="dataset", name="webdataset", node=WebdatasetDataModuleConfig)
    config_store.store(group="dataset", name="dummy_dataset", node=DummyDataModuleConfig)


def register_resolvers(omegaconf):
    omegaconf.register_new_resolver("dataset_prefix", dataset_prefix)

ocl/config/datasets.py
amazon-science/object-centric-multiple-object-tracking
"""Configurations for feature extractors."""
import dataclasses

from hydra_zen import make_custom_builds_fn

from ocl import feature_extractors


@dataclasses.dataclass
class FeatureExtractorConfig:
    """Base class for PyTorch Lightning DataModules.

    This class does not actually do anything but ensures that feature extractors give outputs of
    a defined structure.
    """

    pass


builds_feature_extractor = make_custom_builds_fn(
    populate_full_signature=True,
)

TimmFeatureExtractorConfig = builds_feature_extractor(
    feature_extractors.TimmFeatureExtractor,
    builds_bases=(FeatureExtractorConfig,),
)
SlotAttentionFeatureExtractorConfig = builds_feature_extractor(
    feature_extractors.SlotAttentionFeatureExtractor,
    builds_bases=(FeatureExtractorConfig,),
)
SAViFeatureExtractorConfig = builds_feature_extractor(
    feature_extractors.SAViFeatureExtractor,
    builds_bases=(FeatureExtractorConfig,),
)


def register_configs(config_store):
    config_store.store(group="schemas", name="feature_extractor", node=FeatureExtractorConfig)
    config_store.store(
        group="feature_extractor",
        name="timm_model",
        node=TimmFeatureExtractorConfig,
    )
    config_store.store(
        group="feature_extractor",
        name="slot_attention",
        node=SlotAttentionFeatureExtractorConfig,
    )
    config_store.store(
        group="feature_extractor",
        name="savi",
        node=SAViFeatureExtractorConfig,
    )

ocl/config/feature_extractors.py
amazon-science/object-centric-multiple-object-tracking
"""Perceptual grouping models."""
import dataclasses

from hydra_zen import builds

from ocl import memory


@dataclasses.dataclass
class MemoryConfig:
    """Configuration class of Predictor."""


TransitionConfig = builds(
    memory.SelfSupervisedMemory,
    builds_bases=(MemoryConfig,),
    populate_full_signature=True,
)


def register_configs(config_store):
    config_store.store(group="schemas", name="memory", node=MemoryConfig)
    config_store.store(group="memory", name="mem", node=TransitionConfig)

ocl/config/memory.py
amazon-science/object-centric-multiple-object-tracking
"""Register metric related configs."""
import dataclasses

from hydra_zen import builds, make_custom_builds_fn

from ocl import metrics


@dataclasses.dataclass
class MetricConfig:
    """Base class for metrics."""

    pass


builds_metric = make_custom_builds_fn(
    populate_full_signature=True,
)

TensorStatisticConfig = builds_metric(metrics.TensorStatistic, builds_bases=(MetricConfig,))
TorchmetricsWrapperConfig = builds_metric(metrics.TorchmetricsWrapper, builds_bases=(MetricConfig,))
ARIMetricConfig = builds_metric(metrics.ARIMetric, builds_bases=(MetricConfig,))
PatchARIMetricConfig = builds_metric(
    metrics.PatchARIMetric,
    builds_bases=(MetricConfig,),
)

MOTMetricConfig = builds_metric(
    metrics.MOTMetric,
    builds_bases=(MetricConfig,),
)
ThreeDIOUMetric = builds_metric(
    metrics.ThreeDIOUMetric,
    builds_bases=(MetricConfig,),
)
UnsupervisedMaskIoUMetricConfig = builds_metric(
    metrics.UnsupervisedMaskIoUMetric,
    builds_bases=(MetricConfig,),
)
MaskCorLocMetricConfig = builds_metric(
    metrics.UnsupervisedMaskIoUMetric,
    matching="best_overlap",
    correct_localization=True,
    builds_bases=(MetricConfig,),
)
AverageBestOverlapMetricConfig = builds_metric(
    metrics.UnsupervisedMaskIoUMetric,
    matching="best_overlap",
    builds_bases=(MetricConfig,),
)
BestOverlapObjectRecoveryMetricConfig = builds_metric(
    metrics.UnsupervisedMaskIoUMetric,
    matching="best_overlap",
    compute_discovery_fraction=True,
    builds_bases=(MetricConfig,),
)
UnsupervisedBboxIoUMetricConfig = builds_metric(
    metrics.UnsupervisedBboxIoUMetric,
    builds_bases=(MetricConfig,),
)
BboxCorLocMetricConfig = builds_metric(
    metrics.UnsupervisedBboxIoUMetric,
    matching="best_overlap",
    correct_localization=True,
    builds_bases=(MetricConfig,),
)
BboxRecallMetricConfig = builds_metric(
    metrics.UnsupervisedBboxIoUMetric,
    matching="best_overlap",
    compute_discovery_fraction=True,
    builds_bases=(MetricConfig,),
)

DatasetSemanticMaskIoUMetricConfig = builds_metric(metrics.DatasetSemanticMaskIoUMetric)

SklearnClusteringConfig = builds(
    metrics.SklearnClustering,
    populate_full_signature=True,
)


def register_configs(config_store):
    config_store.store(group="metrics", name="tensor_statistic", node=TensorStatisticConfig)
    config_store.store(group="metrics", name="torchmetric", node=TorchmetricsWrapperConfig)
    config_store.store(group="metrics", name="ari_metric", node=ARIMetricConfig)
    config_store.store(group="metrics", name="mot_metric", node=MOTMetricConfig)
    config_store.store(group="metrics", name="three_d_iou", node=ThreeDIOUMetric)
    config_store.store(group="metrics", name="patch_ari_metric", node=PatchARIMetricConfig)
    config_store.store(
        group="metrics", name="unsupervised_mask_iou_metric", node=UnsupervisedMaskIoUMetricConfig
    )
    config_store.store(group="metrics", name="mask_corloc_metric", node=MaskCorLocMetricConfig)
    config_store.store(
        group="metrics", name="average_best_overlap_metric", node=AverageBestOverlapMetricConfig
    )
    config_store.store(
        group="metrics",
        name="best_overlap_object_recovery_metric",
        node=BestOverlapObjectRecoveryMetricConfig,
    )
    config_store.store(
        group="metrics", name="unsupervised_bbox_iou_metric", node=UnsupervisedBboxIoUMetricConfig
    )
    config_store.store(group="metrics", name="bbox_corloc_metric", node=BboxCorLocMetricConfig)
    config_store.store(group="metrics", name="bbox_recall_metric", node=BboxRecallMetricConfig)

    config_store.store(
        group="metrics",
        name="dataset_semantic_mask_iou",
        node=DatasetSemanticMaskIoUMetricConfig,
    )
    config_store.store(
        group="clustering",
        name="sklearn_clustering",
        node=SklearnClusteringConfig,
    )

ocl/config/metrics.py
amazon-science/object-centric-multiple-object-tracking
"""Configs for neural networks."""
import omegaconf
from hydra_zen import builds

from ocl import neural_networks

MLPBuilderConfig = builds(
    neural_networks.build_mlp,
    features=omegaconf.MISSING,
    zen_partial=True,
    populate_full_signature=True,
)
TransformerEncoderBuilderConfig = builds(
    neural_networks.build_transformer_encoder,
    n_layers=omegaconf.MISSING,
    n_heads=omegaconf.MISSING,
    zen_partial=True,
    populate_full_signature=True,
)
TransformerDecoderBuilderConfig = builds(
    neural_networks.build_transformer_decoder,
    n_layers=omegaconf.MISSING,
    n_heads=omegaconf.MISSING,
    zen_partial=True,
    populate_full_signature=True,
)


def register_configs(config_store):
    config_store.store(group="neural_networks", name="mlp", node=MLPBuilderConfig)
    config_store.store(
        group="neural_networks", name="transformer_encoder", node=TransformerEncoderBuilderConfig
    )
    config_store.store(
        group="neural_networks", name="transformer_decoder", node=TransformerDecoderBuilderConfig
    )

ocl/config/neural_networks.py
amazon-science/object-centric-multiple-object-tracking
"""Pytorch optimizers."""
import dataclasses

import torch.optim
from hydra_zen import make_custom_builds_fn


@dataclasses.dataclass
class OptimizerConfig:
    pass


# TODO(hornmax): We cannot automatically extract type information from the torch SGD implementation,
# thus we define it manually here.
@dataclasses.dataclass
class SGDConfig(OptimizerConfig):
    learning_rate: float
    momentum: float = 0.0
    dampening: float = 0.0
    nestov: bool = False
    _target_: str = "hydra_zen.funcs.zen_processing"
    _zen_target: str = "torch.optim.SGD"
    _zen_partial: bool = True


pbuilds = make_custom_builds_fn(
    zen_partial=True,
    populate_full_signature=True,
)

AdamConfig = pbuilds(torch.optim.Adam, builds_bases=(OptimizerConfig,))
AdamWConfig = pbuilds(torch.optim.AdamW, builds_bases=(OptimizerConfig,))


def register_configs(config_store):
    config_store.store(group="optimizers", name="sgd", node=SGDConfig)
    config_store.store(group="optimizers", name="adam", node=AdamConfig)
    config_store.store(group="optimizers", name="adamw", node=AdamWConfig)

ocl/config/optimizers.py
amazon-science/object-centric-multiple-object-tracking
"""Perceptual grouping models."""
import dataclasses

from hydra_zen import builds

from ocl import perceptual_grouping


@dataclasses.dataclass
class PerceptualGroupingConfig:
    """Configuration class of perceptual grouping models."""


SlotAttentionConfig = builds(
    perceptual_grouping.SlotAttentionGrouping,
    builds_bases=(PerceptualGroupingConfig,),
    populate_full_signature=True,
)
StickBreakingGroupingConfig = builds(
    perceptual_grouping.StickBreakingGrouping,
    builds_bases=(PerceptualGroupingConfig,),
    populate_full_signature=True,
)


def register_configs(config_store):
    config_store.store(group="schemas", name="perceptual_grouping", node=PerceptualGroupingConfig)
    config_store.store(group="perceptual_grouping", name="slot_attention", node=SlotAttentionConfig)
    config_store.store(
        group="perceptual_grouping", name="stick_breaking", node=StickBreakingGroupingConfig
    )

ocl/config/perceptual_groupings.py
amazon-science/object-centric-multiple-object-tracking
"""Configuration of plugins."""
import dataclasses
import functools

import hydra_zen
from hydra_zen import builds
from torch.optim.lr_scheduler import LambdaLR

from ocl import plugins, scheduling
from ocl.config.optimizers import OptimizerConfig


@dataclasses.dataclass
class PluginConfig:
    """Base class for plugin configurations."""

    pass


@dataclasses.dataclass
class LRSchedulerConfig:
    pass


def exponential_decay_with_optional_warmup(
    optimizer, decay_rate: float = 1.0, decay_steps: int = 10000, warmup_steps: int = 0
):
    """Return pytorch lighting optimizer configuration for exponential decay with optional warmup.

    Returns:
        Dict with structure compatible with ptl.  See
        https://pytorch-lightning.readthedocs.io/en/stable/api/pytorch_lightning.core.lightning.html#pytorch_lightning.core.lightning.LightningModule.configure_optimizers
    """
    decay_fn = functools.partial(
        scheduling.exp_decay_with_warmup_fn,
        decay_rate=decay_rate,
        decay_steps=decay_steps,
        warmup_steps=warmup_steps,
    )

    return {"lr_scheduler": {"scheduler": LambdaLR(optimizer, decay_fn), "interval": "step"}}


def cosine_annealing_with_optional_warmup(
    optimizer,
    T_max: int = 100000,
    eta_min: float = 0.0,
    warmup_steps: int = 0,
    error_on_exceeding_steps: bool = True,
):
    """Return pytorch lighting optimizer configuration for cosine annealing with warmup.

    Returns:
        Dict with structure compatible with ptl.  See
        https://pytorch-lightning.readthedocs.io/en/stable/api/pytorch_lightning.core.lightning.html#pytorch_lightning.core.lightning.LightningModule.configure_optimizers
    """
    return {
        "lr_scheduler": {
            "scheduler": scheduling.CosineAnnealingWithWarmup(
                optimizer,
                T_max,
                eta_min=eta_min,
                warmup_steps=warmup_steps,
                error_on_exceeding_steps=error_on_exceeding_steps,
            ),
            "interval": "step",
        }
    }


ExpDecayLR = builds(
    exponential_decay_with_optional_warmup,
    zen_partial=True,
    populate_full_signature=True,
    builds_bases=(LRSchedulerConfig,),
)

CosineAnnealingLR = builds(
    cosine_annealing_with_optional_warmup,
    zen_partial=True,
    populate_full_signature=True,
    builds_bases=(LRSchedulerConfig,),
)


@dataclasses.dataclass
class HPSchedulerConfig:
    """Base class for hyperparameter scheduler configuration."""


LinearHPSchedulerConfig = builds(
    scheduling.LinearHPScheduler,
    builds_bases=(HPSchedulerConfig,),
    populate_full_signature=True,
)

StepHPSchedulerConfig = builds(
    scheduling.StepHPScheduler,
    builds_bases=(HPSchedulerConfig,),
    populate_full_signature=True,
)


builds_plugin = hydra_zen.make_custom_builds_fn(
    populate_full_signature=True,
)
OptimizationConfig = builds_plugin(
    plugins.Optimization,
    optimizer=OptimizerConfig,
    lr_scheduler=LRSchedulerConfig,
    builds_bases=(PluginConfig,),
)
SingleElementPreprocessingConfig = builds_plugin(
    plugins.SingleElementPreprocessing, builds_bases=(PluginConfig,)
)
MultiElementPreprocessingConfig = builds_plugin(
    plugins.MultiElementPreprocessing, builds_bases=(PluginConfig,)
)
DataPreprocessingConfig = builds_plugin(plugins.DataPreprocessing, builds_bases=(PluginConfig,))
SubsetDatasetConfig = builds_plugin(plugins.SubsetDataset, builds_bases=(PluginConfig,))
SampleFramesFromVideoConfig = builds_plugin(
    plugins.SampleFramesFromVideo, builds_bases=(PluginConfig,)
)
SplitConsecutiveFramesConfig = builds_plugin(
    plugins.SplitConsecutiveFrames, builds_bases=(PluginConfig,)
)
RandomStridedWindowConfig = builds_plugin(plugins.RandomStridedWindow, builds_bases=(PluginConfig,))


def register_configs(config_store):
    config_store.store(group="schemas", name="lr_scheduler", node=LRSchedulerConfig)
    config_store.store(group="lr_schedulers", name="exponential_decay", node=ExpDecayLR)
    config_store.store(group="lr_schedulers", name="cosine_annealing", node=CosineAnnealingLR)

    config_store.store(group="schemas", name="hp_scheduler", node=HPSchedulerConfig)
    config_store.store(group="hp_schedulers", name="linear", node=LinearHPSchedulerConfig)
    config_store.store(group="hp_schedulers", name="step", node=StepHPSchedulerConfig)

    config_store.store(group="schemas", name="plugin", node=PluginConfig)
    config_store.store(group="plugins", name="optimization", node=OptimizationConfig)
    config_store.store(
        group="plugins",
        name="single_element_preprocessing",
        node=SingleElementPreprocessingConfig,
    )
    config_store.store(
        group="plugins",
        name="multi_element_preprocessing",
        node=MultiElementPreprocessingConfig,
    )
    config_store.store(
        group="plugins",
        name="data_preprocessing",
        node=DataPreprocessingConfig,
    )
    config_store.store(group="plugins", name="subset_dataset", node=SubsetDatasetConfig)
    config_store.store(
        group="plugins", name="sample_frames_from_video", node=SampleFramesFromVideoConfig
    )
    config_store.store(
        group="plugins", name="split_consecutive_frames", node=SplitConsecutiveFramesConfig
    )
    config_store.store(group="plugins", name="random_strided_window", node=RandomStridedWindowConfig)


def _torchvision_interpolation_mode(mode):
    import torchvision

    return torchvision.transforms.InterpolationMode[mode.upper()]


def register_resolvers(omegaconf):
    omegaconf.register_new_resolver(
        "torchvision_interpolation_mode", _torchvision_interpolation_mode
    )

ocl/config/plugins.py
amazon-science/object-centric-multiple-object-tracking
"""Perceptual grouping models."""
import dataclasses

from hydra_zen import builds

from ocl import predictor


@dataclasses.dataclass
class PredictorConfig:
    """Configuration class of Predictor."""


TransitionConfig = builds(
    predictor.TransformerBlock,
    builds_bases=(PredictorConfig,),
    populate_full_signature=True,
)



def register_configs(config_store):
    config_store.store(group="schemas", name="perceptual_grouping", node=PredictorConfig)
    config_store.store(group="perceptual_grouping", name="slot_attention", node=TransitionConfig)


ocl/config/predictor.py
amazon-science/object-centric-multiple-object-tracking
"""Utility functions useful for configuration."""
import ast
from typing import Any, Callable

from hydra_zen import builds

from ocl.config.feature_extractors import FeatureExtractorConfig
from ocl.config.perceptual_groupings import PerceptualGroupingConfig
from ocl.config.predictor import PredictorConfig
from ocl.distillation import EMASelfDistillation
from ocl.utils import Combined, CreateSlotMask, Recurrent


def lambda_string_to_function(function_string: str) -> Callable[..., Any]:
    """Convert string of the form "lambda x: x" into a callable Python function."""
    # This is a bit hacky but ensures that the syntax of the input is correct and contains
    # a valid lambda function definition without requiring to run `eval`.
    parsed = ast.parse(function_string)
    is_lambda = isinstance(parsed.body[0], ast.Expr) and isinstance(parsed.body[0].value, ast.Lambda)
    if not is_lambda:
        raise ValueError(f"'{function_string}' is not a valid lambda definition.")

    return eval(function_string)


class ConfigDefinedLambda:
    """Lambda function defined in the config.

    This allows lambda functions defined in the config to be pickled.
    """

    def __init__(self, function_string: str):
        self.__setstate__(function_string)

    def __getstate__(self) -> str:
        return self.function_string

    def __setstate__(self, function_string: str):
        self.function_string = function_string
        self._fn = lambda_string_to_function(function_string)

    def __call__(self, *args, **kwargs):
        return self._fn(*args, **kwargs)


def eval_lambda(function_string, *args):
    lambda_fn = lambda_string_to_function(function_string)
    return lambda_fn(*args)


FunctionConfig = builds(ConfigDefinedLambda, populate_full_signature=True)

# Inherit from all so it can be used in place of any module.
CombinedConfig = builds(
    Combined,
    populate_full_signature=True,
    builds_bases=(FeatureExtractorConfig, PerceptualGroupingConfig, PredictorConfig),
)
RecurrentConfig = builds(
    Recurrent,
    populate_full_signature=True,
    builds_bases=(FeatureExtractorConfig, PerceptualGroupingConfig, PredictorConfig),
)
CreateSlotMaskConfig = builds(CreateSlotMask, populate_full_signature=True)


EMASelfDistillationConfig = builds(
    EMASelfDistillation,
    populate_full_signature=True,
    builds_bases=(FeatureExtractorConfig, PerceptualGroupingConfig, PredictorConfig),
)


def register_configs(config_store):
    config_store.store(group="schemas", name="lambda_fn", node=FunctionConfig)
    config_store.store(group="utils", name="combined", node=CombinedConfig)
    config_store.store(group="utils", name="selfdistillation", node=EMASelfDistillationConfig)
    config_store.store(group="utils", name="recurrent", node=RecurrentConfig)
    config_store.store(group="utils", name="create_slot_mask", node=CreateSlotMaskConfig)


def register_resolvers(omegaconf):
    omegaconf.register_new_resolver("lambda_fn", ConfigDefinedLambda)
    omegaconf.register_new_resolver("eval_lambda", eval_lambda)

ocl/config/utils.py
amazon-science/object-centric-multiple-object-tracking
"""Modules to compute the IoU matching cost and solve the corresponding LSAP."""
import numpy as np
import torch
from scipy.optimize import linear_sum_assignment
from torch import nn


class HungarianMatcher(nn.Module):
    """This class computes an assignment between the targets and the predictions of the network."""

    @torch.no_grad()
    def forward(self, mask_preds, mask_targets):
        """Performs the matching.

        Params:
            mask_preds: Tensor of dim [batch_size, n_objects, N, N] with the predicted masks
            mask_targets: Tensor of dim [batch_size, n_objects, N, N]
            with the target masks from another augmentation

        Returns:
            A list of size batch_size, containing tuples of (index_i, index_j) where:
                - index_i is the indices of the selected predictions
                - index_j is the indices of the corresponding selected targets
        """
        bs, n_objects, _, _ = mask_preds.shape
        # Compute the iou cost betwen masks
        cost_iou = -get_iou_matrix(mask_preds, mask_targets)
        cost_iou = cost_iou.reshape(bs, n_objects, bs, n_objects).cpu()
        self.costs = torch.stack([cost_iou[i, :, i, :][None] for i in range(bs)])
        indices = [linear_sum_assignment(c[0]) for c in self.costs]
        return torch.as_tensor(np.array(indices))


def get_iou_matrix(preds, targets):

    bs, n_objects, H, W = targets.shape
    targets = targets.reshape(bs * n_objects, H * W).float()
    preds = preds.reshape(bs * n_objects, H * W).float()

    intersection = torch.matmul(targets, preds.t())
    targets_area = targets.sum(dim=1).view(1, -1)
    preds_area = preds.sum(dim=1).view(1, -1)
    union = (targets_area.t() + preds_area) - intersection

    return torch.where(
        union == 0,
        torch.tensor(0.0, device=targets.device),
        intersection / union,
    )

ocl/consistency.py
amazon-science/object-centric-multiple-object-tracking
"""Implementation of datasets."""
import collections
import logging
import math
import os
from functools import partial
from itertools import chain
from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union

import numpy as np
import pytorch_lightning as pl
import torch
import webdataset
from torch.utils.data._utils import collate as torch_collate

from ocl import base
from ocl.hooks import FakeHooks, hook_implementation

LOGGER = logging.getLogger(__name__)


def filter_keys(d: dict, keys_to_keep=tuple):
    """Filter dict for keys in keys_to_keep."""
    keys_to_keep = ("_",) + keys_to_keep
    return {
        key: value
        for key, value in d.items()
        if any(key.startswith(prefix) for prefix in keys_to_keep)
    }


def combine_keys(list_of_key_tuples: Sequence[Tuple[str]]):
    return tuple(set(chain.from_iterable(list_of_key_tuples)))


class WebdatasetDataModule(pl.LightningDataModule):
    """Imagenet Data Module."""

    def __init__(
        self,
        train_shards: Optional[Union[str, List[str]]] = None,
        val_shards: Optional[Union[str, List[str]]] = None,
        test_shards: Optional[Union[str, List[str]]] = None,
        batch_size: int = 32,
        eval_batch_size: Optional[int] = None,
        num_workers: int = 2,
        train_size: Optional[int] = None,
        val_size: Optional[int] = None,
        test_size: Optional[int] = None,
        shuffle_buffer_size: int = 3000,
        use_autopadding: bool = False,
        continue_on_shard_error: bool = False,
        hooks: Optional[base.PluggyHookRelay] = None,
    ):
        super().__init__()
        if train_shards is None and val_shards is None and test_shards is None:
            raise ValueError("No split was specified. Need to specify at least one split.")
        self.train_shards = train_shards
        self.val_shards = val_shards
        self.test_shards = test_shards
        self.train_size = train_size
        self.val_size = val_size
        self.test_size = test_size
        self.batch_size = batch_size
        self.eval_batch_size = eval_batch_size if eval_batch_size is not None else batch_size
        self.num_workers = num_workers
        self.shuffle_buffer_size = shuffle_buffer_size
        self.continue_on_shard_error = continue_on_shard_error
        self.hooks = hooks if hooks else FakeHooks()

        if use_autopadding:
            self.collate_fn = collate_with_autopadding
        else:
            self.collate_fn = collate_with_batch_size

    @staticmethod
    def _remove_extensions(input_dict):
        def _remove_extension(name: str):
            if name.endswith(".gz"):
                # Webdataset automatically decompresses these, we want to remove two layers of
                # extensions due to that.
                name = os.path.splitext(name)[0]
            return os.path.splitext(name)[0]

        return {_remove_extension(name): value for name, value in input_dict.items()}

    @hook_implementation
    def on_train_epoch_start(self, model) -> None:
        """Set environment variables required for better shuffling."""
        # Required for shuffling of instances across workers, see `epoch_shuffle` parameter of
        # `webdataset.PytorchShardList`.
        os.environ["WDS_EPOCH"] = str(model.current_epoch)

    def _create_webdataset(
        self,
        uri_expression: Union[str, List[str]],
        shuffle=False,
        n_datapoints: Optional[int] = None,
        keys_to_keep: Sequence[str] = tuple(),
        transforms: Sequence[Callable[[webdataset.Processor], webdataset.Processor]] = tuple(),
    ):
        shard_list = webdataset.PytorchShardList(
            uri_expression, shuffle=shuffle, epoch_shuffle=shuffle
        )
        if self.continue_on_shard_error:
            handler = webdataset.warn_and_continue
        else:
            handler = webdataset.reraise_exception
        dataset = webdataset.WebDataset(shard_list, handler=handler)
        # Discard unneeded properties of the elements prior to shuffling and decoding.
        dataset = dataset.map(partial(filter_keys, keys_to_keep=keys_to_keep))

        if shuffle:
            dataset = dataset.shuffle(self.shuffle_buffer_size)

        # Decode files and remove extensions from input as we already decoded the elements. This
        # makes our pipeline invariant to the exact encoding used in the dataset.
        dataset = dataset.decode("rgb8").map(WebdatasetDataModule._remove_extensions)

        # Apply transforms
        for transform in transforms:
            dataset = transform(dataset)
        return dataset.with_length(n_datapoints)

    def _create_dataloader(self, dataset, size, batch_size, partial_batches):
        # Don't return partial batches during training as these give the partial samples a higher
        # weight in the optimization than the other samples of the dataset.
        dataloader = webdataset.WebLoader(
            dataset.batched(
                batch_size,
                partial=partial_batches,
                collation_fn=self.collate_fn,
            ),
            num_workers=self.num_workers,
            batch_size=None,
        )

        if size:
            # This is required for ddp training as we otherwise cannot guarantee that each worker
            # gets the same number of batches.
            equalized_size: int
            if partial_batches:
                # Round up in the case of partial batches.
                equalized_size = int(math.ceil(size / batch_size))
            else:
                equalized_size = size // batch_size

            dataloader = dataloader.ddp_equalize(equalized_size, with_length=True)
        else:
            LOGGER.warning(
                "Size not provided in the construction of webdataset. "
                "This may lead to problems when running distributed training."
            )
        return dataloader

    def train_data_iterator(self):
        if self.train_shards is None:
            raise ValueError("Can not create train_data_iterator. No training split was specified.")
        return self._create_webdataset(
            self.train_shards,
            shuffle=True,
            n_datapoints=self.train_size,
            keys_to_keep=combine_keys(self.hooks.training_fields()),
            transforms=self.hooks.training_transform(),
        )

    def train_dataloader(self):
        return self._create_dataloader(
            self.train_data_iterator(), self.train_size, self.batch_size, partial_batches=False
        )

    def val_data_iterator(self):
        if self.val_shards is None:
            raise ValueError("Can not create val_data_iterator. No val split was specified.")
        return self._create_webdataset(
            self.val_shards,
            shuffle=False,
            n_datapoints=self.val_size,
            keys_to_keep=combine_keys(self.hooks.evaluation_fields()),
            transforms=self.hooks.evaluation_transform(),
        )

    def val_dataloader(self):
        return self._create_dataloader(
            self.val_data_iterator(),
            self.val_size,
            self.eval_batch_size,
            partial_batches=True,
        )

    def test_data_iterator(self):
        if self.test_shards is None:
            raise ValueError("Can not create test_data_iterator. No test split was specified.")
        return self._create_webdataset(
            self.test_shards,
            shuffle=False,
            n_datapoints=self.test_size,
            keys_to_keep=combine_keys(self.hooks.evaluation_fields()),
            transforms=self.hooks.evaluation_transform(),
        )

    def test_dataloader(self):
        return self._create_dataloader(
            self.test_data_iterator(),
            self.test_size,
            self.eval_batch_size,
            partial_batches=True,
        )


class DummyDataModule(pl.LightningDataModule):
    """Dataset providing dummy data for testing."""

    def __init__(
        self,
        data_shapes: Dict[str, List[int]],
        data_types: Dict[str, str],
        hooks: Optional[base.PluggyHookRelay] = None,
        batch_size: int = 4,
        eval_batch_size: Optional[int] = None,
        train_size: Optional[int] = None,
        val_size: Optional[int] = None,
        test_size: Optional[int] = None,
        # Remaining args needed for compatibility with other datamodules
        train_shards: Optional[str] = None,
        val_shards: Optional[str] = None,
        test_shards: Optional[str] = None,
        num_workers: Optional[int] = None,
    ):
        super().__init__()
        self.data_shapes = {key: list(shape) for key, shape in data_shapes.items()}
        self.data_types = data_types
        self.hooks = hooks if hooks else FakeHooks
        self.batch_size = batch_size
        self.eval_batch_size = eval_batch_size if eval_batch_size is not None else batch_size

        self.train_size = train_size
        if self.train_size is None:
            self.train_size = 3 * batch_size + 1
        self.val_size = val_size
        if self.val_size is None:
            self.val_size = 2 * batch_size
        self.test_size = test_size
        if self.test_size is None:
            self.test_size = 2 * batch_size

    @staticmethod
    def _get_random_data_for_dtype(dtype: str, shape: List[int]):
        if dtype == "image":
            return np.random.randint(0, 256, size=shape, dtype=np.uint8)
        elif dtype == "binary":
            return np.random.randint(0, 2, size=shape, dtype=bool)
        elif dtype == "uniform":
            return np.random.rand(*shape).astype(np.float32)
        elif dtype.startswith("categorical_"):
            bounds = [int(b) for b in dtype.split("_")[1:]]
            if len(bounds) == 1:
                lower, upper = 0, bounds[0]
            else:
                lower, upper = bounds
            np_dtype = np.uint8 if upper <= 256 else np.uint64
            return np.random.randint(lower, upper, size=shape, dtype=np_dtype)
        elif dtype.startswith("mask"):
            categories = shape[0]
            np_dtype = np.uint8 if categories <= 256 else np.uint64
            slot_per_pixel = np.random.randint(
                0, categories, size=shape[:1] + shape[2:], dtype=np_dtype
            )
            return (
                np.eye(categories)[slot_per_pixel.reshape(-1)]
                .reshape(shape[:1] + shape[2:] + [categories])
                .transpose([0, 4, 1, 2, 3])
            )
        else:
            raise ValueError(f"Unsupported dtype `{dtype}`")

    def _create_dataset(
        self,
        n_datapoints: int,
        transforms: Sequence[Callable[[Any], Any]],
    ):
        class NumpyDataset(torch.utils.data.IterableDataset):
            def __init__(self, data: Dict[str, np.ndarray], size: int):
                super().__init__()
                self.data = data
                self.size = size

            def __iter__(self):
                for i in range(self.size):
                    elem = {key: value[i] for key, value in self.data.items()}
                    elem["__key__"] = str(i)
                    yield elem

        data = {}
        for key, shape in self.data_shapes.items():
            data[key] = self._get_random_data_for_dtype(self.data_types[key], [n_datapoints] + shape)

        dataset = webdataset.Processor(NumpyDataset(data, n_datapoints), lambda x: x)
        for transform in transforms:
            dataset = transform(dataset)

        return dataset

    def _create_dataloader(self, dataset, batch_size):
        return torch.utils.data.DataLoader(
            dataset, batch_size=batch_size, num_workers=0, collate_fn=collate_with_autopadding
        )

    def train_dataloader(self):
        dataset = self._create_dataset(self.train_size, self.hooks.training_transform())
        return self._create_dataloader(dataset, self.batch_size)

    def val_dataloader(self):
        dataset = self._create_dataset(self.val_size, self.hooks.evaluation_transform())
        return self._create_dataloader(dataset, self.eval_batch_size)

    def test_dataloader(self):
        dataset = self._create_dataset(self.test_size, self.hooks.evaluation_transform())
        return self._create_dataloader(dataset, self.eval_batch_size)


def collate_with_batch_size(batch):
    """Call default pytorch collate function yet for dict type input additionally add batch size."""
    if isinstance(batch[0], collections.abc.Mapping):
        out = torch_collate.default_collate(batch)
        out["batch_size"] = len(batch)
        return out
    return torch_collate.default_collate(batch)


def collate_with_autopadding(batch):
    """Collate function that takes a batch of data and stacks it with a batch dimension.

    In contrast to torch's collate function, this function automatically pads tensors of different
    sizes with zeros such that they can be stacked.

    Adapted from https://github.com/pytorch/pytorch/blob/master/torch/utils/data/_utils/collate.py.
    """
    elem = batch[0]
    elem_type = type(elem)
    if isinstance(elem, torch.Tensor):
        out = None
        # As most tensors will not need padding to be stacked, we first try to stack them normally
        # and pad only if normal padding fails. This avoids explicitly checking whether all tensors
        # have the same shape before stacking.
        try:
            if torch.utils.data.get_worker_info() is not None:
                # If we're in a background process, concatenate directly into a
                # shared memory tensor to avoid an extra copy
                numel = sum(x.numel() for x in batch)
                if len(batch) * elem.numel() != numel:
                    # Check whether resizing will fail because tensors have unequal sizes to avoid
                    # a memory allocation. This is a sufficient but not necessary condition, so it
                    # can happen that this check will not trigger when padding is necessary.
                    raise RuntimeError()
                storage = elem.storage()._new_shared(numel)
                out = elem.new(storage).resize_(len(batch), *elem.shape)
            return torch.stack(batch, 0, out=out)
        except RuntimeError:
            # Stacking did not work. Try to pad tensors to the same dimensionality.
            if not all(x.ndim == elem.ndim for x in batch):
                raise ValueError("Tensors in batch have different number of dimensions.")

            shapes = [x.shape for x in batch]
            max_dims = [max(shape[idx] for shape in shapes) for idx in range(elem.ndim)]

            paddings = []
            for shape in shapes:
                padding = []
                # torch.nn.functional.pad wants padding from last to first dim, so go in reverse
                for idx in reversed(range(len(shape))):
                    padding.append(0)
                    padding.append(max_dims[idx] - shape[idx])
                paddings.append(padding)

            batch_padded = [
                torch.nn.functional.pad(x, pad, mode="constant", value=0.0)
                for x, pad in zip(batch, paddings)
            ]

            if torch.utils.data.get_worker_info() is not None:
                # If we're in a background process, concatenate directly into a
                # shared memory tensor to avoid an extra copy
                numel = sum(x.numel() for x in batch_padded)
                storage = elem.storage()._new_shared(numel)
                out = elem.new(storage).resize_(len(batch_padded), *batch_padded[0].shape)
            return torch.stack(batch_padded, 0, out=out)
    elif (
        elem_type.__module__ == "numpy"
        and elem_type.__name__ != "str_"
        and elem_type.__name__ != "string_"
    ):
        if elem_type.__name__ == "ndarray" or elem_type.__name__ == "memmap":
            # array of string classes and object
            if torch_collate.np_str_obj_array_pattern.search(elem.dtype.str) is not None:
                raise TypeError(torch_collate.default_collate_err_msg_format.format(elem.dtype))

            return collate_with_autopadding([torch.as_tensor(b) for b in batch])
        elif elem.shape == ():  # scalars
            return torch.as_tensor(batch)
    elif isinstance(elem, float):
        return torch.tensor(batch, dtype=torch.float64)
    elif isinstance(elem, int):
        return torch.tensor(batch)
    elif isinstance(elem, str):
        return batch
    elif isinstance(elem, collections.abc.Mapping):
        out = {key: collate_with_autopadding([d[key] for d in batch]) for key in elem}
        out["batch_size"] = len(batch)
        try:
            return elem_type(out)
        except TypeError:
            # The mapping type may not support `__init__(iterable)`.
            return out
    elif isinstance(elem, tuple) and hasattr(elem, "_fields"):  # namedtuple
        return elem_type(*(collate_with_autopadding(samples) for samples in zip(*batch)))
    elif isinstance(elem, collections.abc.Sequence):
        # check to make sure that the elements in batch have consistent size
        it = iter(batch)
        elem_size = len(next(it))
        if not all(len(elem) == elem_size for elem in it):
            raise RuntimeError("each element in list of batch should be of equal size")
        transposed = list(zip(*batch))  # It may be accessed twice, so we use a list.

        if isinstance(elem, tuple):
            return [
                collate_with_autopadding(samples) for samples in transposed
            ]  # Backwards compatibility.
        else:
            try:
                return elem_type([collate_with_autopadding(samples) for samples in transposed])
            except TypeError:
                # The sequence type may not support `__init__(iterable)` (e.g., `range`).
                return [collate_with_autopadding(samples) for samples in transposed]

    raise TypeError(torch_collate.default_collate_err_msg_format.format(elem_type))

ocl/datasets.py
amazon-science/object-centric-multiple-object-tracking
"""Implementation of different types of decoders."""
import dataclasses
import math
from typing import Callable, Optional, Tuple, Union

import torch
from torch import nn
from torch.nn import functional as F
from torchtyping import TensorType
from ocl.base import Instances
from ocl import neural_networks
from ocl.path_defaults import OBJECTS
from ocl.utils import RoutableMixin, SoftPositionEmbed, resize_patches_to_image
from ocl.utils import box_cxcywh_to_xyxy

@dataclasses.dataclass
class BBoxOutput:
    bboxes: TensorType["batch_size", "n_objects", "box_dim"]  # noqa: F821
    classes: TensorType["batch_size", "n_objects", "num_classes"]  # noqa: F821
    ori_res_bboxes: TensorType["batch_size", "n_objects", "box_dim"]  # noqa: F821
    inference_obj_idxes: TensorType["batch_size", "n_objects"]  # noqa: F821


@dataclasses.dataclass
class ReconstructionOutput:
    reconstruction: TensorType["batch_size", "channels", "height", "width"]  # noqa: F821
    object_reconstructions: TensorType[
        "batch_size", "n_objects", "channels", "height", "width"  # noqa: F821
    ]
    masks: TensorType["batch_size", "n_objects", "height", "width"]  # noqa: F821
    masks_vis: TensorType["batch_size", "n_objects", "height", "width"]  # noqa: F821
    masks_eval: TensorType["batch_size", "n_objects", "height", "width"]  # noqa: F821


@dataclasses.dataclass
class PatchReconstructionOutput:
    reconstruction: TensorType["batch_size", "n_patches", "n_patch_features"]  # noqa: F821
    masks: TensorType["batch_size", "n_objects", "n_patches"]  # noqa: F821
    masks_as_image: Optional[
        TensorType["batch_size", "n_objects", "height", "width"]  # noqa: F821
    ] = None
    target: Optional[TensorType["batch_size", "n_patches", "n_patch_features"]] = None  # noqa: F821


@dataclasses.dataclass
class DepthReconstructionOutput(ReconstructionOutput):
    masks_amodal: Optional[
        TensorType["batch_size", "n_objects", "height", "width"]  # noqa: F821
    ] = None
    depth_map: Optional[TensorType["batch_size", "height", "width"]] = None  # noqa: F821
    object_depth_map: Optional[
        TensorType["batch_size", "n_objects", "height", "width"]  # noqa: F821
    ] = None
    densities: Optional[
        TensorType["batch_size", "n_objects", "n_depth", "height", "width"]  # noqa: F821
    ] = None
    colors: Optional[
        TensorType["batch_size", "n_objects", "n_depth", "channels", "height", "width"]  # noqa: F821
    ] = None

@dataclasses.dataclass
class OpticalFlowPredictionTaskOutput:
    predicted_flow: TensorType["batch_size", "channels", "height", "width"]  # noqa: F821
    object_flows: TensorType["batch_size", "n_objects", "channels", "height", "width"]  # noqa: F821
    masks: TensorType["batch_size", "n_objects", "height", "width"]  # noqa: F821


def build_grid_of_positions(resolution):
    """Build grid of positions which can be used to create positions embeddings."""
    ranges = [torch.linspace(0.0, 1.0, steps=res) for res in resolution]
    grid = torch.meshgrid(*ranges, indexing="ij")
    grid = torch.stack(grid, dim=-1)
    grid = torch.reshape(grid, [resolution[0], resolution[1], -1])
    return grid


def get_slotattention_decoder_backbone(object_dim: int, output_dim: int = 4):
    """Get CNN decoder backbone form the original slot attention paper."""
    return nn.Sequential(
        nn.ConvTranspose2d(object_dim, 64, 5, stride=2, padding=2, output_padding=1),
        nn.ReLU(inplace=True),
        nn.ConvTranspose2d(64, 64, 5, stride=2, padding=2, output_padding=1),
        nn.ReLU(inplace=True),
        nn.ConvTranspose2d(64, 64, 5, stride=2, padding=2, output_padding=1),
        nn.ReLU(inplace=True),
        nn.ConvTranspose2d(64, 64, 5, stride=2, padding=2, output_padding=1),
        nn.ReLU(inplace=True),
        nn.ConvTranspose2d(64, 64, 5, stride=1, padding=2, output_padding=0),
        nn.ReLU(inplace=True),
        nn.ConvTranspose2d(64, output_dim, 3, stride=1, padding=1, output_padding=0),
    )


def get_savi_decoder_backbone(
    object_dim: int,
    output_dim: int = 4,
    larger_input_arch: bool = False,
    channel_multiplier: float = 1,
):
    """Get CNN decoder backbone form the slot attention for video paper."""
    channels = int(64 * channel_multiplier)
    if larger_input_arch:
        output_stride = 2
        output_padding = 1
    else:
        output_stride = 1
        output_padding = 0
    return nn.Sequential(
        nn.ConvTranspose2d(object_dim, channels, 5, stride=2, padding=2, output_padding=1),
        nn.ReLU(inplace=True),
        nn.ConvTranspose2d(channels, channels, 5, stride=2, padding=2, output_padding=1),
        nn.ReLU(inplace=True),
        nn.ConvTranspose2d(channels, channels, 5, stride=2, padding=2, output_padding=1),
        nn.ReLU(inplace=True),
        nn.ConvTranspose2d(
            channels, channels, 5, stride=output_stride, padding=2, output_padding=output_padding
        ),
        nn.ReLU(inplace=True),
        nn.ConvTranspose2d(
            channels,
            output_dim,
            1,
            stride=1,
            padding=0,
            output_padding=0,
        ),
    )


class StyleGANv2Decoder(nn.Module):
    """CNN decoder as used in StyleGANv2 and GIRAFFE."""

    def __init__(
        self,
        object_feature_dim: int,
        output_dim: int = 4,
        min_features=32,
        input_size: int = 8,
        output_size: int = 128,
        activation_fn: str = "leaky_relu",
        leaky_relu_slope: float = 0.2,
    ):
        super().__init__()
        input_size_log2 = math.log2(input_size)
        assert math.floor(input_size_log2) == input_size_log2, "Input size needs to be power of 2"

        output_size_log2 = math.log2(output_size)
        assert math.floor(output_size_log2) == output_size_log2, "Output size needs to be power of 2"

        n_blocks = int(output_size_log2 - input_size_log2)

        self.convs = nn.ModuleList()
        cur_dim = object_feature_dim
        for _ in range(n_blocks):
            next_dim = max(cur_dim // 2, min_features)
            self.convs.append(nn.Conv2d(cur_dim, next_dim, 3, stride=1, padding=1))
            cur_dim = next_dim

        self.skip_convs = nn.ModuleList()
        cur_dim = object_feature_dim
        for _ in range(n_blocks + 1):
            self.skip_convs.append(nn.Conv2d(cur_dim, output_dim, 1, stride=1))
            cur_dim = max(cur_dim // 2, min_features)

        nn.init.zeros_(self.skip_convs[-1].bias)

        if activation_fn == "relu":
            self.activation_fn = nn.ReLU(inplace=True)
        elif activation_fn == "leaky_relu":
            self.activation_fn = nn.LeakyReLU(leaky_relu_slope, inplace=True)
        else:
            raise ValueError(f"Unknown activation function {activation_fn}")

    def forward(self, inp):
        output = self.skip_convs[0](inp)

        features = inp
        for conv, skip_conv in zip(self.convs, self.skip_convs[1:]):
            features = F.interpolate(features, scale_factor=2, mode="nearest-exact")
            features = conv(features)
            features = self.activation_fn(features)

            output = F.interpolate(
                output, scale_factor=2, mode="bilinear", align_corners=False, antialias=True
            )
            output = output + skip_conv(features)

        return output


class SlotAttentionDecoder(nn.Module, RoutableMixin):
    """Decoder used in the original slot attention paper."""

    def __init__(
        self,
        decoder: nn.Module,
        final_activation: Union[str, Callable] = "identity",
        positional_embedding: Optional[nn.Module] = None,
        object_features_path: Optional[str] = OBJECTS,
    ):
        nn.Module.__init__(self)
        RoutableMixin.__init__(self, {"object_features": object_features_path})
        self.initial_conv_size = (8, 8)
        self.decoder = decoder
        self.final_activation = neural_networks.get_activation_fn(final_activation)
        self.positional_embedding = positional_embedding
        if positional_embedding:
            self.register_buffer("grid", build_grid_of_positions(self.initial_conv_size))

    def rescale_mask(self, mask):
        max = torch.max(mask)
        min = torch.min(mask)
        mask_new = (mask-min) / (max - min)
        return mask_new

    @RoutableMixin.route
    def forward(self, object_features: torch.Tensor):
        assert object_features.dim() >= 3  # Image or video data.
        initial_shape = object_features.shape[:-1]

        object_features_ori = object_features.clone()

        object_features = object_features.flatten(0, -2)
        object_features = (
            object_features.unsqueeze(-1).unsqueeze(-1).expand(-1, -1, *self.initial_conv_size)
        )
        if self.positional_embedding:
            object_features = self.positional_embedding(object_features, self.grid.unsqueeze(0))

        # Apply deconvolution and restore object dimension.
        output = self.decoder(object_features)
        output = output.unflatten(0, initial_shape)

        # Split out alpha channel and normalize over slots.
        # The decoder is assumed to output tensors in CNN order, i.e. * x C x H x W.
        rgb, alpha = output.split([3, 1], dim=-3)
        rgb = self.final_activation(rgb)
        alpha1 = alpha.softmax(dim=-4)    # visible masks
        alpha2 = alpha.sigmoid()         # amodal masks

        n_objects = alpha1.shape[1]
        masks_argmax = alpha1.argmax(dim=1)[:, None]
        classes = torch.arange(n_objects)[None, :, None, None, None].to(masks_argmax)
        masks_one_hot = masks_argmax == classes

        # remove zero object features for reconstruction
        # rec = []
        # masks = torch.zeros(alpha2.shape).to(alpha2.device)
        masks_vis = torch.zeros(alpha1.shape).to(alpha1.device)
        for b in range(object_features_ori.shape[0]):
            index = torch.sum(object_features_ori[b], dim=-1).nonzero(as_tuple=True)[0]
            # zero_index = (torch.sum(object_features_ori[b], dim=-1) == 0).nonzero(as_tuple=True)[0]
            # masks[b][index] =alpha2[b][index]
            masks_vis[b][index] = alpha1[b][index]
            for i in index:
                masks_vis[b][i] = self.rescale_mask(alpha1[b][i])
            # reconstruction = (rgb[b] * alpha1[b])[index].sum(-4)
            # rec.append(reconstruction.unsqueeze(0))
        # recon = torch.stack(rec, dim=0).squeeze(1)
        return ReconstructionOutput(
            # Combine rgb weighted according to alpha channel.
            reconstruction=(rgb * alpha1).sum(-4),
            object_reconstructions=rgb,
            masks=alpha2.squeeze(-3),
            masks_vis = alpha1.squeeze(-3),
            masks_eval = masks_vis.squeeze(-3),
        )


class SlotAttentionOpticalFlowDecoder(nn.Module, RoutableMixin):
    # TODO(flwenzel): for now use the same decoder as for rbg reconstruction. Might implement
    # improved/specialized decoder.
    # TODO(hornmax): Maybe we can merge this with the RGB decoder and generalize the task outputs.
    # This is something for a later time though.

    def __init__(
        self,
        decoder: nn.Module,
        positional_embedding: Optional[nn.Module] = None,
        object_features_path: Optional[str] = OBJECTS,
    ):
        nn.Module.__init__(self)
        RoutableMixin.__init__(self, {"object_features": object_features_path})
        self.initial_conv_size = (8, 8)
        self.decoder = decoder
        self.positional_embedding = positional_embedding
        if positional_embedding:
            self.register_buffer("grid", build_grid_of_positions(self.initial_conv_size))

    @RoutableMixin.route
    def forward(self, object_features: torch.Tensor):
        assert object_features.dim() >= 3  # Image or video data.
        initial_shape = object_features.shape[:-1]
        object_features = object_features.flatten(0, -2)

        object_features = (
            object_features.unsqueeze(-1).unsqueeze(-1).expand(-1, -1, *self.initial_conv_size)
        )
        if self.positional_embedding:
            object_features = self.positional_embedding(object_features, self.grid.unsqueeze(0))

        # Apply deconvolution and restore object dimension.
        output = self.decoder(object_features)
        output = output.unflatten(0, initial_shape)

        # Split out alpha channel and normalize over slots.
        # The decoder is assumed to output tensors in CNN order, i.e. * x C x H x W.
        flow, alpha = output.split([2, 1], dim=2)  # flow is assumed to be 2-dim.
        alpha = alpha.softmax(dim=-4)

        return OpticalFlowPredictionTaskOutput(
            # Combine rgb weighted according to alpha channel.
            predicted_flow=(flow * alpha).sum(-4),
            object_flows=flow,
            masks=alpha.squeeze(-3),
        )


class PatchDecoder(nn.Module, RoutableMixin):
    """Decoder that takes object representations and reconstructs patches.

    Args:
        object_dim: Dimension of objects representations.
        output_dim: Dimension of each patch.
        num_patches: Number of patches P to reconstruct.
        decoder: Function that returns backbone to use for decoding. Function takes input and output
            dimensions and should return module that takes inputs of shape (B * K), P, N, and produce
            outputs of shape (B * K), P, M, where K is the number of objects, N is the number of
            input dimensions and M the number of output dimensions.
        decoder_input_dim: Input dimension to decoder backbone. If specified, a linear
            transformation from object to decoder dimension is added. If not specified, the object
            dimension is used and no linear transform is added.
    """

    def __init__(
        self,
        object_dim: int,
        output_dim: int,
        num_patches: int,
        decoder: Callable[[int, int], nn.Module],
        decoder_input_dim: Optional[int] = None,
        upsample_target: Optional[float] = None,
        object_features_path: Optional[str] = OBJECTS,
        target_path: Optional[str] = None,
        image_path: Optional[str] = None,
    ):
        nn.Module.__init__(self)
        RoutableMixin.__init__(
            self,
            {"object_features": object_features_path, "target": target_path, "image": image_path},
        )
        self.output_dim = output_dim
        self.num_patches = num_patches
        self.upsample_target = upsample_target

        if decoder_input_dim is not None:
            self.inp_transform = nn.Linear(object_dim, decoder_input_dim, bias=True)
            nn.init.xavier_uniform_(self.inp_transform.weight)
            nn.init.zeros_(self.inp_transform.bias)
        else:
            self.inp_transform = None
            decoder_input_dim = object_dim

        self.decoder = decoder(decoder_input_dim, output_dim + 1)
        self.pos_embed = nn.Parameter(torch.randn(1, num_patches, decoder_input_dim) * 0.02)

    @RoutableMixin.route
    def forward(
        self,
        object_features: torch.Tensor,
        target: Optional[torch.Tensor] = None,
        image: Optional[torch.Tensor] = None,
    ):
        assert object_features.dim() >= 3  # Image or video data.
        if self.upsample_target is not None and target is not None:
            target = (
                resize_patches_to_image(
                    target.detach().transpose(-2, -1),
                    scale_factor=self.upsample_target,
                    resize_mode="bilinear",
                )
                .flatten(-2, -1)
                .transpose(-2, -1)
            )

        initial_shape = object_features.shape[:-1]
        object_features = object_features.flatten(0, -2)

        if self.inp_transform is not None:
            object_features = self.inp_transform(object_features)

        object_features = object_features.unsqueeze(1).expand(-1, self.num_patches, -1)

        # Simple learned additive embedding as in ViT
        object_features = object_features + self.pos_embed

        output = self.decoder(object_features)
        output = output.unflatten(0, initial_shape)

        # Split out alpha channel and normalize over slots.
        decoded_patches, alpha = output.split([self.output_dim, 1], dim=-1)
        alpha = alpha.softmax(dim=-3)

        reconstruction = torch.sum(decoded_patches * alpha, dim=-3)
        masks = alpha.squeeze(-1)

        if image is not None:
            masks_as_image = resize_patches_to_image(
                masks, size=image.shape[-1], resize_mode="bilinear"
            )
        else:
            masks_as_image = None

        return PatchReconstructionOutput(
            reconstruction=reconstruction,
            masks=alpha.squeeze(-1),
            masks_as_image=masks_as_image,
            target=target if target is not None else None,
        )


class AutoregressivePatchDecoder(nn.Module, RoutableMixin):
    """Decoder that takes object representations and reconstructs patches autoregressively.

    Args:
        object_dim: Dimension of objects representations.
        output_dim: Dimension of each patch.
        num_patches: Number of patches P to reconstruct.
        decoder: Function that returns backbone to use for decoding. Function takes input and output
            dimensions and should return module that takes autoregressive targets of shape B, P, M,
            conditioning of shape B, K, N, masks of shape P, P, and produces outputs of shape
            B, P, M, where K is the number of objects, N is the number of input dimensions and M the
            number of output dimensions.
        decoder_cond_dim: Dimension of conditioning input of decoder backbone. If specified, a linear
            transformation from object to decoder dimension is added. If not specified, the object
            dimension is used and no linear transform is added.
    """

    def __init__(
        self,
        object_dim: int,
        output_dim: int,
        num_patches: int,
        decoder: Callable[[int, int], nn.Module],
        decoder_dim: Optional[int] = None,
        decoder_cond_dim: Optional[int] = None,
        upsample_target: Optional[float] = None,
        use_decoder_masks: bool = False,
        use_input_transform: bool = False,
        use_output_transform: bool = False,
        use_positional_embedding: bool = False,
        object_features_path: Optional[str] = OBJECTS,
        masks_path: Optional[str] = "perceptual_grouping.masks",
        target_path: Optional[str] = None,
        image_path: Optional[str] = None,
        empty_objects_path: Optional[str] = None,
    ):
        nn.Module.__init__(self)
        RoutableMixin.__init__(
            self,
            {
                "object_features": object_features_path,
                "target": target_path,
                "image": image_path,
                "masks": masks_path,
                "empty_objects": empty_objects_path,
            },
        )
        self.output_dim = output_dim
        self.num_patches = num_patches
        self.upsample_target = upsample_target
        self.use_decoder_masks = use_decoder_masks

        if decoder_dim is None:
            decoder_dim = output_dim

        self.decoder = decoder(decoder_dim, decoder_dim)
        self.bos_token = nn.Parameter(torch.randn(1, 1, output_dim) * output_dim**-0.5)

        if decoder_cond_dim is not None:
            self.cond_transform = nn.Sequential(
                nn.Linear(object_dim, decoder_cond_dim, bias=False),
                nn.LayerNorm(decoder_cond_dim, eps=1e-5),
            )
            nn.init.xavier_uniform_(self.cond_transform[0].weight)
        else:
            decoder_cond_dim = object_dim
            self.cond_transform = nn.LayerNorm(decoder_cond_dim, eps=1e-5)

        if use_input_transform:
            self.inp_transform = nn.Sequential(
                nn.Linear(output_dim, decoder_dim, bias=False),
                nn.LayerNorm(decoder_dim, eps=1e-5),
            )
            nn.init.xavier_uniform_(self.inp_transform[0].weight)
        else:
            self.inp_transform = None

        if use_output_transform:
            self.outp_transform = nn.Linear(decoder_dim, output_dim)
            nn.init.xavier_uniform_(self.outp_transform.weight)
            nn.init.zeros_(self.outp_transform.bias)
        else:
            self.outp_transform = None

        if use_positional_embedding:
            self.pos_embed = nn.Parameter(
                torch.randn(1, num_patches, decoder_dim) * decoder_dim**-0.5
            )
        else:
            self.pos_embed = None

        mask = torch.triu(torch.full((num_patches, num_patches), float("-inf")), diagonal=1)
        self.register_buffer("mask", mask)

    @RoutableMixin.route
    def forward(
        self,
        object_features: torch.Tensor,
        masks: torch.Tensor,
        target: torch.Tensor,
        image: Optional[torch.Tensor] = None,
        empty_objects: Optional[torch.Tensor] = None,
    ) -> PatchReconstructionOutput:
        assert object_features.dim() >= 3  # Image or video data.
        if self.upsample_target is not None and target is not None:
            target = (
                resize_patches_to_image(
                    target.detach().transpose(-2, -1),
                    scale_factor=self.upsample_target,
                    resize_mode="bilinear",
                )
                .flatten(-2, -1)
                .transpose(-2, -1)
            )
        # Squeeze frames into batch if present.
        object_features = object_features.flatten(0, -3)

        object_features = self.cond_transform(object_features)

        # Squeeze frame into batch size if necessary.
        initial_targets_shape = target.shape[:-2]
        targets = target.flatten(0, -3)

        bs = len(object_features)
        inputs = torch.cat((self.bos_token.expand(bs, -1, -1), targets[:, :-1].detach()), dim=1)

        if self.inp_transform is not None:
            inputs = self.inp_transform(inputs)

        if self.pos_embed is not None:
            # Simple learned additive embedding as in ViT
            inputs = inputs + self.pos_embed

        if empty_objects is not None:
            outputs = self.decoder(
                inputs, object_features, self.mask, memory_key_padding_mask=empty_objects
            )
        else:
            outputs = self.decoder(inputs, object_features, self.mask)

        if self.use_decoder_masks:
            decoded_patches, masks = outputs
        else:
            decoded_patches = outputs

        if self.outp_transform is not None:
            decoded_patches = self.outp_transform(decoded_patches)

        decoded_patches = decoded_patches.unflatten(0, initial_targets_shape)

        if image is not None:
            masks_as_image = resize_patches_to_image(
                masks, size=image.shape[-1], resize_mode="bilinear"
            )
        else:
            masks_as_image = None

        return PatchReconstructionOutput(
            reconstruction=decoded_patches, masks=masks, masks_as_image=masks_as_image, target=target
        )


class DensityPredictingSlotAttentionDecoder(nn.Module, RoutableMixin):
    """Decoder predicting color and densities along a ray into the scene."""

    def __init__(
        self,
        object_dim: int,
        decoder: nn.Module,
        depth_positions: int,
        white_background: bool = False,
        normalize_densities_along_slots: bool = False,
        initial_alpha: float = None,
        object_features_path: Optional[str] = OBJECTS,
    ):
        nn.Module.__init__(self)
        RoutableMixin.__init__(self, {"object_features": object_features_path})
        self.initial_conv_size = (8, 8)
        self.depth_positions = depth_positions
        self.white_background = white_background
        self.normalize_densities_along_slots = normalize_densities_along_slots
        self.register_buffer("grid", build_grid_of_positions(self.initial_conv_size))
        self.pos_embedding = SoftPositionEmbed(2, object_dim, cnn_channel_order=True)

        self.decoder = decoder
        if isinstance(self.decoder, nn.Sequential) and hasattr(self.decoder[-1], "bias"):
            nn.init.zeros_(self.decoder[-1].bias)

        if initial_alpha is not None:
            # Distance between neighboring ray points, currently assumed to be 1
            point_distance = 1
            # Value added to density output of network before softplus activation. If network outputs
            # are approximately zero, the initial mask value per voxel becomes `initial_alpha`. See
            # https://arxiv.org/abs/2111.11215 for a derivation.
            self.initial_density_offset = math.log((1 - initial_alpha) ** (-1 / point_distance) - 1)
        else:
            self.initial_density_offset = 0.0

    def _render_objectwise(self, densities, rgbs):
        """Render objects individually.

        Args:
            densities: Predicted densities of shape (B, S, Z, H, W), where S is the number of slots
                and Z is the number of depth positions.
            rgbs: Predicted color values of shape (B, S, 3, H, W), where S is the number of slots.
            background: Optional background to render on.
        """
        densities_objectwise = densities.flatten(0, 1).unsqueeze(2)
        rgbs_objectwise = rgbs.flatten(0, 1).unsqueeze(1)
        rgbs_objectwise = rgbs_objectwise.expand(-1, densities_objectwise.shape[1], -1, -1, -1)

        if self.white_background:
            background = torch.full_like(rgbs_objectwise[:, 0], 1.0)  # White color, i.e. 0xFFFFFF
        else:
            background = None

        object_reconstructions, _, object_masks_per_depth, p_ray_hits_points = volume_rendering(
            densities_objectwise, rgbs_objectwise, background=background
        )

        object_reconstructions = object_reconstructions.unflatten(0, rgbs.shape[:2])
        object_masks_per_depth = object_masks_per_depth.squeeze(2).unflatten(0, rgbs.shape[:2])
        p_ray_hits_points = p_ray_hits_points.squeeze(2).unflatten(0, rgbs.shape[:2])

        p_ray_hits_points_and_reflects = p_ray_hits_points * object_masks_per_depth
        object_masks, object_depth_map = p_ray_hits_points_and_reflects.max(2)

        return object_reconstructions, object_masks, object_depth_map

    @RoutableMixin.route
    def forward(self, object_features: torch.Tensor):
        # TODO(hornmax): Adapt this for video data.
        # Reshape object dimension into batch dimension and broadcast.
        bs, n_objects, object_feature_dim = object_features.shape
        object_features = object_features.view(bs * n_objects, object_feature_dim, 1, 1).expand(
            -1, -1, *self.initial_conv_size
        )
        object_features = self.pos_embedding(object_features, self.grid.unsqueeze(0))

        # Apply deconvolution and restore object dimension.
        output = self.decoder(object_features)
        output = output.view(bs, n_objects, *output.shape[-3:])

        # Split rgb and density channels and transform to appropriate ranges.
        rgbs, densities = output.split([3, self.depth_positions], dim=2)
        rgbs = torch.sigmoid(rgbs)  # B x S x 3 x H x W
        densities = F.softplus(densities + self.initial_density_offset)  # B x S x Z x H x W

        if self.normalize_densities_along_slots:
            densities_depthwise_sum = torch.einsum("bszhw -> bzhw", densities).unsqueeze(1)
            densities_weighted = densities * F.softmax(densities, dim=1)
            densities_weighted_sum = torch.einsum("bszhw -> bzhw", densities_weighted).unsqueeze(1)
            densities = densities_weighted * densities_depthwise_sum / densities_weighted_sum

        # Combine densities from different slots by summing over slot dimension
        density = torch.einsum("bszhw -> bzhw", densities).unsqueeze(2)
        # Combine colors from different slots by density-weighted mean
        rgb = torch.einsum("bszhw, bschw -> bzchw", densities, rgbs) / density

        if self.white_background:
            background = torch.full_like(rgb[:, 0], 1.0)  # White color, i.e. 0xFFFFFF
        else:
            background = None

        reconstruction, _, _, p_ray_hits_point = volume_rendering(
            density, rgb, background=background
        )

        if self.training:
            # Get object masks by taking the max density over all depth positions
            masks = 1 - torch.exp(-densities.detach().max(dim=2).values)
            object_reconstructions = rgbs.detach() * masks.unsqueeze(2)

            if background is not None:
                masks = torch.cat((masks, p_ray_hits_point[:, -1:, 0]), dim=1)
                object_reconstructions = torch.cat(
                    (object_reconstructions, background[:, None]), dim=1
                )

            return ReconstructionOutput(
                reconstruction=reconstruction,
                object_reconstructions=object_reconstructions,
                masks=masks,
            )
        else:
            object_reconstructions, object_masks, object_depth_map = self._render_objectwise(
                densities, rgbs
            )

            # Joint depth map results from taking minimum depth over objects per pixel, whereas
            # joint mask results from the index of the object with minimum depth
            depth_map, mask_dense = object_depth_map.min(1)

            if background is not None:
                object_reconstructions = torch.cat(
                    (object_reconstructions, background[:, None]), dim=1
                )
                # Assign designated background class wherever the depth map indicates background
                mask_dense[depth_map == self.depth_positions] = n_objects
                n_classes = n_objects + 1
            else:
                n_classes = n_objects

            masks = F.one_hot(mask_dense, num_classes=n_classes)
            masks = masks.squeeze(1).permute(0, 3, 1, 2).contiguous()  # B x C x H x W

            return DepthReconstructionOutput(
                reconstruction=reconstruction,
                object_reconstructions=object_reconstructions,
                masks=masks,
                masks_amodal=object_masks,
                depth_map=depth_map,
                object_depth_map=object_depth_map,
                densities=densities,
                colors=rgbs.unsqueeze(2).expand(-1, -1, self.depth_positions, -1, -1, -1),
            )


def volume_rendering(
    densities: torch.Tensor,
    colors: torch.Tensor,
    distances: Union[float, torch.Tensor] = None,
    background: torch.Tensor = None,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    """Volume render along camera rays (also known as alpha compositing).

    For each ray, assumes input of Z density and C color channels, corresponding to Z points along
    the ray from front to back of the scene.

    Args:
        densities: Tensor of shape (B, Z, 1, ...). Non-negative, real valued density values along
            the ray.
        colors: Tensor of shape (B, Z, C, ...). Color values along the ray.
        distances: Tensor of shape (B, Z, 1, ...). Optional distances between this ray point and
            the next. Can also be a single float value. If not given, distances between all points
            are assumed to be one. The last value corresponds to the distance between the last point
            and the background.
        background: Tensor of shape (B, C, ...). An optional background image that the rendering can
            be put on.

    Returns:
        Tuple of tensors of shape (B, C, ...), (B, Z, C, ...), (B, Z, 1, ...), (B, Z, 1, ...).
        First tensor is the rendered image, second tensor are the rendered images along different
        points of the ray, third tensor the alpha masks for each point of the ray, fourth tensor the
        probabilities of reaching each point of the ray (the transmittances). If background is not
        None, the background is included as the last ray point.
    """
    if distances is None:
        transmittances = torch.exp(-torch.cumsum(densities, dim=1))
        p_ray_reflects = 1.0 - torch.exp(-densities)
    else:
        densities_distance_weighted = densities * distances
        transmittances = torch.exp(-torch.cumsum(densities_distance_weighted, dim=1))
        p_ray_reflects = 1.0 - torch.exp(-densities_distance_weighted)

    # First object has 100% probability of being hit as it cannot be occluded by other objects
    p_ray_hits_point = torch.cat((torch.ones_like(densities[:, :1]), transmittances), dim=1)

    if background is not None:
        background = background.unsqueeze(1)

        # All rays reaching the background reflect
        p_ray_reflects = torch.cat((p_ray_reflects, torch.ones_like(p_ray_reflects[:, :1])), dim=1)
        colors = torch.cat((colors, background), dim=1)
    else:
        p_ray_hits_point = p_ray_hits_point[:, :-1]

    z_images = p_ray_reflects * colors
    image = (p_ray_hits_point * z_images).sum(dim=1)

    return image, z_images, p_ray_reflects, p_ray_hits_point

class BBoxAndClsDecoder(nn.Module, RoutableMixin):
    """Decoder used for multiple object tracking."""

    def __init__(
        self,
        object_dim: int,
        hidden_dim: int,
        num_layers: int,
        num_classes: int,
        img_w: int,
        img_h: int,
        score_thresh: float = 0.5,
        filter_score_thresh: float = 0.3,
        miss_tolerance: float = 5,
        object_features_path: Optional[str] = OBJECTS,
    ):
        nn.Module.__init__(self)
        RoutableMixin.__init__(self, {"object_features": object_features_path})
        self.num_layers = num_layers
        self.num_classes = num_classes
        self.img_w = img_w
        self.img_h = img_h
        self.score_thresh = score_thresh
        self.filter_score_thresh = filter_score_thresh
        self.miss_tolerance = miss_tolerance
        h = [hidden_dim] * (num_layers - 1)
        self.bbox_layers = nn.ModuleList(nn.Linear(n, k) for n, k in zip([object_dim] + h, h + [4]))

        self.cls_layers = nn.ModuleList(
            nn.Linear(n, k) for n, k in zip([object_dim] + h, h + [num_classes])
        )

    def update(self, track_instances: Instances):
        track_instances.disappear_time[track_instances.scores >= self.score_thresh] = 0
        for i in range(len(track_instances)):
            if track_instances.obj_idxes[i] == -1 and track_instances.scores[i] >= self.score_thresh:
                track_instances.obj_idxes[i] = self.max_obj_id
                self.max_obj_id += 1
            elif (
                track_instances.obj_idxes[i] >= 0
                and track_instances.scores[i] < self.filter_score_thresh
            ):
                track_instances.disappear_time[i] += 1
                # When using SAVi structure where no new det query,
                # we do not kill tracks.
                """
                if track_instances.disappear_time[i] >= self.miss_tolerance:
                    # Set the obj_id to -1.
                    # Then this track will be removed by TrackEmbeddingLayer.
                    track_instances.obj_idxes[i] = -1
                """

    # Note, this is duplicated with loss.
    def _generate_empty_tracks(self, num_queries, device):
        track_instances = Instances((1, 1))

        # At init, the number of track_instances is the same as slot number
        track_instances.obj_idxes = torch.full((num_queries,), -1, dtype=torch.long, device=device)
        track_instances.matched_gt_idxes = torch.full(
            (num_queries,), -1, dtype=torch.long, device=device
        )
        track_instances.disappear_time = torch.zeros((num_queries,), dtype=torch.long, device=device)
        track_instances.iou = torch.zeros((num_queries,), dtype=torch.float, device=device)
        track_instances.scores = torch.zeros((num_queries,), dtype=torch.float, device=device)
        track_instances.track_scores = torch.zeros((num_queries,), dtype=torch.float, device=device)
        track_instances.pred_boxes = torch.zeros((num_queries, 4), dtype=torch.float, device=device)
        track_instances.pred_logits = torch.zeros(
            (num_queries, self.num_classes), dtype=torch.float, device=device
        )

        return track_instances.to(device)

    @RoutableMixin.route
    def forward(self, object_features: torch.Tensor):
        assert object_features.dim() >= 3  # Image or video data.
        bbox_x = object_features
        cls_x = object_features

        for i, layer in enumerate(self.bbox_layers):
            bbox_x = F.relu(layer(bbox_x)) if i < self.num_layers - 1 else layer(bbox_x)
        for i, layer in enumerate(self.cls_layers):
            cls_x = F.relu(layer(cls_x)) if i < self.num_layers - 1 else layer(cls_x)
        # call softmax on the objectness or cls
        cls_x = F.sigmoid(cls_x)
        # make sure bbox prediction are positive.
        # since we normalize the box to [0, 1]
        bbox_x = F.sigmoid(bbox_x)

        # Postprocessing
        # cls_x to class idx
        scores, labels = cls_x.max(-1)
        # converting bbox_x to ori image space
        ori_res_bbox_x = box_cxcywh_to_xyxy(bbox_x)
        ori_res_bbox_x[:, :, :, 0::2] = ori_res_bbox_x[:, :, :, 0::2] * self.img_w
        ori_res_bbox_x[:, :, :, 1::2] = ori_res_bbox_x[:, :, :, 1::2] * self.img_h

        batch_size, num_frames, num_queries, _ = object_features.shape
        device = bbox_x.device

        inference_obj_idxes = torch.full(
            (
                batch_size,
                num_frames,
                num_queries,
            ),
            -1,
            dtype=torch.long,
            device=device,
        )

        for cidx in range(batch_size):
            # Init empty prediction tracks
            track_instances = self._generate_empty_tracks(num_queries, device)
            self.max_obj_id = 0
            for fidx in range(num_frames):
                track_instances.scores = scores[cidx, fidx]
                track_instances.pred_boxes = ori_res_bbox_x[cidx, fidx]
                self.update(track_instances)
                inference_obj_idxes[cidx, fidx] = track_instances.obj_idxes

        return BBoxOutput(
            bboxes=bbox_x,
            classes=cls_x,
            ori_res_bboxes=ori_res_bbox_x,
            inference_obj_idxes=inference_obj_idxes,
        )
ocl/decoding.py
amazon-science/object-centric-multiple-object-tracking
import copy
from typing import Any, Dict, List, Optional, Union

import torch
from torch import nn

from ocl import scheduling, tree_utils, utils


class EMASelfDistillation(nn.Module):
    def __init__(
        self,
        student: Union[nn.Module, Dict[str, nn.Module]],
        schedule: scheduling.HPScheduler,
        student_remapping: Optional[Dict[str, str]] = None,
        teacher_remapping: Optional[Dict[str, str]] = None,
    ):
        super().__init__()
        # Do this for convenience to reduce crazy amount of nesting.
        if isinstance(student, dict):
            student = utils.Combined(student)
        if student_remapping is None:
            student_remapping = {}
        if teacher_remapping is None:
            teacher_remapping = {}

        self.student = student
        self.teacher = copy.deepcopy(student)
        self.schedule = schedule
        self.student_remapping = {key: value.split(".") for key, value in student_remapping.items()}
        self.teacher_remapping = {key: value.split(".") for key, value in teacher_remapping.items()}

    def build_input_dict(self, inputs, remapping):
        if not remapping:
            return inputs
        # This allows us to bing the initial input and previous_output into a similar format.
        output_dict = {}
        for output_path, input_path in remapping.items():
            source = tree_utils.get_tree_element(inputs, input_path)

            output_path = output_path.split(".")
            cur_search = output_dict
            for path_part in output_path[:-1]:
                # Iterate along path and create nodes that do not exist yet.
                try:
                    # Get element prior to last.
                    cur_search = tree_utils.get_tree_element(cur_search, [path_part])
                except ValueError:
                    # Element does not yet exist.
                    cur_search[path_part] = {}
                    cur_search = cur_search[path_part]

            cur_search[output_path[-1]] = source
        return output_dict

    def forward(self, inputs: Dict[str, Any]):
        if self.training:
            with torch.no_grad():
                m = self.schedule(inputs["global_step"])  # momentum parameter
                for param_q, param_k in zip(self.student.parameters(), self.teacher.parameters()):
                    param_k.data.mul_(m).add_((1 - m) * param_q.detach().data)

        # prefix variable similar to combined module.
        prefix: List[str]
        if "prefix" in inputs.keys():
            prefix = inputs["prefix"]
        else:
            prefix = []
            inputs["prefix"] = prefix

        outputs = tree_utils.get_tree_element(inputs, prefix)

        # Forward pass student.
        prefix.append("student")
        outputs["student"] = {}
        student_inputs = self.build_input_dict(inputs, self.student_remapping)
        outputs["student"] = self.student(inputs={**inputs, **student_inputs})
        # Teacher and student share the same code, thus paths also need to be the same.  To ensure
        # that we save the student outputs and run the teacher as if it where the student.
        student_output = outputs["student"]

        # Forward pass teacher, but pretending to be student.
        outputs["student"] = {}
        teacher_inputs = self.build_input_dict(inputs, self.teacher_remapping)

        with torch.no_grad():
            outputs["teacher"] = self.teacher(inputs={**inputs, **teacher_inputs})
        prefix.pop()

        # Set correct outputs again.
        outputs["student"] = student_output

        return outputs

ocl/distillation.py
amazon-science/object-centric-multiple-object-tracking
"""Implementation of feature extractors."""
import enum
import itertools
import math
from typing import List, Optional, Union

import torch
from torch import nn

from ocl import base, path_defaults
from ocl.utils import RoutableMixin


def cnn_compute_positions_and_flatten(features: torch.Tensor):
    """Flatten output image CNN output and return it with positions of the features."""
    # todo(hornmax): see how this works with vision transformer based architectures.
    spatial_dims = features.shape[2:]
    positions = torch.cartesian_prod(
        *[torch.linspace(0.0, 1.0, steps=dim, device=features.device) for dim in spatial_dims]
    )
    # reorder into format (batch_size, flattened_spatial_dims, feature_dim).
    flattened = torch.permute(features.view(features.shape[:2] + (-1,)), (0, 2, 1)).contiguous()
    return positions, flattened


def transformer_compute_positions(features: torch.Tensor):
    """Compute positions for Transformer features."""
    n_tokens = features.shape[1]
    image_size = math.sqrt(n_tokens)
    image_size_int = int(image_size)
    assert (
        image_size_int == image_size
    ), "Position computation for Transformers requires square image"

    spatial_dims = (image_size_int, image_size_int)
    positions = torch.cartesian_prod(
        *[torch.linspace(0.0, 1.0, steps=dim, device=features.device) for dim in spatial_dims]
    )

    return positions


class ImageFeatureExtractor(base.FeatureExtractor, RoutableMixin):
    """Feature extractor which operates on images.

    For these we reshape the frame dimension into the batch dimension and process the frames as
    individual images.
    """

    def __init__(self, video_path: Optional[str] = path_defaults.VIDEO):
        base.FeatureExtractor.__init__(self)
        RoutableMixin.__init__(self, {"video": video_path})

    def forward_images(self, images: torch.Tensor):
        pass

    @RoutableMixin.route
    def forward(self, video: torch.Tensor) -> base.FeatureExtractorOutput:
        ndim = video.dim()
        assert ndim == 4 or ndim == 5

        if ndim == 5:
            # Handling video data.
            bs, frames, channels, height, width = video.shape
            images = video.view(bs * frames, channels, height, width).contiguous()
        else:
            images = video

        result = self.forward_images(images)

        if len(result) == 2:
            positions, features = result
            aux_features = None
        elif len(result) == 3:
            positions, features, aux_features = result

        if ndim == 5:
            features = features.unflatten(0, (bs, frames))
            if aux_features is not None:
                aux_features = {k: f.unflatten(0, (bs, frames)) for k, f in aux_features.items()}

        return base.FeatureExtractorOutput(features, positions, aux_features)


class VitFeatureType(enum.Enum):
    BLOCK = 1
    KEY = 2
    VALUE = 3
    QUERY = 4
    CLS = 5


class VitFeatureHook:
    """Auxilliary class used to extract features from timm ViT models.

    Args:
        mode: Type of feature to extract.
        block: Number of block to extract features from. Note that this is not zero-indexed.
    """

    def __init__(self, feature_type: VitFeatureType, block: int, drop_cls_token: bool = True):
        assert isinstance(feature_type, VitFeatureType)
        self.feature_type = feature_type
        self.block = block
        self.drop_cls_token = drop_cls_token
        self.name = f"{feature_type.name.lower()}{block}"
        self.remove_handle = None  # Can be used to remove this hook from the model again

        self._features = None

    @staticmethod
    def create_hook_from_feature_level(feature_level: Union[int, str]):
        feature_level = str(feature_level)
        prefixes = ("key", "query", "value", "block", "cls")
        for prefix in prefixes:
            if feature_level.startswith(prefix):
                _, _, block = feature_level.partition(prefix)
                feature_type = VitFeatureType[prefix.upper()]
                block = int(block)
                break
        else:
            feature_type = VitFeatureType.BLOCK
            try:
                block = int(feature_level)
            except ValueError:
                raise ValueError(f"Can not interpret feature_level '{feature_level}'.")

        return VitFeatureHook(feature_type, block)

    def register_with(self, model):
        import timm

        if not isinstance(model, timm.models.vision_transformer.VisionTransformer):
            raise ValueError(
                "This hook only supports timm.models.vision_transformer.VisionTransformer."
            )

        if self.block > len(model.blocks):
            raise ValueError(
                f"Trying to extract features of block {self.block}, but model only has "
                f"{len(model.blocks)} blocks"
            )

        block = model.blocks[self.block - 1]
        if self.feature_type == VitFeatureType.BLOCK:
            self.remove_handle = block.register_forward_hook(self)
        else:
            if isinstance(block, timm.models.vision_transformer.ParallelBlock):
                raise ValueError(
                    f"ViT with `ParallelBlock` not supported for {self.feature_type} extraction."
                )
            self.remove_handle = block.attn.qkv.register_forward_hook(self)

        return self

    def pop(self) -> torch.Tensor:
        """Remove and return extracted feature from this hook.

        We only allow access to the features this way to not have any lingering references to them.
        """
        assert self._features is not None, "Feature extractor was not called yet!"
        features = self._features
        self._features = None
        return features

    def __call__(self, module, inp, outp):
        if self.feature_type == VitFeatureType.BLOCK:
            features = outp
            if self.drop_cls_token:
                # First token is CLS token.
                features = features[:, 1:]
        elif self.feature_type in {VitFeatureType.KEY, VitFeatureType.QUERY, VitFeatureType.VALUE}:
            # This part is adapted from the timm implementation. Unfortunately, there is no more
            # elegant way to access keys, values, or queries.
            B, N, C = inp[0].shape
            qkv = outp.reshape(B, N, 3, C)  # outp has shape B, N, 3 * H * (C // H)
            q, k, v = qkv.unbind(2)

            if self.feature_type == VitFeatureType.QUERY:
                features = q
            elif self.feature_type == VitFeatureType.KEY:
                features = k
            else:
                features = v
            if self.drop_cls_token:
                # First token is CLS token.
                features = features[:, 1:]
        elif self.feature_type == VitFeatureType.CLS:
            # We ignore self.drop_cls_token in this case as it doesn't make any sense.
            features = outp[:, 0]  # Only get class token.
        else:
            raise ValueError("Invalid VitFeatureType provided.")

        self._features = features


class TimmFeatureExtractor(ImageFeatureExtractor):
    """Feature extractor implementation for timm models.

    Args:
        model_name: Name of model. See `timm.list_models("*")` for available options.
        feature_level: Level of features to return. For CNN-based models, a single integer. For ViT
            models, either a single or a list of feature descriptors. If a list is passed, multiple
            levels of features are extracted and concatenated. A ViT feature descriptor consists of
            the type of feature to extract, followed by an integer indicating the ViT block whose
            features to use. The type of features can be one of "block", "key", "query", "value",
            specifying that the block's output, attention keys, query or value should be used. If
            omitted, assumes "block" as the type. Example: "block1" or ["block1", "value2"].
        aux_features: Features to store as auxilliary features. The format is the same as in the
            `feature_level` argument. Features are stored as a dictionary, using their string
            representation (e.g. "block1") as the key. Only valid for ViT models.
        pretrained: Whether to load pretrained weights.
        freeze: Whether the weights of the feature extractor should be trainable.
        n_blocks_to_unfreeze: Number of blocks that should be trainable, beginning from the last
            block.
        unfreeze_attention: Whether weights of ViT attention layers should be trainable (only valid
            for ViT models). According to http://arxiv.org/abs/2203.09795, finetuning attention
            layers only can yield better results in some cases, while being slightly cheaper in terms
            of computation and memory.
    """

    def __init__(
        self,
        model_name: str,
        feature_level: Optional[Union[int, str, List[Union[int, str]]]] = None,
        aux_features: Optional[Union[int, str, List[Union[int, str]]]] = None,
        pretrained: bool = False,
        freeze: bool = False,
        n_blocks_to_unfreeze: int = 0,
        unfreeze_attention: bool = False,
        video_path: Optional[str] = path_defaults.VIDEO,
    ):
        super().__init__(video_path)
        try:
            import timm
        except ImportError:
            raise Exception("Using timm models requires installation with extra `timm`.")

        register_custom_timm_models()

        self.is_vit = model_name.startswith("vit")

        def feature_level_to_list(feature_level):
            if feature_level is None:
                return []
            elif isinstance(feature_level, (int, str)):
                return [feature_level]
            else:
                return list(feature_level)

        self.feature_levels = feature_level_to_list(feature_level)
        self.aux_features = feature_level_to_list(aux_features)

        if self.is_vit:
            model = timm.create_model(model_name, pretrained=pretrained)
            if hasattr(model, "head"):
                # Classification head adds unused parameters during training, so delete it
                del model.head

            if len(self.feature_levels) > 0 or len(self.aux_features) > 0:
                self._feature_hooks = [
                    VitFeatureHook.create_hook_from_feature_level(level).register_with(model)
                    for level in itertools.chain(self.feature_levels, self.aux_features)
                ]
                if len(self.feature_levels) > 0:
                    feature_dim = model.num_features * len(self.feature_levels)

                    # Remove modules not needed in computation of features
                    max_block = max(hook.block for hook in self._feature_hooks)
                    new_blocks = model.blocks[:max_block]  # Creates a copy
                    del model.blocks
                    model.blocks = new_blocks
                    model.norm = nn.Identity()
                else:
                    feature_dim = model.num_features
            else:
                self._feature_hooks = None
                feature_dim = model.num_features
        else:
            if len(self.feature_levels) == 0:
                raise ValueError(
                    f"Feature extractor {model_name} requires specifying `feature_level`"
                )
            elif len(self.feature_levels) != 1:
                raise ValueError(
                    f"Feature extractor {model_name} only supports a single `feature_level`"
                )
            elif not isinstance(self.feature_levels[0], int):
                raise ValueError("`feature_level` needs to be an integer")

            if len(self.aux_features) > 0:
                raise ValueError("`aux_features` not supported by feature extractor {model_name}")

            model = timm.create_model(
                model_name,
                pretrained=pretrained,
                features_only=True,
                out_indices=self.feature_levels,
            )
            feature_dim = model.feature_info.channels()[0]

        self.model = model
        self._feature_dim = feature_dim

        if freeze:
            self.model.requires_grad_(False)

        if n_blocks_to_unfreeze > 0:
            if not self.is_vit:
                raise NotImplementedError(
                    "`unfreeze_n_blocks` option only implemented for ViT models"
                )
            self.model.blocks[-n_blocks_to_unfreeze:].requires_grad_(True)
            if self.model.norm is not None:
                self.model.norm.requires_grad_(True)

        if unfreeze_attention:
            if not self.is_vit:
                raise ValueError("`unfreeze_attention` option only works with ViT models")
            for module in self.model.modules():
                if isinstance(module, timm.models.vision_transformer.Attention):
                    module.requires_grad_(True)

    @property
    def feature_dim(self):
        return self._feature_dim

    def forward_images(self, images: torch.Tensor):
        if self.is_vit:
            features = self.model.forward_features(images)
            if self._feature_hooks is not None:
                hook_features = [hook.pop() for hook in self._feature_hooks]

            if len(self.feature_levels) == 0:
                # Remove class token when not using hooks.
                features = features[:, 1:]
                positions = transformer_compute_positions(features)
            else:
                features = hook_features[: len(self.feature_levels)]
                positions = transformer_compute_positions(features[0])
                features = torch.cat(features, dim=-1)

            if len(self.aux_features) > 0:
                aux_hooks = self._feature_hooks[len(self.feature_levels) :]
                aux_features = hook_features[len(self.feature_levels) :]
                aux_features = {hook.name: feat for hook, feat in zip(aux_hooks, aux_features)}
            else:
                aux_features = None
        else:
            features = self.model(images)[0]
            positions, features = cnn_compute_positions_and_flatten(features)
            aux_features = None

        return positions, features, aux_features


class SlotAttentionFeatureExtractor(ImageFeatureExtractor):
    """Feature extractor as used in slot attention paper."""

    def __init__(self, video_path: Optional[str] = path_defaults.VIDEO):
        super().__init__(video_path)
        self.layers = nn.Sequential(
            nn.Conv2d(3, out_channels=64, kernel_size=5, padding="same"),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, out_channels=64, kernel_size=5, padding="same"),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, out_channels=64, kernel_size=5, padding="same"),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, out_channels=64, kernel_size=5, padding="same"),
            nn.ReLU(inplace=True),
        )

    @property
    def feature_dim(self):
        return 64

    def forward_images(self, images: torch.Tensor):
        features = self.layers(images)
        positions, flattened = cnn_compute_positions_and_flatten(features)
        return positions, flattened


class SAViFeatureExtractor(ImageFeatureExtractor):
    """Feature extractor as used in the slot attention for video paper."""

    def __init__(self, larger_input_arch=False, video_path: Optional[str] = path_defaults.VIDEO):
        """Feature extractor as used in the slot attention for video paper.

        Args:
            larger_input_arch: Use the architecture for larger image datasets such as MOVi++, which
                contains more a stride in the first layer and a higher number of feature channels in
                the CNN backbone.
            video_path: Path of input video or also image.
        """
        super().__init__()
        self.larger_input_arch = larger_input_arch
        if larger_input_arch:
            self.layers = nn.Sequential(
                nn.Conv2d(3, out_channels=64, kernel_size=5, stride=2, padding="same"),
                nn.ReLU(inplace=True),
                nn.Conv2d(64, out_channels=64, kernel_size=5, padding="same"),
                nn.ReLU(inplace=True),
                nn.Conv2d(64, out_channels=64, kernel_size=5, padding="same"),
                nn.ReLU(inplace=True),
                nn.Conv2d(64, out_channels=64, kernel_size=5, padding="same"),
            )
        else:
            self.layers = nn.Sequential(
                nn.Conv2d(3, out_channels=32, kernel_size=5, padding="same"),
                nn.ReLU(inplace=True),
                nn.Conv2d(32, out_channels=32, kernel_size=5, padding="same"),
                nn.ReLU(inplace=True),
                nn.Conv2d(32, out_channels=32, kernel_size=5, padding="same"),
                nn.ReLU(inplace=True),
                nn.Conv2d(32, out_channels=32, kernel_size=5, padding="same"),
            )

    @property
    def feature_dim(self):
        return 64 if self.larger_input_arch else 32

    def forward_images(self, images: torch.Tensor):
        features = self.layers(images)
        positions, flattened = cnn_compute_positions_and_flatten(features)
        return positions, flattened


def register_custom_timm_models():
    import timm

    @timm.models.registry.register_model
    def resnet34_savi(pretrained=False, **kwargs):
        """ResNet34 as used in SAVi and SAVi++.

        As of now, no official code including the ResNet was released, so we can only guess which of
        the numerous ResNet variants was used. This modifies the basic timm ResNet34 to have 1x1
        strides in the stem, and replaces batch norm with group norm. Gives 16x16 feature maps with
        an input size of 224x224.

        From SAVi:
        > For the modified SAVi (ResNet) model on MOVi++, we replace the convolutional backbone [...]
        > with a ResNet-34 backbone. We use a modified ResNet root block without strides
        > (i.e. 1×1 stride), resulting in 16×16 feature maps after the backbone [w. 128x128 images].
        > We further use group normalization throughout the ResNet backbone.

        From SAVi++:
        > We used a ResNet-34 backbone with modified root convolutional layer that has 1×1 stride.
        > For all layers, we replaced the batch normalization operation by group normalization.
        """
        if pretrained:
            raise ValueError("No pretrained weights available for `savi_resnet34`.")
        from timm.models import layers, resnet

        model_args = dict(
            block=resnet.BasicBlock, layers=[3, 4, 6, 3], norm_layer=layers.GroupNorm, **kwargs
        )
        model = resnet._create_resnet("resnet34", pretrained=pretrained, **model_args)
        model.conv1.stride = (1, 1)
        model.maxpool.stride = (1, 1)
        return model

ocl/feature_extractors.py
amazon-science/object-centric-multiple-object-tracking
from typing import Any, Callable, Dict, Tuple

import webdataset
from pluggy import HookimplMarker, HookspecMarker

from ocl.combined_model import CombinedModel

hook_specification = HookspecMarker("ocl")
hook_implementation = HookimplMarker("ocl")


class FakeHooks:
    """Class that mimics the behavior of the plugin manager hooks property."""

    def __getattr__(self, attribute):
        """Return a fake hook handler for any attribute query."""

        def fake_hook_handler(*args, **kwargs):
            return tuple()

        return fake_hook_handler


# @transform_hooks
# def input_dependencies() -> Tuple[str, ...]:
#     """Provide list of variables that are required for the plugin to function."""
#
#
# @transform_hooks
# def provided_inputs() -> Tuple[str, ...]:
#     """Provide list of variables that are provided by the plugin."""


@hook_specification
def training_transform() -> Callable[[webdataset.Processor], webdataset.Processor]:
    """Provide a transformation which processes a component of a webdataset pipeline."""


@hook_specification
def training_fields() -> Tuple[str]:
    """Provide list of fields that are required to be decoded during training."""


@hook_specification
def evaluation_transform() -> Callable[[webdataset.Processor], webdataset.Processor]:
    """Provide a transformation which processes a component of a webdataset pipeline."""


@hook_specification
def evaluation_fields() -> Tuple[str]:
    """Provide list of fields that are required to be decoded during evaluation."""


@hook_specification
def configure_optimizers(model: CombinedModel) -> Dict[str, Any]:
    """Return optimizers in the format of pytorch lightning."""


@hook_specification
def on_train_epoch_start(model: CombinedModel) -> None:
    """Hook called when starting training epoch."""

ocl/hooks.py
amazon-science/object-centric-multiple-object-tracking
from functools import partial
from typing import Callable, Optional, Union
from ocl.matching import CPUHungarianMatcher
import torch
from torch import nn
from torchvision import transforms
import torch.nn.functional as F
from ocl import base, consistency, path_defaults, scheduling
from ocl.utils import RoutableMixin
from typing import Any
from scipy.optimize import linear_sum_assignment
from ocl.base import Instances
from torchvision.ops import generalized_box_iou
from ocl.utils import box_cxcywh_to_xyxy


def _constant_weight(weight: float, global_step: int):
    return weight


class ReconstructionLoss(nn.Module, RoutableMixin):
    def __init__(
        self,
        loss_type: str,
        weight: Union[Callable, float] = 1.0,
        normalize_target: bool = False,
        input_path: Optional[str] = None,
        target_path: Optional[str] = None,
    ):
        nn.Module.__init__(self)
        RoutableMixin.__init__(
            self,
            {"input": input_path, "target": target_path, "global_step": path_defaults.GLOBAL_STEP},
        )
        if loss_type == "mse":
            self.loss_fn = nn.functional.mse_loss
        elif loss_type == "mse_sum":
            # Used for slot_attention and video slot attention.
            self.loss_fn = (
                lambda x1, x2: nn.functional.mse_loss(x1, x2, reduction="sum") / x1.shape[0]
            )
        elif loss_type == "l1":
            self.loss_name = "l1_loss"
            self.loss_fn = nn.functional.l1_loss
        elif loss_type == "cosine":
            self.loss_name = "cosine_loss"
            self.loss_fn = lambda x1, x2: -nn.functional.cosine_similarity(x1, x2, dim=-1).mean()
        else:
            raise ValueError(f"Unknown loss {loss_type}. Valid choices are (mse, l1, cosine).")
        # If weight is callable use it to determine scheduling otherwise use constant value.
        self.weight = weight if callable(weight) else partial(_constant_weight, weight)
        self.normalize_target = normalize_target

    @RoutableMixin.route
    def forward(self, input: torch.Tensor, target: torch.Tensor, global_step: int):
        target = target.detach()
        if self.normalize_target:
            mean = target.mean(dim=-1, keepdim=True)
            var = target.var(dim=-1, keepdim=True)
            target = (target - mean) / (var + 1.0e-6) ** 0.5
        loss = self.loss_fn(input, target)
        weight = self.weight(global_step)
        return weight * loss


class LatentDupplicateSuppressionLoss(nn.Module, RoutableMixin):
    def __init__(
        self,
        weight: Union[float, scheduling.HPSchedulerT],
        eps: float = 1e-08,
        grouping_path: Optional[str] = "perceptual_grouping",
    ):
        nn.Module.__init__(self)
        RoutableMixin.__init__(
            self, {"grouping": grouping_path, "global_step": path_defaults.GLOBAL_STEP}
        )
        self.weight = weight
        self.similarity = nn.CosineSimilarity(dim=-1, eps=eps)

    @RoutableMixin.route
    def forward(self, grouping: base.PerceptualGroupingOutput, global_step: int):
        if grouping.dim() == 4:
            # Build large tensor of reconstructed video.
            # objects = grouping.objects
            objects = grouping
            bs, n_frames, n_objects, n_features = objects.shape

            off_diag_indices = torch.triu_indices(
                n_objects, n_objects, offset=1, device=objects.device
            )

            sq_similarities = (
                self.similarity(
                    objects[:, :, off_diag_indices[0], :], objects[:, :, off_diag_indices[1], :]
                )
                ** 2
            )

            # if grouping.is_empty is not None:
            #     p_not_empty = 1.0 - grouping.is_empty
            #     # Assume that the probability of of individual objects being present is independent,
            #     # thus the probability of both being present is the product of the individual
            #     # probabilities.
            #     p_pair_present = (
            #         p_not_empty[..., off_diag_indices[0]] * p_not_empty[..., off_diag_indices[1]]
            #     )
            #     # Use average expected penalty as loss for each frame.
            #     losses = (sq_similarities * p_pair_present) / torch.sum(
            #         p_pair_present, dim=-1, keepdim=True
            #     )
            # else:
            losses = sq_similarities.mean(dim=-1)

            weight = self.weight(global_step) if callable(self.weight) else self.weight

            return weight * losses.sum() / (bs * n_frames)
        elif grouping.dim() == 3:
            # Build large tensor of reconstructed image.
            objects = grouping
            bs, n_objects, n_features = objects.shape

            off_diag_indices = torch.triu_indices(
                n_objects, n_objects, offset=1, device=objects.device
            )

            sq_similarities = (
                self.similarity(
                    objects[:, off_diag_indices[0], :], objects[:, off_diag_indices[1], :]
                )
                ** 2
            )

            if grouping.is_empty is not None:
                p_not_empty = 1.0 - grouping.is_empty
                # Assume that the probability of of individual objects being present is independent,
                # thus the probability of both being present is the product of the individual
                # probabilities.
                p_pair_present = (
                    p_not_empty[..., off_diag_indices[0]] * p_not_empty[..., off_diag_indices[1]]
                )
                # Use average expected penalty as loss for each frame.
                losses = (sq_similarities * p_pair_present) / torch.sum(
                    p_pair_present, dim=-1, keepdim=True
                )
            else:
                losses = sq_similarities.mean(dim=-1)

            weight = self.weight(global_step) if callable(self.weight) else self.weight
            return weight * losses.sum() / bs
        else:
            raise ValueError("Incompatible input format.")


class EM_loss(nn.Module, RoutableMixin):
    def __init__(
            self,
            loss_weight: float = 20,
            pred_mask_path: Optional[str] = None,
            rec_path: Optional[str] = None,
            tgt_mask_path: Optional[str] = None,
            img_path: Optional[str] = None,
            tgt_vis_path: Optional[str] = None,
            attn_index_path: Optional[str] = None,
            pred_feat_path: Optional[str] = None,
            gt_feat_path: Optional[str] = None,
    ):
        nn.Module.__init__(self)
        RoutableMixin.__init__(
            self,
            {
                "pred_mask": pred_mask_path,
                "reconstructions": rec_path,
                "tgt_mask": tgt_mask_path,
                "masks_vis": tgt_vis_path,
                "rec_tgt": img_path,
                "pred_feats": pred_feat_path,
                "gt_feats": gt_feat_path,
                "attn_index": attn_index_path,
            },
        )
        self.loss_weight = loss_weight
        self.loss_fn = (
            lambda x1, x2: nn.functional.mse_loss(x1, x2, reduction="none")
        )

    @RoutableMixin.route
    def forward(
            self,
            pred_mask: torch.Tensor,  # rollout_decode.masks
            tgt_mask: torch.Tensor,  # decoder.masks
            reconstructions: torch.Tensor,
            rec_tgt: torch.Tensor,
            masks_vis: torch.Tensor,
            pred_feats: torch.Tensor,
            gt_feats: torch.Tensor,
            attn_index: torch.Tensor,

    ):
        b, f, c, h, w = pred_mask.shape
        _, _, n_slots, n_buffer = attn_index.shape
        dim = pred_feats.shape[-1]

        pred_feats = F.normalize(pred_feats, dim=-1)
        gt_feats = F.normalize(gt_feats, dim=-1)

        pred_feats = pred_feats.reshape(-1, n_buffer, dim).unsqueeze(1).repeat(1, n_slots, 1, 1)
        gt_feats = gt_feats.reshape(-1, n_slots, dim).unsqueeze(2).repeat(1, 1, n_buffer, 1)

        pred_mask = pred_mask.reshape(-1, n_buffer, h, w).unsqueeze(1).repeat(1,n_slots,1,1,1)
        tgt_mask = tgt_mask.reshape(-1, n_slots, h, w).unsqueeze(2).repeat(1,1,n_buffer,1,1)
        tgt_mask = tgt_mask > 0.5
        masks_vis = masks_vis.reshape(-1, n_slots, h, w).unsqueeze(2).unsqueeze(3).repeat(1,1,n_buffer,3,1,1)
        masks_vis = masks_vis > 0.5
        attn_index = attn_index.reshape(-1, n_slots, n_buffer)
        rec_tgt = rec_tgt.reshape(-1,3,h,w).unsqueeze(1).unsqueeze(2).repeat(1,n_slots,n_buffer,1,1,1)
        reconstructions = reconstructions.reshape(-1, n_buffer, 3, h, w).unsqueeze(1).repeat(1,n_slots,1,1,1,1)
        rec_pred = reconstructions * masks_vis
        rec_tgt_ = rec_tgt * masks_vis
        loss = torch.sum(F.binary_cross_entropy(pred_mask, tgt_mask.float(), reduction = 'none'), (-1,-2)) / (h*w) + 0.1 * torch.sum(self.loss_fn(rec_pred, rec_tgt_), (-3,-2,-1))
        #loss = torch.sum(self.loss_fn(pred_feats, gt_feats), -1)
        total_loss = torch.sum(attn_index * loss, (0,1,2)) / (b * f * n_slots * n_buffer)
        return (total_loss) * self.loss_weight






ocl/losses.py
amazon-science/object-centric-multiple-object-tracking
"""Methods for matching between sets of elements."""
from typing import Tuple, Type

import numpy as np
import torch
from scipy.optimize import linear_sum_assignment
from torchtyping import TensorType

# Avoid errors due to flake:
batch_size = None
n_elements = None

CostMatrix = Type[TensorType["batch_size", "n_elements", "n_elements"]]
AssignmentMatrix = Type[TensorType["batch_size", "n_elements", "n_elements"]]
CostVector = Type[TensorType["batch_size"]]


class Matcher(torch.nn.Module):
    """Matcher base class to define consistent interface."""

    def forward(self, C: CostMatrix) -> Tuple[AssignmentMatrix, CostVector]:
        pass


class CPUHungarianMatcher(Matcher):
    """Implementaiton of a cpu hungarian matcher using scipy.optimize.linear_sum_assignment."""

    def forward(self, C: CostMatrix) -> Tuple[AssignmentMatrix, CostVector]:
        X = torch.zeros_like(C)
        C_cpu: np.ndarray = C.detach().cpu().numpy()
        for i, cost_matrix in enumerate(C_cpu):
            row_ind, col_ind = linear_sum_assignment(cost_matrix)
            X[i][row_ind, col_ind] = 1.0
        return X, (C * X).sum(dim=(1, 2))

ocl/matching.py
amazon-science/object-centric-multiple-object-tracking
import torch
import torch.nn as nn
import torch.nn.functional as F
from ocl.utils import RoutableMixin
from ocl.predictor import MLP
import dataclasses
from typing import Optional
from ocl.memory_rollout import GPT
from ocl import base, path_defaults
from ocl.mha import MultiHeadAttention, MultiHeadAttention_for_index


@dataclasses.dataclass
class MemoryOutput:
    rollout: torch.Tensor  # noqa: F821
    object_features: torch.Tensor  # noqa: F821
    mem: torch.Tensor
    eval_mem_features: torch.Tensor
    table: torch.Tensor
    attn_index: torch.Tensor


class SelfSupervisedMemory(nn.Module, RoutableMixin):
    def __init__(self,
                 embed_dim: int = 128,
                 num_objects: int = 20,
                 memory_len: int = 30,
                 stale_number: int = 5,
                 object_features_path: Optional[str] = path_defaults.OBJECTS,
                 conditioning_path: Optional[str] = path_defaults.CONDITIONING,
                 attention_maps_path: Optional[str] = None,
                 frame_features_path: Optional[str] = path_defaults.FEATURES,
                 first_box_path: Optional[str] = None,
                 init_flag_path: Optional[str] = None,
                 matched_idx_path: Optional[str] = None,
                 # eval_flag: Optional[str] = False,
                 ):
        nn.Module.__init__(self)
        RoutableMixin.__init__(self, {
            "observation": object_features_path,
            "conditions": conditioning_path,
            "attention_maps": attention_maps_path,
            "frame_features": frame_features_path,
            "first_frame_boxes": first_box_path,
            "init_flag": init_flag_path,
            "matched_index": matched_idx_path,
        }
                               )
        self.embed_dim = embed_dim
        self.memory_len = memory_len
        self.object_num = num_objects
        self.stale_number = stale_number
        self.num_heads = 4
        self.roll_out_module = GPT(buffer_len=memory_len, n_layer=8, n_head=8, n_embd = embed_dim)
        self.register_buffer("memory", torch.zeros(8, 7, 12, embed_dim))
        self.register_buffer("memory_table", torch.zeros(8, 12))
        self.register_buffer("stale_counter", torch.zeros(8, 12))
        self.MultiHead_1 = MultiHeadAttention_for_index(n_head = 4, d_model = embed_dim, d_k = embed_dim, d_v = embed_dim)
        self.MultiHead_2 = MultiHeadAttention(n_head = 4, d_model = embed_dim, d_k = embed_dim, d_v = embed_dim)

    def remove_duplicated_slot_id_by_box_intersection(self, slot_masks):
        slot_masks = slot_masks>0.5
        n, h, w = slot_masks.shape
        mask_sum = torch.sum(slot_masks.reshape(-1, h * w), dim=-1)
        empty_idx = (mask_sum <= 10).nonzero(as_tuple=True)[0]
        bg_idx = [3]

        dup_idx = []
        mask = slot_masks.unsqueeze(1).to(torch.bool).reshape(n, 1, -1)
        mask_ = slot_masks.unsqueeze(0).to(torch.bool).reshape(1, n, -1)
        intersection = torch.sum(mask & mask_, dim=-1).to(torch.float64)
        union = torch.sum(mask | mask_, dim=-1).to(torch.float64)
        union1 = torch.sum(mask | mask, dim=-1).to(torch.float64)
        # union2 = torch.sum(mask_ | mask_, dim=-1).to(torch.float64)
        pairwise_iou = intersection / union
        self_iou1 = intersection / union1
        # self_iou2 = intersection / union2
        pairwise_iou[union == 0] = 1.0
        self_iou1[union == 0] = 1.0
        # self_iou2[union == 0] = 1.0

        for i in range(n):
            for j in range(i + 1, n):
                # if pairwise_iou[i, j] > 0.9 or self_iou1[i, j]>0.6 or self_iou2[i, j]>0.6:
                if pairwise_iou[i, j] > 0.9 or self_iou1[i, j] > 0.6:
                    dup_idx.append(i)

        valid_idx = []
        invalid_idx = [*set(list(empty_idx) + list(dup_idx)+ bg_idx)]
        for i in range(n):
            if i not in invalid_idx:
                valid_idx.append(i)
        return valid_idx


    def remove_duplicated_slot_id(self, slot_masks, iou_threshold = 0.5):
        slot_masks = slot_masks>0.5
        n, h, w = slot_masks.shape
        # remove background or none masks
        mask_sum = torch.sum(slot_masks.reshape(-1, h * w), dim=-1)
        bg_value = mask_sum[0]
        # bg_idx = (mask_sum == bg_value).nonzero(as_tuple=True)[0]
        bg_idx = (mask_sum > 0.15 * h * w).nonzero(as_tuple=True)[0]
        empty_idx = (mask_sum <= 10).nonzero(as_tuple=True)[0]

        # remove duplicated masks
        mask = slot_masks.unsqueeze(1).to(torch.bool).reshape(n, 1, -1)
        mask_ = slot_masks.unsqueeze(0).to(torch.bool).reshape(1, n, -1)
        intersection = torch.sum(mask & mask_, dim=-1).to(torch.float64)
        union = torch.sum(mask | mask_, dim=-1).to(torch.float64)
        pairwise_iou = intersection / union
        pairwise_iou[union == 0] = 1.0
        dup_idx = []
        for i in range(n):
            for j in range(i + 1, n):
                if pairwise_iou[i, j] > iou_threshold:
                    dup_idx.append(i)
        invalid_idx = [*set(list(bg_idx) + list(empty_idx)+ list(dup_idx))]
        valid_idx = [0]
        for i in range(n):
            if i not in invalid_idx:
                valid_idx.append(i)
        # if len(empty_idx) == n:
        #     valid_idx.append(0)

        return valid_idx



    def initialization(self, conditions, cur_slots, cur_slot_masks, prev_slot_masks, amodal_masks, frame_id):
        if frame_id == 0:
            # For each video, we should initialize the register buffers as zero
            bs = conditions.shape[0]
            memory_shape = (bs, self.memory_len, self.object_num, self.embed_dim)
            memory_table_shape = (bs, self.object_num)
            stale_counter_shape = (bs, self.object_num)
            self.memory = torch.zeros(memory_shape).to(conditions.device)
            self.memory_table = torch.zeros(memory_table_shape).to(conditions.device)
            self.stale_counter = torch.zeros(stale_counter_shape).to(conditions.device)
            for b in range(bs):
                valid_idx = self.remove_duplicated_slot_id_by_box_intersection(amodal_masks[b])
                num_obj = len(valid_idx)
                self.memory[b, 0, 0, :] = conditions[b, 3, :]
                self.memory[b, 0, 1:num_obj+1, :] = conditions[b, valid_idx, :]
                self.memory_table[b, :num_obj+1] += 1

        else:
            '''IoU score to find new objects'''
            bs, n, h, w = prev_slot_masks.shape
            for b in range(bs):
                # self.memory_eval[b, frame_id, -1, :] = ori_slots[b, 0, :]
                cur_valid_idx = self.remove_duplicated_slot_id(cur_slot_masks[b])
                pre_valid_idx = self.remove_duplicated_slot_id(prev_slot_masks[b])

                cur_slot_mask = cur_slot_masks[b][cur_valid_idx] > 0.7
                prev_slot_mask = prev_slot_masks[b][pre_valid_idx] > 0.7

                # calculate pairwise iou
                cur_mask = cur_slot_mask.unsqueeze(1).to(torch.bool).reshape(len(cur_valid_idx), 1, -1)
                prev_mask = prev_slot_mask.unsqueeze(0).to(torch.bool).reshape(1, len(pre_valid_idx), -1)
                intersection = torch.sum(cur_mask & prev_mask, dim=-1).to(torch.float64)
                union = torch.sum(cur_mask | prev_mask, dim=-1).to(torch.float64)
                pairwise_iou = intersection / union
                pairwise_iou[union == 0] = 1.0
                sim, _ = torch.max(pairwise_iou, dim=-1)
                # new_obj_idx = list((sim < 10).nonzero(as_tuple=True)[0])
                new_obj_idx = (sim ==0).nonzero(as_tuple=True)[0].tolist()

                new_obj_idx_ori = [cur_valid_idx[id] for id in new_obj_idx]
                num_new_obj = len(new_obj_idx_ori)

                new_mem_idx = (self.memory_table[b] == 0).nonzero(as_tuple=True)[0].tolist()
                old_mem_idx = (self.memory_table[b] != 0).nonzero(as_tuple=True)[0].tolist()
                if num_new_obj >0 and len(new_mem_idx) > 0:
                    last_pos = old_mem_idx[-1]+1
                    if last_pos + num_new_obj - 1 in new_mem_idx:
                        self.memory[b, 0, last_pos: last_pos + num_new_obj] = cur_slots[b, new_obj_idx_ori]
                        # self.memory_eval[b, 0, last_pos: last_pos + num_new_obj] = cur_slots[b, new_obj_idx_ori]
                        self.memory_table[b, last_pos: last_pos + num_new_obj] += 1


    def buffer_terminate(self):
        bs = self.stale_counter.shape[0]
        for b in range(bs):
            terminate_idx = (self.stale_counter[b] >= self.stale_number).nonzero(as_tuple=True)[0]
            if len(terminate_idx)>0:
                num_dead_buffer = len(list(terminate_idx))
                tmp = torch.zeros((self.memory_len, num_dead_buffer, self.embed_dim)).to(self.memory.device)
                self.memory[b, :, terminate_idx] = tmp
                self.stale_counter[b, terminate_idx] = 0



    def sms_attn(self, observations, predictions, attn_mask, eval_flag):
        attn_o_to_p, attn_o_to_p_weights = self.MultiHead_1(observations, predictions, predictions, mask = attn_mask)


        mask = torch.zeros(attn_o_to_p_weights.shape).to(attn_o_to_p_weights.device)
        b, w, h = mask.shape
        for i in range(b):
            for j in range(w):
                index = torch.argmax(attn_o_to_p_weights[i, j, :])
                mask[i, j, index] = 1



        weights = mask.clone()

        # MultiHead_2 layer
        if not eval_flag:
            attn_o_to_p_weights_trans = torch.transpose(attn_o_to_p_weights, 1, 2)
        else:
            attn_o_to_p_weights_trans = torch.transpose(weights, 1, 2)

        attn_p_to_o, attn_p_to_o_weights = self.MultiHead_2(predictions, observations, observations, mask=attn_o_to_p_weights_trans)

        # replace the attn_p_to_o with predictions if the buffer is not assigned
        if eval_flag:
            b, h, w = weights.shape
            weights_new = torch.zeros((b, h + 1, w)).to(attn_o_to_p.device)  # [b, n+1, m]
            weights_new[:, 0:h, :] = weights
            weights_new[:, h, :] = torch.sum(weights, dim=1)
            weights_new_convert_zero = weights_new[:, h, :].clone()
            weights_new_convert_zero[weights_new[:, h, :] == 0] = 1
            weights_new_convert_zero[weights_new[:, h, :] > 0] = 0
            weights_new[:, h, :] = weights_new_convert_zero
            b_p, h_p, w_p = attn_p_to_o.shape # merged features
            for j in range(b_p):
                index = weights_new[j, h, :].nonzero(as_tuple=True)[0]  # hard index
                if len(index) > 0:
                    # update the buffer that no slots matched with zero embeddings
                    attn_p_to_o[j][index] = torch.zeros((len(index), self.embed_dim)).to(observations.device)
                    # attn_p_to_o[j][index] = predictions[j][index].clone()
        else:
            b_p, h_p, w_p = attn_p_to_o.shape # merged features
            for j in range(b_p):
                # index = weights_new[j, h, :].nonzero(as_tuple=True)[0]  # hard index
                index = (self.memory_table[j] == 0).nonzero(as_tuple=True)[0]
                if len(index) > 0:
                    # update the buffer that no slots matched with zero embeddings
                    attn_p_to_o[j][index] = torch.zeros((len(index), self.embed_dim)).to(observations.device)
        return attn_p_to_o, weights, attn_o_to_p_weights, attn_p_to_o_weights




    def update_sms(self, object_features):
        object_features_ = object_features.clone().detach()

        for b in range(object_features_.shape[0]):
            for i in range(object_features_.shape[1]):
                tmp = torch.sum(object_features_[b, i, :], dim=-1)
                if tmp != 0 and self.memory_table[b, i] != 0:
                    pos = self.memory_table[b, i].cpu().numpy().astype(int)
                    if pos == self.memory_len-2:
                        self.memory[b, :-2, i] = self.memory[b, 1:-1, i]
                        self.memory[b, -2, i] = object_features_[b, i]
                    else:
                        self.memory[b, pos, i] = object_features_[b, i]
                        self.memory_table[b, i] += 1
                else:
                    self.stale_counter[b, i] += 1

        return object_features

    def generate_mem_attnmask(self, memory_table, slots):
        bs, num_buffer = memory_table.shape
        num_slot = slots.shape[1]
        attn_mask = torch.zeros((bs, num_slot, num_buffer)).to(memory_table.device)
        for b in range(bs):
            tmp = torch.sum(slots[b], dim=-1)
            valid_slot_idx = (tmp!=0).nonzero(as_tuple=True)[0].tolist()
            valid_buffer_idx = (self.memory_table[b] != 0).nonzero(as_tuple=True)[0].tolist()
            if len(valid_slot_idx) > 0 and len(valid_buffer_idx) > 0 :
                num_slot_valid = len(valid_slot_idx)
                num_buffer_valid = len(valid_buffer_idx)
                attn_mask[b, :num_slot_valid, :num_buffer_valid] += 1
        return attn_mask


    @RoutableMixin.route
    def forward(self,
                observations: torch.Tensor,
                prev_slot_masks: torch.Tensor,
                cur_slot_masks: torch.Tensor,
                amodal_masks: torch.Tensor,
                conditions: torch.Tensor,
                frame_id: int,
                phase: str,
                ):
        if phase == 'train':
            eval = False
        elif phase == 'val':
            eval = True
        self.initialization(conditions, observations, cur_slot_masks, prev_slot_masks, amodal_masks, frame_id)
        if frame_id == 0:
            predictions = self.memory[:, 0].clone()
            object_features = predictions.clone()
            b, n_slots = observations.shape[:2]
            n_buffer = predictions.shape[1]
            attn_index = torch.zeros((b, n_slots, n_buffer)).to(observations.device)

        else:
            predictions = self.roll_out_module(self.memory, self.memory_table)

            predictions[:, 0] = observations[:, 3] # index 3 denotes the background slot

            attn_mask = self.generate_mem_attnmask(self.memory_table, observations)
            object_features, weights, attn_index, mask_o_p = self.sms_attn(observations, predictions, attn_mask, eval_flag=eval)
            object_features[:, 0] = observations[:, 3]

            # memory update
            _ = self.update_sms(object_features)

        # memory terminate
        # if eval:
        #     self.buffer_terminate()
        return MemoryOutput(
            rollout=predictions,
            object_features=object_features,
            mem = self.memory,
            eval_mem_features = object_features,
            table = self.memory_table,
            attn_index=attn_index,
        )



ocl/memory.py
amazon-science/object-centric-multiple-object-tracking
"""Memory roll-out module, following GPT-2 architecture.

References:
1) minGPT by Andrej Karpathy:
https://github.com/karpathy/minGPT/tree/master/mingpt
2) the official GPT-2 TensorFlow implementation released by OpenAI:
https://github.com/openai/gpt-2/blob/master/src/model.py
"""

import math

import torch
from torch import nn
import torch.nn.functional as F
# -----------------------------------------------------------------------------


class GELU(nn.Module):
    def forward(self, x):
        return (
            0.5
            * x
            * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi) * (x + 0.044715 * torch.pow(x, 3.0))))
        )


class Block(nn.Module):
    """One GPT-2 decoder block, consists of a Masked Self-Attn and a FFN."""

    def __init__(self, n_embd, n_heads, dropout_rate):
        super().__init__()
        self.ln_1 = nn.LayerNorm(n_embd)
        self.attn = nn.MultiheadAttention(n_embd, n_heads, batch_first=True)
        self.ln_2 = nn.LayerNorm(n_embd)
        self.mlp = nn.ModuleDict(
            dict(
                c_fc=nn.Linear(n_embd, 4 * n_embd),
                c_proj=nn.Linear(4 * n_embd, n_embd),
                act=GELU(),
                dropout=nn.Dropout(dropout_rate),
            )
        )
        m = self.mlp
        self.ffn = lambda x: m.dropout(m.c_proj(m.act(m.c_fc(x))))

    def forward(self, x, causal_mask):
        # query: [160,33,128]
        # print(causal_mask.shape) #[640, 33, 33]
        att, att_weights = self.attn(query=self.ln_1(x), key=self.ln_1(x), value=self.ln_1(x), attn_mask=causal_mask)
        # att, att_weights = self.attn(query=self.ln_1(x), key=self.ln_1(x), value=self.ln_1(x))


        x = x + att
        x = x + self.ffn(self.ln_2(x))

        # att_weights[att_weights>0] = 1
        # att_weights = F.softmax(att_weights, dim=-1)
        # att = torch.matmul(att_weights, self.ln_1(x))
        # x = x + att
        # x = x + self.ffn(self.ln_2(x))

        return x, att_weights


class GPT(nn.Module):
    """Memory roll-out GPT."""

    def __init__(
        self, buffer_len, n_layer, n_head, n_embd, embd_pdrop=0.0, resid_pdrop=0.0, attn_pdrop=0.0
    ):
        super().__init__()
        self.buffer_len = buffer_len
        self.n_layer = n_layer
        self.n_head = n_head
        self.n_embd = n_embd
        self.embd_pdrop = embd_pdrop
        self.resid_pdrop = resid_pdrop
        self.attn_pdrop = attn_pdrop

        self.transformer = nn.ModuleDict(
            dict(
                wte=nn.Linear(self.n_embd, self.n_embd, bias=False),
                wpe=nn.Embedding(self.buffer_len, self.n_embd),
                drop=nn.Dropout(self.embd_pdrop),
                h=nn.ModuleList(
                    [Block(self.n_embd, self.n_head, self.resid_pdrop) for _ in range(self.n_layer)]
                ),
                ln_f=nn.LayerNorm(self.n_embd),
            )
        )
        # roll out to the same dimension
        self.roll_out_head = nn.Linear(self.n_embd, self.n_embd, bias=False)

        # init all weights
        self.apply(self._init_weights)
        for pn, p in self.named_parameters():
            if pn.endswith("c_proj.weight"):
                torch.nn.init.normal_(p, mean=0.0, std=0.02 / math.sqrt(2 * self.n_layer))

        # report number of parameters (note we don't count the decoder parameters in lm_head)
        n_params = sum(p.numel() for p in self.transformer.parameters())
        print("number of parameters: %.2fM" % (n_params / 1e6,))

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                torch.nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
        elif isinstance(module, nn.LayerNorm):
            torch.nn.init.zeros_(module.bias)
            torch.nn.init.ones_(module.weight)

    def forward(self, mem, mem_table, targets=None):
        device = mem.device
        b, t, n, d = mem.shape

        # reshape to merge the batch and num_buffer dimensionsni
        mem = mem.permute(0, 2, 1, 3).reshape(b*n, t, d)
        mem_table = mem_table.view(b * n, -1)
        pos = torch.arange(0, t, dtype=torch.long, device=device).unsqueeze(0)

        tok_emb = self.transformer.wte(mem)  # token embeddings of shape (b, t, n_embd)
        pos_emb = self.transformer.wpe(pos)  # position embeddings of shape (1, t, n_embd)
        x = self.transformer.drop(tok_emb + pos_emb)

        # create causal attention masks
        # need to check correctness
        causal_masks = []
        for idx in range(b * n):
            occupied_len = mem_table[idx].cpu().numpy().astype(int)[0]
            if occupied_len == 0:
                occupied_len = 1
            # causal_mask = torch.tril(torch.ones(self.buffer_len, self.buffer_len).to(device)).view(
            #     1, self.buffer_len, self.buffer_len
            # )
            causal_mask = torch.zeros(self.buffer_len, self.buffer_len).to(device).view( 1, self.buffer_len, self.buffer_len)
            causal_mask[:, occupied_len:, occupied_len:] = 1
            causal_mask = causal_mask > 0
            causal_masks.append(causal_mask)
        causal_masks = torch.stack(causal_masks)
        causal_masks = causal_masks.repeat(1,self.n_head,1,1).view(-1, t, t)


        for block in self.transformer.h:
            x, attn_weights = block(x, causal_masks)
        x = self.transformer.ln_f(x)
        x = self.roll_out_head(x) #[b*n, t, d]

        out = torch.zeros((b*n, d)).to(device)

        for idx in range(b * n):
            t_pos = mem_table[idx].cpu().numpy().astype(int)[0]
            if t_pos > 0 and t_pos < t:
                # print(attn_weights[idx, t_pos])
                out[idx] = x[idx, t_pos-1]

        # for idx in range(b * n):
        #     t_pos = mem_table[idx].cpu().numpy().astype(int)[0]
        #     out[idx] = x[idx, t_pos]
        return out.view(b,n,d)

ocl/memory_rollout.py
amazon-science/object-centric-multiple-object-tracking
from typing import Any, Dict, Optional

import numpy as np
import scipy.optimize
import torch
import torchmetrics
import torchvision
import motmetrics as mm
mm.lap.default_solver = 'lap'
import pandas as pd
from ocl.utils import RoutableMixin, resize_patches_to_image, tensor_to_one_hot
import io
from numba import jit

class TensorStatistic(torchmetrics.Metric, RoutableMixin):
    """Metric that computes summary statistic of tensors for logging purposes.

    First dimension of tensor is assumed to be batch dimension. Other dimensions are reduced to a
    scalar by the chosen reduction approach (sum or mean).
    """

    def __init__(self, path: Optional[str], reduction: str = "mean"):
        torchmetrics.Metric.__init__(self)
        RoutableMixin.__init__(self, {"tensor": path})
        if reduction not in ("sum", "mean"):
            raise ValueError(f"Unknown reduction {reduction}")
        self.reduction = reduction
        self.add_state(
            "values", default=torch.tensor(0.0, dtype=torch.float64), dist_reduce_fx="sum"
        )
        self.add_state("total", default=torch.tensor(0), dist_reduce_fx="sum")

    @RoutableMixin.route
    def update(self, tensor: torch.Tensor):
        tensor = torch.atleast_2d(tensor).flatten(1, -1).to(dtype=torch.float64)

        if self.reduction == "mean":
            tensor = torch.mean(tensor, dim=1)
        elif self.reduction == "sum":
            tensor = torch.sum(tensor, dim=1)

        self.values += tensor.sum()
        self.total += len(tensor)

    def compute(self) -> torch.Tensor:
        return self.values / self.total


class TorchmetricsWrapper(torchmetrics.Metric, RoutableMixin):
    """Wrapper for torchmetrics classes that works with routing."""

    def __init__(
        self,
        metric: str,
        prediction_path: str,
        target_path: str,
        metric_kwargs: Optional[Dict[str, Any]] = None,
    ):
        torchmetrics.Metric.__init__(self)
        RoutableMixin.__init__(self, {"prediction": prediction_path, "target": target_path})
        if not hasattr(torchmetrics, metric):
            raise ValueError(f"Metric {metric} does not exist in torchmetrics")
        self.metric = getattr(torchmetrics, metric)(**(metric_kwargs if metric_kwargs else {}))

    @RoutableMixin.route
    def update(self, prediction: torch.Tensor, target: torch.Tensor):
        return self.metric.update(prediction, target)

    def compute(self) -> torch.Tensor:
        return self.metric.compute()


class ARIMetric(torchmetrics.Metric, RoutableMixin):
    """Computes ARI metric."""

    def __init__(
        self,
        prediction_path: str,
        target_path: str,
        foreground: bool = True,
        convert_target_one_hot: bool = False,
    ):
        torchmetrics.Metric.__init__(self)
        RoutableMixin.__init__(self, {"prediction": prediction_path, "target": target_path})
        self.foreground = foreground
        self.convert_target_one_hot = convert_target_one_hot
        self.add_state(
            "values", default=torch.tensor(0.0, dtype=torch.float64), dist_reduce_fx="sum"
        )
        self.add_state("total", default=torch.tensor(0), dist_reduce_fx="sum")

    @RoutableMixin.route
    def update(self, prediction: torch.Tensor, target: torch.Tensor):
        """Update this metric.

        Args:
            prediction: Predicted mask of shape (B, C, H, W) or (B, F, C, H, W), where C is the
                number of classes.
            target: Ground truth mask of shape (B, K, H, W) or (B, F, K, H, W), where K is the
                number of classes.
        """
        if prediction.ndim == 5:
            # Merge frames, height and width to single dimension.
            prediction = prediction.transpose(1, 2).flatten(-3, -1)
            target = target.transpose(1, 2).flatten(-3, -1)
        elif prediction.ndim == 4:
            # Merge height and width to single dimension.
            prediction = prediction.flatten(-2, -1)
            target = target.flatten(-2, -1)
        else:
            raise ValueError(f"Incorrect input shape: f{prediction.shape}")

        # Make channels / gt labels the last dimension.
        prediction = prediction.transpose(-2, -1)
        target = target.transpose(-2, -1)

        if self.convert_target_one_hot:
            target_oh = tensor_to_one_hot(target, dim=2)
            # For empty pixels (all values zero), one-hot assigns 1 to the first class, correct for
            # this (then it is technically not one-hot anymore).
            target_oh[:, :, 0][target.sum(dim=2) == 0] = 0
            target = target_oh

        # Should be either 0 (empty, padding) or 1 (single object).
        assert torch.all(target.sum(dim=-1) < 2), "Issues with target format, mask non-exclusive"

        if self.foreground:
            ari = fg_adjusted_rand_index(prediction, target)
        else:
            ari = adjusted_rand_index(prediction, target)

        self.values += ari.sum()
        self.total += len(ari)

    def compute(self) -> torch.Tensor:
        return self.values / self.total


class PatchARIMetric(ARIMetric):
    """Computes ARI metric assuming patch masks as input."""

    def __init__(
        self,
        prediction_key: str,
        target_key: str,
        foreground=True,
        resize_masks_mode: str = "bilinear",
        **kwargs,
    ):
        super().__init__(prediction_key, target_key, foreground, **kwargs)
        self.resize_masks_mode = resize_masks_mode

    @RoutableMixin.route
    def update(self, prediction: torch.Tensor, target: torch.Tensor):
        """Update this metric.

        Args:
            prediction: Predicted mask of shape (B, C, P) or (B, F, C, P), where C is the
                number of classes and P the number of patches.
            target: Ground truth mask of shape (B, K, H, W) or (B, F, K, H, W), where K is the
                number of classes.
        """
        h, w = target.shape[-2:]
        assert h == w

        prediction_resized = resize_patches_to_image(
            prediction, size=h, resize_mode=self.resize_masks_mode
        )

        return super().update(prediction=prediction_resized, target=target)


def adjusted_rand_index(pred_mask: torch.Tensor, true_mask: torch.Tensor) -> torch.Tensor:
    """Computes adjusted Rand index (ARI), a clustering similarity score.

    This implementation ignores points with no cluster label in `true_mask` (i.e. those points for
    which `true_mask` is a zero vector). In the context of segmentation, that means this function
    can ignore points in an image corresponding to the background (i.e. not to an object).

    Implementation adapted from https://github.com/deepmind/multi_object_datasets and
    https://github.com/google-research/slot-attention-video/blob/main/savi/lib/metrics.py

    Args:
        pred_mask: Predicted cluster assignment encoded as categorical probabilities of shape
            (batch_size, n_points, n_pred_clusters).
        true_mask: True cluster assignment encoded as one-hot of shape (batch_size, n_points,
            n_true_clusters).

    Returns:
        ARI scores of shape (batch_size,).
    """
    n_pred_clusters = pred_mask.shape[-1]
    pred_cluster_ids = torch.argmax(pred_mask, axis=-1)

    # Convert true and predicted clusters to one-hot ('oh') representations. We use float64 here on
    # purpose, otherwise mixed precision training automatically casts to FP16 in some of the
    # operations below, which can create overflows.
    true_mask_oh = true_mask.to(torch.float64)  # already one-hot
    pred_mask_oh = torch.nn.functional.one_hot(pred_cluster_ids, n_pred_clusters).to(torch.float64)

    n_ij = torch.einsum("bnc,bnk->bck", true_mask_oh, pred_mask_oh)
    a = torch.sum(n_ij, axis=-1)
    b = torch.sum(n_ij, axis=-2)
    n_fg_points = torch.sum(a, axis=1)

    rindex = torch.sum(n_ij * (n_ij - 1), axis=(1, 2))
    aindex = torch.sum(a * (a - 1), axis=1)
    bindex = torch.sum(b * (b - 1), axis=1)
    expected_rindex = aindex * bindex / torch.clamp(n_fg_points * (n_fg_points - 1), min=1)
    max_rindex = (aindex + bindex) / 2
    denominator = max_rindex - expected_rindex
    ari = (rindex - expected_rindex) / denominator

    # There are two cases for which the denominator can be zero:
    # 1. If both true_mask and pred_mask assign all pixels to a single cluster.
    #    (max_rindex == expected_rindex == rindex == n_fg_points * (n_fg_points-1))
    # 2. If both true_mask and pred_mask assign max 1 point to each cluster.
    #    (max_rindex == expected_rindex == rindex == 0)
    # In both cases, we want the ARI score to be 1.0:
    return torch.where(denominator > 0, ari, torch.ones_like(ari))


def fg_adjusted_rand_index(
    pred_mask: torch.Tensor, true_mask: torch.Tensor, bg_dim: int = 0
) -> torch.Tensor:
    """Compute adjusted random index using only foreground groups (FG-ARI).

    Args:
        pred_mask: Predicted cluster assignment encoded as categorical probabilities of shape
            (batch_size, n_points, n_pred_clusters).
        true_mask: True cluster assignment encoded as one-hot of shape (batch_size, n_points,
            n_true_clusters).
        bg_dim: Index of background class in true mask.

    Returns:
        ARI scores of shape (batch_size,).
    """
    n_true_clusters = true_mask.shape[-1]
    assert 0 <= bg_dim < n_true_clusters
    if bg_dim == 0:
        true_mask_only_fg = true_mask[..., 1:]
    elif bg_dim == n_true_clusters - 1:
        true_mask_only_fg = true_mask[..., :-1]
    else:
        true_mask_only_fg = torch.cat(
            (true_mask[..., :bg_dim], true_mask[..., bg_dim + 1 :]), dim=-1
        )

    return adjusted_rand_index(pred_mask, true_mask_only_fg)


def _all_equal_masked(values: torch.Tensor, mask: torch.Tensor, dim=-1) -> torch.Tensor:
    """Check if all masked values along a dimension of a tensor are the same.

    All non-masked values are considered as true, i.e. if no value is masked, true is returned
    for this dimension.
    """
    assert mask.dtype == torch.bool
    _, first_non_masked_idx = torch.max(mask, dim=dim)

    comparison_value = values.gather(index=first_non_masked_idx.unsqueeze(dim), dim=dim)

    return torch.logical_or(~mask, values == comparison_value).all(dim=dim)


class UnsupervisedMaskIoUMetric(torchmetrics.Metric, RoutableMixin):
    """Computes IoU metric for segmentation masks when correspondences to ground truth are not known.

    Uses Hungarian matching to compute the assignment between predicted classes and ground truth
    classes.

    Args:
        use_threshold: If `True`, convert predicted class probabilities to mask using a threshold.
            If `False`, class probabilities are turned into mask using a softmax instead.
        threshold: Value to use for thresholding masks.
        matching: Approach to match predicted to ground truth classes. For "hungarian", computes
            assignment that maximizes total IoU between all classes. For "best_overlap", uses the
            predicted class with maximum overlap for each ground truth class. Using "best_overlap"
            leads to the "average best overlap" metric.
        compute_discovery_fraction: Instead of the IoU, compute the fraction of ground truth classes
            that were "discovered", meaning that they have an IoU greater than some threshold.
        correct_localization: Instead of the IoU, compute the fraction of images on which at least
            one ground truth class was correctly localised, meaning that they have an IoU
            greater than some threshold.
        discovery_threshold: Minimum IoU to count a class as discovered/correctly localized.
        ignore_background: If true, assume class at index 0 of ground truth masks is background class
            that is removed before computing IoU.
    """

    def __init__(
        self,
        prediction_path: str,
        target_path: str,
        use_threshold: bool = False,
        threshold: float = 0.5,
        matching: str = "hungarian",
        compute_discovery_fraction: bool = False,
        correct_localization: bool = False,
        discovery_threshold: float = 0.5,
        ignore_background: bool = False,
    ):
        torchmetrics.Metric.__init__(self)
        RoutableMixin.__init__(self, {"prediction": prediction_path, "target": target_path})
        self.use_threshold = use_threshold
        self.threshold = threshold
        self.discovery_threshold = discovery_threshold
        self.compute_discovery_fraction = compute_discovery_fraction
        self.correct_localization = correct_localization
        if compute_discovery_fraction and correct_localization:
            raise ValueError(
                "Only one of `compute_discovery_fraction` and `correct_localization` can be enabled."
            )

        matchings = ("hungarian", "best_overlap")
        if matching not in matchings:
            raise ValueError(f"Unknown matching type {matching}. Valid values are {matchings}.")
        self.matching = matching
        self.ignore_background = ignore_background

        self.add_state(
            "values", default=torch.tensor(0.0, dtype=torch.float64), dist_reduce_fx="sum"
        )
        self.add_state("total", default=torch.tensor(0), dist_reduce_fx="sum")

    @RoutableMixin.route
    def update(self, prediction: torch.Tensor, target: torch.Tensor):
        """Update this metric.

        Args:
            prediction: Predicted mask of shape (B, C, H, W) or (B, F, C, H, W), where C is the
                number of classes. Assumes class probabilities as inputs.
            target: Ground truth mask of shape (B, K, H, W) or (B, F, K, H, W), where K is the
                number of classes.
        """
        if prediction.ndim == 5:
            # Merge frames, height and width to single dimension.
            predictions = prediction.transpose(1, 2).flatten(-3, -1)
            targets = target.transpose(1, 2).flatten(-3, -1)
        elif prediction.ndim == 4:
            # Merge height and width to single dimension.
            predictions = prediction.flatten(-2, -1)
            targets = target.flatten(-2, -1)
        else:
            raise ValueError(f"Incorrect input shape: f{prediction.shape}")

        if self.use_threshold:
            predictions = predictions > self.threshold
        else:
            indices = torch.argmax(predictions, dim=1)
            predictions = torch.nn.functional.one_hot(indices, num_classes=predictions.shape[1])
            predictions = predictions.transpose(1, 2)

        if self.ignore_background:
            target = target[:, 1:]

        for pred, target in zip(predictions, targets):
            nonzero_classes = torch.nonzero(torch.sum(target, dim=-1) > 0).view(-1)
            if len(nonzero_classes) == 0:
                continue  # Skip elements without any target mask
            n_gt_classes = torch.max(nonzero_classes) + 1
            target = target[:n_gt_classes]  # Remove padded classes

            iou_per_class = unsupervised_mask_iou(
                pred, target, matching=self.matching, reduction="none"
            )
            if self.compute_discovery_fraction:
                discovered = iou_per_class > self.discovery_threshold
                self.values += discovered.sum() / len(discovered)
            elif self.correct_localization:
                correctly_localized = torch.any(iou_per_class > self.discovery_threshold)
                self.values += correctly_localized.sum()
            else:
                self.values += iou_per_class.mean()
            self.total += 1

    def compute(self) -> torch.Tensor:
        if self.total == 0:
            return torch.zeros_like(self.values)
        else:
            return self.values / self.total


def unsupervised_mask_iou(
    pred_mask: torch.Tensor,
    true_mask: torch.Tensor,
    matching: str = "hungarian",
    reduction: str = "mean",
) -> torch.Tensor:
    """Compute intersection-over-union (IoU) between masks with unknown class correspondences.

    This metric is also known as Jaccard index. Note that this is a non-batched implementation.

    Args:
        pred_mask: Predicted mask of shape (C, N), where C is the number of predicted classes and
            N is the number of points. Masks are assumed to be binary.
        true_mask: Ground truth mask of shape (K, N), where K is the number of ground truth
            classes and N is the number of points. Masks are assumed to be binary.
        matching: How to match predicted classes to ground truth classes. For "hungarian", computes
            assignment that maximizes total IoU between all classes. For "best_overlap", uses the
            predicted class with maximum overlap for each ground truth class (each predicted class
            can be assigned to multiple ground truth classes). Empty ground truth classes are
            assigned IoU of zero.
        reduction: If "mean", return IoU averaged over classes. If "none", return per-class IoU.

    Returns:
        Mean IoU over classes if reduction is `mean`, tensor of shape (K,) containing per-class IoU
        otherwise.
    """
    n_gt_classes = len(true_mask)
    pred_mask = pred_mask.unsqueeze(1).to(torch.bool)
    true_mask = true_mask.unsqueeze(0).to(torch.bool)

    intersection = torch.sum(pred_mask & true_mask, dim=-1).to(torch.float64)
    union = torch.sum(pred_mask | true_mask, dim=-1).to(torch.float64)
    pairwise_iou = intersection / union
    # Remove NaN from divide-by-zero: class does not occur, and class was not predicted.
    pairwise_iou[union == 0] = 1.0

    if matching == "hungarian":
        pred_idxs, true_idxs = scipy.optimize.linear_sum_assignment(
            pairwise_iou.cpu(), maximize=True
        )
        pred_idxs = torch.as_tensor(pred_idxs, dtype=torch.int64, device=pairwise_iou.device)
        true_idxs = torch.as_tensor(true_idxs, dtype=torch.int64, device=pairwise_iou.device)
    elif matching == "best_overlap":
        non_empty_gt = torch.sum(true_mask.squeeze(0), dim=1) > 0
        pred_idxs = torch.argmax(pairwise_iou, dim=0)[non_empty_gt]
        true_idxs = torch.arange(pairwise_iou.shape[1])[non_empty_gt]
    else:
        raise ValueError(f"Unknown matching {matching}")

    matched_iou = pairwise_iou[pred_idxs, true_idxs]

    iou = torch.zeros(n_gt_classes, dtype=torch.float64, device=pairwise_iou.device)
    iou[true_idxs] = matched_iou

    if reduction == "mean":
        return iou.mean()
    else:
        return iou


class UnsupervisedBboxIoUMetric(torchmetrics.Metric, RoutableMixin):
    """Computes IoU metric for bounding boxes when correspondences to ground truth are not known.

    Currently, assumes segmentation masks as input for both prediction and targets.

    Args:
        target_is_mask: If `True`, assume input is a segmentation mask, in which case the masks are
            converted to bounding boxes before computing IoU. If `False`, assume the input for the
            targets are already bounding boxes.
        use_threshold: If `True`, convert predicted class probabilities to mask using a threshold.
            If `False`, class probabilities are turned into mask using a softmax instead.
        threshold: Value to use for thresholding masks.
        matching: How to match predicted boxes to ground truth boxes. For "hungarian", computes
            assignment that maximizes total IoU between all boxes. For "best_overlap", uses the
            predicted box with maximum overlap for each ground truth box (each predicted box
            can be assigned to multiple ground truth boxes).
        compute_discovery_fraction: Instead of the IoU, compute the fraction of ground truth classes
            that were "discovered", meaning that they have an IoU greater than some threshold. This
            is recall, or sometimes called the detection rate metric.
        correct_localization: Instead of the IoU, compute the fraction of images on which at least
            one ground truth bounding box was correctly localised, meaning that they have an IoU
            greater than some threshold.
        discovery_threshold: Minimum IoU to count a class as discovered/correctly localized.
    """

    def __init__(
        self,
        prediction_path: str,
        target_path: str,
        target_is_mask: bool = False,
        use_threshold: bool = False,
        threshold: float = 0.5,
        matching: str = "hungarian",
        compute_discovery_fraction: bool = False,
        correct_localization: bool = False,
        discovery_threshold: float = 0.5,
    ):
        torchmetrics.Metric.__init__(self)
        RoutableMixin.__init__(self, {"prediction": prediction_path, "target": target_path})
        self.target_is_mask = target_is_mask
        self.use_threshold = use_threshold
        self.threshold = threshold
        self.discovery_threshold = discovery_threshold
        self.compute_discovery_fraction = compute_discovery_fraction
        self.correct_localization = correct_localization
        if compute_discovery_fraction and correct_localization:
            raise ValueError(
                "Only one of `compute_discovery_fraction` and `correct_localization` can be enabled."
            )

        matchings = ("hungarian", "best_overlap")
        if matching not in matchings:
            raise ValueError(f"Unknown matching type {matching}. Valid values are {matchings}.")
        self.matching = matching

        self.add_state(
            "values", default=torch.tensor(0.0, dtype=torch.float64), dist_reduce_fx="sum"
        )
        self.add_state("total", default=torch.tensor(0), dist_reduce_fx="sum")

    @RoutableMixin.route
    def update(self, prediction: torch.Tensor, target: torch.Tensor):
        """Update this metric.

        Args:
            prediction: Predicted mask of shape (B, C, H, W) or (B, F, C, H, W), where C is the
                number of instances. Assumes class probabilities as inputs.
            target: Ground truth mask of shape (B, K, H, W) or (B, F, K, H, W), where K is the
                number of instance, if using masks as input, or bounding boxes of shape (B, K, 4)
                or (B, F, K, 4).
        """
        if prediction.ndim == 5:
            # Merge batch and frame dimensions
            prediction = prediction.flatten(0, 1)   # Not tracklet
            target = target.flatten(0, 1)
        elif prediction.ndim != 4:
            raise ValueError(f"Incorrect input shape: f{prediction.shape}")

        bs, n_pred_classes = prediction.shape[:2]
        n_gt_classes = target.shape[1]

        if self.use_threshold:
            prediction = prediction > self.threshold
        else:
            indices = torch.argmax(prediction, dim=1)
            prediction = torch.nn.functional.one_hot(indices, num_classes=n_pred_classes)
            prediction = prediction.permute(0, 3, 1, 2)

        pred_bboxes = masks_to_bboxes(prediction.flatten(0, 1)).unflatten(0, (bs, n_pred_classes))
        if self.target_is_mask:
            target_bboxes = masks_to_bboxes(target.flatten(0, 1)).unflatten(0, (bs, n_gt_classes))
        else:
            assert target.shape[-1] == 4
            # Convert all-zero boxes added during padding to invalid boxes
            target[torch.all(target == 0.0, dim=-1)] = -1.0
            target_bboxes = target

        for pred, target in zip(pred_bboxes, target_bboxes):
            valid_pred_bboxes = pred[:, 0] != -1.0
            valid_target_bboxes = target[:, 0] != -1.0
            if valid_target_bboxes.sum() == 0:
                continue  # Skip data points without any target bbox

            pred = pred[valid_pred_bboxes]
            target = target[valid_target_bboxes]

            if valid_pred_bboxes.sum() > 0:
                iou_per_bbox = unsupervised_bbox_iou(
                    pred, target, matching=self.matching, reduction="none"
                )
            else:
                iou_per_bbox = torch.zeros_like(valid_target_bboxes, dtype=torch.float32)

            if self.compute_discovery_fraction:
                discovered = iou_per_bbox > self.discovery_threshold
                self.values += discovered.sum() / len(iou_per_bbox)
            elif self.correct_localization:
                correctly_localized = torch.any(iou_per_bbox > self.discovery_threshold)
                self.values += correctly_localized.sum()
            else:
                self.values += iou_per_bbox.mean()
            self.total += 1

    def compute(self) -> torch.Tensor:
        if self.total == 0:
            return torch.zeros_like(self.values)
        else:
            return self.values / self.total

class MOTMetric(torchmetrics.Metric, RoutableMixin):
    def __init__(
        self,
        prediction_path: str,
        target_path: str,
        target_is_mask: bool = True,
        use_threshold: bool = True,
        threshold: float = 0.5,
        ignore_background: bool = False,
        pred_is_mask: bool = True,
    ):
        torchmetrics.Metric.__init__(self)
        RoutableMixin.__init__(self, {"prediction": prediction_path, "target": target_path})
        self.target_is_mask = target_is_mask
        self.use_threshold = use_threshold
        self.threshold = threshold
        self.pred_is_mask = pred_is_mask
        self.reset_accumulator()
        self.accuracy = []
        self.ignore_background = ignore_background
        self.add_state(
            "values", default=torch.tensor(0.0, dtype=torch.float64), dist_reduce_fx="sum"
        )
        self.add_state("total", default=torch.tensor(0), dist_reduce_fx="sum")

    def reset_accumulator(self):
        self.acc = mm.MOTAccumulator(auto_id=True)


    @RoutableMixin.route
    def update(self, prediction: torch.Tensor, target: torch.Tensor):
        """Update this metric.

        Args:
            prediction: Predicted mask of shape  (B, F, C, H, W), where C is the
                number of instances. Assumes class probabilities as inputs.
            target: Ground truth mask of shape  (B, F, K, H, W), where K is the
                number of instance, if using masks as input, or bounding boxes of shape (B, K, 4)
                or (B, F, K, 4).
        """
        # Merge batch and frame dimensions
        B, F = prediction.shape[:2]
        prediction = prediction.flatten(0, 1)
        target = target.flatten(0, 1)

        bs, n_pred_classes = prediction.shape[:2]
        n_gt_classes = target.shape[1]


        if self.pred_is_mask:
            if self.use_threshold:
                prediction = prediction > self.threshold
            pred_bboxes = masks_to_bboxes(prediction.flatten(0, 1)).unflatten(0, (B, F, n_pred_classes))
        else:
            assert prediction.shape[-1] == 4
            # prediction[torch.all(prediction == 0.0, dim=-1)] = -1.0
            pred_bboxes = prediction.unflatten(0, (B, F))
        if self.target_is_mask:
            target_bboxes = masks_to_bboxes(target.flatten(0, 1)).unflatten(0, (B, F, n_gt_classes))
        else:
            assert target.shape[-1] == 4
            # Convert all-zero boxes added during padding to invalid boxes
            target[torch.all(target == 0.0, dim=-1)] = -1.0
            target_bboxes = target

        self.reset_accumulator()
        for preds, targets in zip(pred_bboxes, target_bboxes):
            # seq evaluation
            self.reset_accumulator()
            pred_obj = preds[:, :]
            if not self.ignore_background:
                targets_obj = targets[:, :]
            else:
                targets_obj = targets[:, 1:]
            count = 0
            for pred, target, mask in zip(pred_obj, targets_obj, prediction):
                valid_track_box = pred[:, 0] != -1.0
                valid_target_box = target[:, 0] != -1.0

                track_id = valid_track_box.nonzero()[:, 0].detach().cpu().numpy()
                target_id = valid_target_box.nonzero()[:, 0].detach().cpu().numpy()

                if valid_target_box.sum() == 0:
                    continue  # Skip data points without any target bbox

                pred = pred[valid_track_box].detach().cpu().numpy()
                target = target[valid_target_box].detach().cpu().numpy()
                # frame evaluation
                self.eval_frame(pred, target, track_id, target_id, count)
            self.accuracy.append(self.acc)

        self.total += 1

    def eval_frame(self, trk_tlwhs, tgt_tlwhs, trk_ids, tgt_ids, count):
        # get distance matrix
        trk_tlwhs = np.copy(trk_tlwhs)
        tgt_tlwhs = np.copy(tgt_tlwhs)
        trk_ids = np.copy(trk_ids)
        tgt_ids = np.copy(tgt_ids)
        iou_distance = mm.distances.iou_matrix(tgt_tlwhs, trk_tlwhs, max_iou=0.7)
        # acc
        self.acc.update(tgt_ids, trk_ids, iou_distance)



    def convert_motmetric_to_value(self, res):
        dp = res.replace(' ', ';').replace(';;', ';').replace(';;', ';').replace(';;', ';')
        tmp = list(dp)
        tmp[0] = '-'
        dp = ''.join(tmp)
        return io.StringIO(dp)

    def compute(self) -> torch.Tensor:
        if self.total == 0:
            return torch.zeros_like(self.values)
        else:
            metrics = mm.metrics.motchallenge_metrics
            mh = mm.metrics.create()
            summary = mh.compute_many(
                self.accuracy,
                metrics=metrics,
                names=None,
                generate_overall=True
            )
            strsummary = mm.io.render_summary(
                summary,
                formatters=mh.formatters,
                namemap=mm.io.motchallenge_metric_names
            )

            print(strsummary)
            res = self.convert_motmetric_to_value(strsummary)
            df = pd.read_csv(res, sep=";", engine='python')

            mota = df.iloc[-1]['MOTA']
            self.values = torch.tensor(float(mota[:-1]), dtype=torch.float64).to(self.values.device)
            self.reset_accumulator()
            self.accuracy = []
            return self.values


class ThreeDIOUMetric(torchmetrics.Metric, RoutableMixin):
    def __init__(
            self,
            prediction_path: str,
            target_path: str,
            target_is_mask: bool = True,
            use_threshold: bool = True,
            threshold: float = 0.5,
            ignore_background: bool = False,
            pred_is_mask: bool = True,
    ):
        torchmetrics.Metric.__init__(self)
        RoutableMixin.__init__(self, {"prediction": prediction_path, "target": target_path})
        self.target_is_mask = target_is_mask
        self.use_threshold = use_threshold
        self.threshold = threshold
        self.pred_is_mask = pred_is_mask
        self.reset_accumulator()
        self.accuracy = []
        self.ious_ap = []
        self.iou_thrs = [0.5]
        self.ignore_background = ignore_background
        self.add_state(
            "values", default=torch.tensor(0.0, dtype=torch.float64), dist_reduce_fx="sum"
        )
        self.add_state("total", default=torch.tensor(0), dist_reduce_fx="sum")

    def reset_accumulator(self):
        self.acc = mm.MOTAccumulator(auto_id=True)

    @RoutableMixin.route
    def update(self, prediction: torch.Tensor, target: torch.Tensor):
        """Update this metric.

        Args:
            prediction: Predicted mask of shape  (B, F, C, H, W), where C is the
                number of instances. Assumes class probabilities as inputs.
            target: Ground truth mask of shape  (B, F, K, H, W), where K is the
                number of instance, if using masks as input, or bounding boxes of shape (B, K, 4)
                or (B, F, K, 4).
        """

        # Merge batch and frame dimensions
        B, F = prediction.shape[:2]
        prediction = prediction.flatten(0, 1)
        target = target.flatten(0, 1)

        bs, n_pred_classes = prediction.shape[:2]
        n_gt_classes = target.shape[1]

        if self.pred_is_mask:
            if self.use_threshold:
                prediction = prediction > self.threshold
            pred_bboxes = masks_to_bboxes(prediction.flatten(0, 1)).unflatten(0, (B, F, n_pred_classes))
        else:
            assert prediction.shape[-1] == 4
            # prediction[torch.all(prediction == 0.0, dim=-1)] = -1.0
            pred_bboxes = prediction.unflatten(0, (B, F))
        if self.target_is_mask:
            target_bboxes = masks_to_bboxes(target.flatten(0, 1)).unflatten(0, (B, F, n_gt_classes))
        else:
            assert target.shape[-1] == 4
            # Convert all-zero boxes added during padding to invalid boxes
            target[torch.all(target == 0.0, dim=-1)] = -1.0
            target_bboxes = target

        self.reset_accumulator()
        num_video = 0
        for preds, targets in zip(pred_bboxes, target_bboxes):
            # seq evaluation
            self.reset_accumulator()
            pred_obj = preds[:, :]
            if not self.ignore_background:
                targets_obj = targets[:, :]
            else:
                targets_obj = targets[:, 1:]
            count = 0
            frame = 0
            preds_track_dic_annotation = {}
            targets_track_dic_annotation = {}
            preds_track_id = []
            target_track_id = []
            preds_output = []
            targets_output = []
            import math
            for pred, target, mask in zip(pred_obj, targets_obj, prediction):
                valid_track_box = pred[:, 0] != -1.0
                valid_target_box = target[:, 0] != -1.0

                track_id = valid_track_box.nonzero()[:, 0].detach().cpu().numpy()
                target_id = valid_target_box.nonzero()[:, 0].detach().cpu().numpy()

                # if valid_target_box.sum() == 0:
                #    continue  # Skip data points without any target bbox
                cur_obj_idx = track_id.tolist()
                pred = pred[valid_track_box].detach().cpu().numpy()
                target = target[valid_target_box].detach().cpu().numpy()
                # frame evaluation
                target_obj_idx = valid_target_box.detach().cpu().numpy().nonzero()
                for i in cur_obj_idx:
                    preds_track_id.append(i)
                for i in target_obj_idx[0]:
                    target_track_id.append(i)

                for i in range(len(cur_obj_idx)):
                    x, y, width, height = pred[i]
                    if cur_obj_idx[i] not in preds_track_dic_annotation:
                        preds_track_dic_annotation[cur_obj_idx[i]] = [{
                            "image_id": frame,
                            "category_id": 1,
                            "bbox": [x, y, width, height],
                            "score": 1,
                        }]
                    else:
                        preds_track_dic_annotation[cur_obj_idx[i]].append({
                            "image_id": frame,
                            "category_id": 1,
                            "bbox": [x, y, width, height],
                            "score": 1,
                        })
                for i in range(len(target_obj_idx[0])):
                    x, y, width, height = target[i]
                    if target_obj_idx[0][i] not in targets_track_dic_annotation:
                        targets_track_dic_annotation[target_obj_idx[0][i]] = [{
                            "image_id": frame,
                            "category_id": 1,
                            "bbox": [x, y, width, height],
                            "score": 1,
                        }]
                    else:
                        targets_track_dic_annotation[target_obj_idx[0][i]].append({
                            "image_id": frame,
                            "category_id": 1,
                            "bbox": [x, y, width, height],
                            "score": 1,
                        })
                frame = frame + 1
            preds_track_id = set(preds_track_id)
            target_track_id = set(target_track_id)

            for i in preds_track_id:
                preds_output.append(
                    {'track_id': i, 'annotations': preds_track_dic_annotation[i], 'video_id': num_video})
            for i in target_track_id:
                targets_output.append(
                    {'track_id': i, 'annotations': targets_track_dic_annotation[i], 'video_id': num_video})
            gt = [{g['image_id']: g["bbox"]
                   for g in gt_track['annotations']} for gt_track in preds_output]
            dt = [{d['image_id']: d["bbox"]
                   for d in dt_track['annotations']} for dt_track in targets_output]
            ious = np.zeros([len(dt), len(gt)])
            for i, j in np.ndindex(ious.shape):
                ious[i, j] = compute_track_box_iou(dt[i], gt[j])

            eval_result = accumulate(preds_output, targets_output, ious, iou_thrs=self.iou_thrs)
            result_ap = _summarize(eval_result, 'ap', iou_thrs=np.array(self.iou_thrs))
            self.ious_ap.append(result_ap)
            num_video = num_video + 1

        self.total += 1

    def compute(self) -> torch.Tensor:
        if self.total == 0:
            return torch.zeros_like(self.values)
        else:
            print("self.ious_ap num:", len(self.ious_ap))
            values = np.array(self.ious_ap).mean()
            self.values = torch.tensor(values).to(self.values.device)
            print("self.values:", self.values)
            return self.values

@jit
def bb_intersect_union(d, g):
    """Compute intersection and union separately.
    Inspired by
    <https://github.com/cocodataset/cocoapi/blob/636becdc73d54283b3aac6d4ec363cffbb6f9b20/common/maskApi.c#L109>
    >>> bb_intersect_union([0, 0, 20, 20], [0, 0, 20, 20], False)
    (400, 400)
    >>> bb_intersect_union([0, 0, 20, 20], [0, 0, 10, 10], False)
    (100, 400)
    >>> bb_intersect_union([0, 0, 20, 20], [0, 0, 10, 10], True)
    (100, 400)
    >>> bb_intersect_union([0, 0, 20, 20], [0, 0, 30, 30], True)
    (400, 400)
    >>> bb_intersect_union([10, 20, 10, 10], [10, 20, 5, 5], True)
    (25, 100)
    """
    dx, dy, dw, dh = d
    gx, gy, gw, gh = g

    detection_area = dw * dh
    groundtruth_area = gw * gh

    intersection_left = max(dx, gx)
    intersection_right = min(dx + dw, gx + gw)
    intersection_top = max(dy, gy)
    intersection_bottom = min(dy + dh, gy + gh)

    w = max(intersection_right - intersection_left, 0)
    h = max(intersection_bottom - intersection_top, 0)

    intersect = w * h
    union = detection_area + groundtruth_area - intersect
    return intersect, union



def compute_track_box_iou(dt_track, gt_track):
    """
    Args:
        dt_track (dict): Map image id to [x0, y0, w, h]
        gt_track (dict): Map image id to [x0, y0, w, h]
    """
    i = 0
    u = 0
    image_ids = set(gt_track.keys()) | set(dt_track.keys())
    for image in image_ids:
        g = gt_track.get(image, None)
        d = dt_track.get(image, None)
        if d and g:
            i_, u_ = bb_intersect_union(d, g)
            i += i_
            u += u_
        elif not d and g:
            u += g[2] * g[3]
        elif d and not g:
            u += d[2] * d[3]
    assert i <= u
    return i / u if u > 0 else 0

def evaluate_vid(gt, dt, ious, iou_thrs = [0.9]):
    """Perform evaluation for single category and video."""

    if len(gt) == 0 and len(dt) == 0:
        return None

    # Add another filed _ignore to only consider anns based on area range.
    for g in gt:
        g["_ignore"] = 0

    num_thrs = len(iou_thrs)
    num_gt = len(gt)
    num_dt = len(dt)

    # Array to store the "id" of the matched dt/gt
    gt_m = np.zeros((num_thrs, num_gt)) - 1
    dt_m = np.zeros((num_thrs, num_dt)) - 1

    gt_ig = np.array([g["_ignore"] for g in gt])
    dt_ig = np.zeros((num_thrs, num_dt))

    for iou_thr_idx, iou_thr in enumerate(iou_thrs):
        if len(ious) == 0:
            break

        for dt_idx, _dt in enumerate(dt):
            iou = min([iou_thr, 1 - 1e-10])
            # information about best match so far (m=-1 -> unmatched)
            # store the gt_idx which matched for _dt
            m = -1
            for gt_idx, _ in enumerate(gt):
                # if this gt already matched continue
                if gt_m[iou_thr_idx, gt_idx] > 0:
                    continue
                # if _dt matched to reg gt, and on ignore gt, stop
                if m > -1 and gt_ig[m] == 0 and gt_ig[gt_idx] == 1:
                    break
                # continue to next gt unless better match made
                if ious[dt_idx, gt_idx] < iou:
                    continue
                # if match successful and best so far, store appropriately
                iou = ious[dt_idx, gt_idx]
                m = gt_idx

            # No match found for _dt, go to next _dt
            if m == -1:
                continue

            # if gt to ignore for some reason update dt_ig.
            # Should not be used in evaluation.
            dt_ig[iou_thr_idx, dt_idx] = gt_ig[m]
            # _dt match found, update gt_m, and dt_m with "id"
            dt_m[iou_thr_idx, dt_idx] = gt[m]["track_id"]
            gt_m[iou_thr_idx, m] = _dt["track_id"]

    vid_id = gt[0]["video_id"]
    # store results for given video and category
    return {
        "video_id": vid_id,
        "dt_ids": [d["track_id"] for d in dt],
        "gt_ids": [g["track_id"] for g in gt],
        "dt_matches": dt_m,
        "gt_matches": gt_m,
        "gt_ignore": gt_ig,
        "dt_ignore": dt_ig,
    }




def accumulate(preds_output, targets_output, ious, iou_thrs):
    """Accumulate per video evaluation results and store the result in
    self.eval.
    """

    num_thrs = len(iou_thrs)
    rec_thrs = np.linspace(
        0.0, 1.00, int(np.round((1.00 - 0.0) / 0.01) + 1), endpoint=True
    )
    num_recalls = len(rec_thrs)
    num_cats = 1
    num_vids = 1

    # -1 for absent categories
    precision = -np.ones(
        (num_thrs, num_recalls)
    )
    recall = -np.ones((num_thrs))

    # Initialize dt_pointers
    dt_pointers = {}

    # Per category evaluation
    E = [
        evaluate_vid(preds_output, targets_output, ious, iou_thrs)
    ]
    # Remove elements which are None
    E = [e for e in E if e is not None]
    # if len(E) == 0:
    #    continue

    # Append all scores: shape (N,)
    # dt_scores = np.concatenate([e["dt_scores"] for e in E], axis=0)
    dt_ids = np.concatenate([e["dt_ids"] for e in E], axis=0)
    dt_m = np.concatenate([e["dt_matches"] for e in E],
                          axis=1)
    dt_ig = np.concatenate([e["dt_ignore"] for e in E],
                           axis=1)

    gt_ig = np.concatenate([e["gt_ignore"] for e in E])
    # num gt anns to consider
    num_gt = np.count_nonzero(gt_ig == 0)
    tps = np.logical_and(dt_m != -1, np.logical_not(dt_ig))
    fps = np.logical_and(dt_m == -1, np.logical_not(dt_ig))

    tp_sum = np.cumsum(tps, axis=1).astype(dtype=np.float)
    fp_sum = np.cumsum(fps, axis=1).astype(dtype=np.float)

    dt_pointers = {
        "dt_ids": dt_ids,
        "tps": tps,
        "fps": fps,
    }

    for iou_thr_idx, (tp, fp) in enumerate(zip(tp_sum, fp_sum)):
        tp = np.array(tp)
        fp = np.array(fp)
        num_tp = len(tp)
        rc = tp / num_gt
        if num_tp:
            recall[iou_thr_idx] = rc[-1]
        else:
            recall[iou_thr_idx] = 0

        # np.spacing(1) ~= eps
        pr = tp / (fp + tp + np.spacing(1))
        pr = pr.tolist()

        # Replace each precision value with the maximum precision
        # value to the right of that recall level. This ensures
        # that the  calculated AP value will be less suspectable
        # to small variations in the ranking.
        for i in range(num_tp - 1, 0, -1):
            if pr[i] > pr[i - 1]:
                pr[i - 1] = pr[i]

        rec_thrs_insert_idx = np.searchsorted(
            rc, rec_thrs, side="left"
        )

        pr_at_recall = [0.0] * num_recalls

        try:
            for _idx, pr_idx in enumerate(rec_thrs_insert_idx):
                pr_at_recall[_idx] = pr[pr_idx]
        except:  # noqa: E722
            pass
        precision[iou_thr_idx, :] = (
            np.array(pr_at_recall))

    eval_result = {
        # "params": self.params,
        "counts": [
            num_thrs, num_recalls
        ],
        "precision": precision,
        "recall": recall,
        "dt_pointers": dt_pointers,
    }
    return eval_result


def _summarize(eval_result,
               summary_type,
               iou_thr=None,
               iou_thrs=[0.8, 0.9]):
    if summary_type == 'ap':
        s = eval_result["precision"]

        if iou_thr is not None:
            tidx = np.where(iou_thr == iou_thrs)[0]
            s = s[tidx]

    else:
        s = eval_result["recall"]

        if iou_thr is not None:
            tidx = np.where(iou_thr == iou_thrs)[0]
            s = s[tidx]

    if len(s[s > -1]) == 0:
        mean_s = -1
    else:
        mean_s = np.mean(s[s > -1])
    return mean_s


def unsupervised_bbox_iou(
    pred_bboxes: torch.Tensor,
    true_bboxes: torch.Tensor,
    matching: str = "best_overlap",
    reduction: str = "mean",
) -> torch.Tensor:
    """Compute IoU between two sets of bounding boxes.

    Args:
        pred_bboxes: Predicted bounding boxes of shape N x 4.
        true_bboxes: True bounding boxes of shape M x 4.
        matching: Method to assign predicted to true bounding boxes.
        reduction: Whether to average the computes IoUs per true box.
    """
    n_gt_bboxes = len(true_bboxes)

    pairwise_iou = torchvision.ops.box_iou(pred_bboxes, true_bboxes)

    if matching == "hungarian":
        pred_idxs, true_idxs = scipy.optimize.linear_sum_assignment(
            pairwise_iou.cpu(), maximize=True
        )
        pred_idxs = torch.as_tensor(pred_idxs, dtype=torch.int64, device=pairwise_iou.device)
        true_idxs = torch.as_tensor(true_idxs, dtype=torch.int64, device=pairwise_iou.device)
    elif matching == "best_overlap":
        pred_idxs = torch.argmax(pairwise_iou, dim=0)
        true_idxs = torch.arange(pairwise_iou.shape[1], device=pairwise_iou.device)
    else:
        raise ValueError(f"Unknown matching {matching}")

    matched_iou = pairwise_iou[pred_idxs, true_idxs]

    iou = torch.zeros(n_gt_bboxes, dtype=torch.float32, device=pairwise_iou.device)
    iou[true_idxs] = matched_iou

    if reduction == "mean":
        return iou.mean()
    else:
        return iou


def masks_to_bboxes(masks: torch.Tensor, empty_value: float = -1.0) -> torch.Tensor:
    """Compute bounding boxes around the provided masks.

    Adapted from DETR: https://github.com/facebookresearch/detr/blob/main/util/box_ops.py

    Args:
        masks: Tensor of shape (N, H, W), where N is the number of masks, H and W are the spatial
            dimensions.
        empty_value: Value bounding boxes should contain for empty masks.

    Returns:
        Tensor of shape (N, 4), containing bounding boxes in (x1, y1, x2, y2) format, where (x1, y1)
        is the coordinate of top-left corner and (x2, y2) is the coordinate of the bottom-right
        corner (inclusive) in pixel coordinates. If mask is empty, all coordinates contain
        `empty_value` instead.
    """
    masks = masks.bool()
    if masks.numel() == 0:
        return torch.zeros((0, 4), device=masks.device)

    large_value = 1e8
    inv_mask = ~masks

    h, w = masks.shape[-2:]

    y = torch.arange(0, h, dtype=torch.float, device=masks.device)
    x = torch.arange(0, w, dtype=torch.float, device=masks.device)
    y, x = torch.meshgrid(y, x, indexing="ij")

    x_mask = masks * x.unsqueeze(0)
    x_max = x_mask.flatten(1).max(-1)[0]
    x_min = x_mask.masked_fill(inv_mask, large_value).flatten(1).min(-1)[0]

    y_mask = masks * y.unsqueeze(0)
    y_max = y_mask.flatten(1).max(-1)[0]
    y_min = y_mask.masked_fill(inv_mask, large_value).flatten(1).min(-1)[0]

    bboxes = torch.stack((x_min, y_min, x_max-x_min, y_max-y_min), dim=1) #x1y1wh
    bboxes[x_min == large_value] = empty_value

    return bboxes

def masks_to_bboxes_xyxy(masks: torch.Tensor, empty_value: float = -1.0) -> torch.Tensor:
    """Compute bounding boxes around the provided masks.

    Adapted from DETR: https://github.com/facebookresearch/detr/blob/main/util/box_ops.py

    Args:
        masks: Tensor of shape (N, H, W), where N is the number of masks, H and W are the spatial
            dimensions.
        empty_value: Value bounding boxes should contain for empty masks.

    Returns:
        Tensor of shape (N, 4), containing bounding boxes in (x1, y1, x2, y2) format, where (x1, y1)
        is the coordinate of top-left corner and (x2, y2) is the coordinate of the bottom-right
        corner (inclusive) in pixel coordinates. If mask is empty, all coordinates contain
        `empty_value` instead.
    """
    masks = masks.bool()
    if masks.numel() == 0:
        return torch.zeros((0, 4), device=masks.device)

    large_value = 1e8
    inv_mask = ~masks

    h, w = masks.shape[-2:]

    y = torch.arange(0, h, dtype=torch.float, device=masks.device)
    x = torch.arange(0, w, dtype=torch.float, device=masks.device)
    y, x = torch.meshgrid(y, x, indexing="ij")

    x_mask = masks * x.unsqueeze(0)
    x_max = x_mask.flatten(1).max(-1)[0]
    x_min = x_mask.masked_fill(inv_mask, large_value).flatten(1).min(-1)[0]

    y_mask = masks * y.unsqueeze(0)
    y_max = y_mask.flatten(1).max(-1)[0]
    y_min = y_mask.masked_fill(inv_mask, large_value).flatten(1).min(-1)[0]

    bboxes = torch.stack((x_min, y_min, x_max, y_max), dim=1) #x1y1x2y2
    bboxes[x_min == large_value] = empty_value
    ## filter background

    return bboxes

class DatasetSemanticMaskIoUMetric(torchmetrics.Metric):
    """Unsupervised IoU metric for semantic segmentation using dataset-wide matching of classes.

    The input to this metric is an instance-level mask with objects, and a class id for each object.
    This is required to convert the mask to semantic classes. The number of classes for the
    predictions does not have to match the true number of classes.

    Note that contrary to the other metrics in this module, this metric is not supposed to be added
    in the online metric computation loop, which is why it does not inherit from `RoutableMixin`.

    Args:
        n_predicted_classes: Number of predictable classes, i.e. highest prediction class id that can
            occur.
        n_classes: Total number of classes, i.e. highest class id that can occur.
        threshold: Value to use for thresholding masks.
        use_threshold: If `True`, convert predicted class probabilities to mask using a threshold.
            If `False`, class probabilities are turned into mask using an argmax instead.
        matching: Method to produce matching between clusters and ground truth classes. If
            "hungarian", assigns each class one cluster such that the total IoU is maximized. If
            "majority", assigns each cluster to the class with the highest IoU (each class can be
            assigned multiple clusters).
        ignore_background: If true, pixels labeled as background (class zero) in the ground truth
            are not taken into account when computing IoU.
        use_unmatched_as_background: If true, count predicted classes not selected after Hungarian
            matching as the background predictions.
    """

    def __init__(
        self,
        n_predicted_classes: int,
        n_classes: int,
        use_threshold: bool = False,
        threshold: float = 0.5,
        matching: str = "hungarian",
        ignore_background: bool = False,
        use_unmatched_as_background: bool = False,
    ):
        super().__init__()
        matching_methods = {"hungarian", "majority"}
        if matching not in matching_methods:
            raise ValueError(
                f"Unknown matching method {matching}. Valid values are {matching_methods}."
            )

        self.matching = matching
        self.n_predicted_classes = n_predicted_classes
        self.n_predicted_classes_with_bg = n_predicted_classes + 1
        self.n_classes = n_classes
        self.n_classes_with_bg = n_classes + 1
        self.matching = matching
        self.use_threshold = use_threshold
        self.threshold = threshold
        self.ignore_background = ignore_background
        self.use_unmatched_as_background = use_unmatched_as_background
        if use_unmatched_as_background and ignore_background:
            raise ValueError(
                "Option `use_unmatched_as_background` not compatible with option `ignore_background`"
            )
        if use_unmatched_as_background and matching == "majority":
            raise ValueError(
                "Option `use_unmatched_as_background` not compatible with matching `majority`"
            )

        confusion_mat = torch.zeros(
            self.n_predicted_classes_with_bg, self.n_classes_with_bg, dtype=torch.int64
        )
        self.add_state("confusion_mat", default=confusion_mat, dist_reduce_fx="sum", persistent=True)

    def update(
        self, predictions: torch.Tensor, targets: torch.Tensor, prediction_class_ids: torch.Tensor
    ):
        """Update metric by computing confusion matrix between predicted and target classes.

        Args:
            predictions: Probability mask of shape (B, K, H, W) or (B, F, K, H, W), where K is the
                number of object instances in the image.
            targets: Mask of shape (B, K, H, W) or (B, F, K, H, W), where K is the number of object
                instances in the image. Class ID of objects is encoded as the value, i.e. densely
                represented.
            prediction_class_ids: Tensor of shape (B, K), containing the class id of each predicted
                object instance in the image. Id must be 0 <= id <= n_predicted_classes.
        """
        predictions = self.preprocess_predicted_mask(predictions)
        predictions = _remap_one_hot_mask(
            predictions, prediction_class_ids, self.n_predicted_classes, strip_empty=False
        )
        assert predictions.shape[-1] == self.n_predicted_classes_with_bg

        targets = self.preprocess_ground_truth_mask(targets)
        assert targets.shape[-1] == self.n_classes_with_bg

        # We are doing the multiply in float64 instead of int64 because it proved to be significantly
        # faster on GPU. We need to use 64 bits because we can easily exceed the range of 32 bits
        # if we aggregate over a full dataset.
        confusion_mat = torch.einsum(
            "bpk,bpc->kc", predictions.to(torch.float64), targets.to(torch.float64)
        )
        self.confusion_mat += confusion_mat.to(torch.int64)

    def preprocess_predicted_mask(self, mask: torch.Tensor) -> torch.Tensor:
        """Preprocess predicted masks for metric computation.

        Args:
            mask: Probability mask of shape (B, K, H, W) or (B, F, K, H, W), where K is the number
                of object instances in the prediction.

        Returns:
            Binary tensor of shape (B, P, K), where P is the number of points. If `use_threshold` is
            True, overlapping objects for the same point are possible.
        """
        if mask.ndim == 5:  # Video case
            mask = mask.flatten(0, 1)
        mask = mask.flatten(-2, -1)

        if self.use_threshold:
            mask = mask > self.threshold
            mask = mask.transpose(1, 2)
        else:
            maximum, indices = torch.max(mask, dim=1)
            mask = torch.nn.functional.one_hot(indices, num_classes=mask.shape[1])
            mask[:, :, 0][maximum == 0.0] = 0

        return mask

    def preprocess_ground_truth_mask(self, mask: torch.Tensor) -> torch.Tensor:
        """Preprocess ground truth mask for metric computation.

        Args:
            mask: Mask of shape (B, K, H, W) or (B, F, K, H, W), where K is the number of object
                instances in the image. Class ID of objects is encoded as the value, i.e. densely
                represented.

        Returns:
            One-hot tensor of shape (B, P, J), where J is the number of the classes and P the number
            of points, with object instances with the same class ID merged together. In the case of
            an overlap of classes for a point, the class with the highest ID is assigned to that
            point.
        """
        if mask.ndim == 5:  # Video case
            mask = mask.flatten(0, 1)
        mask = mask.flatten(-2, -1)

        # Pixels which contain no object get assigned the background class 0. This also handles the
        # padding of zero masks which is done in preprocessing for batching.
        mask = torch.nn.functional.one_hot(
            mask.max(dim=1).values.to(torch.long), num_classes=self.n_classes_with_bg
        )

        return mask

    def compute(self):
        """Compute per-class IoU using matching."""
        if self.ignore_background:
            n_classes = self.n_classes
            confusion_mat = self.confusion_mat[:, 1:]
        else:
            n_classes = self.n_classes_with_bg
            confusion_mat = self.confusion_mat

        pairwise_iou, _, _, area_gt = self._compute_iou_from_confusion_mat(confusion_mat)

        if self.use_unmatched_as_background:
            # Match only in foreground
            pairwise_iou = pairwise_iou[1:, 1:]
            confusion_mat = confusion_mat[1:, 1:]
        else:
            # Predicted class zero is not matched against anything
            pairwise_iou = pairwise_iou[1:]
            confusion_mat = confusion_mat[1:]

        if self.matching == "hungarian":
            cluster_idxs, class_idxs = scipy.optimize.linear_sum_assignment(
                pairwise_iou.cpu(), maximize=True
            )
            cluster_idxs = torch.as_tensor(
                cluster_idxs, dtype=torch.int64, device=self.confusion_mat.device
            )
            class_idxs = torch.as_tensor(
                class_idxs, dtype=torch.int64, device=self.confusion_mat.device
            )
            matched_iou = pairwise_iou[cluster_idxs, class_idxs]
            true_pos = confusion_mat[cluster_idxs, class_idxs]

            if self.use_unmatched_as_background:
                cluster_oh = torch.nn.functional.one_hot(
                    cluster_idxs, num_classes=pairwise_iou.shape[0]
                )
                matched_clusters = cluster_oh.max(dim=0).values.to(torch.bool)
                bg_pred = self.confusion_mat[:1]
                bg_pred += self.confusion_mat[1:][~matched_clusters].sum(dim=0)
                bg_iou, _, _, _ = self._compute_iou_from_confusion_mat(bg_pred, area_gt)
                class_idxs = torch.cat((torch.zeros_like(class_idxs[:1]), class_idxs + 1))
                matched_iou = torch.cat((bg_iou[0, :1], matched_iou))
                true_pos = torch.cat((bg_pred[0, :1], true_pos))

        elif self.matching == "majority":
            max_iou, class_idxs = torch.max(pairwise_iou, dim=1)
            # Form new clusters by merging old clusters which are assigned the same ground truth
            # class. After merging, the number of clusters equals the number of classes.
            _, old_to_new_cluster_idx = torch.unique(class_idxs, return_inverse=True)

            confusion_mat_new = torch.zeros(
                n_classes, n_classes, dtype=torch.int64, device=self.confusion_mat.device
            )
            for old_cluster_idx, new_cluster_idx in enumerate(old_to_new_cluster_idx):
                if max_iou[old_cluster_idx] > 0.0:
                    confusion_mat_new[new_cluster_idx] += confusion_mat[old_cluster_idx]

            # Important: use previously computed area_gt because it includes background predictions,
            # whereas the new confusion matrix does not contain the bg predicted class anymore.
            pairwise_iou, _, _, _ = self._compute_iou_from_confusion_mat(confusion_mat_new, area_gt)
            max_iou, class_idxs = torch.max(pairwise_iou, dim=1)
            valid = max_iou > 0.0  # Ignore clusters without any kind of overlap
            class_idxs = class_idxs[valid]
            cluster_idxs = torch.arange(pairwise_iou.shape[1])[valid]
            matched_iou = pairwise_iou[cluster_idxs, class_idxs]
            true_pos = confusion_mat_new[cluster_idxs, class_idxs]

        iou = torch.zeros(n_classes, dtype=torch.float64, device=pairwise_iou.device)
        iou[class_idxs] = matched_iou

        accuracy = true_pos.sum().to(torch.float64) / area_gt.sum()
        empty_classes = area_gt == 0

        return iou, accuracy, empty_classes

    @staticmethod
    def _compute_iou_from_confusion_mat(
        confusion_mat: torch.Tensor, area_gt: Optional[torch.Tensor] = None
    ):
        area_pred = torch.sum(confusion_mat, axis=1)
        if area_gt is None:
            area_gt = torch.sum(confusion_mat, axis=0)
        union = area_pred.unsqueeze(1) + area_gt.unsqueeze(0) - confusion_mat
        pairwise_iou = confusion_mat.to(torch.float64) / union

        # Ignore classes that occured on no image.
        pairwise_iou[union == 0] = 0.0

        return pairwise_iou, union, area_pred, area_gt


def _remap_one_hot_mask(
    mask: torch.Tensor, new_classes: torch.Tensor, n_new_classes: int, strip_empty: bool = False
):
    """Remap classes from binary mask to new classes.

    In the case of an overlap of classes for a point, the new class with the highest ID is
    assigned to that point. If no class is assigned to a point, the point will have no class
    assigned after remapping as well.

    Args:
        mask: Binary mask of shape (B, P, K) where K is the number of old classes and P is the
            number of points.
        new_classes: Tensor of shape (B, K) containing ids of new classes for each old class.
        n_new_classes: Number of classes after remapping, i.e. highest class id that can occur.
        strip_empty: Whether to remove the empty pixels mask

    Returns:
        Tensor of shape (B, P, J), where J is the new number of classes.
    """
    assert new_classes.shape[1] == mask.shape[2]
    mask_dense = (mask * new_classes.unsqueeze(1)).max(dim=-1).values
    mask = torch.nn.functional.one_hot(mask_dense.to(torch.long), num_classes=n_new_classes + 1)

    if strip_empty:
        mask = mask[..., 1:]

    return mask


class SklearnClustering:
    """Wrapper around scikit-learn clustering algorithms.

    Args:
        n_clusters: Number of clusters.
        method: Clustering method to use.
        clustering_kwargs: Dictionary of additional keyword arguments to pass to clustering object.
        use_l2_normalization: Whether to L2 normalize the representations before clustering (but
            after PCA).
        use_pca: Whether to apply PCA before fitting the clusters.
        pca_dimensions: Number of dimensions for PCA dimensionality reduction. If `None`, do not
            reduce dimensions with PCA.
        pca_kwargs: Dictionary of additional keyword arguments to pass to PCA object.
    """

    def __init__(
        self,
        n_clusters: int,
        method: str = "kmeans",
        clustering_kwargs: Optional[Dict[str, Any]] = None,
        use_l2_normalization: bool = False,
        use_pca: bool = False,
        pca_dimensions: Optional[int] = None,
        pca_kwargs: Optional[Dict[str, Any]] = None,
    ):
        methods = ("kmeans", "spectral")
        if method not in methods:
            raise ValueError(f"Unknown clustering method {method}. Valid values are {methods}.")

        self._n_clusters = n_clusters
        self.method = method
        self.clustering_kwargs = clustering_kwargs
        self.use_l2_normalization = use_l2_normalization
        self.use_pca = use_pca
        self.pca_dimensions = pca_dimensions
        self.pca_kwargs = pca_kwargs

        self._clustering = None
        self._pca = None

    @property
    def n_clusters(self):
        return self._n_clusters

    def _init(self):
        from sklearn import cluster, decomposition

        kwargs = self.clustering_kwargs if self.clustering_kwargs is not None else {}
        if self.method == "kmeans":
            self._clustering = cluster.KMeans(n_clusters=self.n_clusters, **kwargs)
        elif self.method == "spectral":
            self._clustering = cluster.SpectralClustering(n_clusters=self.n_clusters, **kwargs)
        else:
            raise NotImplementedError(f"Clustering {self.method} not implemented.")

        if self.use_pca:
            kwargs = self.pca_kwargs if self.pca_kwargs is not None else {}
            self._pca = decomposition.PCA(n_components=self.pca_dimensions, **kwargs)

    def fit_predict(self, features: torch.Tensor):
        self._init()
        features = features.detach().cpu().numpy()
        if self.use_pca:
            features = self._pca.fit_transform(features)
        if self.use_l2_normalization:
            features /= np.maximum(np.linalg.norm(features, ord=2, axis=1, keepdims=True), 1e-8)
        cluster_ids = self._clustering.fit_predict(features).astype(np.int64)
        return torch.from_numpy(cluster_ids)

    def predict(self, features: torch.Tensor) -> torch.Tensor:
        if self._clustering is None:
            raise ValueError("Clustering was not fitted. Call `fit_predict` first.")

        features = features.detach().cpu().numpy()
        if self.use_pca:
            features = self._pca.transform(features)
        if self.use_l2_normalization:
            features /= np.maximum(np.linalg.norm(features, ord=2, axis=1, keepdims=True), 1e-8)
        cluster_ids = self._clustering.predict(features).astype(np.int64)
        return torch.from_numpy(cluster_ids)

ocl/metrics.py
amazon-science/object-centric-multiple-object-tracking
import torch
import torch.nn as nn
import torch.nn.functional as F
class ScaledDotProductAttention(nn.Module):
    ''' Scaled Dot-Product Attention '''

    def __init__(self, temperature, attn_dropout=0.0):
        super().__init__()
        self.temperature = temperature
        self.dropout = nn.Dropout(attn_dropout)

    def forward(self, q, k, v, mask=None):
        attn = torch.matmul(q / self.temperature, k.transpose(2, 3))

        # if mask is not None:
        #     attn = attn.masked_fill(mask == 0, -1e9)
        if mask is not None:
            bias = (1-mask)*(-1e9)
            attn = attn * mask + bias


        attn = F.softmax(attn, dim=-1)
        output = torch.matmul(attn, v)

        return output, attn


class MultiHeadAttention_for_index(nn.Module):
    ''' Multi-Head Attention module '''

    def __init__(self, n_head, d_model, d_k, d_v, dropout=0):
        super().__init__()

        self.n_head = n_head
        self.d_k = d_k
        self.d_v = d_v

        self.w_qs = nn.Linear(d_model, n_head * d_k, bias=False)
        self.w_ks = nn.Linear(d_model, n_head * d_k, bias=False)
        self.w_vs = nn.Linear(d_model, n_head * d_v, bias=False)
        self.fc = nn.Linear(n_head * d_v, d_model, bias=False)

        '''
        nn.init.constant_(self.w_qs.weight, 0)
        nn.init.constant_(self.w_ks.weight, 0)
        nn.init.constant_(self.w_vs.weight, 0)
        '''

        # nn.init.eye_(self.w_qs.weight)
        nn.init.eye_(self.w_ks.weight)
        nn.init.eye_(self.w_vs.weight)
        # nn.init.eye_(self.w_qs.weight)

        # self.w_vs.eval()
        # self.w_ks.eval()
        # self.w_qs.eval()
        # self.w_vs.weight.requires_grad = False
        # self.w_ks.weight.requires_grad = False
        # self.w_qs.weight.requires_grad = False

        # nn.init.eye_(self.fc.weight)

        self.attention = ScaledDotProductAttention(temperature=d_k ** 0.5)  #temperature=d_k ** 0.5

        self.dropout = nn.Dropout(dropout)
        self.layer_norm = nn.LayerNorm(d_model, eps=1e-6)

    def forward(self, q, k, v, mask=None):
        # print("w_qs:", self.w_qs.weight)
        d_k, d_v, n_head = self.d_k, self.d_v, self.n_head
        sz_b, len_q, len_k, len_v = q.size(0), q.size(1), k.size(1), v.size(1)


        # Pass through the pre-attention projection: b x lq x (n*dv)
        # Separate different heads: b x lq x n x dv
        q = self.w_qs(q).view(sz_b, len_q, n_head, d_k)
        k = self.w_ks(k).view(sz_b, len_k, n_head, d_k)
        v = self.w_vs(v).view(sz_b, len_v, n_head, d_v)
        # print (self.w_qs.weight.abs().sum(), self.w_ks.weight.abs().sum(), self.w_vs.weight.abs().sum())
        # print ('kqv', q.abs().sum(), k.abs().sum(), v.abs().sum())

        # Transpose for attention dot product: b x n x lq x dv
        q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)

        if mask is not None:
            mask = mask.unsqueeze(1)  # For head axis broadcasting.

        q, attn2 = self.attention(q, k, v, mask=mask)


        # attn = attn1 + attn2
        # attn = attn1.unsqueeze(1)
        attn = attn2

        # Transpose to move the head dimension back: b x lq x n x dv
        # Combine the last two dimensions to concatenate all the heads together: b x lq x (n*dv)
        q = q.transpose(1, 2).contiguous().view(sz_b, len_q, -1)
        q = self.dropout(self.fc(q))
        # q += residual

        q = self.layer_norm(q)

        attn = torch.mean(attn, 1)
        return q, attn


class MultiHeadAttention(nn.Module):
    ''' Multi-Head Attention module '''

    def __init__(self, n_head, d_model, d_k, d_v, dropout=0):
        super().__init__()

        self.n_head = n_head
        self.d_k = d_k
        self.d_v = d_v

        self.w_qs = nn.Linear(d_model, n_head * d_k, bias=False)
        self.w_ks = nn.Linear(d_model, n_head * d_k, bias=False)
        self.w_vs = nn.Linear(d_model, n_head * d_v, bias=False)
        self.fc = nn.Linear(n_head * d_v, d_model, bias=False)



        nn.init.eye_(self.w_qs.weight)
        nn.init.eye_(self.w_ks.weight)
        nn.init.eye_(self.w_vs.weight)
        nn.init.eye_(self.fc.weight)
        self.attention = ScaledDotProductAttention(temperature=0.5)

        self.dropout = nn.Dropout(dropout)
        self.layer_norm = nn.LayerNorm(d_model, eps=1e-6)


    def forward(self, q, k, v, mask=None):
        #print("w_qs:", self.w_qs.weight)
        d_k, d_v, n_head = self.d_k, self.d_v, self.n_head
        sz_b, len_q, len_k, len_v = q.size(0), q.size(1), k.size(1), v.size(1)

        residual = q

        # Pass through the pre-attention projection: b x lq x (n*dv)
        # Separate different heads: b x lq x n x dv
        q = self.w_qs(q).view(sz_b, len_q, n_head, d_k)
        k = self.w_ks(k).view(sz_b, len_k, n_head, d_k)
        v = self.w_vs(v).view(sz_b, len_v, n_head, d_v)

        # Transpose for attention dot product: b x n x lq x dv
        q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)

        if mask is not None:
            mask = mask.unsqueeze(1)   # For head axis broadcasting.

        q, attn = self.attention(q, k, v, mask=mask)

        # Transpose to move the head dimension back: b x lq x n x dv
        # Combine the last two dimensions to concatenate all the heads together: b x lq x (n*dv)
        q = q.transpose(1, 2).contiguous().view(sz_b, len_q, -1)
        q = self.dropout(self.fc(q))

        # q += residual
        # q = self.fc2(q)
        # do not norm
        # q = self.layer_norm(q)

        attn = torch.mean(attn, 1)
        return q, attn



class MultiHeadAttention_dotversion_merge(nn.Module):
    ''' Multi-Head Attention module '''

    def __init__(self, n_head, d_model, d_k, d_v, dropout=0):
        super().__init__()

        self.n_head = n_head
        self.fc = nn.Linear(n_head * d_v, d_model, bias=False)
        self.n_head = n_head
        self.d_k = d_k
        self.d_v = d_v

        self.attention = ScaledDotProductAttention(temperature=0.5)

        self.dropout = nn.Dropout(dropout)
        self.layer_norm = nn.LayerNorm(d_model, eps=1e-6)


    def forward(self, q, k, v, mask=None):
        #print("w_qs:", self.w_qs.weight)
        d_k, d_v, n_head = self.d_k, self.d_v, self.n_head
        sz_b, len_q, len_k, len_v = q.size(0), q.size(1), k.size(1), v.size(1)

        residual = q

        # Pass through the pre-attention projection: b x lq x (n*dv)
        # Separate different heads: b x lq x n x dv
        q = q.view(sz_b, len_q, n_head, d_k)
        k = k.view(sz_b, len_k, n_head, d_k)
        v = v.view(sz_b, len_v, n_head, d_v)

        # Transpose for attention dot product: b x n x lq x dv
        q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)

        if mask is not None:
            mask = mask.unsqueeze(1)   # For head axis broadcasting.

        q, attn = self.attention(q, k, v, mask=mask)

        # Transpose to move the head dimension back: b x lq x n x dv
        # Combine the last two dimensions to concatenate all the heads together: b x lq x (n*dv)
        q = q.transpose(1, 2).contiguous().view(sz_b, len_q, -1)
        q = self.dropout(self.fc(q))


        attn = torch.mean(attn, 1)
        return q, attn

class MultiHeadAttention_dotversion_index(nn.Module):
    ''' Multi-Head Attention module '''

    def __init__(self, n_head, d_model, d_k, d_v, dropout=0):
        super().__init__()

        self.n_head = n_head

        '''
        nn.init.constant_(self.w_qs.weight, 0)
        nn.init.constant_(self.w_ks.weight, 0)
        nn.init.constant_(self.w_vs.weight, 0)
        '''
        self.n_head = n_head
        self.d_k = d_k
        self.d_v = d_v

        self.attention = ScaledDotProductAttention(temperature=d_k ** 0.5)  #temperature=d_k ** 0.5

        self.dropout = nn.Dropout(dropout)
        self.layer_norm = nn.LayerNorm(d_model, eps=1e-6)

    def forward(self, q, k, v, mask=None):
        # print("w_qs:", self.w_qs.weight)
        d_k, d_v, n_head = self.d_k, self.d_v, self.n_head
        sz_b, len_q, len_k, len_v = q.size(0), q.size(1), k.size(1), v.size(1)


        # Pass through the pre-attention projection: b x lq x (n*dv)
        # Separate different heads: b x lq x n x dv
        q = q.view(sz_b, len_q, n_head, d_k)
        k = k.view(sz_b, len_k, n_head, d_k)
        v = v.view(sz_b, len_v, n_head, d_v)
        # print (self.w_qs.weight.abs().sum(), self.w_ks.weight.abs().sum(), self.w_vs.weight.abs().sum())
        # print ('kqv', q.abs().sum(), k.abs().sum(), v.abs().sum())

        # Transpose for attention dot product: b x n x lq x dv
        q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)

        if mask is not None:
            mask = mask.unsqueeze(1)  # For head axis broadcasting.

        q, attn2 = self.attention(q, k, v, mask=mask)


        # attn = attn1 + attn2
        # attn = attn1.unsqueeze(1)
        attn = attn2

        # Transpose to move the head dimension back: b x lq x n x dv
        # Combine the last two dimensions to concatenate all the heads together: b x lq x (n*dv)
        q = q.transpose(1, 2).contiguous().view(sz_b, len_q, -1)
        q = self.dropout(q)
        # q += residual

        q = self.layer_norm(q)

        attn = torch.mean(attn, 1)
        return q, attn
ocl/mha.py
amazon-science/object-centric-multiple-object-tracking
"""Models defined in code."""
from ocl.models.savi import SAVi
from ocl.models.savi_with_memory import SAVi_mem
__all__ = ["SAVi", "SAVi_mem"]

ocl/models/__init__.py
amazon-science/object-centric-multiple-object-tracking
from functools import partial
from typing import Any, Dict
import copy
import torch
from torch import nn

from ocl.path_defaults import VIDEO, BOX
from ocl.tree_utils import get_tree_element, reduce_tree


class SAVi(nn.Module):
    def __init__(
        self,
        conditioning: nn.Module,
        feature_extractor: nn.Module,
        perceptual_grouping: nn.Module,
        decoder: nn.Module,
        transition_model: nn.Module,
    ):
        super().__init__()
        self.conditioning = conditioning
        self.feature_extractor = feature_extractor
        self.perceptual_grouping = perceptual_grouping
        self.decoder = decoder
        self.transition_model = transition_model
        self.batched_input = None

    def forward(self, inputs: Dict[str, Any], phase = 'train'):
        # if self.batched_input is None:
        #     video = get_tree_element(inputs, VIDEO.split("."))
        #     # if video.shape[1] == 6:
        #     self.batched_input = copy.deepcopy(inputs)
        # else:
        #     print ('use catched')
        #     inputs = self.batched_input

        output = inputs
        video = get_tree_element(inputs, VIDEO.split("."))
        box = get_tree_element(inputs, BOX.split("."))
        batch_size = video.shape[0]

        features = self.feature_extractor(video=video)
        output["feature_extractor"] = features
        conditioning = self.conditioning(batch_size=batch_size)
        # conditioning = self.conditioning(batch_size=batch_size)
        output["initial_conditioning"] = conditioning

        # Loop over time.
        perceptual_grouping_outputs = []
        decoder_outputs = []
        transition_model_outputs = []
        trackers = []
        for frame_features in features:
            perceptual_grouping_output = self.perceptual_grouping(
                extracted_features=frame_features, conditioning=conditioning
            )
            slots = perceptual_grouping_output.objects
            decoder_output = self.decoder(object_features=slots)

            # remove background
            masks = decoder_output.masks_eval
            valid_idx = [0,1,2,4,5,6,7,8,9,10]
            masks_obj = masks[:, valid_idx]

            conditioning = self.transition_model(slots)
            # Store outputs.
            perceptual_grouping_outputs.append(slots)
            decoder_outputs.append(decoder_output)
            transition_model_outputs.append(conditioning)
            trackers.append(masks_obj)

        # Stack all recurrent outputs.
        stacking_fn = partial(torch.stack, dim=1)
        output["perceptual_grouping"] = reduce_tree(perceptual_grouping_outputs, stacking_fn)
        output["decoder"] = reduce_tree(decoder_outputs, stacking_fn)
        output["transition_model"] = reduce_tree(transition_model_outputs, stacking_fn)
        output["tracks"] = reduce_tree(trackers, stacking_fn)
        return output


ocl/models/savi.py
amazon-science/object-centric-multiple-object-tracking
from functools import partial
from typing import Any, Dict
import copy
import torch
from torch import nn

from ocl.path_defaults import VIDEO, BOX, MASK
from ocl.tree_utils import get_tree_element, reduce_tree
from ocl.metrics import masks_to_bboxes

class SAVi_mem(nn.Module):
    def __init__(
        self,
        conditioning: nn.Module,
        feature_extractor: nn.Module,
        perceptual_grouping: nn.Module,
        decoder: nn.Module,
        transition_model: nn.Module,
        memory: nn.Module,
        freeze = False,
    ):
        super().__init__()
        self.conditioning = conditioning
        self.feature_extractor = feature_extractor
        self.perceptual_grouping = perceptual_grouping
        self.decoder = decoder
        self.transition_model = transition_model
        self.memory = memory
        if freeze:
            self.conditioning.eval()
            self.feature_extractor.eval()
            self.perceptual_grouping.eval()
            self.decoder.eval()
            self.transition_model.eval()
            # freeze params
            for param in self.conditioning.parameters():
                param.requires_grad = False
            for param in self.feature_extractor.parameters():
                param.requires_grad = False
            for param in self.perceptual_grouping.parameters():
                param.requires_grad = False
            for param in self.decoder.parameters():
                param.requires_grad = False
            for param in self.transition_model.parameters():
                param.requires_grad = False
        self.batched_input = None

    def remove_bg_id(self, slot_masks):
        slot_masks = slot_masks > 0.7
        n, h, w = slot_masks.shape
        # remove background or none masks
        mask_sum = torch.sum(slot_masks.reshape(-1, h * w), dim=-1)
        empty_idx = (mask_sum <= 10).nonzero(as_tuple=True)[0].tolist()
        # bg_value = mask_sum[3]
        # bg_idx = (mask_sum == bg_value).nonzero(as_tuple=True)[0]
        # bg_idx = torch.cat([bg_idx, empty_idx], dim=0)
        fg_idxs = (mask_sum>20).nonzero(as_tuple=True)[0].tolist()
        bg_idx = 3
        fg_idx = 0
        for i in fg_idxs:
            if i not in empty_idx:
                if fg_idx != bg_idx:
                    fg_idx = i
                break
        empty_idx = list(empty_idx)
        if bg_idx in empty_idx:
            del empty_idx[empty_idx.index(3)]
        return empty_idx, fg_idx, bg_idx

    def remove_duplicated_slot_id(self, slot_masks):
        slot_masks = slot_masks > 0.7
        n, h, w = slot_masks.shape
        # remove background or none masks
        mask_sum = torch.sum(slot_masks.reshape(-1, h * w), dim=-1)
        empty_idx = (mask_sum <= 10).nonzero(as_tuple=True)[0]

        # remove duplicated masks
        mask = slot_masks.unsqueeze(1).to(torch.bool).reshape(n, 1, -1)
        mask_ = slot_masks.unsqueeze(0).to(torch.bool).reshape(1, n, -1)
        intersection = torch.sum(mask & mask_, dim=-1).to(torch.float64)
        union = torch.sum(mask | mask_, dim=-1).to(torch.float64)
        union1 = torch.sum(mask | mask, dim=-1).to(torch.float64)
        # union2 = torch.sum(mask_ | mask_, dim=-1).to(torch.float64)
        pairwise_iou = intersection / union
        self_iou1 = intersection / union1
        # self_iou2 = intersection / union2
        pairwise_iou[union == 0] = 1.0
        dup_idx = []
        for i in range(n):
            for j in range(i + 1, n):
                # if pairwise_iou[i, j] > 0.9 or self_iou1[i, j]>0.6 or self_iou2[i, j]>0.6:
                if pairwise_iou[i, j] > 0.9 or self_iou1[i, j] > 0.6:
                    dup_idx.append(i)

        invalid_idx = [*set(list(empty_idx) + list(dup_idx) +[3])]
        valid_idx = [3]
        for i in range(n):
            if i not in invalid_idx:
                valid_idx.append(i)
        return valid_idx

    def forward(self, inputs: Dict[str, Any], phase: str):

        output = inputs
        video = get_tree_element(inputs, VIDEO.split("."))
        # box = get_tree_element(inputs, BOX.split("."))

        batch_size = video.shape[0]
        features = self.feature_extractor(video=video)
        output["feature_extractor"] = features
        conditioning = self.conditioning(batch_size=batch_size)
        output["initial_conditioning"] = conditioning

        # Loop over time.
        perceptual_grouping_outputs = []
        decoder_outputs = []
        memory_outputs = []
        rollout_outputs = []
        object_features = []
        object_eval_masks= []
        objects = []
        attn_index = []
        tables = []
        slots_feats = []
        for (frame_id, frame_features) in enumerate(features):
            if frame_id == 0:
                query = conditioning
            else:
                query = conditioning

            perceptual_grouping_output = self.perceptual_grouping(
                extracted_features=frame_features, conditioning=query
            )
            slots = perceptual_grouping_output.objects
            # slots_ori = slots.clone()
            conditioning = self.transition_model(slots).clone()

            decoder_output = self.decoder(object_features=slots)
            cur_slot_masks = decoder_output.masks_eval
            bs, n = cur_slot_masks.shape[:2]
            for b in range(bs):
                empty_id, fg_id, bg_id = self.remove_bg_id(cur_slot_masks[b])
                if len(list(empty_id)) > 0:
                    tmp = slots[b,fg_id].unsqueeze(0).repeat(len(list(empty_id)),1)
                    slots[b, empty_id] = tmp

            decoder_output_new = self.decoder(object_features=slots)

            cur_slot_masks = decoder_output_new.masks_eval
            cur_slot_amodal_masks = decoder_output_new.masks

            if frame_id == 0:
                prev_slot_masks = cur_slot_masks.clone()
            memory_output = self.memory(observations=slots , prev_slot_masks = prev_slot_masks, amodal_masks = cur_slot_amodal_masks,
                cur_slot_masks = cur_slot_masks, conditions=slots, frame_id=frame_id, phase = phase)
            # prev_slot_masks = cur_slot_masks.clone()

            rollout_output = self.decoder(object_features= memory_output.rollout)
            prev_slot_masks = rollout_output.masks_eval
            object_feature_masks = self.decoder(object_features=memory_output.object_features)

            # remove background slots for evaluation
            mem_masks = object_feature_masks.masks_eval
            mem_masks_eval = torch.zeros(mem_masks.shape).to(mem_masks.device)
            num_buffer, w, h = mem_masks.shape[1:]
            for b in range(bs):
                masks = mem_masks[b] > 0.7
                boxes = masks_to_bboxes(masks)
                for i in range(num_buffer):
                    if not(boxes[i, 2] > 0.4 * w or boxes[i, 3] > 0.4*h):
                        mem_masks_eval[b, i] = mem_masks[b, i]


            # Store outputs.
            perceptual_grouping_outputs.append(perceptual_grouping_output)
            decoder_outputs.append(decoder_output_new) # decoder_output_new for memory training
            memory_outputs.append(memory_output.rollout)
            rollout_outputs.append(rollout_output)
            object_features.append(object_feature_masks)
            objects.append(memory_output.object_features)
            attn_index.append(memory_output.attn_index)
            slots_feats.append(slots)
            object_eval_masks.append(mem_masks_eval)
            tables.append(memory_output.table)

        # Stack all recurrent outputs.
        stacking_fn = partial(torch.stack, dim=1)
        output["perceptual_grouping"] = reduce_tree(perceptual_grouping_outputs, stacking_fn)
        output["decoder"] = reduce_tree(decoder_outputs, stacking_fn)
        output["memory"] = reduce_tree(memory_outputs, stacking_fn)
        output["rollout_decode"] = reduce_tree(rollout_outputs, stacking_fn)
        output["mem_masks"] = reduce_tree(object_features, stacking_fn)
        output['objects'] = reduce_tree(objects, stacking_fn)
        output["attn_index"] = reduce_tree(attn_index, stacking_fn)
        output["slots"] = reduce_tree(slots_feats, stacking_fn)
        output["tracks"] = reduce_tree(object_eval_masks, stacking_fn)
        output["table"] = reduce_tree(tables, stacking_fn)


        return output


ocl/models/savi_with_memory.py
amazon-science/object-centric-multiple-object-tracking
"""Neural network backbones."""
from typing import Callable, List, Optional, Union

import torch
from torch import nn

from ocl.utils import Residual


class ReLUSquared(nn.Module):
    def __init__(self, inplace=False):
        super().__init__()
        self.inplace = inplace

    def forward(self, x):
        return nn.functional.relu(x, inplace=self.inplace) ** 2


def get_activation_fn(name: str, inplace: bool = True, leaky_relu_slope: Optional[float] = None):
    if callable(name):
        return name

    name = name.lower()
    if name == "relu":
        return nn.ReLU(inplace=inplace)
    elif name == "relu_squared":
        return ReLUSquared(inplace=inplace)
    elif name == "leaky_relu":
        if leaky_relu_slope is None:
            raise ValueError("Slope of leaky ReLU was not defined")
        return nn.LeakyReLU(leaky_relu_slope, inplace=inplace)
    elif name == "tanh":
        return nn.Tanh()
    elif name == "sigmoid":
        return nn.Sigmoid()
    elif name == "identity":
        return nn.Identity()
    else:
        raise ValueError(f"Unknown activation function {name}")


def build_mlp(
    input_dim: int,
    output_dim: int,
    features: List[int],
    activation_fn: Union[str, Callable] = "relu",
    final_activation_fn: Optional[Union[str, Callable]] = None,
    initial_layer_norm: bool = False,
    residual: bool = False,
) -> nn.Sequential:
    layers = []
    current_dim = input_dim
    if initial_layer_norm:
        layers.append(nn.LayerNorm(current_dim))

    for n_features in features:
        layers.append(nn.Linear(current_dim, n_features))
        nn.init.zeros_(layers[-1].bias)
        layers.append(get_activation_fn(activation_fn))
        current_dim = n_features

    layers.append(nn.Linear(current_dim, output_dim))
    nn.init.zeros_(layers[-1].bias)
    if final_activation_fn is not None:
        layers.append(get_activation_fn(final_activation_fn))

    if residual:
        return Residual(nn.Sequential(*layers))
    return nn.Sequential(*layers)


def build_two_layer_mlp(
    input_dim, output_dim, hidden_dim, initial_layer_norm: bool = False, residual: bool = False
):
    """Build a two layer MLP, with optional initial layer norm.

    Separate class as this type of construction is used very often for slot attention and
    transformers.
    """
    return build_mlp(
        input_dim, output_dim, [hidden_dim], initial_layer_norm=initial_layer_norm, residual=residual
    )


def build_transformer_encoder(
    input_dim: int,
    output_dim: int,
    n_layers: int,
    n_heads: int,
    hidden_dim: Optional[int] = None,
    dropout: float = 0.0,
    activation_fn: Union[str, Callable] = "relu",
    layer_norm_eps: float = 1e-5,
    use_output_transform: bool = True,
):
    if hidden_dim is None:
        hidden_dim = 4 * input_dim

    layers = []
    for _ in range(n_layers):
        layers.append(
            nn.TransformerEncoderLayer(
                d_model=input_dim,
                nhead=n_heads,
                dim_feedforward=hidden_dim,
                dropout=dropout,
                activation=activation_fn,
                layer_norm_eps=layer_norm_eps,
                batch_first=True,
                norm_first=True,
            )
        )

    if use_output_transform:
        layers.append(nn.LayerNorm(input_dim, eps=layer_norm_eps))
        output_transform = nn.Linear(input_dim, output_dim, bias=True)
        nn.init.xavier_uniform_(output_transform.weight)
        nn.init.zeros_(output_transform.bias)
        layers.append(output_transform)

    return nn.Sequential(*layers)


def build_transformer_decoder(
    input_dim: int,
    output_dim: int,
    n_layers: int,
    n_heads: int,
    hidden_dim: Optional[int] = None,
    dropout: float = 0.0,
    activation_fn: Union[str, Callable] = "relu",
    layer_norm_eps: float = 1e-5,
    return_attention_weights: bool = False,
    attention_weight_type: Union[int, str] = -1,
):
    if hidden_dim is None:
        hidden_dim = 4 * input_dim

    decoder_layer = nn.TransformerDecoderLayer(
        d_model=input_dim,
        nhead=n_heads,
        dim_feedforward=hidden_dim,
        dropout=dropout,
        activation=activation_fn,
        layer_norm_eps=layer_norm_eps,
        batch_first=True,
        norm_first=True,
    )

    if return_attention_weights:
        return TransformerDecoder(
            decoder_layer,
            n_layers,
            return_attention_weights=True,
            attention_weight_type=attention_weight_type,
        )
    else:
        return nn.TransformerDecoder(decoder_layer, n_layers)


class TransformerDecoder(nn.TransformerDecoder):
    """Modified nn.TransformerDecoder class that returns attention weights over memory."""

    def __init__(
        self,
        decoder_layer,
        num_layers,
        norm=None,
        return_attention_weights=False,
        attention_weight_type: Union[int, str] = "mean",
    ):
        super(TransformerDecoder, self).__init__(decoder_layer, num_layers, norm)

        if return_attention_weights:
            self.attention_hooks = []
            for layer in self.layers:
                self.attention_hooks.append(self._prepare_layer(layer))
        else:
            self.attention_hooks = None

        if isinstance(attention_weight_type, int):
            if attention_weight_type >= num_layers or attention_weight_type < -num_layers:
                raise ValueError(
                    f"Index {attention_weight_type} exceeds number of layers {num_layers}"
                )
        elif attention_weight_type != "mean":
            raise ValueError("`weights` needs to be a number or 'mean'.")
        self.weights = attention_weight_type

    def _prepare_layer(self, layer):
        assert isinstance(layer, nn.TransformerDecoderLayer)

        def _mha_block(self, x, mem, attn_mask, key_padding_mask):
            x = self.multihead_attn(
                x,
                mem,
                mem,
                attn_mask=attn_mask,
                key_padding_mask=key_padding_mask,
                need_weights=True,
            )[0]
            return self.dropout2(x)

        # Patch _mha_block method to compute attention weights
        layer._mha_block = _mha_block.__get__(layer, nn.TransformerDecoderLayer)

        class AttentionHook:
            def __init__(self):
                self._attention = None

            def pop(self) -> torch.Tensor:
                assert self._attention is not None, "Forward was not called yet!"
                attention = self._attention
                self._attention = None
                return attention

            def __call__(self, module, inp, outp):
                self._attention = outp[1]

        hook = AttentionHook()
        layer.multihead_attn.register_forward_hook(hook)
        return hook

    def forward(
        self,
        tgt: torch.Tensor,
        memory: torch.Tensor,
        tgt_mask: Optional[torch.Tensor] = None,
        memory_mask: Optional[torch.Tensor] = None,
        tgt_key_padding_mask: Optional[torch.Tensor] = None,
        memory_key_padding_mask: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        output = tgt

        for mod in self.layers:
            output = mod(
                output,
                memory,
                tgt_mask=tgt_mask,
                memory_mask=memory_mask,
                tgt_key_padding_mask=tgt_key_padding_mask,
                memory_key_padding_mask=memory_key_padding_mask,
            )

        if self.norm is not None:
            output = self.norm(output)

        if self.attention_hooks is not None:
            attentions = []
            for hook in self.attention_hooks:
                attentions.append(hook.pop())

            if self.weights == "mean":
                attentions = torch.stack(attentions, dim=-1)
                # Take mean over all layers
                attention = attentions.mean(dim=-1)
            else:
                attention = attentions[self.weights]

            return output, attention.transpose(1, 2)
        else:
            return output

ocl/neural_networks.py
amazon-science/object-centric-multiple-object-tracking
"""Default paths for different types of inputs.

These are only defined for convenience and can also be overwritten using the appropriate *_path
constructor variables of RoutableMixin subclasses.
"""
INPUT = "input"
VIDEO = f"{INPUT}.image"
BOX = f"{INPUT}.instance_bbox"
MASK = f"{INPUT}.mask"
ID = f"{INPUT}.instance_id"
BATCH_SIZE = f"{INPUT}.batch_size"
GLOBAL_STEP = "global_step"
FEATURES = "feature_extractor"
CONDITIONING = "conditioning"
# TODO(hornmax): Currently decoders are nested in the task and accept PerceptualGroupingOutput as
# input. In the future this will change and decoders should just be regular parts of the model.
OBJECTS = "perceptual_grouping.objects"
FEATURE_ATTRIBUTIONS = "perceptual_grouping.feature_attributions"
OBJECT_DECODER = "object_decoder"

ocl/path_defaults.py
amazon-science/object-centric-multiple-object-tracking
"""Implementations of perceptual grouping algorithms."""
import math
from typing import Dict, Optional

import torch
from torch import nn

from ocl import base, path_defaults
from ocl.utils import RoutableMixin
from ocl.predictor import TransformerBlock

class SlotAttention(nn.Module):
    """Implementation of SlotAttention.

    Based on the slot attention implementation of Phil Wang available at:
    https://github.com/lucidrains/slot-attention
    """

    def __init__(
        self,
        dim: int,
        feature_dim: int,
        kvq_dim: Optional[int] = None,
        n_heads: int = 1,
        iters: int = 3,
        eps: float = 1e-8,
        ff_mlp: Optional[nn.Module] = None,
        use_projection_bias: bool = False,
        use_implicit_differentiation: bool = False,
    ):
        super().__init__()
        self.dim = dim
        self.n_heads = n_heads
        self.iters = iters
        self.eps = eps
        self.use_implicit_differentiation = use_implicit_differentiation

        if kvq_dim is None:
            self.kvq_dim = dim
        else:
            self.kvq_dim = kvq_dim

        if self.kvq_dim % self.n_heads != 0:
            raise ValueError("Key, value, query dimensions must be divisible by number of heads.")
        self.dims_per_head = self.kvq_dim // self.n_heads
        self.scale = self.dims_per_head**-0.5

        self.to_q = nn.Linear(dim, self.kvq_dim, bias=use_projection_bias)
        self.to_k = nn.Linear(feature_dim, self.kvq_dim, bias=use_projection_bias)
        self.to_v = nn.Linear(feature_dim, self.kvq_dim, bias=use_projection_bias)

        self.gru = nn.GRUCell(self.kvq_dim, dim)

        self.norm_input = nn.LayerNorm(feature_dim)
        self.norm_slots = nn.LayerNorm(dim)
        self.ff_mlp = ff_mlp

    def step(self, slots, k, v, masks=None):
        bs, n_slots, _ = slots.shape
        slots_prev = slots

        slots = self.norm_slots(slots)
        q = self.to_q(slots).view(bs, n_slots, self.n_heads, self.dims_per_head)

        dots = torch.einsum("bihd,bjhd->bihj", q, k) * self.scale
        if masks is not None:
            # Masked slots should not take part in the competition for features. By replacing their
            # dot-products with -inf, their attention values will become zero within the softmax.
            dots.masked_fill_(masks.to(torch.bool).view(bs, n_slots, 1, 1), float("-inf"))

        attn = dots.flatten(1, 2).softmax(dim=1)  # Take softmax over slots and heads
        attn = attn.view(bs, n_slots, self.n_heads, -1)
        attn_before_reweighting = attn
        attn = attn + self.eps
        attn = attn / attn.sum(dim=-1, keepdim=True)

        updates = torch.einsum("bjhd,bihj->bihd", v, attn)

        slots = self.gru(updates.reshape(-1, self.kvq_dim), slots_prev.reshape(-1, self.dim))

        slots = slots.reshape(bs, -1, self.dim)

        if self.ff_mlp:
            slots = self.ff_mlp(slots)

        return slots, attn_before_reweighting.mean(dim=2)

    def iterate(self, slots, k, v, masks=None):
        for _ in range(self.iters):
            slots, attn = self.step(slots, k, v, masks)
        return slots, attn

    def forward(
        self, inputs: torch.Tensor, conditioning: torch.Tensor, masks: Optional[torch.Tensor] = None
    ):
        b, n, d = inputs.shape
        slots = conditioning

        inputs = self.norm_input(inputs)

        k = self.to_k(inputs).view(b, n, self.n_heads, self.dims_per_head)
        v = self.to_v(inputs).view(b, n, self.n_heads, self.dims_per_head)

        if self.use_implicit_differentiation:
            slots, attn = self.iterate(slots, k, v, masks)
            slots, attn = self.step(slots.detach(), k, v, masks)
        else:
            slots, attn = self.iterate(slots, k, v, masks)

        return slots, attn, v


class SlotAttentionGrouping(base.PerceptualGrouping, RoutableMixin):
    """Implementation of SlotAttention for perceptual grouping.

    Args:
        feature_dim: Dimensionality of features to slot attention (after positional encoding).
        object_dim: Dimensionality of slots.
        kvq_dim: Dimensionality after projecting to keys, values, and queries. If `None`,
            `object_dim` is used.
        n_heads: Number of heads slot attention uses.
        iters: Number of slot attention iterations.
        eps: Epsilon in slot attention.
        ff_mlp: Optional module applied slot-wise after GRU update.
        positional_embedding: Optional module applied to the features before slot attention, adding
            positional encoding.
        use_projection_bias: Whether to use biases in key, value, query projections.
        use_implicit_differentiation: Whether to use implicit differentiation trick. If true,
            performs one more iteration of slot attention that is used for the gradient step after
            `iters` iterations of slot attention without gradients. Faster and more memory efficient
            than the standard version, but can not backpropagate gradients to the conditioning input.
        input_dim: Dimensionality of features before positional encoding is applied. Specifying this
            is optional but can be convenient to structure configurations.
    """

    def __init__(
        self,
        feature_dim: int,
        object_dim: int,
        kvq_dim: Optional[int] = None,
        n_heads: int = 1,
        iters: int = 3,
        eps: float = 1e-8,
        ff_mlp: Optional[nn.Module] = None,
        positional_embedding: Optional[nn.Module] = None,
        use_projection_bias: bool = False,
        use_implicit_differentiation: bool = False,
        use_empty_slot_for_masked_slots: bool = False,
        input_dim: Optional[int] = None,
        feature_path: Optional[str] = path_defaults.FEATURES,
        conditioning_path: Optional[str] = path_defaults.CONDITIONING,
        slot_mask_path: Optional[str] = None,
    ):
        base.PerceptualGrouping.__init__(self)
        RoutableMixin.__init__(
            self,
            {
                "extracted_features": feature_path,
                "conditioning": conditioning_path,
                "slot_masks": slot_mask_path,
            },
        )

        self._object_dim = object_dim
        self.slot_attention = SlotAttention(
            dim=object_dim,
            feature_dim=feature_dim,
            kvq_dim=kvq_dim,
            n_heads=n_heads,
            iters=iters,
            eps=eps,
            ff_mlp=ff_mlp,
            use_projection_bias=use_projection_bias,
            use_implicit_differentiation=use_implicit_differentiation,
        )

        self.positional_embedding = positional_embedding

        if use_empty_slot_for_masked_slots:
            if slot_mask_path is None:
                raise ValueError("Need `slot_mask_path` for `use_empty_slot_for_masked_slots`")
            self.empty_slot = nn.Parameter(torch.randn(object_dim) * object_dim**-0.5)
        else:
            self.empty_slot = None

    @property
    def object_dim(self):
        return self._object_dim

    @RoutableMixin.route
    def forward(
        self,
        extracted_features: base.FeatureExtractorOutput,
        conditioning: base.ConditioningOutput,
        slot_masks: Optional[torch.Tensor] = None,
    ):
        if self.positional_embedding:
            features = self.positional_embedding(
                extracted_features.features, extracted_features.positions
            )
        else:
            features = extracted_features.features
        slots, attn, v = self.slot_attention(features, conditioning, slot_masks)

        if slot_masks is not None and self.empty_slot is not None:
            slots[slot_masks] = self.empty_slot.to(dtype=slots.dtype)

        return base.PerceptualGroupingOutput(slots, feature_attributions=attn, is_empty=slot_masks, value=v)


class StickBreakingGrouping(base.Conditioning, RoutableMixin):
    """Perceptual grouping based on a stick-breaking process.

    The idea is to pick a random feature from a yet unexplained part of the feature map, then see
    which parts of the feature map are "explained" by this feature using a kernel distance. This
    process is iterated until some termination criterion is reached. In principle, this process
    allows to extract a variable number of slots per image.

    This is based on Engelcke et al, GENESIS-V2: Inferring Unordered Object Representations without
    Iterative Refinement, http://arxiv.org/abs/2104.09958. Our implementation here differs a bit from
    the one described there:

    - It only implements one kernel distance, the Gaussian kernel
    - It does not take features positions into account when computing the kernel distances
    - It L2-normalises the input features to get comparable scales of the kernel distance
    - It has multiple termination criteria, namely termination based on fraction explained, mean
      mask value, and min-max mask value. GENESIS-V2 implements termination based on mean mask
      value, but does not mention it in the paper. Note that by default, all termination criteria
      are disabled.

    Args:
        object_dim: Dimensionality of extracted slots.
        feature_dim: Dimensionality of features to operate on.
        n_slots: Maximum number of slots.
        kernel_var: Variance in Gaussian kernel.
        learn_kernel_var: Whether kernel variance should be included as trainable parameter.
        max_unexplained: If fraction of unexplained features drops under this value, drop the slot.
        min_slot_mask: If slot mask has lower average value than this value, drop the slot.
        min_max_mask_value: If slot mask's maximum value is lower than this value, drop the slot.
        early_termination: If true, all slots after the first dropped slot are also dropped.
        add_unexplained: If true, add a slot that covers all unexplained parts at the point when the
            first slot was dropped.
        eps: Minimum value for masks.
        detach_features: If true, detach input features such that no gradient flows through this
            operation.
    """

    def __init__(
        self,
        object_dim: int,
        feature_dim: int,
        n_slots: int,
        kernel_var: float = 1.0,
        learn_kernel_var: bool = False,
        max_unexplained: float = 0.0,
        min_slot_mask: float = 0.0,
        min_max_mask_value: float = 0.0,
        early_termination: bool = False,
        add_unexplained: bool = False,
        eps: float = 1e-8,
        detach_features: bool = False,
        use_input_layernorm: bool = False,
        features_path: Optional[str] = path_defaults.FEATURES,
    ):
        base.Conditioning.__init__(self)
        RoutableMixin.__init__(self, {"features": features_path})
        self.n_slots = n_slots
        self.object_dim = object_dim

        assert kernel_var > 0.0
        if learn_kernel_var:
            self.kernel_logvar = nn.Parameter(torch.tensor(math.log(kernel_var)))
        else:
            self.register_buffer("kernel_logvar", torch.tensor(math.log(kernel_var)))

        assert 0.0 <= max_unexplained < 1.0
        self.max_unexplained = max_unexplained
        assert 0.0 <= min_slot_mask < 1.0
        self.min_slot_mask = min_slot_mask
        assert 0.0 <= min_max_mask_value < 1.0
        self.min_max_mask_value = min_max_mask_value

        self.early_termination = early_termination
        self.add_unexplained = add_unexplained
        if add_unexplained and not early_termination:
            raise ValueError("`add_unexplained=True` only works with `early_termination=True`")

        self.eps = eps
        self.log_eps = math.log(eps)
        self.detach_features = detach_features

        if use_input_layernorm:
            self.in_proj = nn.Sequential(
                nn.LayerNorm(feature_dim), nn.Linear(feature_dim, feature_dim)
            )
            torch.nn.init.xavier_uniform_(self.in_proj[-1].weight)
            torch.nn.init.zeros_(self.in_proj[-1].bias)
        else:
            self.in_proj = nn.Linear(feature_dim, feature_dim)
            torch.nn.init.xavier_uniform_(self.in_proj.weight)
            torch.nn.init.zeros_(self.in_proj.bias)

        self.out_proj = nn.Linear(feature_dim, object_dim)
        torch.nn.init.xavier_uniform_(self.out_proj.weight)
        torch.nn.init.zeros_(self.out_proj.bias)

    @RoutableMixin.route
    def forward(self, features: base.FeatureExtractorOutput) -> Dict[str, torch.Tensor]:
        features = features.features
        bs, n_features, feature_dim = features.shape
        if self.detach_features:
            features = features.detach()

        proj_features = torch.nn.functional.normalize(self.in_proj(features), dim=-1)

        # The scope keep tracks of the unexplained parts of the feature map
        log_scope = torch.zeros_like(features[:, :, 0])
        # Seeds are used for random sampling of features
        log_seeds = torch.rand_like(log_scope).clamp_min(self.eps).log()

        slot_masks = []
        log_scopes = []

        # Always iterate for `n_iters` steps for batching reasons. Termination is modeled afterwards.
        n_iters = self.n_slots - 1 if self.add_unexplained else self.n_slots
        for _ in range(n_iters):
            log_scopes.append(log_scope)

            # Sample random features from unexplained parts of the feature map
            rand_idxs = torch.argmax(log_scope + log_seeds, dim=1)
            cur_centers = proj_features.gather(
                1, rand_idxs.view(bs, 1, 1).expand(-1, -1, feature_dim)
            )

            # Compute similarity between selected features and other features. alpha can be
            # considered an attention mask.
            dists = torch.sum((cur_centers - proj_features) ** 2, dim=-1)
            log_alpha = (-dists / self.kernel_logvar.exp()).clamp_min(self.log_eps)

            # To get the slot mask, we subtract already explained parts from alpha using the scope
            mask = (log_scope + log_alpha).exp()
            slot_masks.append(mask)

            # Update scope by masking out parts explained by the current iteration
            log_1m_alpha = (1 - log_alpha.exp()).clamp_min(self.eps).log()
            log_scope = log_scope + log_1m_alpha

        if self.add_unexplained:
            slot_masks.append(log_scope.exp())
            log_scopes.append(log_scope)

        slot_masks = torch.stack(slot_masks, dim=1)
        scopes = torch.stack(log_scopes, dim=1).exp()

        # Compute criteria for ignoring slots
        empty_slots = torch.zeros_like(slot_masks[:, :, 0], dtype=torch.bool)
        # When fraction of unexplained features drops under threshold, ignore slot,
        empty_slots |= scopes.mean(dim=-1) < self.max_unexplained
        # or when slot's mean mask is under threshold, ignore slot,
        empty_slots |= slot_masks.mean(dim=-1) < self.min_slot_mask
        # or when slot's masks maximum value is under threshold, ignore slot.
        empty_slots |= slot_masks.max(dim=-1).values < self.min_max_mask_value

        if self.early_termination:
            # Simulate early termination by marking all slots after the first empty slot as empty
            empty_slots = torch.cummax(empty_slots, dim=1).values
            if self.add_unexplained:
                # After termination, add one more slot using the unexplained parts at that point
                first_empty = torch.argmax(empty_slots.to(torch.int32), dim=1).unsqueeze(-1)
                empty_slots.scatter_(1, first_empty, torch.zeros_like(first_empty, dtype=torch.bool))

                idxs = first_empty.view(bs, 1, 1).expand(-1, -1, n_features)
                unexplained = scopes.gather(1, idxs)
                slot_masks.scatter_(1, idxs, unexplained)

        # Create slot representations as weighted average of feature map
        slots = torch.einsum("bkp,bpd->bkd", slot_masks, features)
        slots = slots / slot_masks.sum(dim=-1, keepdim=True).clamp_min(self.eps)
        slots = self.out_proj(slots)

        # Zero-out masked slots
        slots.masked_fill_(empty_slots.view(bs, slots.shape[1], 1), 0.0)

        return base.PerceptualGroupingOutput(
            slots, feature_attributions=slot_masks, is_empty=empty_slots
        )



ocl/perceptual_grouping.py
amazon-science/object-centric-multiple-object-tracking
import functools
import random
from typing import Any, Callable, Dict, Iterable, List, Optional, Sequence

import numpy as np
import torch
import webdataset

from ocl import hooks


class Plugin:
    """A plugin which defines a set of hooks to be called by the code."""


class Optimization(Plugin):
    """Optimize (a subset of) the parameters using a optimizer and a LR scheduler."""

    def __init__(
        self, optimizer, lr_scheduler=None, parameter_groups: Optional[List[Dict[str, Any]]] = None
    ):
        self.optimizer = optimizer
        self.lr_scheduler = lr_scheduler
        self.parameter_group_specs = parameter_groups
        if self.parameter_group_specs:
            for idx, param_group_spec in enumerate(self.parameter_group_specs):
                if "params" not in param_group_spec:
                    raise ValueError(f'Parameter group {idx + 1} does not contain key "params"')
                param_spec = param_group_spec["params"]
                if isinstance(param_spec, str):
                    param_group_spec["params"] = [param_spec]
                elif isinstance(param_spec, Iterable):
                    param_group_spec["params"] = list(param_spec)
                else:
                    raise ValueError(
                        f'"params" for parameter group {idx + 1} is not of type str or iterable'
                    )

                if "predicate" in param_group_spec:
                    if not callable(param_group_spec["predicate"]):
                        raise ValueError(
                            f'"predicate" for parameter group {idx + 1} is not a callable'
                        )

    def _get_parameter_groups(self, model):
        """Build parameter groups from specification."""
        parameter_groups = []
        for param_group_spec in self.parameter_group_specs:
            param_spec = param_group_spec["params"]
            # Default predicate includes all parameters
            predicate = param_group_spec.get("predicate", lambda name, param: True)

            parameters = []
            for parameter_path in param_spec:
                root = model
                for child in parameter_path.split("."):
                    root = getattr(root, child)
                parameters.extend(
                    param for name, param in root.named_parameters() if predicate(name, param)
                )

            param_group = {
                k: v for k, v in param_group_spec.items() if k not in ("params", "predicate")
            }
            param_group["params"] = parameters
            parameter_groups.append(param_group)

        return parameter_groups

    @hooks.hook_implementation
    def configure_optimizers(self, model):
        if self.parameter_group_specs:
            params_or_param_groups = self._get_parameter_groups(model)
        else:
            params_or_param_groups = model.parameters()

        optimizer = self.optimizer(params_or_param_groups)
        output = {"optimizer": optimizer}
        if self.lr_scheduler:
            output.update(self.lr_scheduler(optimizer))
        return output


def transform_with_duplicate(elements: dict, *, transform, element_key: str, duplicate_key: str):
    """Utility function to fix issues with pickling."""
    element = transform(elements[element_key])
    elements[element_key] = element
    elements[duplicate_key] = element
    return elements


class SingleElementPreprocessing(Plugin):
    """Preprocessing of a single element in the input data.

    This is useful to build preprocessing pipelines based on existing element transformations such
    as those provided by torchvision. The element can optionally be duplicated and stored under a
    different key after the transformation by specifying `duplicate_key`. This is useful to further
    preprocess this element in different ways afterwards.
    """

    def __init__(
        self,
        training_transform: Callable,
        evaluation_transform: Callable,
        element_key: str = "image",
        duplicate_key: Optional[str] = None,
    ):
        self._training_transform = training_transform
        self._evaluation_transform = evaluation_transform
        self.element_key = element_key
        self.duplicate_key = duplicate_key

    @hooks.hook_implementation
    def training_fields(self):
        return (self.element_key,)

    @hooks.hook_implementation
    def training_transform(self):
        if self._training_transform:

            if self.duplicate_key is None:

                def transform(pipeline: webdataset.Processor):
                    return pipeline.map_dict(**{self.element_key: self._training_transform})

            else:

                def transform(pipeline: webdataset.Processor):
                    transform_func = functools.partial(
                        transform_with_duplicate,
                        transform=self._training_transform,
                        element_key=self.element_key,
                        duplicate_key=self.duplicate_key,
                    )
                    return pipeline.map(transform_func)

            return transform
        else:
            return None

    @hooks.hook_implementation
    def evaluation_fields(self):
        return (self.element_key,)

    @hooks.hook_implementation
    def evaluation_transform(self):
        if self._evaluation_transform:

            if self.duplicate_key is None:

                def transform(pipeline: webdataset.Processor):
                    return pipeline.map_dict(**{self.element_key: self._evaluation_transform})

            else:

                def transform(pipeline: webdataset.Processor):
                    transform_func = functools.partial(
                        transform_with_duplicate,
                        transform=self._evaluation_transform,
                        element_key=self.element_key,
                        duplicate_key=self.duplicate_key,
                    )

                    return pipeline.map(transform_func)

            return transform
        else:
            return None


class MultiElementPreprocessing(Plugin):
    """Preprocessing of multiple elements in the input data.

    This is useful preprocessing pipelines based on existing element transformations such as those
    provided by torchvision.
    """

    def __init__(
        self,
        training_transforms: Optional[Dict[str, Any]] = None,
        evaluation_transforms: Optional[Dict[str, Any]] = None,
    ):
        if training_transforms is None:
            training_transforms = {}
        self.training_keys = tuple(training_transforms)
        self._training_transforms = {
            key: transf for key, transf in training_transforms.items() if transf is not None
        }

        if evaluation_transforms is None:
            evaluation_transforms = {}
        self.evaluation_keys = tuple(evaluation_transforms)
        self._evaluation_transforms = {
            key: transf for key, transf in evaluation_transforms.items() if transf is not None
        }

    @hooks.hook_implementation
    def training_fields(self):
        return self.training_keys

    @hooks.hook_implementation
    def training_transform(self):
        if self._training_transforms:

            def transform(pipeline: webdataset.Processor):
                return pipeline.map_dict(**self._training_transforms)

            return transform
        else:
            return None

    @hooks.hook_implementation
    def evaluation_fields(self):
        return self.evaluation_keys

    @hooks.hook_implementation
    def evaluation_transform(self):
        if self._evaluation_transforms:

            def transform(pipeline: webdataset.Processor):
                return pipeline.map_dict(**self._evaluation_transforms)

            return transform
        else:
            return None


class DataPreprocessing(Plugin):
    """Arbitrary preprocessing of input data.

    The transform takes in a dictionary of elements and should return a dictionary of elements.
    Plugin must specify the elements that should be included in the dictionary using
    `training_fields` and `evaluation_fields` arguments.
    """

    def __init__(
        self,
        training_transform: Optional[Callable] = None,
        evaluation_transform: Optional[Callable] = None,
        training_fields: Optional[Sequence[str]] = None,
        evaluation_fields: Optional[Sequence[str]] = None,
    ):
        if training_transform is not None and training_fields is None:
            raise ValueError(
                "If passing `training_transform`, `training_fields` must also be specified."
            )
        if evaluation_transform is not None and evaluation_fields is None:
            raise ValueError(
                "If passing `evaluation_transform`, `evaluation_fields` must also be specified."
            )

        self._training_transform = training_transform
        self._evaluation_transform = evaluation_transform
        self._training_fields = tuple(training_fields) if training_fields else tuple()
        self._evaluation_fields = tuple(evaluation_fields) if evaluation_fields else tuple()

    @hooks.hook_implementation
    def training_fields(self):
        return self._training_fields

    @hooks.hook_implementation
    def training_transform(self):
        if self._training_transform:

            def transform(pipeline: webdataset.Processor):
                return pipeline.map(self._training_transform)

            return transform
        else:
            return None

    @hooks.hook_implementation
    def evaluation_fields(self):
        return self._evaluation_fields

    @hooks.hook_implementation
    def evaluation_transform(self):
        if self._evaluation_transform:

            def transform(pipeline: webdataset.Processor):
                return pipeline.map(self._evaluation_transform)

            return transform
        else:
            return None


class SubsetDataset(Plugin):
    """Create a subset of a dataset by discarding samples."""

    def __init__(self, predicate, fields: Sequence[str]):
        """Plugin to create a subset of a dataset by discarding samples.

        Args:
            predicate: Function which determines if elements should be kept (return value is True)
                or discarded (return value is False). The function is only provided with the fields
                specified in the `fields` parameter.
            fields (Sequence[str]): The fields from the input which should be passed on to the
                predicate for evaluation.
        """
        self.predicate = predicate
        self.fields = tuple(fields)

    def _get_transform_function(self):
        def wrapped_predicate(d: dict):
            return self.predicate(*(d[field] for field in self.fields))

        def select(pipeline: webdataset.Processor):
            return pipeline.select(wrapped_predicate)

        return select

    @hooks.hook_implementation
    def training_fields(self):
        return self.fields

    @hooks.hook_implementation(tryfirst=True)
    def training_transform(self):
        return self._get_transform_function()

    @hooks.hook_implementation
    def evaluation_fields(self):
        return self.fields

    @hooks.hook_implementation(tryfirst=True)
    def evaluation_transform(self):
        return self._get_transform_function()


class SampleFramesFromVideo(Plugin):
    def __init__(
        self,
        n_frames_per_video: int,
        training_fields: Sequence[str],
        evaluation_fields: Sequence[str],
        dim: int = 0,
        seed: int = 39480234,
        shuffle_buffer_size: int = 1000,
    ):
        """Sample frames from input tensors.

        Args:
            n_frames_per_video: Number of frames per video to sample. -1 indicates that all frames
                should be sampled.
            training_fields: The fields that should be considered video data and thus sliced
                according to the frame sampling during training.
            evaluation_fields: The fields that should be considered video data and thus sliced
                according to the frame sampling during evaluation.
            dim: The dimension along which to slice the tensors.
            seed: Random number generator seed to deterministic sampling during evaluation.
            shuffle_buffer_size: Size of shuffle buffer used during training. An additional
                shuffling step ensures each batch contains a diverse set of images and not only
                images from the same video.
        """
        self.n_frames_per_video = n_frames_per_video
        self._training_fields = tuple(training_fields)
        self._evaluation_fields = tuple(evaluation_fields)
        self.dim = dim
        self.seed = seed
        self.shuffle_buffer_size = shuffle_buffer_size

    def slice_data(self, data, index: int):
        """Small utility method to slice a numpy array along a specified axis."""
        n_dims_before = self.dim
        n_dims_after = data.ndim - 1 - self.dim
        slices = (slice(None),) * n_dims_before + (index,) + (slice(None),) * n_dims_after
        return data[slices]

    def sample_frames_using_key(self, data, fields, seed):
        """Sample frames deterministically from generator of videos using the __key__ field."""
        for sample in data:
            # Initialize random number generator dependent on instance key. This should make the
            # sampling process deterministic, which is useful when sampling frames for the
            # validation/test data.
            key = sample["__key__"]
            # TODO (hornmax): We assume all fields to have the same size. I do not want to check
            # this here as it seems a bit verbose.
            n_frames = sample[fields[0]].shape[self.dim]
            rand = random.Random(int(key) + seed)
            frames_per_video = self.n_frames_per_video if self.n_frames_per_video != -1 else n_frames
            selected_frames = rand.sample(range(n_frames), k=frames_per_video)

            for frame in selected_frames:
                # Slice the fields according to the frame.
                sliced_fields = {field: self.slice_data(sample[field], frame) for field in fields}
                # Leave all fields besides the sliced ones as before, augment the __key__ field to
                # include the frame number.
                sliced_fields["__key__"] = f"{key}_{frame}"
                yield {**sample, **sliced_fields}

    @hooks.hook_implementation
    def training_fields(self):
        return self._training_fields

    @hooks.hook_implementation
    def training_transform(self):
        def apply_deterministic_sampling(pipeline: webdataset.Processor):
            if len(self._training_fields) > 0:
                return pipeline.then(
                    functools.partial(
                        self.sample_frames_using_key, fields=self._training_fields, seed=self.seed
                    )
                ).shuffle(self.shuffle_buffer_size)
            else:
                return pipeline

        return apply_deterministic_sampling

    @hooks.hook_implementation
    def evaluation_fields(self):
        return self._evaluation_fields

    @hooks.hook_implementation
    def evaluation_transform(self):
        def apply_deterministic_sampling(pipeline: webdataset.Processor):
            if len(self._evaluation_fields) > 0:
                return pipeline.then(
                    functools.partial(
                        self.sample_frames_using_key,
                        fields=self._evaluation_fields,
                        seed=self.seed + 1,
                    )
                )
            else:
                return pipeline

        return apply_deterministic_sampling


class SplitConsecutiveFrames(Plugin):
    def __init__(
        self,
        n_consecutive_frames: int,
        training_fields: Sequence[str],
        evaluation_fields: Sequence[str],
        dim: int = 0,
        shuffle_buffer_size: int = 1000,
        drop_last: bool = True,
    ):
        self.n_consecutive_frames = n_consecutive_frames
        self._training_fields = tuple(training_fields)
        self._evaluation_fields = tuple(evaluation_fields)
        self.dim = dim
        self.shuffle_buffer_size = shuffle_buffer_size
        self.drop_last = drop_last

    @hooks.hook_implementation
    def training_fields(self):
        return self._training_fields

    def split_to_consecutive_frames(self, data, fields):
        """Sample frames deterministically from generator of videos using the __key__ field."""
        for sample in data:
            key = sample["__key__"]
            n_frames = sample[fields[0]].shape[self.dim]

            splitted_fields = [
                np.array_split(
                    sample[field],
                    range(self.n_consecutive_frames, n_frames, self.n_consecutive_frames),
                    axis=self.dim,
                )
                for field in fields
            ]

            for i, slices in enumerate(zip(*splitted_fields)):
                if self.drop_last and slices[0].shape[self.dim] < self.n_consecutive_frames:
                    # Last slice of not equally divisible input, discard.
                    continue

                sliced_fields = dict(zip(fields, slices))
                sliced_fields["__key__"] = f"{key}_{i}"
                yield {**sample, **sliced_fields}

    @hooks.hook_implementation
    def training_transform(self):
        def apply_deterministic_sampling(pipeline: webdataset.Processor):
            if len(self._training_fields) > 0:
                return pipeline.then(
                    functools.partial(self.split_to_consecutive_frames, fields=self._training_fields)
                ).shuffle(self.shuffle_buffer_size)
            else:
                return pipeline

        return apply_deterministic_sampling

    @hooks.hook_implementation
    def evaluation_fields(self):
        return self._evaluation_fields

    @hooks.hook_implementation
    def evaluation_transform(self):
        def apply_deterministic_sampling(pipeline: webdataset.Processor):
            if len(self._evaluation_fields) > 0:
                return pipeline.then(
                    functools.partial(
                        self.split_to_consecutive_frames,
                        fields=self._evaluation_fields,
                    )
                )
            else:
                return pipeline

        return apply_deterministic_sampling


class RandomStridedWindow(Plugin):
    """Select a random consecutive subsequence of frames in a strided manner.

    Given a sequence of [1, 2, 3, 4, 5, 6, 7, 8, 9] this will return one of
    [1, 2, 3] [4, 5, 6] [7, 8, 9].
    """

    def __init__(
        self,
        n_consecutive_frames: int,
        training_fields: Sequence[str],
        evaluation_fields: Sequence[str],
        dim: int = 0,
    ):
        self.n_consecutive_frames = n_consecutive_frames
        self._training_fields = tuple(training_fields)
        self._evaluation_fields = tuple(evaluation_fields)
        self.dim = dim
        self._random = None

    @hooks.hook_implementation
    def training_fields(self):
        return self._training_fields

    @property
    def random(self):
        if not self._random:
            worker_info = torch.utils.data.get_worker_info()
            if worker_info:
                self._random = random.Random(worker_info.seed)
            else:
                self._random = random.Random(torch.initial_seed())

        return self._random

    def split_to_consecutive_frames(self, data, fields):
        """Sample frames deterministically from generator of videos using the __key__ field."""
        for sample in data:
            key = sample["__key__"]
            n_frames = sample[fields[0]].shape[self.dim]
            splitted_fields = []
            for field in fields:
                splitted_field = []
                for step in range(n_frames // self.n_consecutive_frames):
                    ind = []
                    clip_len = (self.n_consecutive_frames-1)*step + self.n_consecutive_frames-1
                    end = n_frames-1-clip_len
                    start_id = self.random.randint(0, end)
                    end_id  = start_id + (self.n_consecutive_frames-1)*step + self.n_consecutive_frames -1
                    for i in range(start_id, end_id+1, step+1):
                        ind.append(i)
                    # ind = self.random.sample(range(0, (step+1)*self.n_consecutive_frames), self.n_consecutive_frames)
                    splitted_field.append(sample[field][sorted(ind)])
                splitted_fields.append(splitted_field)
            # splitted_fields = [
            #     np.array_split(
            #         sample[field],
            #         range(self.n_consecutive_frames, n_frames, self.n_consecutive_frames),
            #         axis=self.dim,
            #     )
            #     for field in fields
            # ]


            n_fragments = len(splitted_fields[0])

            if len(splitted_fields[0][-1] < self.n_consecutive_frames):
                # Discard last fragment if too short.
                n_fragments -= 1

            # fragment_id = self.random.randint(0, n_fragments - 1)
            fragment_id = self.random.randint(0, n_frames // self.n_consecutive_frames-1)
            # fragment_id = 0
            sliced_fields = {
                field_name: splitted_field[fragment_id]
                for field_name, splitted_field in zip(fields, splitted_fields)
            }
            sliced_fields["__key__"] = f"{key}_{fragment_id}"
            yield {**sample, **sliced_fields}

    @hooks.hook_implementation
    def training_transform(self):
        def apply_deterministic_sampling(pipeline: webdataset.Processor):
            if len(self._training_fields) > 0:
                return pipeline.then(
                    functools.partial(self.split_to_consecutive_frames, fields=self._training_fields)
                )
            else:
                return pipeline

        return apply_deterministic_sampling

    @hooks.hook_implementation
    def evaluation_fields(self):
        return self._evaluation_fields

    @hooks.hook_implementation
    def evaluation_transform(self):
        def apply_deterministic_sampling(pipeline: webdataset.Processor):
            if len(self._evaluation_fields) > 0:
                return pipeline.then(
                    functools.partial(
                        self.split_to_consecutive_frames,
                        fields=self._evaluation_fields,
                    )
                )
            else:
                return pipeline

        return apply_deterministic_sampling

ocl/plugins.py
amazon-science/object-centric-multiple-object-tracking
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from ocl.utils import RoutableMixin, init_fn
from typing import Callable, Optional, Tuple, Union, Any
from ocl.path_defaults import OBJECTS
DType = Any
import dataclasses
from torchtyping import TensorType
from ocl import base, path_defaults

class MLP(nn.Module):
    def __init__(self,
                 input_size: int,  # FIXME: added because or else can't instantiate submodules
                 hidden_size: int,
                 output_size: int,  # if not given, should be inputs.shape[-1] at forward
                 num_hidden_layers: int = 1,
                 activation_fn: nn.Module = nn.ReLU,
                 layernorm: Optional[str] = None,
                 activate_output: bool = False,
                 residual: bool = False,
                 weight_init=None
                 ):
        super().__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.output_size = output_size
        self.num_hidden_layers = num_hidden_layers
        self.activation_fn = activation_fn
        self.layernorm = layernorm
        self.activate_output = activate_output
        self.residual = residual
        self.weight_init = weight_init
        if self.layernorm == "pre":
            self.layernorm_module = nn.LayerNorm(input_size, eps=1e-6)
        elif self.layernorm == "post":
            self.layernorm_module = nn.LayerNorm(output_size, eps=1e-6)
        ## mlp
        self.model = nn.ModuleList()
        self.model.add_module("dense_mlp_0", nn.Linear(self.input_size, self.hidden_size))
        self.model.add_module("dense_mlp_0_act", self.activation_fn())
        for i in range(1, self.num_hidden_layers):
            self.model.add_module(f"den_mlp_{i}", nn.Linear(self.hidden_size, self.hidden_size))
            self.model.add_module(f"dense_mlp_{i}_act", self.activation_fn())
        self.model.add_module(f"dense_mlp_{self.num_hidden_layers}", nn.Linear(self.hidden_size, self.output_size))
        if self.activate_output:
            self.model.add_module(f"dense_mlp_{self.num_hidden_layers}_act", self.activation_fn())
        for name, module in self.model.named_children():
            if 'act' not in name:
                nn.init.xavier_uniform_(module.weight)
                # init_fn[weight_init['linear_w']](module.weight)
                # init_fn[weight_init['linear_b']](module.bias)

    def forward(self, inputs: torch.Tensor, train: bool = False) -> torch.Tensor:
        del train  # Unused

        x = inputs
        if self.layernorm == "pre":
            x = self.layernorm_module(x)
        for layer in self.model:
            x = layer(x)
        if self.residual:
            x = x + inputs
        if self.layernorm == "post":
            x = self.layernorm_module(x)
        return x




class GeneralizedDotProductAttention(nn.Module):
    """Multi-head dot-product attention with customizable normalization axis.
    This module supports logging of attention weights in a variable collection.
    """

    def __init__(self,
                 dtype: DType = torch.float32,
                 # precision: Optional[] # not used
                 epsilon: float = 1e-8,
                 inverted_attn: bool = False,
                 renormalize_keys: bool = False,
                 attn_weights_only: bool = False
                ):
        super().__init__()

        self.dtype = dtype
        self.epsilon = epsilon
        self.inverted_attn = inverted_attn
        self.renormalize_keys = renormalize_keys
        self.attn_weights_only = attn_weights_only

    def forward(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor,
                train: bool = False, **kwargs) -> torch.Tensor:
        """Computes multi-head dot-product attention given query, key, and value.
        Args:
            query: Queries with shape of `[batch..., q_num, num_heads, qk_features]`.
            key: Keys with shape of `[batch..., kv_num, num_heads, qk_features]`.
            value: Values with shape of `[batch..., kv_num, num_heads, v_features]`.
            train: Indicating whether we're training or evaluating.
            **kwargs: Additional keyword arguments are required when used as attention
                function in nn.MultiHeadDotPRoductAttention, but they will be ignored here.
        Returns:
            Output of shape `[batch..., q_num, num_heads, v_features]`.
        """
        del train # Unused.

        assert query.ndim == key.ndim == value.ndim, (
            "Queries, keys, and values must have the same rank.")
        assert query.shape[:-3] == key.shape[:-3] == value.shape[:-3], (
            "Query, key, and value batch dimensions must match.")
        assert query.shape[-2] == key.shape[-2] == value.shape[-2], (
            "Query, key, and value num_heads dimensions must match.")
        assert key.shape[-3] == value.shape[-3], (
            "Key and value cardinality dimensions must match.")
        assert query.shape[-1] == key.shape[-1], (
            "Query and key feature dimensions must match.")

        if kwargs.get("bias") is not None:
            raise NotImplementedError(
                "Support for masked attention is not yet implemented.")

        if "dropout_rate" in kwargs:
            if kwargs["dropout_rate"] > 0.:
                raise NotImplementedError("Support for dropout is not yet implemented.")

        # Temperature normalization.
        qk_features = query.shape[-1]
        query = query / (qk_features ** 0.5) # torch.sqrt(qk_features)

        # attn.shape = (batch..., num_heads, q_num, kv_num)
        attn = torch.matmul(query.permute(0, 2, 1, 3), key.permute(0, 2, 3, 1)) # bhqd @ bhdk -> bhqk

        if self.inverted_attn:
            attention_dim = -2 # Query dim
        else:
            attention_dim = -1 # Key dim

        # Softmax normalization (by default over key dim)
        attn = torch.softmax(attn, dim=attention_dim, dtype=self.dtype)

        if self.renormalize_keys:
            # Corresponds to value aggregation via weighted mean (as opposed to sum).
            normalizer = torch.sum(attn, axis=-1, keepdim=True) + self.epsilon
            attn_n = attn / normalizer
        else:
            attn_n = attn

        if self.attn_weights_only:
            return attn_n

        # Aggregate values using a weighted sum with weights provided by `attn`
        updates = torch.einsum("bhqk,bkhd->bqhd", attn_n, value)

        return updates, attn # FIXME: return attention too, as no option for intermediate storing in module in torch.

class TransformerBlock(nn.Module, RoutableMixin):
    def __init__(self,
                 embed_dim: int = 128,
                 num_heads: int = 4,
                 qkv_size: int = 128,
                 mlp_size: int = 256,
                 pre_norm: bool = False,
                 weight_init=None,
                 # object_features_path: Optional[str] = OBJECTS,
                 ):
        nn.Module.__init__(self)
        # RoutableMixin.__init__(self, {"object_features": object_features_path})

        self.embed_dim = embed_dim
        self.qkv_size = qkv_size
        self.mlp_size = mlp_size
        self.num_heads = num_heads
        self.pre_norm = pre_norm
        self.weight_init = weight_init

        assert num_heads >= 1
        assert qkv_size % num_heads == 0, "embed dim must be divisible by num_heads"
        self.head_dim = qkv_size // num_heads

        # submodules
        ## MHA #
        self.attn = GeneralizedDotProductAttention()
        ## mlps
        self.mlp = MLP(
            input_size=embed_dim, hidden_size=mlp_size,
            output_size=embed_dim, weight_init=weight_init)
        ## layernorms
        self.layernorm_query = nn.LayerNorm(embed_dim, eps=1e-6)
        self.layernorm_mlp = nn.LayerNorm(embed_dim, eps=1e-6)
        ## weights
        self.dense_q = nn.Linear(embed_dim, qkv_size)
        self.dense_k = nn.Linear(embed_dim, qkv_size)
        self.dense_v = nn.Linear(embed_dim, qkv_size)
        # init_fn[weight_init['linear_w']](self.dense_q.weight)
        # init_fn[weight_init['linear_b']](self.dense_q.bias)
        # init_fn[weight_init['linear_w']](self.dense_k.weight)
        # init_fn[weight_init['linear_b']](self.dense_k.bias)
        # init_fn[weight_init['linear_w']](self.dense_v.weight)
        # init_fn[weight_init['linear_b']](self.dense_v.bias)
        if self.num_heads > 1:
            self.dense_o = nn.Linear(qkv_size, embed_dim)
            # nn.init.xavier_uniform_(self.w_o.weight)
            # init_fn[weight_init['linear_w']](self.dense_o.weight)
            # init_fn[weight_init['linear_b']](self.dense_o.bias)
            self.multi_head = True
        else:
            self.multi_head = False


    # @RoutableMixin.route
    def forward(self, object_features: torch.Tensor):  # TODO: add general attention for q, k, v, not just for x = qkv
        assert object_features.ndim == 3
        B, L, _ = object_features.shape
        head_dim = self.embed_dim // self.num_heads

        if self.pre_norm:
            # Self-attention.
            x = self.layernorm_query(object_features)
            q = self.dense_q(x).view(B, L, self.num_heads, head_dim)
            k = self.dense_k(x).view(B, L, self.num_heads, head_dim)
            v = self.dense_v(x).view(B, L, self.num_heads, head_dim)
            x, _ = self.attn(query=q, key=k, value=v)
            if self.multi_head:
                x = self.dense_o(x.reshape(B, L, self.qkv_size)).view(B, L, self.embed_dim)
            else:
                x = x.squeeze(-2)
            x = x + object_features

            y = x

            # MLP
            z = self.layernorm_mlp(y)
            z = self.mlp(z)
            z = z + y
        else:
            # Self-attention on queries.
            x = object_features
            # q = self.dense_q(x).view(B, L, self.num_heads, head_dim)
            # k = self.dense_k(x).view(B, L, self.num_heads, head_dim)
            # v = self.dense_v(x).view(B, L, self.num_heads, head_dim)
            # x, _ = self.attn(query=q, key=k, value=v)
            # if self.multi_head:
            #     x = self.dense_o(x.reshape(B, L, self.qkv_size)).view(B, L, self.embed_dim)
            # else:
            #     x = x.squeeze(-2)
            # x = x + object_features
            # x = self.layernorm_query(x)

            y = x

            # MLP
            z = self.mlp(y)
            z = z + y
            z = self.layernorm_mlp(z)
        return z






@dataclasses.dataclass
class PredictorOutput:
    objects: TensorType["batch_size", "n_objects", "object_dim"]  # noqa: F821
    is_empty: Optional[TensorType["batch_size", "n_objects"]] = None  # noqa: F821
    feature_attributions: Optional[
        TensorType["batch_size", "n_objects", "n_spatial_features"]  # noqa: F821
    ] = None
ocl/predictor.py
amazon-science/object-centric-multiple-object-tracking
"""Data preprocessing functions."""
from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
from torchvision.ops import masks_to_boxes
import numpy
import torch
from torchvision import transforms


class DropEntries:
    """Drop entries from data dictionary."""

    def __init__(self, keys: List[str]):
        self.keys = tuple(keys)

    def __call__(self, data: Dict[str, Any]):
        return {k: v for k, v in data.items() if k not in self.keys}


class MaskToTensor:
    """Convert a segmentation mask numpy array to a tensor.

    Mask is assumed to be of shape (..., K, H, W, 1), i.e. one-hot encoded with K classes and any
    number of leading dimensions. Returned tensor is of shape (..., K, H, W), containing binary
    entries.
    """

    def __call__(self, mask: numpy.ndarray) -> torch.Tensor:
        assert mask.shape[-1] == 1
        mask_binary = mask > 0.0
        return torch.from_numpy(mask_binary).squeeze(-1).to(torch.float32)

    def __repr__(self) -> str:
        return f"{self.__class__.__name__}()"


class DenseMaskToTensor:
    """Convert a dense segmentation mask numpy array to a tensor.

    Mask is assumed to be of shape (..., K, H, W, 1), i.e. densely encoded with K classes and any
    number of leading dimensions. Returned tensor is of shape (..., K, H, W).
    """

    def __call__(self, mask: numpy.ndarray) -> torch.Tensor:
        assert mask.shape[-1] == 1
        return torch.from_numpy(mask).squeeze(-1).to(torch.uint8)

    def __repr__(self) -> str:
        return f"{self.__class__.__name__}()"

class BBoxToTensor:
    """Convert a bbox numpy array to a tensor.

    BBox is assumed to be of shape (..., 4), i.e. the four dimensions represent [x,y,w,h].
    Returned tensor is of shape (..., 4).
    """

    def __call__(self, bbox: numpy.ndarray) -> torch.Tensor:
        assert bbox.shape[-1] == 4
        return torch.from_numpy(bbox).squeeze(-1).to(torch.float32)

    def __repr__(self) -> str:
        return f"{self.__class__.__name__}()"

class ClsToTensor:
    """Convert a category numpy array to a tensor.

    Category is assumed to be of shape (..., 1).
    Returned tensor is of shape (..., 1).
    """

    def __call__(self, instance_cls: numpy.ndarray) -> torch.Tensor:
        assert instance_cls.shape[-1] == 1
        return torch.from_numpy(instance_cls).squeeze(-1).to(torch.long)

    def __repr__(self) -> str:
        return f"{self.__class__.__name__}()"

class IDToTensor:
    """Convert an ID numpy array to a tensor.

    ID is assumed to be of shape (..., 1).
    Returned tensor is of shape (..., 1).
    """

    def __call__(self, instance_id: numpy.ndarray) -> torch.Tensor:
        assert instance_id.shape[-1] == 1
        return torch.from_numpy(instance_id).squeeze(-1).to(torch.long)

    def __repr__(self) -> str:
        return f"{self.__class__.__name__}()"

class MultiMaskToTensor:
    """Discretize mask, where multiple objects are partially masked into an exclusive binary mask."""

    def __init__(self, axis: int = -4):
        self.axis = axis

    def __call__(self, mask: numpy.ndarray) -> torch.Tensor:
        int_mask = numpy.argmax(mask, axis=self.axis).squeeze(-1)
        out_mask = torch.nn.functional.one_hot(torch.from_numpy(int_mask), mask.shape[self.axis])
        # Ensure the object axis is again at the same location.
        # We operate on the shape prior to squeezing for axis to be consistent.
        last_index = len(out_mask.shape) - 1
        indices = list(range(len(out_mask.shape) + 1))
        indices.insert(self.axis, last_index)
        indices = indices[:-2]  # Remove last indices as they are squeezed or inserted.
        out_mask = out_mask.permute(*indices).to(torch.float32)
        return out_mask


class IntegerToOneHotMask:
    """Convert an integer mask to a one-hot mask.

    Integer masks are masks where the instance ID is written into the mask.
    This transform expands them to a one-hot encoding.

    Args:
        ignore_typical_background: Ignore pixels where the mask is zero or 255.
        This often corresponds to the background or to the segmentation boundary.
    """

    def __init__(self, ignore_typical_background=True, output_axis=-4):
        self.ignore_typical_background = ignore_typical_background
        self.output_axis = output_axis

    def __call__(self, array: numpy.array):
        max_instance = array.max()
        if self.ignore_typical_background:
            if max_instance == 255:
                # Replace 255 with zero, both are ignored.
                array[array == 255] = 0
            max_instance = array.max()
            to_one_hot = numpy.concatenate(
                [
                    numpy.zeros((1, max_instance), dtype=numpy.uint8),
                    numpy.eye(max_instance, dtype=numpy.uint8),
                ],
                axis=0,
            )
        else:
            to_one_hot = numpy.eye(max_instance + 1, dtype=numpy.uint8)
        return numpy.moveaxis(to_one_hot[array], -1, self.output_axis)


class VOCInstanceMasksToDenseMasks:
    """Convert a segmentation mask with integer encoding into a one-hot segmentation mask.

    We use this transform as Pascal VOC segmentatation and object annotations seems to not
    be aligned.
    """

    def __init__(
        self,
        instance_mask_key: str = "segmentation-instance",
        class_mask_key: str = "segmentation-class",
        classes_key: str = "instance_category",
        instance_axis: int = -4,
    ):
        self.instance_mask_key = instance_mask_key
        self.class_mask_key = class_mask_key
        self.classes_key = classes_key
        self.instance_axis = instance_axis

    def __call__(self, data: Dict[str, Any]):
        expanded_segmentation_mask = data[self.instance_mask_key] * numpy.expand_dims(
            data[self.class_mask_key], axis=self.instance_axis
        )
        assert expanded_segmentation_mask.max() != 255
        data[self.instance_mask_key] = expanded_segmentation_mask
        classes = []
        for instance_slice in numpy.rollaxis(expanded_segmentation_mask, self.instance_axis):
            unique_values = numpy.unique(instance_slice)
            assert len(unique_values) == 2  # Should contain 0 and class id.
            classes.append(unique_values[1])
        data[self.classes_key] = numpy.array(classes)
        return data


class AddImageSize:
    """Add height and width of image as data entry.

    Args:
        key: Key of image.
        target_key: Key under which to store size.
    """

    def __init__(self, key: str = "image", target_key: str = "image_size"):
        self.key = key
        self.target_key = target_key

    def __call__(self, data: Dict[str, Any]):
        height, width, _ = data[self.key].shape
        data[self.target_key] = numpy.array([height, width], dtype=numpy.int64)
        return data


class AddEmptyMasks:
    """Add empty masks to data if the data does not include them already.

    Args:
        keys: One or several keys of empty masks to be added.
        take_size_from: Key of element whose height and width is used to create mask. Element is
            assumed to have shape of (H, W, C).
    """

    def __init__(self, mask_keys: Union[str, Sequence[str]], take_size_from: str = "image"):
        if isinstance(mask_keys, str):
            self.mask_keys = (mask_keys,)
        else:
            self.mask_keys = tuple(mask_keys)
        self.source_key = take_size_from

    def __call__(self, data: Dict[str, Any]):
        height, width, _ = data[self.source_key].shape
        for key in self.mask_keys:
            if key not in data:
                data[key] = numpy.zeros((1, height, width, 1), dtype=numpy.uint8)

        return data


class AddEmptyBboxes:
    """Add empty bounding boxes to data if the data does not include them already.

    Args:
        keys: One or several keys of empty boxes to be added.
        empty_value: Value of the empty box at all coordinates.
    """

    def __init__(self, keys: Union[str, Sequence[str]] = "instance_bbox", empty_value: float = -1.0):
        if isinstance(keys, str):
            self.keys = (keys,)
        else:
            self.keys = tuple(keys)
        self.empty_value = empty_value

    def __call__(self, data: Dict[str, Any]):
        for key in self.keys:
            if key not in data:
                data[key] = numpy.ones((1, 4), dtype=numpy.float32) * self.empty_value

        return data


class CanonicalizeBboxes:
    """Convert bounding boxes to canonical (x1, y1, x2, y2) format.

    Args:
        key: Key of bounding box, assumed to have shape K x 4.
        format: Format of bounding boxes. Either "xywh" or "yxyx".
    """

    def __init__(self, key: str = "instance_bbox", format: str = "xywh"):
        self.key = key

        self.format_xywh = False
        self.format_yxyx = False
        if format == "xywh":
            self.format_xywh = True
        elif format == "yxyx":
            self.format_yxyx = True
        else:
            raise ValueError(f"Unknown input format `{format}`")

    def __call__(self, data: Dict[str, Any]):
        if self.key not in data:
            return data

        bboxes = data[self.key]
        if self.format_xywh:
            x1, y1, w, h = numpy.split(bboxes, 4, axis=1)
            x2 = x1 + w
            y2 = y1 + h
        elif self.format_yxyx:
            y1, x1, y2, x2 = numpy.split(bboxes, 4, axis=1)

        data[self.key] = numpy.concatenate((x1, y1, x2, y2), axis=1)

        return data


class RescaleBboxes:
    """Rescale bounding boxes by size taken from data.

    Bounding boxes are assumed to have format (x1, y1, x2, y2). The rescaled box is
        (x1 * width, y1 * height, x2 * width, y2 * height).

    Args:
        key: Key of bounding box, assumed to have shape K x 4.
        take_size_from: Key of element to take the size for rescaling from, assumed to have shape
            H x W x C.
    """

    def __init__(self, key: str = "instance_bbox", take_size_from: str = "image"):
        self.key = key
        self.take_size_from = take_size_from

    def __call__(self, data: Dict[str, Any]):
        if self.key not in data:
            return data

        height, width, _ = data[self.take_size_from].shape
        scaling = numpy.array([[width, height, width, height]], dtype=numpy.float32)
        data[self.key] = data[self.key] * scaling

        return data


def expand_dense_mask(mask: numpy.ndarray) -> numpy.ndarray:
    """Convert dense segmentation mask to one where each class occupies one dimension.

    Args:
        mask: Densely encoded segmentation mask of shape 1 x H x W x 1.

    Returns: Densely encoded segmentation mask of shape K x H x W x 1, where K is the
        number of classes in the mask. Zero is taken to indicate an unoccupied pixel.
    """
    classes = numpy.unique(mask)[:, None, None, None]
    mask = (classes == mask) * classes

    if classes[0].squeeze() == 0:  # Strip empty class
        mask = mask[1:]

    return mask


class AddSegmentationMaskFromInstanceMask:
    """Convert instance to segmentation masks by joining instances with the same category.

    Overlaps of instances of different classes are resolved by taking the class with the higher class
    id.
    """

    def __init__(
        self,
        instance_mask_key: str = "instance_mask",
        target_key: str = "segmentation_mask",
    ):
        self.instance_mask_key = instance_mask_key
        self.target_key = target_key

    @staticmethod
    def convert(instance_mask: numpy.ndarray) -> numpy.ndarray:
        """Convert instance to segmentation mask.

        Args:
            instance_mask: Densely encoded instance masks of shape I x H x W x 1, where I is the
                number of instances.
        """
        # Reduce instance mask to single dimension
        instance_mask = instance_mask.max(axis=0, keepdims=True)

        return expand_dense_mask(instance_mask)

    def __call__(self, data: Dict[str, Any]):
        if self.instance_mask_key not in data:
            return data

        data[self.target_key] = self.convert(data[self.instance_mask_key])

        return data

class AddBBoxFromInstanceMasks:
    """Convert instance mask to bounding box.

    Args:
        instance_mask_key: mask key name.
        target_key: target key name.
    """

    def __init__(
        self,
        instance_mask_key: str = "mask",
        video_id_key: str = "__key__",  # not quite sure if this is the best key
        target_box_key: str = "instance_bbox",
        target_cls_key: str = "instance_cls",
        target_id_key: str = "instance_id",
    ):
        self.instance_mask_key = instance_mask_key
        self.video_id_key = video_id_key
        self.target_box_key = target_box_key
        self.target_cls_key = target_cls_key
        self.target_id_key = target_id_key

    @staticmethod
    def convert(instance_mask: numpy.ndarray, video_id: numpy.ndarray) -> numpy.ndarray:
        num_frame, num_instance, height, width, _ = instance_mask.shape

        # Convert to binary mask
        binary_mask = instance_mask > 0
        # Filter background. TODO: now we assume the first mask for each video is background.
        # Might not apply to every dataset
        binary_mask = binary_mask[:, 1:]
        num_instance -= 1
        binary_mask = (
            torch.tensor(binary_mask).squeeze().reshape(num_frame * num_instance, height, width)
        )
        # Filter empty masks
        non_empty_mask_idx = torch.where(binary_mask.sum(-1).sum(-1) > 0)[0]
        empty_mask_idx = torch.where(binary_mask.sum(-1).sum(-1) == 0)[0]
        non_empty_binary_mask = binary_mask[non_empty_mask_idx]
        non_empty_bboxes = masks_to_boxes(non_empty_binary_mask)

        # Turn box into cxcyhw
        bboxes = torch.zeros(num_frame * num_instance, 4)
        bboxes[non_empty_mask_idx, 0] = non_empty_bboxes[:, 0]
        bboxes[non_empty_mask_idx, 1] = non_empty_bboxes[:, 1]
        bboxes[non_empty_mask_idx, 2] = non_empty_bboxes[:, 2]
        bboxes[non_empty_mask_idx, 3] = non_empty_bboxes[:, 3]


        # normalized to 0,1
        # Make sure width and height are correct
        bboxes[:, 0::2] = bboxes[:, 0::2] / width
        bboxes[:, 1::2] = bboxes[:, 1::2] / height

        # Switch back to numpy, though will turn to tensor again in later phase.
        # I don't know if it is necessary to keep the output of this phase as numpy.
        bboxes = bboxes.view(num_frame, num_instance, 4).numpy()

        # class
        # -1 is background or no object, 0 is the first object class
        instance_cls = torch.ones(num_frame * num_instance, 1) * -1
        instance_cls[non_empty_mask_idx] = 0
        instance_cls = instance_cls.view(num_frame, num_instance, 1).numpy()

        # ID
        instance_id = torch.zeros(num_frame, num_instance, 1)
        # instance_id[:] = torch.range(0, num_instance - 1).unsqueeze(-1)
        # instance_id = instance_id + int(video_id) * 100
        # instance_id = instance_id.view(num_frame * num_instance, 1)
        # instance_id[empty_mask_idx] = -1
        instance_id = instance_id.view(num_frame, num_instance, 1).numpy()

        return bboxes, instance_cls, instance_id

    def __call__(self, data: Dict[str, Any]):
        if self.instance_mask_key not in data:
            return data

        bboxes, instance_cls, instance_id = self.convert(
            data[self.instance_mask_key], data[self.video_id_key]
        )
        data[self.target_box_key] = bboxes
        data[self.target_cls_key] = instance_cls
        data[self.target_id_key] = instance_id
        return data

class RenameFields:
    def __init__(self, mapping: Dict):
        self.mapping = mapping

    def __call__(self, d: Dict):
        # Create shallow copy to avoid issues target key is already used.
        out = d.copy()
        for source, target in self.mapping.items():
            out[target] = d[source]
        return out


class InstanceMasksToDenseMasks:
    """Convert binary instance masks to dense masks, i.e. where the mask value encodes the class id.

    Class ids are taken from a list containing a class id per instance.
    """

    def __init__(
        self,
        instance_mask_key: str = "instance_mask",
        category_key: str = "instance_category",
    ):
        self.instance_mask_key = instance_mask_key
        self.category_key = category_key

    @staticmethod
    def convert(instance_mask: numpy.ndarray, categories: numpy.ndarray) -> numpy.ndarray:
        if numpy.min(categories) <= 0:
            raise ValueError("Detected category smaller equal than 0 in instance masks.")
        if numpy.max(categories) > 255:
            raise ValueError(
                "Detected category greater than 255 in instance masks. This does not fit in uint8."
            )

        categories = categories[:, None, None, None]
        return (instance_mask * categories).astype(numpy.uint8)

    def __call__(self, data: Dict[str, Any]):
        if self.instance_mask_key not in data:
            return data

        data[self.instance_mask_key] = self.convert(
            data[self.instance_mask_key], data[self.category_key]
        )

        return data


class MergeCocoThingsAndStuff:
    """Merge COCO things and stuff segmentation masks.

    Args:
        things_key: Key to things instance mask. Mask is assumed to be densely encoded, i.e.
            the mask value encodes the class id, of shape I x H x W x 1, where I is the number of
            things instances.
        stuff_key: Key to stuff segmentation mask. Mask is assumed to be densely encoded, i.e.
            the mask value encodes the class id, of shape K x H x W x 1, where K is the number stuff
            classes.
        output_key: Key under which the merged mask is stored. Returns mask of shape L x H x W x 1,
            where K <= L <= K + I.
        include_crowd: Whether to include pixels marked as crowd with their class, or with class
            zero.
    """

    def __init__(
        self,
        output_key: str,
        things_key: str = "instance_mask",
        stuff_key: str = "stuff_mask",
        include_crowd: bool = False,
    ):
        self.things_key = things_key
        self.stuff_key = stuff_key
        self.output_key = output_key
        self.include_crowd = include_crowd

    def __call__(self, data: Dict[str, Any]):
        if self.things_key in data:
            things_instance_mask = data[self.things_key]
            things_mask = things_instance_mask.max(axis=0, keepdims=True)
        else:
            things_mask = None

        stuff_mask = data[self.stuff_key]
        merged_mask = stuff_mask.max(axis=0, keepdims=True)

        # In stuff annotations, thing pixels are encoded as class 183.
        use_thing_mask = merged_mask == 183

        if things_mask is not None:
            if self.include_crowd:
                # In the stuff annotations, things marked with the "crowd" label are NOT encoded as
                # class 183, but as class 0. We can take the value of the things mask for those
                # pixels.
                use_thing_mask |= merged_mask == 0
            merged_mask[use_thing_mask] = things_mask[use_thing_mask]
        else:
            # No pixel should have value 183 if the things_mask does not exist, but convert it to
            # zero anyways just to be sure.
            merged_mask[use_thing_mask] = 0

        data[self.output_key] = expand_dense_mask(merged_mask)

        return data


class VideoToTensor:
    """Convert a video numpy array of shape (T, H, W, C) to a torch tensor of shape (T, C, H, W)."""

    def __call__(self, video):
        """Convert a numpy array of a video into a torch tensor.

        Assumes input is a numpy array of shape T x H x W x C (or T x H x W for monochrome videos)
        and convert it into torch tensor of shape T x C x H x W in order to allow application of
        Conv3D operations.
        """
        if isinstance(video, numpy.ndarray):
            # Monochrome video such as mask
            if video.ndim == 3:
                video = video[..., None]

            video = torch.from_numpy(video.transpose((0, 3, 1, 2))).contiguous()
            # backward compatibility
            if isinstance(video, torch.ByteTensor):
                return video.to(dtype=torch.get_default_dtype()).div(255)
            else:
                return video


class ToSingleFrameVideo:
    """Convert image in tensor format to video format by adding frame dimension with single element.

    Converts C x H x W tensors into tensors of shape 1 x C x H x W.
    """

    def __call__(self, image):
        return image.unsqueeze(0)


class NormalizeVideo:
    """Normalize a video tensor of shape (T, C, H, W)."""

    def __init__(self, mean, std):
        self.mean = torch.tensor(mean)[None, :, None, None]
        self.std = torch.tensor(std)[None, :, None, None]

    def __call__(self, video):
        return (video - self.mean) / self.std


class Denormalize(torch.nn.Module):
    """Denormalize a tensor of shape (..., C, H, W) with any number of leading dimensions."""

    def __init__(self, mean, std):
        super().__init__()
        self.register_buffer("mean", torch.tensor(mean)[:, None, None])
        self.register_buffer("std", torch.tensor(std)[:, None, None])

    def __call__(self, tensor):
        return tensor * self.std + self.mean


class ResizeNearestExact:
    """Resize a tensor using mode nearest-exact.

    This mode is not available in torchvision.transforms.Resize as of v0.12. This class was adapted
    from torchvision.transforms.functional_tensor.resize.
    """

    def __init__(self, size: Union[int, List[int]], max_size: Optional[int] = None):
        self.size = size
        self.max_size = max_size

    @staticmethod
    def _cast_squeeze_in(
        img: torch.Tensor, req_dtypes: List[torch.dtype]
    ) -> Tuple[torch.Tensor, bool, bool, torch.dtype]:
        need_squeeze = False
        # make image NCHW
        if img.ndim < 4:
            img = img.unsqueeze(dim=0)
            need_squeeze = True

        out_dtype = img.dtype
        need_cast = False
        if out_dtype not in req_dtypes:
            need_cast = True
            req_dtype = req_dtypes[0]
            img = img.to(req_dtype)
        return img, need_cast, need_squeeze, out_dtype

    @staticmethod
    def _cast_squeeze_out(
        img: torch.Tensor, need_cast: bool, need_squeeze: bool, out_dtype: torch.dtype
    ) -> torch.Tensor:
        if need_squeeze:
            img = img.squeeze(dim=0)

        if need_cast:
            if out_dtype in (torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64):
                # it is better to round before cast
                img = torch.round(img)
            img = img.to(out_dtype)

        return img

    @staticmethod
    def resize(img: torch.Tensor, size: Union[int, List[int]], max_size: Optional[int] = None):
        h, w = img.shape[-2:]

        if isinstance(size, int) or len(size) == 1:  # specified size only for the smallest edge
            short, long = (w, h) if w <= h else (h, w)
            requested_new_short = size if isinstance(size, int) else size[0]

            new_short, new_long = requested_new_short, int(requested_new_short * long / short)

            if max_size is not None:
                if max_size <= requested_new_short:
                    raise ValueError(
                        f"max_size = {max_size} must be strictly greater than the requested "
                        f"size for the smaller edge size = {size}"
                    )
                if new_long > max_size:
                    new_short, new_long = int(max_size * new_short / new_long), max_size

            new_w, new_h = (new_short, new_long) if w <= h else (new_long, new_short)

            if (w, h) == (new_w, new_h):
                return img
        else:  # specified both h and w
            new_w, new_h = size[1], size[0]

        img, need_cast, need_squeeze, out_dtype = ResizeNearestExact._cast_squeeze_in(
            img, (torch.float32, torch.float64)
        )

        img = torch.nn.functional.interpolate(img, size=[new_h, new_w], mode="nearest-exact")

        img = ResizeNearestExact._cast_squeeze_out(
            img, need_cast=need_cast, need_squeeze=need_squeeze, out_dtype=out_dtype
        )

        return img

    def __call__(self, img: torch.Tensor) -> torch.Tensor:
        return ResizeNearestExact.resize(img, self.size, self.max_size)


class ConvertToCocoSuperclasses:
    """Convert segmentation mask from COCO classes (183) to COCO superclasses (27)."""

    ID_TO_SUPERCLASS_AND_NAME = {
        0: ("unlabeled", "unlabeled"),
        1: ("person", "person"),
        2: ("vehicle", "bicycle"),
        3: ("vehicle", "car"),
        4: ("vehicle", "motorcycle"),
        5: ("vehicle", "airplane"),
        6: ("vehicle", "bus"),
        7: ("vehicle", "train"),
        8: ("vehicle", "truck"),
        9: ("vehicle", "boat"),
        10: ("outdoor", "traffic light"),
        11: ("outdoor", "fire hydrant"),
        13: ("outdoor", "stop sign"),
        14: ("outdoor", "parking meter"),
        15: ("outdoor", "bench"),
        16: ("animal", "bird"),
        17: ("animal", "cat"),
        18: ("animal", "dog"),
        19: ("animal", "horse"),
        20: ("animal", "sheep"),
        21: ("animal", "cow"),
        22: ("animal", "elephant"),
        23: ("animal", "bear"),
        24: ("animal", "zebra"),
        25: ("animal", "giraffe"),
        27: ("accessory", "backpack"),
        28: ("accessory", "umbrella"),
        31: ("accessory", "handbag"),
        32: ("accessory", "tie"),
        33: ("accessory", "suitcase"),
        34: ("sports", "frisbee"),
        35: ("sports", "skis"),
        36: ("sports", "snowboard"),
        37: ("sports", "sports ball"),
        38: ("sports", "kite"),
        39: ("sports", "baseball bat"),
        40: ("sports", "baseball glove"),
        41: ("sports", "skateboard"),
        42: ("sports", "surfboard"),
        43: ("sports", "tennis racket"),
        44: ("kitchen", "bottle"),
        46: ("kitchen", "wine glass"),
        47: ("kitchen", "cup"),
        48: ("kitchen", "fork"),
        49: ("kitchen", "knife"),
        50: ("kitchen", "spoon"),
        51: ("kitchen", "bowl"),
        52: ("food", "banana"),
        53: ("food", "apple"),
        54: ("food", "sandwich"),
        55: ("food", "orange"),
        56: ("food", "broccoli"),
        57: ("food", "carrot"),
        58: ("food", "hot dog"),
        59: ("food", "pizza"),
        60: ("food", "donut"),
        61: ("food", "cake"),
        62: ("furniture", "chair"),
        63: ("furniture", "couch"),
        64: ("furniture", "potted plant"),
        65: ("furniture", "bed"),
        67: ("furniture", "dining table"),
        70: ("furniture", "toilet"),
        72: ("electronic", "tv"),
        73: ("electronic", "laptop"),
        74: ("electronic", "mouse"),
        75: ("electronic", "remote"),
        76: ("electronic", "keyboard"),
        77: ("electronic", "cell phone"),
        78: ("appliance", "microwave"),
        79: ("appliance", "oven"),
        80: ("appliance", "toaster"),
        81: ("appliance", "sink"),
        82: ("appliance", "refrigerator"),
        84: ("indoor", "book"),
        85: ("indoor", "clock"),
        86: ("indoor", "vase"),
        87: ("indoor", "scissors"),
        88: ("indoor", "teddy bear"),
        89: ("indoor", "hair drier"),
        90: ("indoor", "toothbrush"),
        92: ("textile", "banner"),
        93: ("textile", "blanket"),
        94: ("plant", "branch"),
        95: ("building", "bridge"),
        96: ("building", "building-other"),
        97: ("plant", "bush"),
        98: ("furniture-stuff", "cabinet"),
        99: ("structural", "cage"),
        100: ("raw-material", "cardboard"),
        101: ("floor", "carpet"),
        102: ("ceiling", "ceiling-other"),
        103: ("ceiling", "ceiling-tile"),
        104: ("textile", "cloth"),
        105: ("textile", "clothes"),
        106: ("sky", "clouds"),
        107: ("furniture-stuff", "counter"),
        108: ("furniture-stuff", "cupboard"),
        109: ("textile", "curtain"),
        110: ("furniture-stuff", "desk-stuff"),
        111: ("ground", "dirt"),
        112: ("furniture-stuff", "door-stuff"),
        113: ("structural", "fence"),
        114: ("floor", "floor-marble"),
        115: ("floor", "floor-other"),
        116: ("floor", "floor-stone"),
        117: ("floor", "floor-tile"),
        118: ("floor", "floor-wood"),
        119: ("plant", "flower"),
        120: ("water", "fog"),
        121: ("food-stuff", "food-other"),
        122: ("food-stuff", "fruit"),
        123: ("furniture-stuff", "furniture-other"),
        124: ("plant", "grass"),
        125: ("ground", "gravel"),
        126: ("ground", "ground-other"),
        127: ("solid", "hill"),
        128: ("building", "house"),
        129: ("plant", "leaves"),
        130: ("furniture-stuff", "light"),
        131: ("textile", "mat"),
        132: ("raw-material", "metal"),
        133: ("furniture-stuff", "mirror-stuff"),
        134: ("plant", "moss"),
        135: ("solid", "mountain"),
        136: ("ground", "mud"),
        137: ("textile", "napkin"),
        138: ("structural", "net"),
        139: ("raw-material", "paper"),
        140: ("ground", "pavement"),
        141: ("textile", "pillow"),
        142: ("plant", "plant-other"),
        143: ("raw-material", "plastic"),
        144: ("ground", "platform"),
        145: ("ground", "playingfield"),
        146: ("structural", "railing"),
        147: ("ground", "railroad"),
        148: ("water", "river"),
        149: ("ground", "road"),
        150: ("solid", "rock"),
        151: ("building", "roof"),
        152: ("textile", "rug"),
        153: ("food-stuff", "salad"),
        154: ("ground", "sand"),
        155: ("water", "sea"),
        156: ("furniture-stuff", "shelf"),
        157: ("sky", "sky-other"),
        158: ("building", "skyscraper"),
        159: ("ground", "snow"),
        160: ("solid", "solid-other"),
        161: ("furniture-stuff", "stairs"),
        162: ("solid", "stone"),
        163: ("plant", "straw"),
        164: ("structural", "structural-other"),
        165: ("furniture-stuff", "table"),
        166: ("building", "tent"),
        167: ("textile", "textile-other"),
        168: ("textile", "towel"),
        169: ("plant", "tree"),
        170: ("food-stuff", "vegetable"),
        171: ("wall", "wall-brick"),
        172: ("wall", "wall-concrete"),
        173: ("wall", "wall-other"),
        174: ("wall", "wall-panel"),
        175: ("wall", "wall-stone"),
        176: ("wall", "wall-tile"),
        177: ("wall", "wall-wood"),
        178: ("water", "water-other"),
        179: ("water", "waterdrops"),
        180: ("window", "window-blind"),
        181: ("window", "window-other"),
        182: ("solid", "wood"),
        183: ("other", "other"),
    }

    SUPERCLASS_TO_ID = {
        "unlabeled": 0,
        "person": 1,
        "vehicle": 2,
        "outdoor": 3,
        "animal": 4,
        "accessory": 5,
        "sports": 6,
        "kitchen": 7,
        "food": 8,
        "furniture": 9,
        "electronic": 10,
        "appliance": 11,
        "indoor": 12,
        "textile": 13,
        "plant": 14,
        "building": 15,
        "furniture-stuff": 16,
        "structural": 17,
        "raw-material": 18,
        "floor": 19,
        "ceiling": 20,
        "sky": 21,
        "ground": 22,
        "water": 23,
        "food-stuff": 24,
        "wall": 25,
        "window": 26,
        "solid": 27,
        "other": 28,
    }

    def __init__(self):
        max_class = max(ConvertToCocoSuperclasses.ID_TO_SUPERCLASS_AND_NAME.keys())
        class_to_superclass = numpy.zeros((max_class + 1,), dtype=numpy.uint8)
        for class_id, (supclass, _) in ConvertToCocoSuperclasses.ID_TO_SUPERCLASS_AND_NAME.items():
            class_to_superclass[class_id] = ConvertToCocoSuperclasses.SUPERCLASS_TO_ID[supclass]
        self.class_to_superclass = class_to_superclass

    def __call__(self, mask: numpy.ndarray) -> numpy.ndarray:
        """Convert mask to superclasses.

        Args:
            mask: Densely encoded segmentation mask of shape K x H x W x 1.

        Returns: Segmentation mask of shape C x H x W x 1, where C is the new set of classes.
        """
        classes = mask.reshape(len(mask), -1).max(axis=-1)
        superclasses = self.class_to_superclass[classes]
        mask = (mask > 0) * superclasses[:, None, None, None]

        return expand_dense_mask(mask.max(axis=0, keepdims=True))


class JointRandomResizedCropwithParameters(transforms.RandomResizedCrop):
    def __init__(
        self,
        size,
        scale=(0.08, 1.0),
        ratio=(3.0 / 4.0, 4.0 / 3.0),
        interpolation=transforms.functional.InterpolationMode.BILINEAR,
    ):
        super().__init__(size, scale, ratio, interpolation)
        self.mask_to_tensor = DenseMaskToTensor()
        self.mask_resize = ResizeNearestExact((size, size))

    def forward(self, img, masks=None):
        """Returns parameters of the resize in addition to the crop.

        Args:
            img (PIL Image or Tensor): Image to be cropped and resized.

        Returns:
            PIL Image or Tensor: Randomly cropped and resized image.
        """
        params = self.get_params(img, self.scale, self.ratio)
        img = transforms.functional.resized_crop(img, *params, self.size, self.interpolation)

        for mask_key, mask in masks.items():
            if not isinstance(mask, torch.Tensor):
                mask = self.mask_to_tensor(mask)
            mask = transforms.functional.crop(mask, *params)
            mask = self.mask_resize(mask)
            masks[mask_key] = mask
        return img, masks, params


class MultiCrop(object):
    def __init__(
        self,
        size: int = 224,
        input_key: str = "image",
        teacher_key: str = "teacher",
        student_key: str = "student",
        global_scale: Tuple[float] = (0.8, 1.0),
        local_scale: Tuple[float] = (0.7, 1.0),
        ratio: Tuple[float] = (3.0 / 4.0, 4.0 / 3.0),
        mask_keys: Optional[Tuple[str]] = None,
    ):
        self.ratio = ratio
        self.teacher_key = teacher_key
        self.student_key = student_key
        self.global_crop = JointRandomResizedCropwithParameters(size, global_scale, ratio)
        self.local_crop = JointRandomResizedCropwithParameters(size, local_scale, ratio)
        self.input_key = input_key
        self.mask_keys = tuple(mask_keys) if mask_keys is not None else tuple()

    def __call__(self, data):
        if self.input_key not in data:
            raise ValueError(f"Wrong input key {self.input_key}")
        img = transforms.functional.to_tensor(data[self.input_key])
        masks = {mask_key: data[mask_key] for mask_key in self.mask_keys}
        teacher_view, global_masks, params = self.global_crop(img, masks)
        data[self.teacher_key] = teacher_view
        for k, mask in global_masks.items():
            data[f"{self.teacher_key}_{k}"] = mask

        student_view, local_masks, params = self.local_crop(teacher_view, global_masks)
        data[self.student_key] = student_view
        for k, mask in local_masks.items():
            data[f"{self.student_key}_{k}"] = mask
        data["params"] = torch.Tensor(numpy.array(params))
        return data

ocl/preprocessing.py
amazon-science/object-centric-multiple-object-tracking
"""Scheduling of learning rate and hyperparameters."""
import abc
import math
import warnings
from typing import Callable

from torch.optim.lr_scheduler import _LRScheduler


def exp_decay_with_warmup_fn(
    step: int, decay_rate: float, decay_steps: int, warmup_steps: int
) -> float:
    """Decay function for exponential decay with learning rate warmup.

    Maps the step to a factor for rescaling the learning rate.
    """
    if warmup_steps:
        factor = min(1.0, step / warmup_steps)
    else:
        factor = 1.0

    return factor * (decay_rate ** (step / decay_steps))


class CosineAnnealingWithWarmup(_LRScheduler):
    """Cosine annealing with warmup."""

    def __init__(
        self,
        optimizer,
        T_max: int,
        warmup_steps: int = 0,
        eta_min: float = 0.0,
        last_epoch: int = -1,
        error_on_exceeding_steps: bool = True,
        verbose: bool = False,
    ):
        self.T_max = T_max
        self.warmup_steps = warmup_steps
        self.eta_min = eta_min
        self.error_on_exceeding_steps = error_on_exceeding_steps
        super().__init__(optimizer, last_epoch, verbose)

    def _linear_lr_warmup(self, base_lr, step_num):
        return base_lr * ((step_num + 0.5) / self.warmup_steps)

    def _cosine_annealing(self, base_lr, step_num):
        fraction_of_steps = (step_num - self.warmup_steps) / (self.T_max - self.warmup_steps - 1)
        return self.eta_min + 1 / 2 * (base_lr - self.eta_min) * (
            1 + math.cos(math.pi * fraction_of_steps)
        )

    def get_lr(self):
        if not self._get_lr_called_within_step:
            warnings.warn(
                "To get the last learning rate computed by the scheduler, "
                "please use `get_last_lr()`."
            )
        step_num = self.last_epoch
        self.T_max = 400005

        if step_num < self.warmup_steps:
            # Warmup.
            return [self._linear_lr_warmup(base_lr, step_num) for base_lr in self.base_lrs]
        elif step_num < self.T_max:
            # Cosine annealing.
            return [self._cosine_annealing(base_lr, step_num) for base_lr in self.base_lrs]
        else:
            if self.error_on_exceeding_steps:
                raise ValueError(
                    "Tried to step {} times. The specified number of total steps is {}".format(
                        step_num + 1, self.T_max
                    )
                )
            else:
                return [self.eta_min for _ in self.base_lrs]


HPSchedulerT = Callable[[int], float]  # Type for function signatures.


class HPScheduler(metaclass=abc.ABCMeta):
    """Base class for scheduling of scalar hyperparameters based on the number of training steps."""

    @abc.abstractmethod
    def __call__(self, step: int) -> float:
        """Return current value of hyperparameter based on global step."""
        pass


class LinearHPScheduler(HPScheduler):
    def __init__(
        self, end_value: float, end_step: int, start_value: float = 0.0, start_step: int = 0
    ):
        super().__init__()
        if start_step > end_step:
            raise ValueError("`start_step` needs to be smaller equal to `end_step`.")

        self.start_value = start_value
        self.end_value = end_value
        self.start_step = start_step
        self.end_step = end_step

    def __call__(self, step: int) -> float:
        if step < self.start_step:
            return self.start_value
        elif step > self.end_step:
            return self.end_value
        else:
            t = step - self.start_step
            T = self.end_step - self.start_step
            return self.start_value + t * (self.end_value - self.start_value) / T


class StepHPScheduler(HPScheduler):
    def __init__(self, end_value: float, switch_step: int, start_value: float = 0.0):
        super().__init__()
        self.start_value = start_value
        self.end_value = end_value
        self.switch_step = switch_step

    def __call__(self, step: int) -> float:
        if step < self.switch_step:
            return self.start_value
        elif step >= self.switch_step:
            return self.end_value

ocl/scheduling.py
amazon-science/object-centric-multiple-object-tracking
"""Utilities for working with our own version of PyTrees which focus on torch tensors.

PyTrees are any nested structure of dictionaries, lists, tuples, namedtuples or dataclasses.
"""
import copy
import dataclasses
from collections import OrderedDict, abc
from typing import Any, Callable, Dict, List, Mapping, Sequence, Tuple, Union

import torch

Tree = Union[Dict, List, Tuple]


def is_tensor_or_module(t: Any):
    """Check if input is a torch.Tensor or a torch.nn.Module."""
    return isinstance(t, (torch.Tensor, torch.nn.Module, dict))  # update


def is_namedtuple(obj) -> bool:
    """Check if input is a named tuple."""
    return isinstance(obj, tuple) and hasattr(obj, "_asdict") and hasattr(obj, "_fields")


def get_tree_element(d: Tree, path: List[str]):
    """Get element of a tree."""
    next_element = d

    for next_element_name in path:
        if isinstance(next_element, abc.Mapping) and next_element_name in next_element:
            next_element = next_element[next_element_name]
        elif hasattr(next_element, next_element_name):
            next_element = getattr(next_element, next_element_name)
        elif isinstance(next_element, (list, tuple)) and next_element_name.isnumeric():
            next_element = next_element[int(next_element_name)]
        else:
            msg = f"Trying to access path {'.'.join(path)}, "
            if isinstance(next_element, abc.Mapping):
                msg += f"but element {next_element_name} is not among keys {next_element.keys()}"
            elif isinstance(next_element, (list, tuple)):
                msg += f"but cannot index into list with {next_element_name}"
            else:

                msg += (
                    f"but element {next_element_name} cannot be used to access attribute of object "
                    f"of type {type(next_element)}"
                )
            raise ValueError(msg)
    return next_element


def _build_walk_path(previous_element, new_element):
    return previous_element + [new_element]


def walk_tree_with_paths(next_element, path=None, instance_check=is_tensor_or_module):
    """Walk over all tensors + modules and their paths in a nested structure.

    This could lead to an infinite loop.
    """
    if path is None:
        path = []

    if instance_check(next_element):
        yield path, next_element
    elif isinstance(next_element, str):
        # Special handling for strings, as even a single element slice is a sequence. This leads to
        # infinite nesting.
        pass
    elif isinstance(next_element, (dict, Mapping)):
        for key, value in next_element.items():
            yield from walk_tree_with_paths(
                value, path=_build_walk_path(path, key), instance_check=instance_check
            )
    elif dataclasses.is_dataclass(next_element):
        for field in dataclasses.fields(next_element):
            yield from walk_tree_with_paths(
                getattr(next_element, field.name),
                path=_build_walk_path(path, field.name),
                instance_check=instance_check,
            )
    elif is_namedtuple(next_element):
        for field_name in next_element._fields:
            yield from walk_tree_with_paths(
                getattr(next_element, field_name),
                path=_build_walk_path(path, field_name),
                instance_check=instance_check,
            )
    elif isinstance(next_element, (List, Sequence, tuple)):
        for index, el in enumerate(next_element):
            yield from walk_tree_with_paths(
                el, path=_build_walk_path(path, index), instance_check=instance_check
            )


def reduce_tree(outputs: List[Dict[str, Any]], fn: Callable[[List[torch.Tensor]], torch.Tensor]):
    """Apply reduction function to a list of nested dicts.

    This only considers tensors at the moment, for other data types are simply copied from the first
    element.
    """
    id_to_reduced_tensor = {}
    for path, tensor in walk_tree_with_paths(outputs[0]):
        stacked_tensor = fn([tensor] + [get_tree_element(output, path) for output in outputs[1:]])
        id_to_reduced_tensor[id(tensor)] = stacked_tensor

    # Replace all tensors with their stacked versions.
    return copy.deepcopy(outputs[0], memo=id_to_reduced_tensor)


def split_tree(d: Tree, split_paths: List[List[str]], split_axis: int, chunk_size: int):
    # We essentially need a deep copy of the input dict that we then update with splitted
    # references. To avoid copies of tensors and thus memory duplication we want to use shallow
    # copies for tensors instead. We do this by defining the memo parameter used in deepcopy for
    # all tensors in the dict. This way deepcopy thinks that these where already copied and uses
    # the provided objects instead. We can further use this trick to replace the original
    # tensors with splitted counterparts when running deepcopy.

    # Create memo containing all tensors to avoid data duplication.
    memo = {id(tensor): tensor for path, tensor in walk_tree_with_paths(d)}

    # Gather tensors that should be replaced and note their id.
    tensors_to_split = [get_tree_element(d, path) for path in split_paths]
    splitted_memos = OrderedDict(
        (id(tensor), torch.split(tensor, chunk_size, dim=split_axis)) for tensor in tensors_to_split
    )

    for tensor_slices in zip(*splitted_memos.values()):
        # Replace entires in memo dict with splitted counterparts.
        if chunk_size == 1:
            # Additionally squeeze the input.
            memo_override = {
                orig_id: tensor_slice.squeeze(split_axis)
                for orig_id, tensor_slice in zip(splitted_memos.keys(), tensor_slices)
            }
        else:
            memo_override = {
                orig_id: tensor_slice
                for orig_id, tensor_slice in zip(splitted_memos.keys(), tensor_slices)
            }
        yield copy.deepcopy(d, {**memo, **memo_override})

ocl/tree_utils.py
amazon-science/object-centric-multiple-object-tracking
"""Utility modules and functions."""
from __future__ import annotations

import dataclasses
import functools
import inspect
import math
from collections import defaultdict
from typing import Any, Dict, List, Mapping, Optional, Sequence, Tuple, Union

import torch
from torch import nn

from ocl import tree_utils
from torchvision.ops.boxes import box_area

class SoftPositionEmbed(nn.Module):
    """Embeding of positions using convex combination of learnable tensors.

    This assumes that the input positions are between 0 and 1.
    """

    def __init__(
        self, n_spatial_dims: int, feature_dim: int, cnn_channel_order=False, savi_style=False
    ):
        """__init__.

        Args:
            n_spatial_dims (int): Number of spatial dimensions.
            feature_dim (int): Dimensionality of the input features.
            cnn_channel_order (bool): Assume features are in CNN channel order (i.e. C x H x W).
            savi_style (bool): Use savi style positional encoding, where positions are normalized
                between -1 and 1 and a single dense layer is used for embedding.
        """
        super().__init__()
        self.savi_style = savi_style
        n_features = n_spatial_dims if savi_style else 2 * n_spatial_dims
        self.dense = nn.Linear(in_features=n_features, out_features=feature_dim)
        self.cnn_channel_order = cnn_channel_order

    def forward(self, inputs: torch.Tensor, positions: torch.Tensor):
        if self.savi_style:
            # Rescale positional encoding to -1 to 1
            positions = (positions - 0.5) * 2
        else:
            positions = torch.cat([positions, 1 - positions], axis=-1)
        emb_proj = self.dense(positions)
        if self.cnn_channel_order:
            emb_proj = emb_proj.permute(*range(inputs.ndim - 3), -1, -3, -2)
        return inputs + emb_proj


class DummyPositionEmbed(nn.Module):
    """Embedding that just passes through inputs without adding any positional embeddings."""

    def __init__(self):
        super().__init__()

    def forward(self, inputs: torch.Tensor, positions: torch.Tensor):
        return inputs


def tensor_to_one_hot(tensor: torch.Tensor, dim: int) -> torch.Tensor:
    """Convert tensor to one-hot encoding by using maximum across dimension as one-hot element."""
    assert 0 <= dim
    max_idxs = torch.argmax(tensor, dim=dim, keepdim=True)
    shape = [1] * dim + [-1] + [1] * (tensor.ndim - dim - 1)
    one_hot = max_idxs == torch.arange(tensor.shape[dim], device=tensor.device).view(*shape)
    return one_hot.to(torch.long)


def resize_patches_to_image(
    patches: torch.Tensor,
    size: Optional[int] = None,
    scale_factor: Optional[float] = None,
    resize_mode: str = "bilinear",
) -> torch.Tensor:
    """Convert and resize a tensor of patches to image shape.

    This method requires that the patches can be converted to a square image.

    Args:
        patches: Patches to be converted of shape (..., C, P), where C is the number of channels and
            P the number of patches.
        size: Image size to resize to.
        scale_factor: Scale factor by which to resize the patches. Can be specified alternatively to
            `size`.
        resize_mode: Method to resize with. Valid options are "nearest", "nearest-exact", "bilinear",
            "bicubic".

    Returns: Tensor of shape (..., C, S, S) where S is the image size.
    """
    has_size = size is None
    has_scale = scale_factor is None
    if has_size == has_scale:
        raise ValueError("Exactly one of `size` or `scale_factor` must be specified.")

    n_channels = patches.shape[-2]
    n_patches = patches.shape[-1]
    patch_size_float = math.sqrt(n_patches)
    patch_size = int(math.sqrt(n_patches))
    if patch_size_float != patch_size:
        raise ValueError("The number of patches needs to be a perfect square.")

    image = torch.nn.functional.interpolate(
        patches.view(-1, n_channels, patch_size, patch_size),
        size=size,
        scale_factor=scale_factor,
        mode=resize_mode,
    )

    return image.view(*patches.shape[:-1], image.shape[-2], image.shape[-1])


def flatten_dict_of_dataclasses(input_dict: Dict[str, Any]) -> Dict[str, Any]:
    """Flatten a dictionary containing dataclasses as values.

    Fields of dataclass are stored with the key "{dataclass_key}.{field_name}" in the flattened dict.
    Entries that are no dataclass are stored with the same key in the flattened dict. Flattening is
    not recursive, only the first level of nesting is considered. Then the first level is a list it
    is also flattened in a similar fashion.

    Example:
    ```
    @dataclasses.dataclass
    class MyDataclass:
        value: int

    dictionary = {"item1": MyDataclass(value=5), "another_item": "no_dataclass"}
    assert _flatten_dict_of_dataclasses(dictionary) == {"item1.value": 5, "item2": "no_dataclass"}
    ```
    """
    flattened_dict = {}
    for key, value in input_dict.items():
        if dataclasses.is_dataclass(value):
            for field in dataclasses.fields(value):
                flattened_dict[f"{key}.{field.name}"] = getattr(value, field.name)
        elif isinstance(value, Sequence):
            sequence_output = defaultdict(list)
            for element in value:
                if dataclasses.is_dataclass(element):
                    for field in dataclasses.fields(element):
                        sequence_output[f"{key}.{field.name}"].append(getattr(element, field.name))
                else:
                    sequence_output[key].append(value)

            for nest_key, nest_value in sequence_output.items():
                if isinstance(nest_value[0], torch.Tensor):
                    flattened_dict[nest_key] = torch.stack(nest_value, dim=1)
                else:
                    flattened_dict[nest_key] = nest_value
        else:
            flattened_dict[key] = value

    return flattened_dict


class Residual(nn.Module):
    def __init__(self, module: nn.Module):
        super().__init__()
        self.module = module

    def forward(self, inputs):
        return inputs + self.module(inputs)


class RoutableMixin:
    """Mixin class that allows to connect any element of a (nested) dict with a module input."""

    def __init__(self, input_mapping: Mapping[str, Optional[str]]):
        self.input_mapping = {
            key: value.split(".") for key, value in input_mapping.items() if value is not None
        }

    def _route(method, filter_parameters=True):
        """Pass arguments to a function based on the mapping defined in `self.input_mapping`.

        This method supports both filtering for parameters that match the arguments of the wrapped
        method and passing all arguments defined in `input_mapping`.  If a non-optional argument is
        missing this will raise an exception.  Additional arguments can also be passed to the method
        to override entries in the input dict.  Non-keyword arguments are always directly passed to
        the method.

        Args:
            method: The method to pass the arguments to.
            filter_parameters: Only pass arguments to wrapped method that match the methods
                signature.  This is practical if different methods require different types of input.

        """
        # Run inspection here to reduce compute time when calling method.
        signature = inspect.signature(method)
        valid_parameters = list(signature.parameters)  # Returns the parameter names.
        valid_parameters = valid_parameters[1:]  # Discard "self".
        # Keep track of default parameters. For these we should not fail if they are not in
        # the input dict.
        with_defaults = [
            name
            for name, param in signature.parameters.items()
            if param.default is not inspect.Parameter.empty
        ]

        @functools.wraps(method)
        def method_with_routing(self, *args, inputs=None, **kwargs):
            if not inputs:
                inputs = {}
            if self.input_mapping:
                if not inputs:  # Empty dict.
                    inputs = kwargs

                routed_inputs = {}
                for input_field, input_path in self.input_mapping.items():
                    if filter_parameters and input_field not in valid_parameters:
                        # Skip parameters that are not the function signature.
                        continue
                    if input_field in kwargs.keys():
                        # Skip parameters that are directly provided as kwargs.
                        continue
                    try:
                        element = tree_utils.get_tree_element(inputs, input_path)
                        routed_inputs[input_field] = element
                    except ValueError as e:
                        if input_field in with_defaults:
                            continue
                        else:
                            raise e
                # Support for additional parameters passed via keyword arguments.
                # TODO(hornmax): This is not ideal as it mixes routing args from the input dict
                # and explicitly passed kwargs and thus could lead to collisions.
                for name, element in kwargs.items():
                    if filter_parameters and name not in valid_parameters:
                        continue
                    else:
                        routed_inputs[name] = element
                return method(self, *args, **routed_inputs)
            else:
                return method(self, *args, **kwargs)

        return method_with_routing

    # This is needed in order to allow the decorator to be used in child classes. The documentation
    # looks a bit hacky but I didn't find an alternative approach on how to do it.
    route = staticmethod(functools.partial(_route, filter_parameters=True))
    route.__doc__ = (
        """Route input arguments according to input_mapping and filter non-matching arguments."""
    )
    route_unfiltered = staticmethod(functools.partial(_route, filter_parameters=False))
    route_unfiltered.__doc__ = """Route all input arguments according to input_mapping."""


class DataRouter(nn.Module, RoutableMixin):
    """Data router for modules that don't support the RoutableMixin.

    This allows the usage of modules without RoutableMixin support in the dynamic information flow
    pattern of the code.
    """

    def __init__(self, module: nn.Module, input_mapping: Mapping[str, str]):
        nn.Module.__init__(self)
        RoutableMixin.__init__(self, input_mapping)
        self.module = module
        self._cached_valid_parameters = None

    @RoutableMixin.route_unfiltered
    def forward(self, *args, **kwargs):
        # We need to filter parameters at runtime as we cannot know them prior to initialization.
        if not self._cached_valid_parameters:
            try:
                signature = inspect.signature(self.module.forward)
            except AttributeError:
                if callable(self.module):
                    signature = inspect.signature(self.module.__call__)
                else:
                    signature = inspect.signature(self.module)

            self._cached_valid_parameters = list(signature.parameters)

        kwargs = {
            name: param for name, param in kwargs.items() if name in self._cached_valid_parameters
        }
        return self.module(*args, **kwargs)


class Combined(nn.ModuleDict):
    """Module to combine multiple modules and store their outputs.

    A combined module groups together multiple model components and allows them to access any
    information that was returned in processing steps prior to their own application.

    It functions similarly to `nn.ModuleDict` yet for modules of type `RoutableMixin` and
    additionally implements a forward routine which will return a dict of the outputs of the
    submodules.

    """

    def __init__(self, modules: Dict[str, Union[RoutableMixin, Combined, Recurrent]]):
        super().__init__(modules)

    def forward(self, inputs: Dict[str, Any]):
        # The combined module does not know where it is positioned and thus also does not know in
        # which sub-path results should be written. As we want different modules of a combined
        # module to be able access previous outputs using their global path in the dictionary, we
        # need to somehow keep track of the nesting level and then directly write results into the
        # input dict at the right path.  The prefix variable keeps track of the nesting level.
        prefix: List[str]
        if "prefix" in inputs.keys():
            prefix = inputs["prefix"]
        else:
            prefix = []
            inputs["prefix"] = prefix

        outputs = tree_utils.get_tree_element(inputs, prefix)
        for name, module in self.items():
            # Update prefix state such that nested calls of combined return dict in the correct
            # location.
            prefix.append(name)
            outputs[name] = {}
            # If module is a Combined module, it will return the same dict as set above. If not the
            # dict will be overwritten with the output of the module.
            outputs[name] = module(inputs=inputs)
            # Remove last component of prefix after execution.
            prefix.pop()
        return outputs


class Recurrent(nn.Module):
    """Module to apply another module in a recurrent fashion over a axis.

    This module takes a set of input tensors and applies a module recurrent over them.  The output
    of the previous iteration is kept in the `previous_output` key of input dict and thus can be
    accessed using data routing. After applying the module to the input slices, the outputs are
    stacked along the same axis as the inputs where split.

    Args:
        module: The module that should be applied recurrently along input tensors.
        inputs_to_split: List of paths that should be split for recurrent application.
        initial_input_mapping: Mapping that constructs the first `previous_output` element.  If
            `previous_output` should just be a tensor, use a mapping of the format
            `{"": "input_path"}`.
        split_axis: Axis along which to split the tensors defined by inputs_to_split.
        chunk_size: The size of each slice, when set to 1, the slice dimension is squeezed prior to
            passing to the module.

    """

    def __init__(
        self,
        module,
        inputs_to_split: List[str],
        initial_input_mapping: Dict[str, str],
        split_axis: int = 1,
        chunk_size: int = 1,
    ):
        super().__init__()
        self.module = module
        self.inputs_to_split = [path.split(".") for path in inputs_to_split]
        self.initial_input_mapping = {
            output: input.split(".") for output, input in initial_input_mapping.items()
        }
        self.split_axis = split_axis
        self.chunk_size = chunk_size

    def _build_initial_dict(self, inputs):
        # This allows us to bing the initial input and previous_output into a similar format.
        output_dict = {}
        for output_path, input_path in self.initial_input_mapping.items():
            source = tree_utils.get_tree_element(inputs, input_path)
            if output_path == "":
                # Just the object itself, no dict nesting.
                return source

            output_path = output_path.split(".")
            cur_search = output_dict
            for path_part in output_path[:-1]:
                # Iterate along path and create nodes that do not exist yet.
                try:
                    # Get element prior to last.
                    cur_search = tree_utils.get_tree_element(cur_search, [path_part])
                except ValueError:
                    # Element does not yet exist.
                    cur_search[path_part] = {}
                    cur_search = cur_search[path_part]

            cur_search[output_path[-1]] = source
        return output_dict

    def forward(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
        # TODO: Come up with a better way of handling the initial input without putting restrictions
        # on modules being run recurrently.
        outputs = [self._build_initial_dict(inputs)]
        for split_dict in tree_utils.split_tree(
            inputs, self.inputs_to_split, self.split_axis, self.chunk_size
        ):
            split_dict["previous_output"] = outputs[-1]
            outputs.append(self.module(inputs=split_dict))

        # TODO: When chunk size is larger than 1 then this should be cat and not stack. Otherwise an
        # additional axis would be added. Evtl. this should be configurable.
        stack_fn = functools.partial(torch.stack, dim=self.split_axis)
        # Ignore initial input.
        return tree_utils.reduce_tree(outputs[1:], stack_fn)


class Sequential(nn.Module):
    """Extended sequential module that supports multiple inputs and outputs to layers.

    This allows a stack of layers where for example the first layer takes two inputs and only has
    a single output or where a layer has multiple outputs and the downstream layer takes multiple
    inputs.
    """

    def __init__(self, *layers):
        super().__init__()
        self.layers = nn.ModuleList(layers)

    def forward(self, *inputs):
        outputs = inputs
        for layer in self.layers:
            if isinstance(outputs, (tuple, list)):
                outputs = layer(*outputs)
            else:
                outputs = layer(outputs)
        return outputs


class CreateSlotMask(nn.Module, RoutableMixin):
    """Module intended to create a mask that marks empty slots.

    Module takes a tensor holding the number of slots per batch entry, and returns a binary mask of
    shape (batch_size, max_slots) where entries exceeding the number of slots are masked out.
    """

    def __init__(self, max_slots: int, n_slots_path: str):
        nn.Module.__init__(self)
        RoutableMixin.__init__(self, {"n_slots": n_slots_path})
        self.max_slots = max_slots

    @RoutableMixin.route
    def forward(self, n_slots: torch.Tensor) -> torch.Tensor:
        (batch_size,) = n_slots.shape

        # Create mask of shape B x K where the first n_slots entries per-row are false, the rest true
        indices = torch.arange(self.max_slots, device=n_slots.device)
        masks = indices.unsqueeze(0).expand(batch_size, -1) >= n_slots.unsqueeze(1)

        return masks

def box_cxcywh_to_xyxy(x):
    x_c, y_c, w, h = x.unbind(-1)
    b = [(x_c - 0.5 * w), (y_c - 0.5 * h), (x_c + 0.5 * w), (y_c + 0.5 * h)]
    return torch.stack(b, dim=-1)


def box_xyxy_to_cxcywh(x):
    x0, y0, x1, y1 = x.unbind(-1)
    b = [(x0 + x1) / 2, (y0 + y1) / 2, (x1 - x0), (y1 - y0)]
    return torch.stack(b, dim=-1)


# modified from torchvision to also return the union
def box_iou(boxes1, boxes2):
    area1 = box_area(boxes1)
    area2 = box_area(boxes2)

    lt = torch.max(boxes1[:, None, :2], boxes2[:, :2])  # [N,M,2]
    rb = torch.min(boxes1[:, None, 2:], boxes2[:, 2:])  # [N,M,2]

    wh = (rb - lt).clamp(min=0)  # [N,M,2]
    inter = wh[:, :, 0] * wh[:, :, 1]  # [N,M]

    union = area1[:, None] + area2 - inter

    iou = inter / union
    return iou, union


class Resize(nn.Module, RoutableMixin):
    """Module resizing tensors."""

    MODES = {"nearest", "linear", "bilinear", "bicubic", "trilinear", "area", "nearest-exact"}

    def __init__(
        self,
        input_path: str,
        size: Optional[Union[int, Tuple[int, int]]] = None,
        take_size_from: Optional[str] = None,
        resize_mode: str = "bilinear",
        patch_mode: bool = False,
        channels_last: bool = False,
    ):
        nn.Module.__init__(self)
        RoutableMixin.__init__(self, {"tensor": input_path, "size_tensor": take_size_from})

        if size is not None and take_size_from is not None:
            raise ValueError("`size` and `take_size_from` can not be set at the same time")
        self.size = size

        if resize_mode not in Resize.MODES:
            raise ValueError(f"`mode` must be one of {Resize.MODES}")
        self.resize_mode = resize_mode
        self.patch_mode = patch_mode
        self.channels_last = channels_last
        self.expected_dims = 3 if patch_mode else 4

    @RoutableMixin.route
    def forward(
        self, tensor: torch.Tensor, size_tensor: Optional[torch.Tensor] = None
    ) -> torch.Tensor:
        """Resize tensor.

        Args:
            tensor: Tensor to resize. If `patch_mode=False`, assumed to be of shape (..., C, H, W).
                If `patch_mode=True`, assumed to be of shape (..., C, P), where P is the number of
                patches. Patches are assumed to be viewable as a perfect square image. If
                `channels_last=True`, channel dimension is assumed to be the last dimension instead.
            size_tensor: Tensor which size to resize to. If tensor has <=2 dimensions and the last
                dimension of this tensor has length 2, the two entries are taken as height and width.
                Otherwise, the size of the last two dimensions of this tensor are used as height
                and width.

        Returns: Tensor of shape (..., C, H, W), where height and width are either specified by
            `size` or `size_tensor`.
        """
        dims_to_flatten = tensor.ndim - self.expected_dims
        if dims_to_flatten > 0:
            flattened_dims = tensor.shape[: dims_to_flatten + 1]
            tensor = tensor.flatten(0, dims_to_flatten)
        elif dims_to_flatten < 0:
            raise ValueError(
                f"Tensor needs at least {self.expected_dims} dimensions, but only has {tensor.ndim}"
            )

        if self.patch_mode:
            if self.channels_last:
                tensor = tensor.transpose(-2, -1)
            n_channels, n_patches = tensor.shape[-2:]
            patch_size_float = math.sqrt(n_patches)
            patch_size = int(math.sqrt(n_patches))
            if patch_size_float != patch_size:
                raise ValueError(
                    f"The number of patches needs to be a perfect square, but is {n_patches}."
                )
            tensor = tensor.view(-1, n_channels, patch_size, patch_size)
        else:
            if self.channels_last:
                tensor = tensor.permute(0, 3, 1, 2)

        if self.size is None:
            if size_tensor is None:
                raise ValueError("`size` is `None` but no `size_tensor` was passed.")
            if size_tensor.ndim <= 2 and size_tensor.shape[-1] == 2:
                height, width = size_tensor.unbind(-1)
                height = torch.atleast_1d(height)[0].squeeze().detach().cpu()
                width = torch.atleast_1d(width)[0].squeeze().detach().cpu()
                size = (int(height), int(width))
            else:
                size = size_tensor.shape[-2:]
        else:
            size = self.size

        tensor = torch.nn.functional.interpolate(
            tensor,
            size=size,
            mode=self.resize_mode,
        )

        if dims_to_flatten > 0:
            tensor = tensor.unflatten(0, flattened_dims)

        return tensor

def lecun_uniform_(tensor, gain=1.):
    fan_in, fan_out = nn.init._calculate_fan_in_and_fan_out(tensor)
    var = gain / float(fan_in)
    a = math.sqrt(3 * var)
    return nn.init._no_grad_uniform_(tensor, -a, a)


def lecun_normal_(tensor, gain=1., mode="fan_in"):
    fan_in, fan_out = nn.init._calculate_fan_in_and_fan_out(tensor)
    if mode == "fan_in":
        scale_mode = fan_in
    elif mode == "fan_out":
        scale_mode = fan_out
    else:
        raise NotImplementedError
    var = gain / float(scale_mode)
    # constant is stddev of standard normal truncated to (-2, 2)
    std = math.sqrt(var) / .87962566103423978
    # return nn.init._no_grad_normal_(tensor, 0., std)
    kernel = torch.nn.init._no_grad_trunc_normal_(tensor, 0, 1, -2, 2) * std
    with torch.no_grad():
        tensor[:] = kernel[:]
    return tensor


def lecun_normal_fan_out_(tensor, gain=1.):
    return lecun_normal_(tensor, gain=gain, mode="fan_out")

def lecun_normal_convtranspose_(tensor, gain=1.):
    # for some reason, the convtranspose weights are [in_channels, out_channels, kernel, kernel]
    # but the _calculate_fan_in_and_fan_out treats dim 1 as fan_in and dim 0 as fan_out.
    # so, for convolution weights, have to use fan_out instead of fan_in
    # which is actually using fan_in instead of fan_out
    return lecun_normal_fan_out_(tensor, gain=gain)


init_fn = {
    'xavier_uniform': nn.init.xavier_uniform_,
    'xavier_normal': nn.init.xavier_normal_,
    'kaiming_uniform': nn.init.kaiming_uniform_,
    'kaiming_normal': nn.init.kaiming_normal_,
    'lecun_uniform': lecun_uniform_,
    'lecun_normal': lecun_normal_,
    'lecun_normal_fan_out': lecun_normal_fan_out_,
    'ones': nn.init.ones_,
    'zeros': nn.init.zeros_,
    'default': lambda x: x}


ocl/utils.py
amazon-science/object-centric-multiple-object-tracking
"""Classes for handling different types of visualizations."""
import dataclasses
from typing import Any, List, Optional, Union

import matplotlib.pyplot
import torch
from torch.utils.tensorboard import SummaryWriter
from torchtyping import TensorType


def dataclass_to_dict(d):
    return {field.name: getattr(d, field.name) for field in dataclasses.fields(d)}


@dataclasses.dataclass
class Visualization:
    def add_to_experiment(self, experiment: SummaryWriter, tag: str, global_step: int):
        pass


@dataclasses.dataclass
class Figure(Visualization):
    """Matplotlib figure."""

    figure: matplotlib.pyplot.figure
    close: bool = True

    def add_to_experiment(self, experiment: SummaryWriter, tag: str, global_step: int):
        experiment.add_figure(**dataclass_to_dict(self), tag=tag, global_step=global_step)


@dataclasses.dataclass
class Image(Visualization):
    """Single image."""

    img_tensor: torch.Tensor
    dataformats: str = "CHW"

    def add_to_experiment(self, experiment: SummaryWriter, tag: str, global_step: int):
        experiment.add_image(**dataclass_to_dict(self), tag=tag, global_step=global_step)


@dataclasses.dataclass
class Images(Visualization):
    """Batch of images."""

    img_tensor: torch.Tensor
    dataformats: str = "NCHW"

    def add_to_experiment(self, experiment: SummaryWriter, tag: str, global_step: int):
        experiment.add_images(**dataclass_to_dict(self), tag=tag, global_step=global_step)


@dataclasses.dataclass
class Video(Visualization):
    """Batch of videos."""

    vid_tensor: TensorType["batch_size", "frames", "channels", "height", "width"]  # noqa: F821
    fps: Union[int, float] = 4

    def add_to_experiment(self, experiment: SummaryWriter, tag: str, global_step: int):
        experiment.add_video(**dataclass_to_dict(self), tag=tag, global_step=global_step)


class Embedding(Visualization):
    """Batch of embeddings."""

    mat: TensorType["batch_size", "feature_dim"]  # noqa: F821
    metadata: Optional[List[Any]] = None
    label_img: Optional[TensorType["batch_size", "channels", "height", "width"]] = None  # noqa: F821
    metadata_header: Optional[List[str]] = None

    def add_to_experiment(self, experiment: SummaryWriter, tag: str, global_step: int):
        experiment.add_embedding(**dataclass_to_dict(self), tag=tag, global_step=global_step)

ocl/visualization_types.py
amazon-science/object-centric-multiple-object-tracking
from typing import Callable, Dict, List, Optional, Tuple

import torch
from torchvision import transforms
from torchvision.utils import draw_segmentation_masks, make_grid, draw_bounding_boxes
import numpy as np
from ocl import consistency, visualization_types
from ocl.utils import RoutableMixin, box_cxcywh_to_xyxy
from ocl.metrics import masks_to_bboxes_xyxy
from torchvision.ops import masks_to_boxes

def _nop(arg):
    return arg


class Image(RoutableMixin):
    def __init__(
        self,
        n_instances: int = 8,
        n_row: int = 8,
        denormalization: Optional[Callable[[torch.Tensor], torch.Tensor]] = None,
        as_grid: bool = True,
        image_path: Optional[str] = None,
    ):
        super().__init__({"image": image_path})
        self.n_instances = n_instances
        self.n_row = n_row
        self.denormalization = denormalization if denormalization else _nop
        self.as_grid = as_grid

    @RoutableMixin.route
    def __call__(self, image: torch.Tensor):
        image = self.denormalization(image[: self.n_instances].cpu())
        if self.as_grid:
            return visualization_types.Image(make_grid(image, nrow=self.n_row))
        else:
            return visualization_types.Images(image)


class Video(RoutableMixin):
    def __init__(
        self,
        n_instances: int = 8,
        n_row: int = 8,
        denormalization: Optional[Callable[[torch.Tensor], torch.Tensor]] = None,
        as_grid: bool = True,
        video_path: Optional[str] = None,
    ):
        super().__init__({"video": video_path})
        self.n_instances = n_instances
        self.n_row = n_row
        self.denormalization = denormalization if denormalization else _nop
        self.as_grid = as_grid

    @RoutableMixin.route
    def __call__(self, video: torch.Tensor):
        video = video[: self.n_instances].cpu()
        if self.as_grid:
            video = torch.stack(
                [
                    make_grid(self.denormalization(frame.unsqueeze(1)).squeeze(1), nrow=self.n_row)
                    for frame in torch.unbind(video, 1)
                ],
                dim=0,
            ).unsqueeze(0)
        return visualization_types.Video(video)


class Mask(RoutableMixin):
    def __init__(
        self,
        n_instances: int = 8,
        mask_path: Optional[str] = None,
    ):
        super().__init__({"masks": mask_path})
        self.n_instances = n_instances

    @RoutableMixin.route
    def __call__(self, masks):
        masks = masks[: self.n_instances].cpu().contiguous()
        image_shape = masks.shape[-2:]
        n_objects = masks.shape[-3]

        if masks.dim() == 5:
            # Handling video data.
            # bs x frames x objects x H x W
            mask_vis = masks.transpose(1, 2).contiguous()
            flattened_masks = mask_vis.flatten(0, 1).unsqueeze(2)

            # Draw masks inverted as they are easier to print.
            mask_vis = torch.stack(
                [
                    make_grid(1.0 - masks, nrow=n_objects)
                    for masks in torch.unbind(flattened_masks, 1)
                ],
                dim=0,
            )
            mask_vis = mask_vis.unsqueeze(0)
            return visualization_types.Video(mask_vis)
        elif masks.dim() == 4:
            # Handling image data.
            # bs x objects x H x W
            # Monochrome image with single channel.
            masks = masks.view(-1, 1, *image_shape)
            # Draw masks inverted as they are easier to print.
            return visualization_types.Image(make_grid(1.0 - masks, nrow=n_objects))


class VisualObject(RoutableMixin):
    def __init__(
        self,
        n_instances: int = 8,
        denormalization: Optional[Callable[[torch.Tensor], torch.Tensor]] = None,
        object_path: Optional[str] = None,
        mask_path: Optional[str] = None,
    ):
        super().__init__({"object_reconstructions": object_path, "object_masks": mask_path})
        self.n_instances = n_instances
        self.denormalization = denormalization if denormalization else _nop

    @RoutableMixin.route
    def __call__(self, object_reconstructions, object_masks):
        objects = object_reconstructions[: self.n_instances].cpu()
        masks = object_masks[: self.n_instances].cpu().contiguous()
        masks = masks > 0.5

        image_shape = objects.shape[-3:]
        n_objects = objects.shape[-4]

        if objects.dim() == 6:
            # Handling video data.
            # bs x frames x objects x C x H x W

            # We need to denormalize prior to constructing the grid, yet the denormalization
            # method assumes video input. We thus convert a frame into a single frame video and
            # remove the additional dimension prior to make_grid.
            # Switch object and frame dimension.
            object_vis = objects.transpose(1, 2).contiguous()
            mask_vis = masks.transpose(1, 2).contiguous()
            flattened_masks = mask_vis.flatten(0, 1).unsqueeze(2).float()
            object_vis = self.denormalization(object_vis.flatten(0, 1))
            # Keep object pixels and apply white background to non-objects parts.
            object_vis = object_vis * flattened_masks + (1.0 - flattened_masks)

            # object_vis = object_vis * flattened_masks
            object_vis = torch.stack(
                [
                    make_grid(
                        object_vis_frame,
                        nrow=n_objects,
                    )
                    for object_vis_frame in torch.unbind(object_vis, 1)
                ],
                dim=0,
            )
            # Add batch dimension as this is required for video input.
            object_vis = object_vis.unsqueeze(0)

            # Draw masks inverted as they are easier to print.
            flattened_masks_ = flattened_masks > 0.5
            mask_vis = torch.stack(
                [
                    # make_grid(1.0 - masks, nrow=n_objects)
                    make_grid(masks, nrow=n_objects)
                    for masks in torch.unbind(flattened_masks_, 1)
                ],
                dim=0,
            )
            mask_vis = mask_vis.unsqueeze(0)
            return {
                "reconstruction": visualization_types.Video(object_vis),
                "mask": visualization_types.Video(mask_vis),
            }
        elif objects.dim() == 5:
            # Handling image data.
            # bs x objects x C x H x W
            object_reconstructions = self.denormalization(objects.view(-1, *image_shape))
            # Monochrome image with single channel.
            masks = masks.view(-1, 1, *image_shape[1:])
            # Save object reconstructions as RGBA image. make_grid does not support RGBA input, thus
            # we combine the channels later.  For the masks we need to pad with 1 as we want the
            # borders between images to remain visible (i.e. alpha value of 1.)
            masks_grid = make_grid(masks, nrow=n_objects, pad_value=1.0)
            object_grid = make_grid(object_reconstructions, nrow=n_objects)
            # masks_grid expands the image to three channels, which we don't need. Only keep one, and
            # use it as the alpha channel. After make_grid the tensor has the shape C X W x H.
            object_grid = torch.cat((object_grid, masks_grid[:1]), dim=0)

            return {
                "reconstruction": visualization_types.Image(object_grid),
                # Draw masks inverted as they are easier to print.
                "mask": visualization_types.Image(make_grid(1.0 - masks, nrow=n_objects)),
            }


class ConsistencyMask(RoutableMixin):
    def __init__(
        self,
        matcher: consistency.HungarianMatcher,
        mask_path: Optional[str] = None,
        mask_target_path: Optional[str] = None,
        params_path: Optional[str] = None,
    ):
        super().__init__(
            {"mask": mask_path, "mask_target": mask_target_path, "cropping_params": params_path}
        )
        self.matcher = matcher

    @RoutableMixin.route
    def __call__(self, mask: torch.Tensor, mask_target: torch.Tensor, cropping_params: torch.Tensor):
        _, _, size, _ = mask.shape
        mask_one_hot = self._to_binary_mask(mask)
        mask_target = self.crop_views(mask_target, cropping_params, size)
        mask_target_one_hot = self._to_binary_mask(mask_target)
        _ = self.matcher(mask_one_hot, mask_target_one_hot)
        return {
            "costs": visualization_types.Image(
                make_grid(-self.matcher.costs, nrow=8, pad_value=0.9)
            ),
        }

    @staticmethod
    def _to_binary_mask(masks: torch.Tensor):
        _, n_objects, _, _ = masks.shape
        m_lables = masks.argmax(dim=1)
        mask_one_hot = torch.nn.functional.one_hot(m_lables, n_objects)
        return mask_one_hot.permute(0, 3, 1, 2)

    def crop_views(self, view: torch.Tensor, param: torch.Tensor, size: int):
        return torch.cat([self.crop_maping(v, p, size) for v, p in zip(view, param)])

    @staticmethod
    def crop_maping(view: torch.Tensor, p: torch.Tensor, size: int):
        p = tuple(p.cpu().numpy().astype(int))
        return transforms.functional.resized_crop(view, *p, size=(size, size))[None]


class Segmentation(RoutableMixin):
    def __init__(
        self,
        n_instances: int = 8,
        denormalization: Optional[Callable[[torch.Tensor], torch.Tensor]] = None,
        image_path: Optional[str] = None,
        mask_path: Optional[str] = None,
    ):
        super().__init__({"image": image_path, "mask": mask_path})
        self.n_instances = n_instances
        self.denormalization = denormalization if denormalization else _nop
        self._cmap_cache: Dict[int, List[Tuple[int, int, int]]] = {}

    def _get_cmap(self, num_classes: int) -> List[Tuple[int, int, int]]:
        if num_classes in self._cmap_cache:
            return self._cmap_cache[num_classes]

        from matplotlib import cm

        if num_classes <= 20:
            mpl_cmap = cm.get_cmap("tab20", num_classes)(range(num_classes))
        else:
            mpl_cmap = cm.get_cmap("turbo", num_classes)(range(num_classes))

        cmap = [tuple((255 * cl[:3]).astype(int)) for cl in mpl_cmap]
        self._cmap_cache[num_classes] = cmap
        return cmap

    @RoutableMixin.route
    def __call__(
        self, image: torch.Tensor, mask: torch.Tensor
    ) -> Optional[visualization_types.Visualization]:
        image = image[: self.n_instances].cpu()
        mask = mask[: self.n_instances].cpu().contiguous()
        if image.dim() == 4:  # Only support image data at the moment.
            input_image = self.denormalization(image)
            n_objects = mask.shape[1]

            masks_argmax = mask.argmax(dim=1)[:, None]
            classes = torch.arange(n_objects)[None, :, None, None].to(masks_argmax)
            masks_one_hot = masks_argmax == classes

            cmap = self._get_cmap(n_objects)
            masks_on_image = torch.stack(
                [
                    draw_segmentation_masks(
                        (255 * img).to(torch.uint8), mask, alpha=0.75, colors=cmap
                    )
                    for img, mask in zip(input_image.to("cpu"), masks_one_hot.to("cpu"))
                ]
            )

            return visualization_types.Image(make_grid(masks_on_image, nrow=8))
        return None

class ObjectMOT(RoutableMixin):
    def __init__(
        self,
        n_clips: int = 3,
        n_row: int = 8,
        denormalization: Optional[Callable[[torch.Tensor], torch.Tensor]] = None,
        video_path: Optional[str] = None,
        mask_path: Optional[str] = None,
        pred_is_box: bool = False,
    ):
        super().__init__({"video": video_path, "object_masks": mask_path})
        self.n_clips = n_clips
        self.n_row = n_row
        self.denormalization = denormalization if denormalization else _nop
        self.pred_is_box = pred_is_box

    def generate_color_list(self, track_num):
        import random
        color_list = []
        for i in range(track_num):
            hexadecimal = ["#" + ''.join([random.choice('ABCDEF0123456789') for i in range(6)])]
            color_list.append(hexadecimal[0])
        return color_list


    @RoutableMixin.route
    def __call__(
        self,
        video: torch.Tensor,
        object_masks: torch.Tensor,
    ) -> Optional[visualization_types.Visualization]:
        video = video[: self.n_clips].cpu()
        num_frames = video.shape[1]

        if not self.pred_is_box:
            masks = object_masks[: self.n_clips].cpu().contiguous()
            B, F, C, h, w = masks.shape  # [5, 6, 11, 64, 64]
            masks = masks.flatten(0, 1)
            masks = masks > 0.7
            bbox = masks_to_bboxes_xyxy(masks.flatten(0, 1)).unflatten(0, (B, F, C))
        else:
            bbox = object_masks[: self.n_clips].cpu().contiguous()
            bbox[:,:,:,2] += bbox[:,:,:,0]
            bbox[:, :, :, 3] += bbox[:, :, :, 1]

        rendered_video = torch.zeros_like(video)

        n_colors = 500
        color_list = self.generate_color_list(n_colors)

        for cidx in range(self.n_clips):
            for fidx in range(num_frames):
                cur_obj_box = bbox[cidx, fidx][:, 0] != -1.0
                cur_obj_idx = cur_obj_box.nonzero()[:, 0].detach().cpu().numpy()
                idx = cur_obj_idx.tolist()

                cur_obj_idx = np.array(idx)
                cur_color_list = [color_list[obj_idx] for obj_idx in idx]
                frame = (video[cidx, fidx] * 256).to(torch.uint8)
                frame = draw_bounding_boxes(
                    frame, bbox[cidx, fidx][cur_obj_idx], colors=cur_color_list
                )
                rendered_video[cidx, fidx] = frame

        rendered_video = (
            torch.stack(
                [
                    make_grid(self.denormalization(frame.unsqueeze(1)).squeeze(1), nrow=self.n_row)
                    for frame in torch.unbind(rendered_video, 1)
                ],
                dim=0,
            )
            .unsqueeze(0)
            .to(torch.float32)
            / 256
        )

        return visualization_types.Video(rendered_video)

ocl/visualizations.py
amazon-science/object-centric-multiple-object-tracking