Taming Contrast Maximization for Learning Sequential Event-based Optical Flow Estimation
import numpy as np
import random
import torch
import yaml


class YAMLParser:
    """
    YAML parser for the config files.
    """

    def __init__(self, config):
        self.reset_config()
        self.parse_config(config)
        self.get_device()
        if self._config["loader"]["seed"] is not None:
            self.init_seeds()

    def parse_config(self, file):
        """
        Load and parse the config file.
        """
        with open(file) as fid:
            yaml_config = yaml.load(fid, Loader=yaml.FullLoader)
        self.parse_dict(yaml_config)

    @property
    def config(self):
        return self._config

    @property
    def device(self):
        return self._device

    @property
    def loader_kwargs(self):
        return self._loader_kwargs

    def reset_config(self):
        self._config = {}

        # MLFlow experiment name
        self._config["experiment"] = "Default"

        # input data mode
        self._config["data"] = {}
        self._config["data"]["mode"] = "events"
        self._config["data"]["window"] = 5000

        # data loader
        self._config["loader"] = {}
        self._config["loader"]["resolution"] = [180, 240]
        self._config["loader"]["batch_size"] = 1
        self._config["loader"]["augment"] = []
        self._config["loader"]["gpu"] = 0
        self._config["loader"]["seed"] = 42

        # model
        self._config["model"] = {}

        # visualization
        self._config["vis"] = {}
        self._config["vis"]["bars"] = False

    def update(self, config):
        """
        Updates the config with the given config.
        :param config: dictionary containing a config to update with
        """
        self.reset_config()
        self.parse_config(config)

    def parse_dict(self, input_dict, parent=None):
        """
        Augments self._config with the given dictionary.
        :param input_dict: dictionary to parse and use to update self._config
        :param parent: parent dictionary to be updated
        """
        if parent is None:
            parent = self._config
        for key, val in input_dict.items():
            if isinstance(val, dict):
                if key not in parent.keys():
                    parent[key] = {}
                self.parse_dict(val, parent[key])
            else:
                parent[key] = val

    def get_device(self):
        """
        Get the device to use in the pipeline.
        """
        cuda = torch.cuda.is_available()
        self._device = torch.device("cuda:" + str(self._config["loader"]["gpu"]) if cuda else "cpu")
        self._loader_kwargs = {"num_workers": 0, "pin_memory": True} if cuda else {}

    @staticmethod
    # TODO: not using multiple workers anymore, enable it
    def worker_init_fn(worker_id):
        np.random.seed(np.random.get_state()[1][0] + worker_id)

    def init_seeds(self):
        """
        Initialize random seeds.
        """
        torch.manual_seed(self._config["loader"]["seed"])
        if torch.cuda.is_available():
            torch.cuda.manual_seed(self._config["loader"]["seed"])
            torch.cuda.manual_seed_all(self._config["loader"]["seed"])
        np.random.seed(self._config["loader"]["seed"])
        random.seed(self._config["loader"]["seed"])

    def merge_configs(self, run):
        """
        Overwrites mlflow metadata with configs.
        :param run: mlflow run object
        """
        # parse mlflow settings
        config = {}
        for key in run.keys():
            if len(run[key]) > 0 and run[key][0] == "{":  # assume dictionary
                config[key] = eval(run[key])
            else:  # string
                config[key] = run[key]

        # overwrite with config settings
        self.parse_dict(self._config, config)

        return config

    @staticmethod
    def combine_entries(config):
        """
        Combines entries that had to be split because of MLFlow's max character limit.
        :param config: dictionary to combine entries in
        """
        return config

configs/parser.py
tudelft/taming_event_flow
from abc import abstractmethod

import cv2
import numpy as np
import random
import torch

from .encodings import events_to_voxel, events_to_channels


class BaseDataLoader(torch.utils.data.Dataset):
    """
    Base class for dataloader.
    """

    def __init__(self, config):
        self.config = config
        self.epoch = 0
        self.seq_num = 0
        self.samples = 0
        self.new_seq = False
        self.rectify = False
        self.device = self.config["loader"]["device"]
        self.res = self.config["loader"]["resolution"]
        self.batch_size = self.config["loader"]["batch_size"]

        # batch-specific data augmentation mechanisms
        self.batch_augmentation = {}
        for mechanism in self.config["loader"]["augment"]:
            self.batch_augmentation[mechanism] = [False for i in range(self.config["loader"]["batch_size"])]

        for i, mechanism in enumerate(self.config["loader"]["augment"]):
            for batch in range(self.config["loader"]["batch_size"]):
                if np.random.random() < self.config["loader"]["augment_prob"][i]:
                    self.batch_augmentation[mechanism][batch] = True

    @abstractmethod
    def __getitem__(self, index):
        raise NotImplementedError

    @abstractmethod
    def get_events(self, history):
        raise NotImplementedError

    def reset_sequence(self, batch):
        """
        Reset sequence-specific variables.
        :param batch: batch index
        """

        self.seq_num += 1

        # data augmentation
        for i, mechanism in enumerate(self.config["loader"]["augment"]):
            if np.random.random() < self.config["loader"]["augment_prob"][i]:
                self.batch_augmentation[mechanism][batch] = True
            else:
                self.batch_augmentation[mechanism][batch] = False

    def rectification_mapping(self, batch):
        """
        Compute the backward rectification map for the input representations.
        See https://github.com/uzh-rpg/DSEC/issues/14 for details.
        :param batch: batch index
        :return K_rect: intrinsic matrix of rectified image
        :return mapping: rectification map
        :return Q_rect: scaling matrix to convert disparity to depth
        """

        # distorted image
        K_dist = eval(self.open_files[batch]["calibration/intrinsics"][()])["cam0"]["camera_matrix"]

        # rectified image
        K_rect = eval(self.open_files[batch]["calibration/intrinsics"][()])["camRect0"]["camera_matrix"]
        R_rect = eval(self.open_files[batch]["calibration/extrinsics"][()])["R_rect0"]
        dist_coeffs = eval(self.open_files[batch]["calibration/intrinsics"][()])["cam0"]["distortion_coeffs"]

        # formatting
        K_dist = np.array([[K_dist[0], 0, K_dist[2]], [0, K_dist[1], K_dist[3]], [0, 0, 1]])
        K_rect = np.array([[K_rect[0], 0, K_rect[2]], [0, K_rect[1], K_rect[3]], [0, 0, 1]])
        R_rect = np.array(
            [
                [R_rect[0][0], R_rect[0][1], R_rect[0][2]],
                [R_rect[1][0], R_rect[1][1], R_rect[1][2]],
                [R_rect[2][0], R_rect[2][1], R_rect[2][2]],
            ]
        )
        dist_coeffs = np.array([dist_coeffs[0], dist_coeffs[1], dist_coeffs[2], dist_coeffs[3]])

        # backward mapping
        mapping = cv2.initUndistortRectifyMap(
            K_dist,
            dist_coeffs,
            R_rect,
            K_rect,
            (self.res[1], self.res[0]),
            cv2.CV_32FC2,
        )[0]

        # disparity to depth (onyl used for evaluation)
        Q_rect = eval(self.open_files[batch]["calibration/disparity_to_depth"][()])["cams_03"]
        Q_rect = np.array(
            [
                [Q_rect[0][0], Q_rect[0][1], Q_rect[0][2], Q_rect[0][3]],
                [Q_rect[1][0], Q_rect[1][1], Q_rect[1][2], Q_rect[1][3]],
                [Q_rect[2][0], Q_rect[2][1], Q_rect[2][2], Q_rect[2][3]],
                [Q_rect[3][0], Q_rect[3][1], Q_rect[3][2], Q_rect[3][3]],
            ]
        ).astype(np.float32)

        for _, mechanism in enumerate(self.config["loader"]["augment"]):

            if mechanism == "Horizontal":
                if self.batch_augmentation["Horizontal"][batch]:
                    K_rect[0, 2] = self.res[1] - 1 - K_rect[0, 2]
                    mapping[:, :, 0] = self.res[1] - 1 - mapping[:, :, 0]
                    mapping = np.flip(mapping, axis=1)
                    Q_rect[0, 3] = -K_rect[0, 2]

            elif mechanism == "Vertical":
                if self.batch_augmentation["Vertical"][batch]:
                    K_rect[1, 2] = self.res[0] - 1 - K_rect[1, 2]
                    mapping[:, :, 1] = self.res[0] - 1 - mapping[:, :, 1]
                    mapping = np.flip(mapping, axis=0)
                    Q_rect[1, 3] = -K_rect[1, 2]

        return K_rect, mapping, Q_rect

    @staticmethod
    def format_intrinsics(K_rect):
        """
        Format camera matrices.
        :param K_rect: [3 x 3] intrinsic matrix (numpy) of rectified image
        :return K_rect: [4 x 4] intrinsic matrix (tensor) of rectified image
        :return inv_K_rect: [4 x 4] inverse of the intrinsic matrix (tensor) of rectified image
        """

        K_rect = np.c_[K_rect, np.zeros(3)]
        K_rect = np.concatenate((K_rect, np.array([[0, 0, 0, 1]])), axis=0)
        inv_K_rect = np.linalg.pinv(K_rect)

        K_rect = torch.from_numpy(K_rect.astype(np.float32))
        inv_K_rect = torch.from_numpy(inv_K_rect.astype(np.float32))

        return K_rect, inv_K_rect

    def event_formatting(self, xs, ys, ts, ps):
        """
        Format input events as torch tensors.
        :param xs: [N] numpy array with event x location
        :param ys: [N] numpy array with event y location
        :param ts: [N] numpy array with event timestamp
        :param ps: [N] numpy array with event polarity ([-1, 1])
        :return rectified_xs: [N] numpy array with rectified event x location
        :return rectified_ys: [N] numpy array with rectified event y location
        :return xs: [N] tensor with event x location
        :return ys: [N] tensor with event y location
        :return ts: [N] tensor with normalized event timestamp
        :return ps: [N] tensor with event polarity ([-1, 1])
        """

        assert len(xs) == len(ys) == len(ts) == len(ps)

        xs = torch.from_numpy(xs.astype(np.float32)).to(self.device)
        ys = torch.from_numpy(ys.astype(np.float32)).to(self.device)
        ts = torch.from_numpy(ts.astype(np.float32)).to(self.device)
        ps = torch.from_numpy(ps.astype(np.float32)).to(self.device) * 2 - 1
        if ts.shape[0] > 0:
            ts = (ts - ts[0]) / (ts[-1] - ts[0])

        return xs, ys, ts, ps

    @staticmethod
    def rectify_events(rectify_map, xs, ys):
        """
        Rectify (and undistort) input events.
        :param rectify_map: map used to rectify events
        :param xs: [N] numpy array with event x location
        :param ys: [N] numpy array with event y location
        :return rectified_xs: [N] numpy array with rectified event x location
        :return rectified_ys: [N] numpy array with rectified event y location
        """

        rectified_events = rectify_map[ys.long(), xs.long()]
        rectified_xs = rectified_events[:, 0]
        rectified_ys = rectified_events[:, 1]

        return rectified_xs, rectified_ys

    def augment_events(self, xs, ys, ps, rec_xs, rec_ys, batch):
        """
        Augment event sequence with horizontal, vertical, and polarity flips.
        :param xs: [N] tensor with event x location
        :param ys: [N] tensor with event y location
        :param ps: [N] tensor with event polarity ([-1, 1])
        :param rec_xs: [N] tensor with rectified event x location
        :param rec_ys: [N] tensor with rectified event y location
        :param batch: batch index
        :return xs: [N] tensor with augmented event x location
        :return ys: [N] tensor with augmented event y location
        :return ps: [N] tensor with augmented event polarity ([-1, 1])
        :return rec_xs: [N] tensor with augmented rectified event x location
        :return rec_ys: [N] tensor with augmented rectified event y location
        """

        for _, mechanism in enumerate(self.config["loader"]["augment"]):

            if mechanism == "Horizontal":
                if self.batch_augmentation["Horizontal"][batch]:
                    xs = self.res[1] - 1 - xs
                    if rec_xs is not None:
                        rec_xs = self.res[1] - 1 - rec_xs

            elif mechanism == "Vertical":
                if self.batch_augmentation["Vertical"][batch]:
                    ys = self.res[0] - 1 - ys
                    if rec_ys is not None:
                        rec_ys = self.res[0] - 1 - rec_ys

            elif mechanism == "Polarity":
                if self.batch_augmentation["Polarity"][batch]:
                    ps *= -1

        return xs, ys, ps, rec_xs, rec_ys

    def augment_gt(self, gt, batch):
        """
        Augment ground truth data with horizontal and vertical.
        :param gt: dictionary containing ground truth data
        :param batch: batch index
        """

        for _, mechanism in enumerate(self.config["loader"]["augment"]):

            if mechanism == "Horizontal":
                if self.batch_augmentation["Horizontal"][batch]:
                    for key in gt.keys():
                        gt[key] = torch.flip(gt[key], dims=[2])
                        if key == "gtflow":
                            gt[key][0, ...] *= -1

            elif mechanism == "Vertical":
                if self.batch_augmentation["Vertical"][batch]:
                    for key in gt.keys():
                        gt[key] = torch.flip(gt[key], dims=[1])
                        if key == "gtflow":
                            gt[key][1, ...] *= -1

        return gt

    @staticmethod
    def create_list_encoding(xs, ys, ts, ps):
        """
        Creates a four channel tensor with all the events in the input partition.
        :param xs: [N] tensor with event x location
        :param ys: [N] tensor with event y location
        :param ts: [N] tensor with normalized event timestamp
        :param ps: [N] tensor with event polarity ([-1, 1])
        :return [N x 4] list event representation
        """

        return torch.stack([ts, ys, xs, ps])

    @staticmethod
    def create_polarity_mask(ps):
        """
        Creates a two channel tensor that acts as a mask for the input event list.
        :param ps: [N] tensor with event polarity ([-1, 1])
        :return [N x 2] polarity list event representation
        """

        event_list_pol_mask = torch.stack([ps, ps])
        event_list_pol_mask[0, :][event_list_pol_mask[0, :] < 0] = 0
        event_list_pol_mask[0, :][event_list_pol_mask[0, :] > 0] = 1
        event_list_pol_mask[1, :][event_list_pol_mask[1, :] < 0] = -1
        event_list_pol_mask[1, :][event_list_pol_mask[1, :] > 0] = 0
        event_list_pol_mask[1, :] *= -1
        return event_list_pol_mask

    def create_cnt_encoding(self, xs, ys, ps, rect_mapping):
        """
        Creates a per-pixel and per-polarity event count representation.
        :param xs: [N] tensor with event x location
        :param ys: [N] tensor with event y location
        :param ps: [N] tensor with event polarity ([-1, 1])
        :param rect_mapping: map used to rectify events
        :return [2 x H x W] rectified event count representation
        """

        # create event count representation and rectify it using backward mapping
        event_cnt = events_to_channels(xs, ys, ps, sensor_size=self.res)
        if rect_mapping is not None:
            event_cnt = event_cnt.permute(1, 2, 0)
            event_cnt = cv2.remap(event_cnt.cpu().numpy(), rect_mapping, None, cv2.INTER_NEAREST)
            event_cnt = torch.from_numpy(event_cnt.astype(np.float32)).to(self.device)
            event_cnt = event_cnt.permute(2, 0, 1)

        return event_cnt

    @staticmethod
    def create_mask_encoding(event_cnt):
        """
        Creates per-pixel event mask based on event count.
        :param event_cnt: [2 x H x W] event count
        :return [H x W] rectified event mask representation
        """

        event_mask = event_cnt.clone()
        event_mask = torch.sum(event_mask, dim=0, keepdim=True)
        event_mask[event_mask > 0.0] = 1.0

        return event_mask

    def create_voxel_encoding(self, xs, ys, ts, ps, rect_mapping, num_bins=5):
        """
        Creates a spatiotemporal voxel grid tensor representation with a certain number of bins,
        as described in Section 3.1 of the paper 'Unsupervised Event-based Learning of Optical Flow,
        Depth, and Egomotion', Zhu et al., CVPR'19..
        Events are distributed to the spatiotemporal closest bins through bilinear interpolation.
        Positive events are added as +1, while negative as -1.
        :param xs: [N] tensor with event x location
        :param ys: [N] tensor with event y location
        :param ts: [N] tensor with normalized event timestamp
        :param ps: [N] tensor with event polarity ([-1, 1])
        :param rect_mapping: map used to rectify events
        :param num_bins: number of bins in the voxel grid
        :return [B x H x W] rectified voxel grid event representation
        """

        event_voxel = events_to_voxel(
            xs,
            ys,
            ts,
            ps,
            num_bins,
            sensor_size=self.res,
        )

        if rect_mapping is not None:
            event_voxel = event_voxel.permute(1, 2, 0)
            event_voxel = cv2.remap(event_voxel.cpu().numpy(), rect_mapping, None, cv2.INTER_NEAREST)
            event_voxel = torch.from_numpy(event_voxel.astype(np.float32)).to(self.device)
            event_voxel = event_voxel.permute(2, 0, 1)

        return event_voxel

    @staticmethod
    def split_event_list(event_list, event_list_pol_mask, max_num_grad_events):
        """
        Splits the event list into two lists, one of them (with max. length) to be used for backprop.
        This helps reducing (VRAM) memory consumption.
        :param event_list: [4 x N] list event representation
        :param event_list_pol_mask: [2 x N] polarity list event representation
        :param max_num_grad_events: maximum number of events to be used for backprop
        :return event_list: [4 x N] list event representation to be used for backprop
        :return event_list_pol_mask: [2 x N] polarity list event representation to be used for backprop
        :return d_event_list: [4 x N] list event representation
        :return d_event_list_pol_mask: [2 x N] polarity list event representation
        """

        d_event_list = torch.zeros((4, 0))
        d_event_list_pol_mask = torch.zeros((2, 0))
        if max_num_grad_events is not None and event_list.shape[1] > max_num_grad_events:
            probs = torch.ones(event_list.shape[1], dtype=torch.float32) / event_list.shape[1]
            sampled_indices = probs.multinomial(
                max_num_grad_events, replacement=False
            )  # sample indices with equal prob.

            unsampled_indices = torch.ones(event_list.shape[1], dtype=torch.bool)
            unsampled_indices[sampled_indices] = False
            d_event_list = event_list[:, unsampled_indices]
            d_event_list_pol_mask = event_list_pol_mask[:, unsampled_indices]

            event_list = event_list[:, sampled_indices]
            event_list_pol_mask = event_list_pol_mask[:, sampled_indices]

        return event_list, event_list_pol_mask, d_event_list, d_event_list_pol_mask

    def __len__(self):
        return 1000  # not used

    def shuffle(self, flag=True):
        """
        Shuffles the training data.
        :param flag: if true, shuffles the data
        """

        if flag:
            random.shuffle(self.files)

    @staticmethod
    def custom_collate(batch):
        """
        Collects the different event representations and stores them together in a dictionary.
        :param batch: batch index
        :return batch_dict: dictionary with the output of a dataloader iteration
        """

        # create dictionary
        batch_dict = {}
        for key in batch[0].keys():
            batch_dict[key] = []

        # collect data
        for entry in batch:
            for key in entry.keys():
                batch_dict[key].append(entry[key])

        # create batches
        for key in batch_dict.keys():

            if batch_dict[key][0] is not None:

                # pad entries of different size
                N = 0
                if key in ["event_list", "event_list_pol_mask", "d_event_list", "d_event_list_pol_mask"]:
                    for i in range(len(batch_dict[key])):
                        if N < batch_dict[key][i].shape[1]:
                            N = batch_dict[key][i].shape[1]

                    for i in range(len(batch_dict[key])):
                        zeros = torch.zeros((batch_dict[key][i].shape[0], N - batch_dict[key][i].shape[1]))
                        batch_dict[key][i] = torch.cat((batch_dict[key][i], zeros), dim=1)

                # create tensor
                item = torch.stack(batch_dict[key])
                if len(item.shape) == 3:
                    item = item.transpose(2, 1)
                batch_dict[key] = item

            else:
                batch_dict[key] = None

        return batch_dict

dataloader/base.py
tudelft/taming_event_flow
import os
import hdf5plugin
import h5py
import numpy as np
import torch
import yaml


class CacheDataset:
    """
    Utility class to "cache" the output of the dataloader in hdf5 files
    for a more efficient dataloader.
    """

    def __init__(self, config, dir, mode="train"):
        self.keys = {}

        data_keys = ["path", "mode", "window", "voxel"]
        for key in data_keys:
            self.keys[key] = config["data"][key]

        loader_keys = ["resolution", "augment", "augment_prob"]
        for key in loader_keys:
            self.keys[key] = config["loader"][key]

        if not os.path.exists(dir):
            os.system("mkdir " + dir)

        self.dir = dir + "cache_" + mode
        dict_file = self.dir + "/dataset_keys.yml"
        if os.path.exists(self.dir):
            if os.path.isfile(dict_file):
                tmp_keys = self.read_yaml(dict_file)
                if self.keys != tmp_keys:  # there are keys, but diff from current settings
                    print("Deleting cache dir:", self.dir)
                    os.system("rm -rf " + self.dir + "/*")
                    self.write_yaml(dict_file, self.keys)
            else:
                self.write_yaml(dict_file, self.keys)  # no keys, write them
        else:
            os.system("mkdir " + self.dir)
            self.write_yaml(dict_file, self.keys)  # no cache, create it

    @staticmethod
    def read_yaml(file):
        with open(file, "r") as f:
            tmp_keys = yaml.load(f, Loader=yaml.FullLoader)
        return tmp_keys

    @staticmethod
    def write_yaml(file, keys):
        with open(file, "w") as outfile:
            yaml.dump(keys, outfile, default_flow_style=False)

    def update(self, filename, dict):
        """
        Update the cached dataset with the given data.
        :param filename: name of the file to update
        :param dict: dictionary with the data to update
        """

        filename = self.dir + "/" + filename.split("/")[-1]
        if not os.path.isfile(filename):
            file = h5py.File(filename, "w")
            file.attrs["idx"] = 0
        else:
            file = h5py.File(filename, "a")
            file.attrs["idx"] += 1

        for key in dict:
            file.create_dataset(
                key + "/{:09d}".format(file.attrs["idx"]),
                data=dict[key].numpy(),
                dtype=np.dtype(np.float32),
                **hdf5plugin.Zstd()
            )

    def load(self, filename, idx):
        """
        Load the cached dataset with the given data.
        :param filename: name of the file to load
        :param idx: index of the data to load
        :return data: dictionary with the loaded data
        :return success: flag indicating if the data was loaded correctly
        """

        filename = self.dir + "/" + filename.split("/")[-1]
        if not os.path.isfile(filename):
            return {}, False

        file = h5py.File(filename, "r")

        data = {}
        success = True
        entry = "{:09d}".format(idx)
        for key in file.keys():
            if entry in file[key].keys():
                data[key] = torch.from_numpy(file[key + "/" + entry][:])
            else:
                success = False
                break

        file.close()

        return data, success

dataloader/cache.py
tudelft/taming_event_flow
"""
Adapted from Monash University https://github.com/TimoStoff/events_contrast_maximization
"""

import torch


def events_to_image(xs, ys, ps, sensor_size=(180, 240), accumulate=True):
    """
    Accumulate events into an image.
    :param xs: event x coordinates
    :param ys: event y coordinates
    :param ps: event polarity
    :param sensor_size: sensor size
    :param accumulate: flag indicating whether to accumulate events into the image
    :return img: image containing per-pixel event counts
    """

    device = xs.device
    img_size = list(sensor_size)
    img = torch.zeros(img_size, device=device)

    if xs.dtype is not torch.long:
        xs = xs.long().to(device)
    if ys.dtype is not torch.long:
        ys = ys.long().to(device)
    img.index_put_((ys, xs), ps, accumulate=accumulate)

    return img


def events_to_voxel(xs, ys, ts, ps, num_bins, sensor_size=(180, 240)):
    """
    Generate a voxel grid from input events using temporal bilinear interpolation.
    :param xs: event x coordinates
    :param ys: event y coordinates
    :param ts: event timestamps
    :param ps: event polarity
    :param num_bins: number of bins in the voxel grid
    :param sensor_size: sensor size
    :return: voxel grid representation
    """

    assert len(xs) == len(ys) and len(ys) == len(ts) and len(ts) == len(ps)

    voxel = []
    ts = ts * (num_bins - 1)
    device = xs.device

    zeros = torch.zeros(ts.size(), device=device)
    for b_idx in range(num_bins):
        weights = torch.max(zeros, 1.0 - torch.abs(ts - b_idx))
        voxel_bin = events_to_image(xs, ys, ps * weights, sensor_size=sensor_size)
        voxel.append(voxel_bin)

    return torch.stack(voxel)


def events_to_channels(xs, ys, ps, sensor_size=(180, 240)):
    """
    Generate a two-channel event image containing per-pixel event counters.
    :param xs: event x coordinates
    :param ys: event y coordinates
    :param ps: event polarity
    :param sensor_size: sensor size
    :return: event image containing per-pixel and per-polarity event counts
    """

    assert len(xs) == len(ys) and len(ys) == len(ps)

    mask_pos = ps.clone()
    mask_neg = ps.clone()
    mask_pos[ps < 0] = 0
    mask_neg[ps > 0] = 0
    mask_pos[ps > 0] = 1
    mask_neg[ps < 0] = -1

    pos_cnt = events_to_image(xs, ys, ps * mask_pos, sensor_size=sensor_size)
    neg_cnt = events_to_image(xs, ys, ps * mask_neg, sensor_size=sensor_size)

    return torch.stack([pos_cnt, neg_cnt])

dataloader/encodings.py
tudelft/taming_event_flow
import os
import sys
import cv2
import hdf5plugin
import h5py
import numpy as np

import torch

from .base import BaseDataLoader
from .cache import CacheDataset
from .utils import ProgressBar

parent_dir_name = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
sys.path.append(parent_dir_name)

from utils.utils import binary_search_array


class FlowMaps:
    """
    Utility class for reading the ground truth optical flow maps encoded in the HDF5 files.
    """

    def __init__(self):
        self.ts_from = []
        self.ts_to = []
        self.names = []

    def __call__(self, name, h5obj):
        if hasattr(h5obj, "dtype") and name not in self.names:
            self.names += [name]
            self.ts_from += [h5obj.attrs["timestamp_from"]]
            self.ts_to += [h5obj.attrs["timestamp_to"]]


class H5Loader(BaseDataLoader):
    def __init__(self, config, shuffle=False, path_cache=""):
        super().__init__(config)
        self.input_window = self.config["data"]["window"]
        if self.config["data"]["mode"] in ["gtflow"] and self.input_window > 1:
            print("DataLoader error: Ground truth data mode cannot be used with window > 1.")
            raise AttributeError

        self.ts_jump = False
        self.ts_jump_reset = False  # used in the inference loops to reset model states

        self.gt_avg_dt = None
        self.gt_avg_idx = 0
        self.last_proc_timestamp = 0

        # "memory" that goes from forward pass to the next
        self.batch_idx = [i for i in range(self.batch_size)]  # event sequence
        self.batch_row = [0 for i in range(self.batch_size)]  # event_idx / time_idx
        self.batch_pass = [0 for i in range(self.batch_size)]  # forward passes

        # input event sequences
        self.files = []
        for root, dirs, files in os.walk(config["data"]["path"]):
            for file in files:
                if file.endswith(".h5"):
                    self.files.append(os.path.join(root, file))

        # shuffle files
        if shuffle:
            self.shuffle()

        # initialize cache
        if self.config["data"]["cache"]:
            self.cache = CacheDataset(config, path_cache)

        # open first files
        self.open_files = []
        self.batch_rectify_map = []
        self.batch_K_rect = []
        self.batch_Q_rect = []
        self.batch_rect_mapping = []
        for batch in range(self.config["loader"]["batch_size"]):
            self.open_files.append(h5py.File(self.files[self.batch_idx[batch] % len(self.files)], "r"))
            if "rectification" in self.open_files[-1].keys():
                self.batch_rectify_map.append(self.open_files[-1]["rectification/rectify_map"][:])
                self.batch_rectify_map[-1] = torch.from_numpy(self.batch_rectify_map[-1]).float().to(self.device)

                K_rect, mapping, Q_rect = self.rectification_mapping(-1)
                self.batch_K_rect.append(K_rect)
                self.batch_Q_rect.append(Q_rect)
                self.batch_rect_mapping.append(mapping)
                self.rectify = True
            else:
                self.batch_rect_mapping.append(None)

        # load GT optical flow maps from open files
        self.open_files_flowmaps = []
        if config["data"]["mode"] == "gtflow":
            for batch in range(self.batch_size):
                flowmaps = FlowMaps()
                if "flow" in self.open_files[batch].keys():
                    self.open_files[batch]["flow"].visititems(flowmaps)
                self.open_files_flowmaps.append(flowmaps)

        # progress bars
        if self.config["vis"]["bars"]:
            self.open_files_bar = []
            for batch in range(self.config["loader"]["batch_size"]):
                max_iters = self.get_iters(batch)
                self.open_files_bar.append(ProgressBar(self.files[batch].split("/")[-1], max=max_iters))

    def get_iters(self, batch):
        """
        Compute the number of forward passes given a sequence and an input mode and window.
        :param batch: batch index
        :return: number of forward passes
        """

        if self.config["data"]["mode"] == "events":
            max_iters = len(self.open_files[batch]["events/xs"])
        elif self.config["data"]["mode"] == "time":
            max_iters = self.open_files[batch].attrs["duration"]
        elif self.config["data"]["mode"] == "gtflow":
            max_iters = len(self.open_files_flowmaps[batch].ts_to) - 1
        else:
            print("DataLoader error: Unknown mode.")
            raise AttributeError

        return max_iters // self.input_window

    def get_events(self, file, idx0, idx1):
        """
        Get all the events in between two indices.
        :param file: file to read from
        :param idx0: start index
        :param idx1: end index
        :return xs: [N] numpy array with event x location
        :return ys: [N] numpy array with event y location
        :return ts: [N] numpy array with event timestamp
        :return ps: [N] numpy array with event polarity ([-1, 1])
        """

        xs = file["events/xs"][idx0:idx1]
        ys = file["events/ys"][idx0:idx1]
        ts = file["events/ts"][idx0:idx1]
        ps = file["events/ps"][idx0:idx1]
        ts -= file.attrs["t0"]  # sequence starting at t0 = 0

        # check if temporal discontinuity in gt data modes
        self.ts_jump = False
        if self.config["data"]["mode"] in ["gtflow"]:
            dt = ts[-1] - self.last_proc_timestamp
            if self.gt_avg_dt is None:
                self.gt_avg_dt = dt
                self.gt_avg_idx += 1

            if dt >= 2 * self.gt_avg_dt / self.gt_avg_idx:
                self.ts_jump = True
                self.ts_jump_reset = True
            else:
                self.gt_avg_dt += dt
                self.gt_avg_idx += 1

        if ts.shape[0] > 0:
            self.last_proc_timestamp = ts[-1]
        return xs, ys, ts, ps

    def get_event_index(self, batch, window=0):
        """
        Get all the event indices to be used for reading.
        :param batch: batch index
        :param window: input window
        :return event_idx0: event index (from)
        :return event_idx1: event index (to)
        """

        restart = False
        event_idx0 = None
        event_idx1 = None
        if self.config["data"]["mode"] == "events":
            event_idx0 = self.batch_row[batch]
            event_idx1 = self.batch_row[batch] + window

        elif self.config["data"]["mode"] == "time":
            event_idx0 = self.find_ts_index(
                self.open_files[batch], self.batch_row[batch] + self.open_files[batch].attrs["t0"]
            )
            event_idx1 = self.find_ts_index(
                self.open_files[batch], self.batch_row[batch] + self.open_files[batch].attrs["t0"] + window
            )

        elif self.config["data"]["mode"] == "gtflow":
            idx1 = int(np.ceil(self.batch_row[batch] + window))
            if np.isclose(self.batch_row[batch] + window, idx1 - 1):
                idx1 -= 1
            event_idx0 = self.find_ts_index(self.open_files[batch], self.open_files_flowmaps[batch].ts_from[idx1])
            event_idx1 = self.find_ts_index(self.open_files[batch], self.open_files_flowmaps[batch].ts_to[idx1])
            if self.open_files_flowmaps[batch].ts_to[idx1] > self.open_files[batch].attrs["tk"]:
                restart = True

        else:
            print("DataLoader error: Unknown mode.")
            raise AttributeError

        return event_idx0, event_idx1, restart

    def find_ts_index(self, file, timestamp, dataset="events/ts"):
        """
        Find closest event index for a given timestamp through binary search.
        :param file: file to read from
        :param timestamp: timestamp to find
        :param dataset: dataset to search in
        :return: event index
        """

        return binary_search_array(file[dataset], timestamp)

    def open_new_h5(self, batch):
        """
        Open new H5 event sequence.
        :param batch: batch index
        """

        self.ts_jump = False
        self.ts_jump_reset = False

        self.gt_avg_dt = None
        self.gt_avg_idx = 0
        self.last_proc_timestamp = 0

        self.open_files[batch] = h5py.File(self.files[self.batch_idx[batch] % len(self.files)], "r+")

        if self.rectify:
            self.batch_rectify_map[batch] = self.open_files[batch]["rectification/rectify_map"][:]
            self.batch_rectify_map[batch] = torch.from_numpy(self.batch_rectify_map[batch]).float().to(self.device)

            K_rect, mapping, Q_rect = self.rectification_mapping(batch)
            self.batch_K_rect[batch] = K_rect
            self.batch_Q_rect[batch] = Q_rect
            self.batch_rect_mapping[batch] = mapping

        if self.config["data"]["mode"] == "gtflow":
            flowmaps = FlowMaps()
            if "flow" in self.open_files[batch].keys():
                self.open_files[batch]["flow"].visititems(flowmaps)
            self.open_files_flowmaps[batch] = flowmaps

        if self.config["vis"]["bars"]:
            self.open_files_bar[batch].finish()
            max_iters = self.get_iters(batch)
            self.open_files_bar[batch] = ProgressBar(
                self.files[self.batch_idx[batch] % len(self.files)].split("/")[-1], max=max_iters
            )

        if "Playback" in self.batch_augmentation.keys() and self.batch_augmentation["Playback"][batch]:
            file = self.open_files[batch]
            xs = np.flip(file["events/xs"][:])
            ys = np.flip(file["events/ys"][:])
            ps = np.flip(file["events/ps"][:])

            ts = np.flip(file["events/ts"][:])
            min_ts = ts[-1]
            max_ts = ts[0]
            ts = np.absolute((ts - min_ts) / (max_ts - min_ts) - 1)
            ts = ts * (max_ts - min_ts) + min_ts

            file["events/xs"][:] = xs
            file["events/ys"][:] = ys
            file["events/ts"][:] = ts
            file["events/ps"][:] = ps

    def __getitem__(self, index):
        while True:
            batch = index % self.config["loader"]["batch_size"]

            # try loading cached data
            if self.config["data"]["cache"]:
                output, success = self.cache.load(
                    self.files[self.batch_idx[batch] % len(self.files)], self.batch_pass[batch]
                )
                if success:
                    self.batch_row[batch] += self.input_window
                    self.batch_pass[batch] += 1
                    return output

            # trigger sequence change
            len_frames = 0
            restart = False
            if self.config["data"]["mode"] == "gtflow":
                len_frames = len(self.open_files_flowmaps[batch].ts_to)
                if int(np.ceil(self.batch_row[batch] + self.input_window)) >= len_frames:
                    restart = True

            # load events
            xs = np.zeros((0))
            ys = np.zeros((0))
            ts = np.zeros((0))
            ps = np.zeros((0))
            if not restart:
                idx0, idx1, restart = self.get_event_index(batch, window=self.input_window)

                if self.config["data"]["mode"] in ["gtflow"] and self.input_window < 1.0:
                    floor_row = int(np.floor(self.batch_row[batch]))
                    ceil_row = int(np.ceil(self.batch_row[batch] + self.input_window))

                    if np.isclose(self.batch_row[batch], floor_row + 1):
                        floor_row += 1
                    if np.isclose(self.batch_row[batch] + self.input_window, ceil_row - 1):
                        ceil_row -= 1

                    idx0_change = self.batch_row[batch] - floor_row
                    idx1_change = self.batch_row[batch] + self.input_window - floor_row

                    delta_idx = idx1 - idx0
                    idx1 = int(idx0 + idx1_change * delta_idx)
                    idx0 = int(idx0 + idx0_change * delta_idx)

                if not restart:
                    xs, ys, ts, ps = self.get_events(self.open_files[batch], idx0, idx1)

            # skip gt sample if temporal discontinuity in gt
            if self.config["data"]["mode"] in ["gtflow"] and self.ts_jump:
                self.batch_row[batch] += self.input_window
                self.batch_pass[batch] += 1
                continue

            # trigger sequence change
            if (self.config["data"]["mode"] == "events" and xs.shape[0] < self.input_window) or (
                self.config["data"]["mode"] == "time"
                and self.batch_row[batch] + self.input_window >= self.open_files[batch].attrs["duration"]
            ):
                restart = True

            # reset sequence if not enough input events
            if restart:
                self.new_seq = True
                self.reset_sequence(batch)
                self.batch_row[batch] = 0
                self.batch_idx[batch] = max(self.batch_idx) + 1
                self.batch_pass[batch] = 0
                self.open_files[batch].close()
                self.open_new_h5(batch)
                continue

            # handle case with very few events
            if xs.shape[0] <= 10:
                xs = np.empty([0])
                ys = np.empty([0])
                ts = np.empty([0])
                ps = np.empty([0])

            # event formatting and timestamp normalization
            xs, ys, ts, ps = self.event_formatting(xs, ys, ts, ps)

            # rectify input events
            rec_xs, rec_ys = None, None
            if self.rectify:
                rec_xs, rec_ys = self.rectify_events(self.batch_rectify_map[batch], xs, ys)

            # data augmentation
            xs, ys, ps, rec_xs, rec_ys = self.augment_events(xs, ys, ps, rec_xs, rec_ys, batch)

            # events to lists
            if self.rectify:
                event_list = self.create_list_encoding(rec_xs, rec_ys, ts, ps)
            else:
                event_list = self.create_list_encoding(xs, ys, ts, ps)
            event_list_pol_mask = self.create_polarity_mask(ps)

            # create event representations
            event_cnt = self.create_cnt_encoding(xs, ys, ps, self.batch_rect_mapping[batch])
            event_mask = self.create_mask_encoding(event_cnt)
            if self.config["data"]["voxel"] is not None:
                event_voxel = self.create_voxel_encoding(
                    xs, ys, ts, ps, self.batch_rect_mapping[batch], num_bins=self.config["data"]["voxel"]
                )

            # voxel is the preferred representation for the network's input
            if self.config["data"]["voxel"] is None:
                net_input = event_cnt.clone()
            else:
                net_input = event_voxel.clone()

            # load (and augment) GT maps when required
            gt = {}
            if self.config["data"]["mode"] == "gtflow":
                idx = int(np.ceil(self.batch_row[batch] + self.input_window))
                if np.isclose(self.batch_row[batch] + self.input_window, idx - 1):
                    idx -= 1
                flowmap = self.open_files[batch]["flow"][self.open_files_flowmaps[batch].names[idx]][:]
                flowmap = flowmap.astype(np.float32)
                flowmap = torch.from_numpy(flowmap).permute(2, 0, 1)
                gt["gtflow"] = flowmap
                gt["gtflow_dt"] = (
                    self.open_files_flowmaps[batch].ts_to[idx] - self.open_files_flowmaps[batch].ts_from[idx]
                )
                gt["gtflow_dt"] = torch.from_numpy(np.asarray(gt["gtflow_dt"])).float()

            gt = self.augment_gt(gt, batch)

            # update window
            self.batch_row[batch] += self.input_window
            self.batch_pass[batch] += 1

            # break while loop if everything went well
            break

        # camera matrix for rectified and cropped events
        if self.rectify:
            K_rect, inv_K_rect = self.format_intrinsics(self.batch_K_rect[batch].copy())

        # split event list (events with and without gradients)
        event_list, event_list_pol_mask, d_event_list, d_event_list_pol_mask = self.split_event_list(
            event_list, event_list_pol_mask, self.config["loader"]["max_num_grad_events"]
        )

        # prepare output
        output = {}
        output["net_input"] = net_input.cpu()
        output["event_cnt"] = event_cnt.cpu()
        output["event_mask"] = event_mask.cpu()
        output["event_list"] = event_list.cpu()
        output["event_list_pol_mask"] = event_list_pol_mask.cpu()
        output["d_event_list"] = d_event_list.cpu()
        output["d_event_list_pol_mask"] = d_event_list_pol_mask.cpu()
        if self.rectify:
            output["K_rect"] = K_rect
            output["inv_K_rect"] = inv_K_rect
        for key in gt.keys():
            output[key] = gt[key]

        if self.config["data"]["cache"]:
            self.cache.update(self.files[self.batch_idx[batch] % len(self.files)], output)

        return output

dataloader/h5.py
tudelft/taming_event_flow
from progress.bar import Bar


class ProgressBar(Bar):
    suffix = "%(percent).1f%%, ETA: %(eta)ds, %(frequency)fHz"

    @property
    def frequency(self):
        """
        :return: frequency of the processing
        """
        return 1 / self.avg

dataloader/utils.py
tudelft/taming_event_flow
import argparse

import mlflow
import torch

from configs.parser import YAMLParser
from dataloader.h5 import H5Loader
from loss.flow_val import *
from models.model import *
from utils.iwe import compute_pol_iwe
from utils.utils import load_model, create_model_dir, initialize_quant_results
from utils.mlflow import log_config, log_results
from utils.visualization import Visualization


def test(args, config_parser):
    """
    Main function of the evaluation pipeline for event-based optical flow estimation.
    :param args: arguments of the script
    :param config_parser: YAMLParser object with config data
    """

    mlflow.set_tracking_uri(args.path_mlflow)

    run = mlflow.get_run(args.runid)
    config = config_parser.merge_configs(run.data.params)
    config = config_parser.combine_entries(config)

    # configs
    config["loader"]["batch_size"] = 1

    # create directory for inference results
    path_results = create_model_dir(args.path_results, args.runid)

    # store validation settings
    eval_id = log_config(path_results, args.runid, config)

    # initialize settings
    device = config_parser.device
    kwargs = config_parser.loader_kwargs
    config["loader"]["device"] = device

    # visualization tool
    vis = Visualization(config, eval_id=eval_id, path_results=path_results)

    # data loader
    data = H5Loader(config, shuffle=True)
    dataloader = torch.utils.data.DataLoader(
        data,
        drop_last=True,
        batch_size=config["loader"]["batch_size"],
        collate_fn=data.custom_collate,
        worker_init_fn=config_parser.worker_init_fn,
        **kwargs,
    )

    # model initialization and settings
    num_bins = 2 if config["data"]["voxel"] is None else config["data"]["voxel"]
    model = eval(config["model"]["name"])(config["model"].copy(), num_bins)
    model = model.to(device)
    model, _ = load_model(args.runid, model, device)
    model.eval()

    # validation metric
    criteria = eval(config["metrics"]["warping"])(config, device)
    val_results = {}

    # inference loop
    end_test = False
    with torch.no_grad():
        while not end_test:
            for inputs in dataloader:
                sequence = data.files[data.batch_idx[0] % len(data.files)].split("/")[-1].split(".")[0]

                if data.new_seq:
                    data.new_seq = False
                    model.reset_states()
                    criteria.reset()

                if config["data"]["mode"] in ["gtflow"] and data.ts_jump_reset:
                    data.ts_jump_reset = False
                    model.reset_states()

                # finish inference loop
                if data.seq_num >= len(data.files):
                    end_test = True
                    break

                # forward pass
                x = model(inputs["net_input"].to(device))
                for i in range(len(x["flow"])):
                    x["flow"][i] = x["flow"][i] * config["loss"]["flow_scaling"]

                # mask flow for visualization
                flow_vis = x["flow"][-1].clone()
                if config["vis"]["mask_output"]:
                    flow_vis *= inputs["event_mask"].to(device)

                # image of warped events
                iwe = None
                if (config["vis"]["enabled"] or config["vis"]["store"]) and (
                    config["vis"]["show"] is None or "iwe" in config["vis"]["show"]
                ):
                    iwe = compute_pol_iwe(
                        flow_vis,
                        inputs["event_list"].to(device),
                        config["loader"]["resolution"],
                        inputs["event_list_pol_mask"].to(device),
                        round_idx=False,
                        round_flow=False,
                    )

                # update validation criteria
                criteria.update(
                    x["flow"],
                    inputs["event_list"].to(device),
                    inputs["event_list_pol_mask"].to(device),
                    inputs["event_mask"].to(device),
                )

                # prepare for visualization
                if config["vis"]["enabled"] or config["vis"]["store"]:

                    # dynamic windows
                    if config["data"]["passes_loss"] > 1 and config["vis"]["dynamic"]:
                        vis.data["events_dynamic"] = criteria.window_events()
                        vis.data["iwe_fw_dynamic"] = criteria.window_iwe(mode="forward")
                        vis.data["iwe_bw_dynamic"] = criteria.window_iwe(mode="backward")
                        vis.data["flow_dynamic"] = criteria.window_flow(mode="forward")

                    # accumulated windows
                    if criteria.num_passes > 1 and criteria.num_passes == config["data"]["passes_loss"]:
                        vis.data["events_window"] = criteria.window_events()
                        vis.data["iwe_fw_window"] = criteria.window_iwe(mode="forward")
                        vis.data["iwe_bw_window"] = criteria.window_iwe(mode="backward")
                        vis.data["flow_window"] = criteria.window_flow(mode="forward")

                # compute error metrics
                vis.data["flow_bw"] = None
                val_results = initialize_quant_results(val_results, sequence, config["metrics"]["name"])
                if criteria.num_passes == config["data"]["passes_loss"]:

                    compute_metrics = True
                    if "eval_time" in config["metrics"].keys():
                        if (
                            data.last_proc_timestamp < config["metrics"]["eval_time"][0]
                            or data.last_proc_timestamp > config["metrics"]["eval_time"][1]
                        ):
                            compute_metrics = False

                    if compute_metrics:

                        # AEE
                        if config["data"]["mode"] == "gtflow" and "AEE" in config["metrics"]["name"]:
                            mask_aee = None
                            if "mask_aee" in config["metrics"].keys() and config["metrics"]["mask_aee"]:
                                mask_aee = criteria.window_events().clone().to(device)

                            vis.data["flow_bw"] = (
                                criteria.window_flow(mode="backward", mask=False) * config["data"]["passes_loss"]
                            )
                            aee = criteria.compute_aee(vis.data["flow_bw"], inputs["gtflow"].to(device), mask=mask_aee)
                            val_results[sequence]["AEE"]["it"] += 1
                            val_results[sequence]["AEE"]["metric"] += aee.cpu().numpy()

                        # deblurring metrics
                        for metric in config["metrics"]["name"]:
                            if metric == "RSAT":
                                rsat = criteria.rsat()
                                val_results[sequence][metric]["metric"] += rsat[0].cpu().numpy()
                                val_results[sequence][metric]["it"] += 1

                            elif metric == "FWL":
                                fwl = criteria.fwl()
                                val_results[sequence][metric]["metric"] += fwl.cpu().numpy()
                                val_results[sequence][metric]["it"] += 1

                    # reset criteria
                    criteria.reset()

                # visualization
                if config["vis"]["bars"]:
                    for bar in data.open_files_bar:
                        bar.next()
                if config["vis"]["enabled"] or config["vis"]["store"]:
                    vis.data["iwe"] = iwe
                    vis.data["flow"] = flow_vis
                    vis.step(
                        inputs,
                        sequence=sequence,
                        ts=data.last_proc_timestamp,
                        show=config["vis"]["show"],
                    )

    if config["vis"]["bars"]:
        for bar in data.open_files_bar:
            bar.finish()

    # store validation config and results
    results = {}
    for metric in config["metrics"]["name"]:
        results[metric] = {}
        for key in val_results.keys():
            if val_results[key][metric]["it"] > 0:
                results[metric][key] = str(val_results[key][metric]["metric"] / val_results[key][metric]["it"])
        log_results(args.runid, results, path_results, eval_id)
    print(results)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("runid", help="mlflow run")
    parser.add_argument(
        "--config",
        default="configs/eval_flow.yml",
        help="config file, overwrites mlflow settings",
    )
    parser.add_argument(
        "--path_mlflow",
        default="",
        help="location of the mlflow ui",
    )
    parser.add_argument("--path_results", default="results_inference/")
    args = parser.parse_args()

    # launch testing
    test(args, YAMLParser(args.config))

eval_flow.py
tudelft/taming_event_flow
from abc import abstractmethod
import os
import sys

import math
import torch

parent_dir_name = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
sys.path.append(parent_dir_name)

from utils.iwe import event_propagation, get_event_flow, purge_unfeasible, get_interpolation, interpolate


class BaseEventWarping(torch.nn.Module):
    """
    Base class for the contrast maximization loss.
    """

    def __init__(self, config, device, loss_scaling=True, border_compensation=True):
        super(BaseEventWarping, self).__init__()
        self.device = device
        self.config = config
        self.loss_scaling = loss_scaling
        self.border_compensation = border_compensation
        self.res = config["loader"]["resolution"]
        self.batch_size = config["loader"]["batch_size"]
        self.flow_spat_smooth_weight = config["loss"]["flow_spat_smooth_weight"]
        self.flow_temp_smooth_weight = config["loss"]["flow_temp_smooth_weight"]

        self._passes = 0
        self._num_flows = None
        self._flow_maps_x = None
        self._flow_maps_y = None

        # warping indices (for temporal consistency)
        my, mx = torch.meshgrid(torch.arange(self.res[0]), torch.arange(self.res[1]))
        indices = torch.stack([my, mx], dim=0).unsqueeze(0)
        self.indices = indices.float().to(device)
        self.indices_mask = torch.ones(self.indices.shape).to(device)

        # timescales for loss computation
        self.passes_loss = []
        for s in range(config["data"]["scales_loss"]):
            self.passes_loss.append(config["data"]["passes_loss"] // (2**s))

    def update_base(self, flow_list):
        """
        Initialize/Update container lists of events and flow maps for forward warping.
        :param flow_list: [[batch_size x 2 x H x W]] list of optical flow (x, y) maps
        """

        if self._num_flows is None:
            self._num_flows = len(flow_list)

        # update optical flow maps
        if self._flow_maps_x is None:
            self._flow_maps_x = []
            self._flow_maps_y = []

        for i, flow in enumerate(flow_list):
            if i == len(self._flow_maps_x):
                self._flow_maps_x.append(flow[:, 0:1, :, :])
                self._flow_maps_y.append(flow[:, 1:2, :, :])
            else:
                self._flow_maps_x[i] = torch.cat([self._flow_maps_x[i], flow[:, 0:1, :, :]], dim=1)
                self._flow_maps_y[i] = torch.cat([self._flow_maps_y[i], flow[:, 1:2, :, :]], dim=1)

    def reset_base(self):
        """
        Reset lists.
        """

        self._passes = 0
        self._flow_maps_x = None
        self._flow_maps_y = None

    @property
    def num_passes(self):
        return self._passes

    def iwe_formatting(self, warped_events, pol_mask, ts_list, tref, ts_scaling, interp_zeros=None, iwe_zeros=None):
        """
        Compues the images of warped events with event count and (accumulated) event timestamp information.
        :param warped_events: [batch_size x N x 2] warped events
        :param pol_mask: [batch_size x 4*N x 2] polarity mask of the (forward) warped events
        :param ts_list: [batch_size x N x 1] event timestamp [0, max_ts]
        :param tref: reference timestamp
        :param ts_scaling: with of the timestamp normalization window
        :return iwe: image of warped events
        :return iwe_ts: image of averaged timestamps
        """

        # normalize timestamps (the ref timestamp is 1, the rest decreases linearly towards 0 on the extremes)
        norm_ts = ts_list.clone()
        norm_ts = 1 - torch.abs(tref - norm_ts) / ts_scaling

        # interpolate forward
        idx, weights = get_interpolation(warped_events, self.res, zeros=interp_zeros)

        # per-polarity image of (forward) warped events
        iwe_pos = interpolate(idx, weights, self.res, polarity_mask=pol_mask[:, :, 0:1], zeros=iwe_zeros)
        iwe_neg = interpolate(idx, weights, self.res, polarity_mask=pol_mask[:, :, 1:2], zeros=iwe_zeros)
        iwe = torch.cat([iwe_pos, iwe_neg], dim=1)

        # image of (forward) warped averaged timestamps
        iwe_pos_ts = interpolate(idx, weights * norm_ts, self.res, polarity_mask=pol_mask[:, :, 0:1], zeros=iwe_zeros)
        iwe_neg_ts = interpolate(idx, weights * norm_ts, self.res, polarity_mask=pol_mask[:, :, 1:2], zeros=iwe_zeros)
        iwe_ts = torch.cat([iwe_pos_ts, iwe_neg_ts], dim=1)

        return iwe, iwe_ts

    def focus_loss(self, iwe, iwe_ts):
        """
        Scaling of the loss function based on the number of events in the image space.
        See "Self-Supervised Learning of Event-Based Optical Flow with Spiking Neural Networks",
        Hagenaars and Paredes-Valles et al., NeurIPS 2021
        :param iwe: [batch_size x N x 2] image of warped events
        :param iwe_ts: [batch_size x N x 2] image of averaged timestamps
        :return loss: loss value
        """

        iwe_ts = iwe_ts.view(iwe_ts.shape[0], 2, -1)
        loss = torch.sum(iwe_ts[:, 0, :] ** 2, dim=1) + torch.sum(iwe_ts[:, 1, :] ** 2, dim=1)
        if self.loss_scaling:
            nonzero_px = torch.sum(iwe, dim=1, keepdim=True).bool()
            nonzero_px = nonzero_px.view(nonzero_px.shape[0], -1)
            loss /= torch.sum(nonzero_px, dim=1) + 1e-9

        return torch.sum(loss)

    def flow_temporal_smoothing(self):
        """
        (Temporal) Scaled Charbonnier smoothness prior on the estimated optical flow vectors.
        :return loss: smoothing loss value
        """

        loss = 0
        for i in range(self._num_flows):
            for j in range(self._flow_maps_x[i].shape[1] - 1):

                # compute (backward) warping indices
                flow = torch.stack([self._flow_maps_y[i][:, j, ...], self._flow_maps_x[i][:, j, ...]], dim=1)
                warped_indices = self.indices + flow
                warped_indices = warped_indices.view(self.batch_size, 2, -1).permute(0, 2, 1)

                # ignore pixels that go out of the image space
                warping_mask = (
                    (warped_indices[..., 0] >= 0)
                    * (warped_indices[..., 0] <= self.res[0] - 1.0)
                    * (warped_indices[..., 1] >= 0)
                    * (warped_indices[..., 1] <= self.res[1] - 1.0)
                )

                # (backward) warp the next flow
                warped_flow = get_event_flow(
                    self._flow_maps_x[i][:, j + 1, ...], self._flow_maps_y[i][:, j + 1, ...], warped_indices
                )
                warped_flow = warped_flow.permute(0, 2, 1).view(self.batch_size, 2, self.res[0], self.res[1])

                # compute flow temporal consistency (charbonnier)
                flow_dt = torch.sqrt((flow - warped_flow) ** 2 + 1e-9)
                flow_dt = torch.sum(flow_dt, dim=1, keepdim=True).view(self.batch_size, -1)
                loss += torch.sum(flow_dt * warping_mask, dim=1) / (torch.sum(warping_mask, dim=1) + 1e-9)

        loss /= self._num_flows
        loss /= self._passes - 1

        return self.flow_temp_smooth_weight * loss.sum()

    def flow_spatial_smoothing(self):
        """
        (Spatial) Scaled Charbonnier smoothness prior on the estimated optical flow vectors.
        :return loss: smoothing loss value
        """

        loss = 0
        for i in range(self._num_flows):

            # forward differences (horizontal, vertical, and diagonals)
            flow_x_dx = self._flow_maps_x[i][:, :, :, :-1] - self._flow_maps_x[i][:, :, :, 1:]
            flow_y_dx = self._flow_maps_y[i][:, :, :, :-1] - self._flow_maps_y[i][:, :, :, 1:]
            flow_x_dy = self._flow_maps_x[i][:, :, :-1, :] - self._flow_maps_x[i][:, :, 1:, :]
            flow_y_dy = self._flow_maps_y[i][:, :, :-1, :] - self._flow_maps_y[i][:, :, 1:, :]
            flow_x_dxdy_dr = self._flow_maps_x[i][:, :, :-1, :-1] - self._flow_maps_x[i][:, :, 1:, 1:]
            flow_y_dxdy_dr = self._flow_maps_y[i][:, :, :-1, :-1] - self._flow_maps_y[i][:, :, 1:, 1:]
            flow_x_dxdy_ur = self._flow_maps_x[i][:, :, 1:, :-1] - self._flow_maps_x[i][:, :, :-1, 1:]
            flow_y_dxdy_ur = self._flow_maps_y[i][:, :, 1:, :-1] - self._flow_maps_y[i][:, :, :-1, 1:]

            # compute flow spatial consistency (charbonnier)
            flow_dx = torch.sqrt((flow_x_dx) ** 2 + 1e-6) + torch.sqrt((flow_y_dx) ** 2 + 1e-6)
            flow_dy = torch.sqrt((flow_x_dy) ** 2 + 1e-6) + torch.sqrt((flow_y_dy) ** 2 + 1e-6)
            flow_dxdy_dr = torch.sqrt((flow_x_dxdy_dr) ** 2 + 1e-6) + torch.sqrt((flow_y_dxdy_dr) ** 2 + 1e-6)
            flow_dxdy_ur = torch.sqrt((flow_x_dxdy_ur) ** 2 + 1e-6) + torch.sqrt((flow_y_dxdy_ur) ** 2 + 1e-6)

            flow_dx = flow_dx.view(self.batch_size, self._passes, -1)
            flow_dy = flow_dy.view(self.batch_size, self._passes, -1)
            flow_dxdy_dr = flow_dxdy_dr.view(self.batch_size, self._passes, -1)
            flow_dxdy_ur = flow_dxdy_ur.view(self.batch_size, self._passes, -1)

            flow_dx = flow_dx.mean(2).mean(1)
            flow_dy = flow_dy.mean(2).mean(1)
            flow_dxdy_ur = flow_dxdy_ur.mean(2).mean(1)
            flow_dxdy_dr = flow_dxdy_dr.mean(2).mean(1)

            loss += (flow_dx + flow_dy + flow_dxdy_dr + flow_dxdy_ur) / 4

        loss /= self._num_flows

        return self.flow_spat_smooth_weight * loss.sum()

    @abstractmethod
    def forward(self):
        raise NotImplementedError


class Linear(BaseEventWarping):
    """
    Contrast maximization loss from Hagenaars and Paredes-Valles et al. (NeurIPS 2021).
    """

    def __init__(self, config, device, loss_scaling=True):
        super().__init__(config, device, loss_scaling=loss_scaling)
        self._event_ts = []
        self._event_loc = []
        self._event_flow = []
        self._event_pol_mask = []

        self._d_event_ts = []
        self._d_event_loc = []
        self._d_event_flow = []
        self._d_event_pol_mask = []

    def update(self, flow_list, event_list, pol_mask, d_event_list, d_pol_mask):
        """
        Initialize/Update container lists of events and flow maps for forward warping.
        :param flow_list: [[batch_size x 2 x H x W]] list of optical flow (x, y) maps
        :param event_list: [batch_size x N x 4] input events (ts, y, x, p)
        :param pol_mask: [batch_size x N x 2] polarity mask (pos, neg)
        :param d_event_list: [batch_size x N x 4] detached input events (ts, y, x, p)
        :param d_pol_mask: [batch_size x N x 2] detached polarity mask (pos, neg)
        """

        # update base lists (flow maps)
        self.update_base(flow_list)

        # update event timestamps
        event_list[:, :, 0:1] += self._passes
        d_event_list[:, :, 0:1] += self._passes
        event_ts = event_list[:, :, 0:1].clone()
        d_event_ts = d_event_list[:, :, 0:1].clone()
        if self.config["loss"]["round_ts"]:
            event_ts[...] = event_ts.min() + 0.5
            d_event_ts[...] = d_event_ts.min() + 0.5
        self._event_ts.append(event_ts)
        self._d_event_ts.append(d_event_ts)

        # update event location
        self._event_loc.append(event_list[:, :, 1:3].clone())
        self._d_event_loc.append(d_event_list[:, :, 1:3].clone())

        # update polarity mask
        self._event_pol_mask.append(pol_mask.clone())
        self._d_event_pol_mask.append(d_pol_mask.clone())

        # update per-event flow vector
        event_flow = []
        d_event_flow = []
        for i in range(self._num_flows):
            event_flow.append(
                get_event_flow(
                    self._flow_maps_x[i][:, -1, ...],
                    self._flow_maps_y[i][:, -1, ...],
                    event_list[:, :, 1:3],
                )
            )
            with torch.no_grad():
                d_event_flow.append(
                    get_event_flow(
                        self._flow_maps_x[i][:, -1, ...],
                        self._flow_maps_y[i][:, -1, ...],
                        d_event_list[:, :, 1:3],
                    )
                )
        self._event_flow.append(event_flow)
        self._d_event_flow.append(d_event_flow)

        # update timestamp index
        self._passes += 1

    def reset(self):
        """
        Reset lists.
        """

        self.reset_base()
        self._event_ts = []
        self._event_loc = []
        self._event_flow = []
        self._event_pol_mask = []

        self._d_event_ts = []
        self._d_event_loc = []
        self._d_event_flow = []
        self._d_event_pol_mask = []

    def forward(self):

        loss = 0
        for s, scale in enumerate(self.passes_loss):

            loss_update = 0
            for w in range(2**s):
                low_pass = w * scale
                high_pass = (w + 1) * scale

                event_ts = torch.cat(self._event_ts[low_pass:high_pass], dim=1)
                event_loc = torch.cat(self._event_loc[low_pass:high_pass], dim=1)
                ts_list = torch.cat([event_ts for _ in range(4)], dim=1)

                d_event_ts = torch.cat(self._d_event_ts[low_pass:high_pass], dim=1)
                d_event_loc = torch.cat(self._d_event_loc[low_pass:high_pass], dim=1)
                d_ts_list = torch.cat([d_event_ts for _ in range(4)], dim=1)

                if not self.border_compensation:
                    event_pol_mask = torch.cat(self._event_pol_mask[low_pass:high_pass], dim=1)
                    event_pol_mask = torch.cat([event_pol_mask for _ in range(4)], dim=1)
                    d_event_pol_mask = torch.cat(self._d_event_pol_mask[low_pass:high_pass], dim=1)
                    d_event_pol_mask = torch.cat([d_event_pol_mask for _ in range(4)], dim=1)

                for i in range(self._num_flows):
                    if self.border_compensation:
                        event_pol_mask = torch.cat(self._event_pol_mask[low_pass:high_pass], dim=1)
                        d_event_pol_mask = torch.cat(self._d_event_pol_mask[low_pass:high_pass], dim=1)

                    # event propagation (with grads)
                    event_flow = torch.cat([flow[i] for flow in self._event_flow[low_pass:high_pass]], dim=1)
                    fw_events = event_propagation(event_ts, event_loc, event_flow, high_pass)
                    bw_events = event_propagation(event_ts, event_loc, event_flow, low_pass)

                    if self.border_compensation:
                        fw_events, event_pol_mask = purge_unfeasible(fw_events, event_pol_mask, self.res)
                        bw_events, event_pol_mask = purge_unfeasible(bw_events, event_pol_mask, self.res)
                        event_pol_mask = torch.cat([event_pol_mask for _ in range(4)], dim=1)

                    fw_iwe, fw_iwe_ts = self.iwe_formatting(
                        fw_events,
                        event_pol_mask,
                        ts_list,
                        high_pass,
                        scale,
                    )
                    bw_iwe, bw_iwe_ts = self.iwe_formatting(
                        bw_events,
                        event_pol_mask,
                        ts_list,
                        low_pass,
                        scale,
                    )

                    # event propagation (without grads)
                    d_event_flow = torch.cat([flow[i] for flow in self._d_event_flow[low_pass:high_pass]], dim=1)
                    d_fw_events = event_propagation(d_event_ts, d_event_loc, d_event_flow, high_pass)
                    d_bw_events = event_propagation(d_event_ts, d_event_loc, d_event_flow, low_pass)
                    if self.border_compensation:
                        d_fw_events, d_event_pol_mask = purge_unfeasible(d_fw_events, d_event_pol_mask, self.res)
                        d_bw_events, d_event_pol_mask = purge_unfeasible(d_bw_events, d_event_pol_mask, self.res)
                        d_event_pol_mask = torch.cat([d_event_pol_mask for _ in range(4)], dim=1)

                    d_fw_iwe, d_fw_iwe_ts = self.iwe_formatting(
                        d_fw_events,
                        d_event_pol_mask,
                        d_ts_list,
                        high_pass,
                        scale,
                    )
                    d_bw_iwe, d_bw_iwe_ts = self.iwe_formatting(
                        d_bw_events,
                        d_event_pol_mask,
                        d_ts_list,
                        low_pass,
                        scale,
                    )

                    # compute loss (forward)
                    fw_iwe = fw_iwe + d_fw_iwe
                    fw_iwe_ts = fw_iwe_ts + d_fw_iwe_ts
                    fw_iwe_ts = fw_iwe_ts / (fw_iwe + 1e-9)  # per-pixel and per-polarity timestamps
                    loss_update += self.focus_loss(fw_iwe, fw_iwe_ts)

                    # compute loss (backward)
                    bw_iwe = bw_iwe + d_bw_iwe
                    bw_iwe_ts = bw_iwe_ts + d_bw_iwe_ts
                    bw_iwe_ts = bw_iwe_ts / (bw_iwe + 1e-9)  # per-pixel and per-polarity timestamps
                    loss_update += self.focus_loss(bw_iwe, bw_iwe_ts)

            loss_update /= 2**s  # number of deblurring windows for a given scale
            loss_update /= 2  # number of deblurring points per deblurring window
            loss += loss_update

        # average loss over all flow predictions
        loss /= self.config["data"]["scales_loss"]
        loss /= self._num_flows

        # spatial smoothing of predicted flow vectors
        if self.flow_spat_smooth_weight is not None:
            loss += self.flow_spatial_smoothing()

        # temporal consistency of predicted flow vectors
        if self.flow_temp_smooth_weight is not None and self._passes > 1:
            loss += self.flow_temporal_smoothing()

        return loss


class Iterative(BaseEventWarping):
    """
    Contrast maximization loss from Hagenaars and Paredes-Valles et al. (NeurIPS 2021) but augmented
    with iterative event warping, loss computation at all intermediate (time) points, and multiple temporal scales.
    """

    def __init__(self, config, device, loss_scaling=True):
        if config["loss"]["iterative_mode"] == "four":
            config["data"]["passes_loss"] *= 2

        super().__init__(config, device, loss_scaling=loss_scaling)
        self._event_ts = []
        self._event_loc = []
        self._event_pol_mask = []

        self._d_event_ts = []
        self._d_event_loc = []
        self._d_event_pol_mask = []

        self.delta_passes = []
        for passes in self.passes_loss:
            if config["loss"]["iterative_mode"] == "one":
                self.delta_passes.append(passes // 1)
            elif config["loss"]["iterative_mode"] == "two":
                self.delta_passes.append(passes // 2)
            elif config["loss"]["iterative_mode"] == "four":
                self.delta_passes.append(passes // 4)

    def update(self, flow_list, event_list, pol_mask, d_event_list, d_pol_mask):
        """
        Initialize/Update container lists of events and flow maps for forward warping.
        :param flow_list: [[batch_size x 2 x H x W]] list of optical flow (x, y) maps
        :param event_list: [batch_size x N x 4] input events (ts, y, x, p)
        :param pol_mask: [batch_size x N x 2] polarity mask (pos, neg)
        :param d_event_list: [batch_size x N x 4] detached input events (ts, y, x, p)
        :param d_pol_mask: [batch_size x N x 2] detached polarity mask (pos, neg)
        """

        # update base lists (event data, flow maps, event masks)
        self.update_base(flow_list)

        # update event timestamps
        event_list[:, :, 0:1] += self._passes
        d_event_list[:, :, 0:1] += self._passes
        event_ts = event_list[:, :, 0:1].clone()
        d_event_ts = d_event_list[:, :, 0:1].clone()
        if self.config["loss"]["round_ts"]:
            event_ts[...] = event_ts.min() + 0.5
            d_event_ts[...] = d_event_ts.min() + 0.5
        self._event_ts.append(event_ts)
        self._d_event_ts.append(d_event_ts)

        # update event locations
        self._event_loc.append(event_list[:, :, 1:3].clone())
        self._d_event_loc.append(d_event_list[:, :, 1:3].clone())

        # update event polarity masks
        self._event_pol_mask.append(pol_mask.clone())
        self._d_event_pol_mask.append(d_pol_mask.clone())

        # update timestamp index
        self._passes += 1

    def reset(self):
        """
        Reset lists.
        """

        self.reset_base()
        self._event_ts = []
        self._event_loc = []
        self._event_pol_mask = []

        self._d_event_ts = []
        self._d_event_loc = []
        self._d_event_pol_mask = []

    def update_warping_indices(self, tref, t, cnt, mode="forward"):
        """
        Updates indices for forward/backward iterative event warping.
        :param tref: reference timestamp for deblurring
        :param t: current (starting) timestamp
        :param cnt: counter that counts the number of warping passes a set of events have undergone
        :param mode: "forward" or "backward" warping mode
        :return cnt: updated counter
        :return break_flag: flag indicating whether the warping should be performed
        :return sampling_idx: indices for flow sampling
        :return warping_ts: reference time for current warping
        """

        warping_ts = t + cnt
        sampling_idx = t + cnt
        if mode == "forward":
            warping_ts += 1
            break_flag = t + cnt < tref
            cnt += 1

        elif mode == "backward":
            break_flag = t + cnt >= tref
            cnt -= 1

        else:
            raise ValueError("Unknown warping mode: {}".format(mode))

        return cnt, break_flag, sampling_idx, warping_ts

    def event_warping(
        self,
        tref,
        t,
        i,
        mode,
        buffer_event_loc,
        buffer_event_ts,
        buffer_event_pol_mask,
        warped_events,
        warped_events_ts,
        warped_events_mask,
    ):
        """
        Perform forward/backward iterative event warping.
        :param tref: reference timestamp for deblurring
        :param t: current (starting) timestamp
        :param i: index indicating which optical flow map to use
        :param mode: "forward" or "backward" warping mode
        :param buffer_event_loc: [batch_size x N x 2] buffer for event locations
        :param buffer_event_ts: [batch_size x N x 1] buffer for event timestamps
        :param buffer_event_pol_mask: [batch_size x N x 2] buffer for event polarity masks
        :param warped_events: [[batch_size x N x 2]] list containing warped events at all trefs
        :param warped_events_ts: [[batch_size x N x 1]] list containing warped events timestamp at all trefs
        :param warped_events_mask: [[batch_size x N x 2]] list containing polarity masks for warped events at all trefs
        :return warped_events: [[batch_size x N x 2]] updated list containing warped events at all trefs
        :return warped_events_ts: [[batch_size x N x 1]] updated list containing warped events timestamp at all trefs
        :return warped_events_mask: [[batch_size x N x 2]] updated list containing polarity masks for warped events at all trefs
        """

        cnt = 0
        event_loc = buffer_event_loc[t].clone()
        event_warp_ts = buffer_event_ts[t].clone()
        event_pol_mask = buffer_event_pol_mask[t].clone()
        while True:
            cnt, break_flag, sampling_idx, warping_ts = self.update_warping_indices(tref, t, cnt, mode=mode)
            if not break_flag:
                break

            # sample optical flow
            event_flow = get_event_flow(
                self._flow_maps_x[i][:, sampling_idx, ...],
                self._flow_maps_y[i][:, sampling_idx, ...],
                event_loc,
            )

            # event warping
            event_loc = event_propagation(
                event_warp_ts,
                event_loc,
                event_flow,
                warping_ts,
            )
            event_warp_ts[...] = warping_ts
            event_loc, event_pol_mask = purge_unfeasible(
                event_loc,
                event_pol_mask,
                self.res,
            )

            # update warping information (when in range)
            warped_events[warping_ts][t] = event_loc.clone()
            warped_events_ts[warping_ts][t] = buffer_event_ts[t].clone()
            warped_events_mask[warping_ts][t] = event_pol_mask.clone()

        return warped_events, warped_events_ts, warped_events_mask

    def forward(self):

        loss = 0
        max_passes = max(self.passes_loss)
        for i in range(self._num_flows):
            none_list = [None for _ in range(max_passes)]

            # iterative event warping
            event_ts = [none_list.copy() for _ in range(max_passes + 1)]
            event_loc = [none_list.copy() for _ in range(max_passes + 1)]
            event_pol_mask = [none_list.copy() for _ in range(max_passes + 1)]
            for t in range(max_passes):
                event_loc, event_ts, event_pol_mask = self.event_warping(
                    max_passes,
                    t,
                    i,
                    "forward",
                    self._event_loc,
                    self._event_ts,
                    self._event_pol_mask,
                    event_loc,
                    event_ts,
                    event_pol_mask,
                )
                event_loc, event_ts, event_pol_mask = self.event_warping(
                    0,
                    t,
                    i,
                    "backward",
                    self._event_loc,
                    self._event_ts,
                    self._event_pol_mask,
                    event_loc,
                    event_ts,
                    event_pol_mask,
                )

            # detached iterative event warping
            d_event_ts = [none_list.copy() for _ in range(max_passes + 1)]
            d_event_loc = [none_list.copy() for _ in range(max_passes + 1)]
            d_event_pol_mask = [none_list.copy() for _ in range(max_passes + 1)]
            with torch.no_grad():
                for t in range(max_passes):
                    d_event_loc, d_event_ts, d_event_pol_mask = self.event_warping(
                        max_passes,
                        t,
                        i,
                        "forward",
                        self._d_event_loc,
                        self._d_event_ts,
                        self._d_event_pol_mask,
                        d_event_loc,
                        d_event_ts,
                        d_event_pol_mask,
                    )
                    d_event_loc, d_event_ts, d_event_pol_mask = self.event_warping(
                        0,
                        t,
                        i,
                        "backward",
                        self._d_event_loc,
                        self._d_event_ts,
                        self._d_event_pol_mask,
                        d_event_loc,
                        d_event_ts,
                        d_event_pol_mask,
                    )

            # learning from multiple temporal scales (i.e., different amounts of blur)
            for s, scale in enumerate(self.passes_loss):

                loss_update = 0
                for w in range(2**s):
                    low_pass = w * scale
                    high_pass = (w + 1) * scale

                    low_tref = low_pass
                    high_tref = high_pass + 1
                    if self.config["loss"]["iterative_mode"] == "four":
                        low_tref = low_pass + self.delta_passes[s]
                        high_tref = low_pass + 3 * self.delta_passes[s] + 1

                    # merge event masks (to deal with partially-observable edges)
                    if self.border_compensation:
                        shared_event_pol_mask = none_list.copy()
                        shared_d_event_pol_mask = none_list.copy()
                        for t in range(low_tref, high_tref - 1):
                            tmp_event_pol_mask = event_pol_mask[low_tref][t].clone()
                            tmp_d_event_pol_mask = d_event_pol_mask[low_tref][t].clone()
                            for tref in range(low_tref + 1, high_tref):
                                tmp_event_pol_mask *= event_pol_mask[tref][t].clone()
                                tmp_d_event_pol_mask *= d_event_pol_mask[tref][t].clone()
                            shared_event_pol_mask[t] = tmp_event_pol_mask
                            shared_d_event_pol_mask[t] = tmp_d_event_pol_mask

                    # compute loss at all intermediate points
                    for tref in range(low_tref, high_tref):
                        low_extreme = max(low_pass, tref - self.delta_passes[s])
                        high_extreme = min(high_pass, tref + self.delta_passes[s])

                        # image of warped events (with grads)
                        ts_list = torch.cat(event_ts[tref][low_extreme:high_extreme], dim=1)
                        warped_events = torch.cat(event_loc[tref][low_extreme:high_extreme], dim=1)
                        if self.border_compensation:
                            warped_pol_mask = torch.cat(shared_event_pol_mask[low_extreme:high_extreme], dim=1)
                        else:
                            warped_pol_mask = torch.cat(event_pol_mask[tref][low_extreme:high_extreme], dim=1)

                        ts_list = torch.cat([ts_list for _ in range(4)], dim=1)
                        warped_pol_mask = torch.cat([warped_pol_mask for _ in range(4)], dim=1)
                        iwe, iwe_ts = self.iwe_formatting(
                            warped_events,
                            warped_pol_mask,
                            ts_list,
                            tref,
                            self.delta_passes[s],
                        )

                        # image of warped events (without grads)
                        d_ts_list = torch.cat(d_event_ts[tref][low_extreme:high_extreme], dim=1)
                        d_warped_events = torch.cat(d_event_loc[tref][low_extreme:high_extreme], dim=1)
                        if self.border_compensation:
                            d_warped_pol_mask = torch.cat(shared_d_event_pol_mask[low_extreme:high_extreme], dim=1)
                        else:
                            d_warped_pol_mask = torch.cat(d_event_pol_mask[tref][low_extreme:high_extreme], dim=1)

                        d_ts_list = torch.cat([d_ts_list for _ in range(4)], dim=1)
                        d_warped_pol_mask = torch.cat([d_warped_pol_mask for _ in range(4)], dim=1)
                        d_iwe, d_iwe_ts = self.iwe_formatting(
                            d_warped_events,
                            d_warped_pol_mask,
                            d_ts_list,
                            tref,
                            self.delta_passes[s],
                        )

                        # combination of images of warped events (with and without grads)
                        iwe = iwe + d_iwe
                        iwe_ts = iwe_ts + d_iwe_ts
                        iwe_ts = iwe_ts / (iwe + 1e-9)  # per-pixel and per-polarity timestamps
                        loss_update += self.focus_loss(iwe, iwe_ts)

                loss_update /= 2**s  # number of deblurring windows for a given scale
                loss_update /= 2 * self.delta_passes[s] + 1  # number of deblurring points per deblurring window
                loss += loss_update

        # average loss over all flow predictions and deblurring points
        loss /= self.config["data"]["scales_loss"]
        loss /= self._num_flows

        # spatial smoothing of predicted flow vectors
        if self.flow_spat_smooth_weight is not None:
            loss += self.flow_spatial_smoothing()

        # temporal consistency of predicted flow vectors
        if self.flow_temp_smooth_weight is not None and self._passes > 1:
            loss += self.flow_temporal_smoothing()

        return loss

loss/flow.py
tudelft/taming_event_flow
import os
import sys

import torch

parent_dir_name = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
sys.path.append(parent_dir_name)

from utils.iwe import event_propagation, get_event_flow, purge_unfeasible, get_interpolation, interpolate


class BaseValidation(torch.nn.Module):
    """
    Base class for validation metrics.
    """

    def __init__(self, config, device):
        super(BaseValidation, self).__init__()
        self.res = config["loader"]["resolution"]
        self.device = device
        self.config = config

        self._passes = 0
        self._event_ts = None
        self._event_loc = None
        self._event_pol_mask = None

        self._flow_maps_x = None
        self._flow_maps_y = None
        self._event_mask = None

        # warping indices (for forward-propagated flow)
        my, mx = torch.meshgrid(torch.arange(self.res[0]), torch.arange(self.res[1]))
        indices = torch.stack([my, mx], dim=0).unsqueeze(0)
        self.indices = indices.float().view(1, 2, -1).permute(0, 2, 1).to(device)
        self.indices_map = self.indices.clone().permute(0, 2, 1).view(1, 2, self.res[0], self.res[1])
        self.indices_mask = torch.ones((1, self.res[0] * self.res[1], 1)).to(device)

    @property
    def num_passes(self):
        return self._passes

    def forward_prop_flow(self, i, tref, flow_maps_x, flow_maps_y):
        """
        Forward propagation of the estimated optical flow using bilinear interpolation.
        :param i: time at which the flow map to be warped is defined
        :param tref: reference time for the forward propagation
        :return warped_flow_x: [[batch_size x 1 x H x W]] warped, horizontal optical flow map
        :return warped_flow_y: [batch_size x 1 x H x W] warped, vertical optical flow map
        """

        # sample per-pixel optical flow
        indices_mask = self.indices_mask.clone()
        indices_flow = get_event_flow(flow_maps_x[:, i, ...], flow_maps_y[:, i, ...], self.indices)

        # optical flow (forward) propagation
        warped_indices = event_propagation(i, self.indices, indices_flow, tref)
        warped_indices, indices_mask = purge_unfeasible(warped_indices, indices_mask, self.res)

        # (bilinearly) interpolate forward
        indices_mask = torch.cat([indices_mask for _ in range(4)], dim=1)
        indices_flow = torch.cat([indices_flow for _ in range(4)], dim=1)
        interp_warped_indices, interp_weights = get_interpolation(warped_indices, self.res)
        warped_weights = interpolate(interp_warped_indices, interp_weights, self.res, polarity_mask=indices_mask)
        warped_flow_y = interpolate(
            interp_warped_indices, interp_weights * indices_flow[..., 0:1], self.res, polarity_mask=indices_mask
        )
        warped_flow_x = interpolate(
            interp_warped_indices, interp_weights * indices_flow[..., 1:2], self.res, polarity_mask=indices_mask
        )
        warped_flow_y /= warped_weights + 1e-9
        warped_flow_x /= warped_weights + 1e-9

        return warped_flow_x, warped_flow_y

    def update_base(self, flow_list, event_list, pol_mask, event_mask):
        """
        Initialize/Update container lists of events and flow maps for forward warping.
        :param flow_list: [[batch_size x 2 x H x W]] list of optical flow (x, y) maps
        :param event_list: [batch_size x N x 4] input events (ts, y, x, p)
        :param pol_mask: [batch_size x N x 2] polarity mask (pos, neg)
        :param event_mask: [batch_size x 1 x H x W] event mask
        """

        # update event timestamps
        event_list[:, :, 0:1] += self._passes  # only nonzero second time
        event_ts = event_list[:, :, 0:1].clone()
        if self.config["loss"]["round_ts"]:
            event_ts[...] = event_ts.min() + 0.5

        if self._event_ts is None:
            self._event_ts = event_ts
            self._event_loc = event_list[:, :, 1:3].clone()
            self._event_pol_mask = pol_mask.clone()
        else:
            self._event_ts = torch.cat([self._event_ts, event_ts], dim=1)
            self._event_loc = torch.cat([self._event_loc, event_list[:, :, 1:3].clone()], dim=1)
            self._event_pol_mask = torch.cat([self._event_pol_mask, pol_mask.clone()], dim=1)

        # update optical flow maps
        flow = flow_list[-1]  # only highest resolution flow
        if self._flow_maps_x is None:
            self._flow_maps_x = flow[:, 0:1, :, :]
            self._flow_maps_y = flow[:, 1:2, :, :]
        else:
            self._flow_maps_x = torch.cat([self._flow_maps_x, flow[:, 0:1, :, :]], dim=1)
            self._flow_maps_y = torch.cat([self._flow_maps_y, flow[:, 1:2, :, :]], dim=1)

        # update internal smoothing mask
        if self._event_mask is None:
            self._event_mask = event_mask
        else:
            self._event_mask = torch.cat([self._event_mask, event_mask], dim=1)

    def reset_base(self):
        """
        Reset lists.
        """

        self._passes = 0
        self._event_ts = None
        self._event_loc = None
        self._event_pol_mask = None

        self._flow_maps_x = None
        self._flow_maps_y = None
        self._event_mask = None

    def window_events_base(self, round_idx=False):
        """
        :param round_idx: if True, round the event coordinates to the nearest integer.
        :return: image-like representation of all the events in the validation time/event window.
        """

        pol_mask_list = self._event_pol_mask
        if not round_idx:
            pol_mask_list = torch.cat([pol_mask_list for i in range(4)], dim=1)

        fw_idx, fw_weights = get_interpolation(self._event_loc, self.res, round_idx=round_idx)
        events_pos = interpolate(fw_idx.long(), fw_weights, self.res, polarity_mask=pol_mask_list[:, :, 0:1])
        events_neg = interpolate(fw_idx.long(), fw_weights, self.res, polarity_mask=pol_mask_list[:, :, 1:2])

        return torch.cat([events_pos, events_neg], dim=1)

    def window_flow_base(self, flow_maps_x, flow_maps_y, mask=False):
        """
        :param flow_maps_x: [batch_size x num_passes x H x W] horizontal flow maps to be averaged
        :param flow_maps_y: [batch_size x num_passes x H x W] vertical flow maps to be averaged
        :return avg_flow: image-like representation of the per-pixel average flow in the validation time/event window.
        """

        flow_x = flow_maps_x[:, 0:1, :, :]
        flow_y = flow_maps_y[:, 0:1, :, :]
        avg_flow = torch.cat([flow_x, flow_y], dim=1)
        flow_mask = (flow_x != 0.0) + (flow_y != 0.0)
        cnt = flow_mask.float()

        for i in range(1, flow_maps_x.shape[1]):
            flow_x = flow_maps_x[:, i : i + 1, :, :]
            flow_y = flow_maps_y[:, i : i + 1, :, :]
            avg_flow += torch.cat([flow_x, flow_y], dim=1)
            flow_mask = (flow_x != 0.0) + (flow_y != 0.0)
            cnt += flow_mask.float()

        if mask:
            mask = torch.sum(self._event_mask, dim=1, keepdim=True) > 0.0
            avg_flow *= mask.float()

        return avg_flow / (cnt + 1e-9)

    def window_iwe_base(self, round_idx=False):
        """
        Assumption: events have NOT been previously warped in a forward fashion in the update() method.
        :param round_idx: if True, round the event coordinates to the nearest integer.
        :return: image-like representation of the IWE of all the events in the validation time/event window.
        """

        pol_mask_list = self._event_pol_mask
        if not round_idx:
            pol_mask_list = torch.cat([pol_mask_list for i in range(4)], dim=1)

        fw_events = event_propagation(self._event_ts, self._event_loc, self._event_flow, self._passes)
        fw_idx, fw_weights = get_interpolation(fw_events, self.res, round_idx=round_idx)
        fw_iwe_pos = interpolate(fw_idx.long(), fw_weights, self.res, polarity_mask=pol_mask_list[:, :, 0:1])
        fw_iwe_neg = interpolate(fw_idx.long(), fw_weights, self.res, polarity_mask=pol_mask_list[:, :, 1:2])

        return torch.cat([fw_iwe_pos, fw_iwe_neg], dim=1)

    def compute_fwl(self, fw_events, zero_events, fw_pol_mask, zero_pol_mask):
        """
        The Flow Warp Loss (FWL) metric is the ratio of the variance of the image of warped events
        and that of the image of (non-warped) events; hence, the higher the value of this metric,
        the better the optical flow estimate.
        See 'Reducing the Sim-to-Real Gap for Event Cameras',
        Stoffregen et al., ECCV 2020.
        """

        # interpolate forward
        fw_idx, fw_weights = get_interpolation(fw_events, self.res, round_idx=True)

        # image of (forward) warped averaged timestamps
        fw_iwe_pos = interpolate(fw_idx, fw_weights, self.res, polarity_mask=fw_pol_mask[:, :, 0:1])
        fw_iwe_neg = interpolate(fw_idx, fw_weights, self.res, polarity_mask=fw_pol_mask[:, :, 1:2])
        fw_iwe = fw_iwe_pos + fw_iwe_neg

        # image of non-warped averaged timestamps
        zero_idx, zero_weights = get_interpolation(zero_events, self.res, round_idx=True)
        zero_iwe_pos = interpolate(zero_idx, zero_weights, self.res, polarity_mask=zero_pol_mask[:, :, 0:1])
        zero_iwe_neg = interpolate(zero_idx, zero_weights, self.res, polarity_mask=zero_pol_mask[:, :, 1:2])
        zero_iwe = zero_iwe_pos + zero_iwe_neg

        return fw_iwe.var() / zero_iwe.var()

    def compute_rsat(self, fw_events, zero_events, fw_pol_mask, zero_pol_mask, ts_list):
        """
        The Ratio of the Squared Averaged Timestamps (RSAT) metric is the ratio of the squared sum of the per-pixel and
        per-polarity average timestamp of the image of warped events and that of the image of (non-warped) events; hence,
        the lower the value of this metric, the better the optical flow estimate.
        See 'Self-Supervised Learning of Event-Based Optical Flow with Spiking Neural Networks',
        Hagenaars and Paredes-Valles et al., NeurIPS 2021.
        """

        # interpolate forward
        fw_idx, fw_weights = get_interpolation(fw_events, self.res, round_idx=True)

        # image of (forward) warped averaged timestamps
        fw_iwe_pos = interpolate(fw_idx, fw_weights, self.res, polarity_mask=fw_pol_mask[:, :, 0:1])
        fw_iwe_neg = interpolate(fw_idx, fw_weights, self.res, polarity_mask=fw_pol_mask[:, :, 1:2])
        fw_iwe_pos_ts = interpolate(fw_idx, fw_weights * ts_list, self.res, polarity_mask=fw_pol_mask[:, :, 0:1])
        fw_iwe_neg_ts = interpolate(fw_idx, fw_weights * ts_list, self.res, polarity_mask=fw_pol_mask[:, :, 1:2])
        fw_iwe_pos_ts /= fw_iwe_pos + 1e-9
        fw_iwe_neg_ts /= fw_iwe_neg + 1e-9
        fw_iwe_pos_ts = fw_iwe_pos_ts / self._passes
        fw_iwe_neg_ts = fw_iwe_neg_ts / self._passes

        # image of non-warped averaged timestamps
        zero_idx, zero_weights = get_interpolation(zero_events, self.res, round_idx=True)
        zero_iwe_pos = interpolate(zero_idx, zero_weights, self.res, polarity_mask=zero_pol_mask[:, :, 0:1])
        zero_iwe_neg = interpolate(zero_idx, zero_weights, self.res, polarity_mask=zero_pol_mask[:, :, 1:2])
        zero_iwe_pos_ts = interpolate(
            zero_idx, zero_weights * ts_list, self.res, polarity_mask=zero_pol_mask[:, :, 0:1]
        )
        zero_iwe_neg_ts = interpolate(
            zero_idx, zero_weights * ts_list, self.res, polarity_mask=zero_pol_mask[:, :, 1:2]
        )
        zero_iwe_pos_ts /= zero_iwe_pos + 1e-9
        zero_iwe_neg_ts /= zero_iwe_neg + 1e-9
        zero_iwe_pos_ts = zero_iwe_pos_ts / self._passes
        zero_iwe_neg_ts = zero_iwe_neg_ts / self._passes

        # (scaled) sum of the squares of the per-pixel and per-polarity average timestamps
        fw_iwe_pos_ts = fw_iwe_pos_ts.view(fw_iwe_pos_ts.shape[0], -1)
        fw_iwe_neg_ts = fw_iwe_neg_ts.view(fw_iwe_neg_ts.shape[0], -1)
        fw_iwe_pos_ts = torch.sum(fw_iwe_pos_ts**2, dim=1)
        fw_iwe_neg_ts = torch.sum(fw_iwe_neg_ts**2, dim=1)
        fw_ts_sum = fw_iwe_pos_ts + fw_iwe_neg_ts

        fw_nonzero_px = fw_iwe_pos + fw_iwe_neg
        fw_nonzero_px[fw_nonzero_px > 0] = 1
        fw_nonzero_px = fw_nonzero_px.view(fw_nonzero_px.shape[0], -1)
        fw_ts_sum /= torch.sum(fw_nonzero_px, dim=1)

        zero_iwe_pos_ts = zero_iwe_pos_ts.view(zero_iwe_pos_ts.shape[0], -1)
        zero_iwe_neg_ts = zero_iwe_neg_ts.view(zero_iwe_neg_ts.shape[0], -1)
        zero_iwe_pos_ts = torch.sum(zero_iwe_pos_ts**2, dim=1)
        zero_iwe_neg_ts = torch.sum(zero_iwe_neg_ts**2, dim=1)
        zero_ts_sum = zero_iwe_pos_ts + zero_iwe_neg_ts

        zero_nonzero_px = zero_iwe_pos + zero_iwe_neg
        zero_nonzero_px[zero_nonzero_px > 0] = 1
        zero_nonzero_px = zero_nonzero_px.view(zero_nonzero_px.shape[0], -1)
        zero_ts_sum /= torch.sum(zero_nonzero_px, dim=1)

        return fw_ts_sum / zero_ts_sum

    def compute_aee(self, pred, gt, mask=None):
        """
        Average endpoint error (i.e., Euclidean distance).
        """

        # compute AEE
        batch_size = pred.shape[0]
        error = (pred - gt).pow(2).sum(1).sqrt()

        # AEE not computed in pixels without valid ground truth
        gtflow_mask = (gt[:, 0, :, :] == 0.0) * (gt[:, 1, :, :] == 0.0)
        gtflow_mask = ~gtflow_mask

        # AEE not computed in pixels without input events (MVSEC)
        if mask is not None:
            mask = torch.sum(mask, axis=1)
            mask = mask > 0

            if "res_aee" in self.config["metrics"].keys():
                yoff = (self.res[0] - self.config["metrics"]["res_aee"][0]) // 2
                xoff = (self.res[1] - self.config["metrics"]["res_aee"][1]) // 2
                mask = mask[:, yoff:-yoff, xoff:-xoff].contiguous()
                error = error[:, yoff:-yoff, xoff:-xoff].contiguous()
                gtflow_mask = gtflow_mask[:, yoff:-yoff, xoff:-xoff].contiguous()

            if "vertical_crop_aee" in self.config["metrics"].keys():
                mask = mask[:, : self.config["metrics"]["vertical_crop_aee"], :]
                error = error[:, : self.config["metrics"]["vertical_crop_aee"], :]
                gtflow_mask = gtflow_mask[:, : self.config["metrics"]["vertical_crop_aee"], :]

            gtflow_mask = gtflow_mask * mask

        # compute AEE
        error = error.view(batch_size, -1)
        gtflow_mask = gtflow_mask.view(batch_size, -1)
        error = error[gtflow_mask]
        aee = torch.mean(error, dim=0)

        return aee


class Linear(BaseValidation):
    """
    Linear event warping validation class.
    """

    def __init__(self, config, device):
        super().__init__(config, device)
        self._event_flow = None

    def update(self, flow_list, event_list, pol_mask, event_mask):
        """
        Initialize/Update container lists of events and flow maps for forward warping.
        :param flow_list: [[batch_size x 2 x H x W]] list of optical flow (x, y) maps
        :param event_list: [batch_size x N x 4] input events (ts, y, x, p)
        :param pol_mask: [batch_size x N x 2] polarity mask (pos, neg)
        :param event_mask: [batch_size x 1 x H x W] event mask
        """

        # update base lists (event data, flow maps, event masks)
        self.update_base(flow_list, event_list, pol_mask, event_mask)

        # get flow for every event in the list
        event_flow = get_event_flow(
            self._flow_maps_x[:, -1, ...],
            self._flow_maps_y[:, -1, ...],
            event_list[:, :, 1:3],
        )

        if self._event_flow is None:
            self._event_flow = event_flow
        else:
            self._event_flow = torch.cat([self._event_flow, event_flow], dim=1)

        # update timestamp index
        self._passes += 1

    def reset(self):
        """
        Reset lists.
        """

        self.reset_base()
        self._event_flow = None

    def window_events(self, round_idx=False):
        """
        :param round_idx: if True, round the event coordinates to the nearest integer.
        :return: image-like representation of all the events in the validation time/event window.
        """

        return self.window_events_base(round_idx)

    def window_flow(self, mode=None, mask=None):
        """
        :return avg_flow: image-like representation of the per-pixel average flow in the validation time/event window.
        """

        if mask is None:
            mask = self.config["vis"]["mask_output"]

        # copy flow tensors to prevent overwriting
        flow_maps_x = self._flow_maps_x.clone()
        flow_maps_y = self._flow_maps_y.clone()

        # forward propagation of the estimated optical flow
        for i in range(self._passes - 1):
            warped_flow_x, warped_flow_y = self.forward_prop_flow(
                i, self._passes - 1, self._flow_maps_x, self._flow_maps_y
            )

            # update lists
            flow_maps_x[:, i : i + 1, ...] = warped_flow_x
            flow_maps_y[:, i : i + 1, ...] = warped_flow_y

        return self.window_flow_base(flow_maps_x, flow_maps_y, mask=mask)

    def window_iwe(self, mode=None, round_idx=False):
        """
        Assumption: events have NOT been previously warped in a forward fashion in the update() method.
        :param round_idx: if True, round the event coordinates to the nearest integer.
        :return: image-like representation of the IWE of all the events in the validation time/event window.
        """

        return self.window_iwe_base(round_idx)

    def rsat(self):
        """
        :return rsat: deblur metric for validation of the estimated optical flow.
        """

        fw_events = event_propagation(self._event_ts, self._event_loc, self._event_flow, self._passes)
        return self.compute_rsat(fw_events, self._event_loc, self._event_pol_mask, self._event_pol_mask, self._event_ts)

    def fwl(self):
        """
        :return fwl: deblur metric for validation of the estimated optical flow (Stoffregen et al, ECCV 2020).
        """

        fw_events = event_propagation(self._event_ts, self._event_loc, self._event_flow, self._passes)
        return self.compute_fwl(fw_events, self._event_loc, self._event_pol_mask, self._event_pol_mask)


class Iterative(BaseValidation):
    """
    Iterative event warping validation class.
    """

    def __init__(self, config, device):
        super().__init__(config, device)
        self._fw_event_loc = None
        self._fw_event_warp_ts = None
        self._fw_event_pol_mask = None

        self._bw_event_loc = None
        self._bw_event_pol_mask = None

        self._fw_prop_flow_maps_x = None
        self._fw_prop_flow_maps_y = None

        self._accum_flow_map_x = None
        self._accum_flow_map_y = None
        self._flow_warping_indices = None
        self._flow_out_mask = torch.zeros(1, 1, self.res[0], self.res[1]).to(device)

    def update_fw_event_lists(self, event_list, event_pol_mask):
        """
        Initialize/Update container lists of events to be udpated during foward warping
        :param event_list: [batch_size x N x 4] input events (ts, y, x, p)
        :param event_pol_mask: [batch_size x N x 2] event polarity mask
        """

        event_ts = event_list[:, :, 0:1].clone()
        if self.config["loss"]["round_ts"]:
            event_ts[...] = event_ts.min() + 0.5

        if self._fw_event_loc is None:
            self._fw_event_warp_ts = event_ts
            self._fw_event_loc = event_list[:, :, 1:3].clone()
            self._fw_event_pol_mask = event_pol_mask.clone()

        else:
            self._fw_event_warp_ts = torch.cat([self._fw_event_warp_ts, event_ts], dim=1)
            self._fw_event_loc = torch.cat([self._fw_event_loc, event_list[:, :, 1:3].clone()], dim=1)
            self._fw_event_pol_mask = torch.cat([self._fw_event_pol_mask, event_pol_mask.clone()], dim=1)

    def update_bw_event_lists(self, event_loc, event_pol_mask):
        """
        Initialize/Update container lists of events to be udpated during foward warping
        :param event_list: [batch_size x N x 2] input events (ts, y, x, p)
        :param event_pol_mask: [batch_size x N x 2] event polarity mask
        """

        if self._bw_event_loc is None:
            self._bw_event_loc = event_loc.clone()
            self._bw_event_pol_mask = event_pol_mask.clone()

        else:
            self._bw_event_loc = torch.cat([self._bw_event_loc, event_loc.clone()], dim=1)
            self._bw_event_pol_mask = torch.cat([self._bw_event_pol_mask, event_pol_mask.clone()], dim=1)

    def update(self, flow_list, event_list, pol_mask, event_mask):
        """
        Initialize/Update container lists of events and flow maps for forward warping.
        :param flow_list: [[batch_size x 2 x H x W]] list of optical flow (x, y) maps
        :param event_list: [batch_size x N x 4] input events (ts, y, x, p)
        :param pol_mask: [batch_size x N x 2] polarity mask (pos, neg)
        :param event_mask: [batch_size x 1 x H x W] event mask
        """

        # update base lists (event data, flow maps, event masks)
        self.update_base(flow_list, event_list, pol_mask, event_mask)

        ############
        # FORWARD WARPING
        ############

        # initialize and update event lists for fw warping
        self.update_fw_event_lists(event_list, pol_mask)

        # sample optical flow
        fw_event_flow = get_event_flow(
            self._flow_maps_x[:, -1, ...],
            self._flow_maps_y[:, -1, ...],
            self._fw_event_loc,
        )

        # event warping process
        self._fw_event_loc = event_propagation(
            self._fw_event_warp_ts,
            self._fw_event_loc,
            fw_event_flow,
            self._passes + 1,
        )
        self._fw_event_loc, self._fw_event_pol_mask = purge_unfeasible(
            self._fw_event_loc,
            self._fw_event_pol_mask,
            self.res,
        )

        # update warping times
        self._fw_event_warp_ts[...] = self._passes + 1

        ############
        # BACKWARD WARPING
        ############

        bw_event_loc = event_list[:, :, 1:3].clone()
        bw_event_pol_mask = pol_mask.clone()
        bw_event_warp_ts = event_list[:, :, 0:1].clone()
        if self.config["loss"]["round_ts"]:
            bw_event_warp_ts[...] = bw_event_warp_ts.min() + 0.5

        cnt = 0
        while self._passes + cnt >= 0:

            # sample optical flow
            bw_event_flow = get_event_flow(
                self._flow_maps_x[:, self._passes + cnt, ...],
                self._flow_maps_y[:, self._passes + cnt, ...],
                bw_event_loc,
            )

            # event warping process
            bw_event_loc = event_propagation(
                bw_event_warp_ts,
                bw_event_loc,
                bw_event_flow,
                self._passes + cnt,
            )
            bw_event_loc, bw_event_pol_mask = purge_unfeasible(
                bw_event_loc,
                bw_event_pol_mask,
                self.res,
            )

            # update warping times
            bw_event_warp_ts[...] = self._passes + cnt
            cnt -= 1

        self.update_bw_event_lists(bw_event_loc, bw_event_pol_mask)

        ########################
        # FORWARD-PROPAGATED FLOW
        ########################

        # forward propagation of the estimated optical flow
        flow = flow_list[-1]  # only highest resolution flow
        if self._fw_prop_flow_maps_x is None:
            self._fw_prop_flow_maps_x = flow[:, 0:1, :, :]
            self._fw_prop_flow_maps_y = flow[:, 1:2, :, :]
        else:
            self._fw_prop_flow_maps_x = torch.cat([self._fw_prop_flow_maps_x, flow[:, 0:1, :, :]], dim=1)
            self._fw_prop_flow_maps_y = torch.cat([self._fw_prop_flow_maps_y, flow[:, 1:2, :, :]], dim=1)

        for i in range(self._passes):
            warped_flow_x, warped_flow_y = self.forward_prop_flow(
                i, i + 1, self._fw_prop_flow_maps_x, self._fw_prop_flow_maps_y
            )
            self._fw_prop_flow_maps_x[:, i : i + 1, ...] = warped_flow_x
            self._fw_prop_flow_maps_y[:, i : i + 1, ...] = warped_flow_y

        ########################
        # ACCUMULATED FLOW (BACKWARD WARPING)
        ########################

        indices = self.indices_map.clone()
        if self._flow_warping_indices is not None:
            indices = self._flow_warping_indices.clone()

        mask_valid = (
            (indices[:, 0:1, ...] >= 0)
            * (indices[:, 0:1, ...] <= self.res[0] - 1.0)
            * (indices[:, 1:2, ...] >= 0)
            * (indices[:, 1:2, ...] <= self.res[1] - 1.0)
        )
        self._flow_out_mask += mask_valid.float()

        curr_flow = get_event_flow(
            self._flow_maps_x[:, -1, ...],
            self._flow_maps_y[:, -1, ...],
            indices.view(1, 2, -1).permute(0, 2, 1),
        )
        curr_flow = curr_flow.permute(0, 2, 1).view(1, 2, self.res[0], self.res[1])

        warped_indices = indices + curr_flow * mask_valid.float()
        self._accum_flow_map_x = warped_indices[:, 1:2, :, :] - self.indices_map[:, 1:2, :, :]
        self._accum_flow_map_y = warped_indices[:, 0:1, :, :] - self.indices_map[:, 0:1, :, :]
        self._flow_warping_indices = warped_indices

        # update timestamp index
        self._passes += 1

    def reset(self):
        """
        Reset lists.
        """

        self.reset_base()
        self._fw_event_loc = None
        self._fw_event_warp_ts = None
        self._fw_event_pol_mask = None

        self._bw_event_loc = None
        self._bw_event_pol_mask = None

        self._fw_prop_flow_maps_x = None
        self._fw_prop_flow_maps_y = None

        self._accum_flow_map_x = None
        self._accum_flow_map_y = None
        self._flow_warping_indices = None
        self._flow_out_mask = torch.zeros(1, 1, self.res[0], self.res[1]).to(self.device)

    def window_events(self, round_idx=False):
        """
        :param round_idx: if True, round the event coordinates to the nearest integer.
        :return: image-like representation of all the events in the validation time/event window.
        """

        return self.window_events_base(round_idx)

    def window_flow(self, mode=None, mask=None):
        """
        :return avg_flow: image-like representation of the per-pixel average flow in the validation time/event window.
        """

        if mask is None:
            mask = self.config["vis"]["mask_output"]

        if mode == "forward":
            return self.window_flow_base(self._fw_prop_flow_maps_x, self._fw_prop_flow_maps_y, mask=mask)
        elif mode == "backward":
            return self.window_flow_base(
                self._accum_flow_map_x / self._flow_out_mask, self._accum_flow_map_y / self._flow_out_mask, mask=mask
            )
        else:
            return self.window_flow_base(self._flow_maps_x, self._flow_maps_y, mask=mask)

    def window_iwe(self, mode="forward", round_idx=False):
        """
        Assumption: events have been warped in a forward fashion in the update() method.
        :param round_idx: if True, round the event coordinates to the nearest integer.
        :return: image-like representation of the IWE of all the events in the validation time/event window.
        """

        if mode == "forward":
            event_loc = self._fw_event_loc
            pol_mask_list = self._fw_event_pol_mask
        elif mode == "backward":
            event_loc = self._bw_event_loc
            pol_mask_list = self._bw_event_pol_mask
        else:
            raise ValueError("Invalid IWE mode: {}".format(mode))

        if not round_idx:
            pol_mask_list = torch.cat([pol_mask_list for _ in range(4)], dim=1)

        idx, weights = get_interpolation(event_loc, self.res, round_idx=round_idx)
        iwe_pos = interpolate(idx.long(), weights, self.res, polarity_mask=pol_mask_list[:, :, 0:1])
        iwe_neg = interpolate(idx.long(), weights, self.res, polarity_mask=pol_mask_list[:, :, 1:2])

        return torch.cat([iwe_pos, iwe_neg], dim=1)

    def rsat(self):
        """
        :return: deblur metric for validation of the estimated optical flow.
        """

        return self.compute_rsat(
            self._fw_event_loc, self._event_loc, self._fw_event_pol_mask, self._event_pol_mask, self._event_ts
        )

    def fwl(self):
        """
        :return fwl: deblur metric for validation of the estimated optical flow (Stoffregen et al, ECCV 2020).
        """

        return self.compute_fwl(self._fw_event_loc, self._event_loc, self._fw_event_pol_mask, self._event_pol_mask)

loss/flow_val.py
tudelft/taming_event_flow
import torch

from .submodules import *


class BaseUNet(nn.Module):
    """
    Base class for conventional UNet architecture.
    Symmetric, skip connections on every encoding layer.
    """

    ff_type = ConvLayer
    res_type = ResidualBlock
    upsample_type = UpsampleConvLayer
    transpose_type = TransposedConvLayer

    def __init__(
        self,
        num_bins,
        base_channels,
        num_encoders,
        num_residual_blocks,
        num_output_channels,
        skip_type,
        norm,
        use_upsample_conv=True,
        kernel_size=3,
        encoder_stride=2,
        channel_multiplier=2,
        activations=["relu", None],
        final_activation=None,
        final_bias=True,
        final_w_scale=None,
        recurrent_block_type=None,
    ):
        super(BaseUNet, self).__init__()
        self.base_channels = base_channels
        self.num_encoders = num_encoders
        self.num_residual_blocks = num_residual_blocks
        self.num_output_channels = num_output_channels
        self.norm = norm
        self.num_bins = num_bins
        self.recurrent_block_type = recurrent_block_type
        self.kernel_size = kernel_size
        self.encoder_stride = encoder_stride
        self.channel_multiplier = channel_multiplier
        self.ff_act, self.rec_act = activations
        self.final_activation = final_activation
        self.final_bias = final_bias
        self.final_w_scale = final_w_scale

        self.skip_type = skip_type
        assert self.skip_type is None or self.skip_type in ["sum", "concat"]

        if use_upsample_conv:
            self.up_type = self.upsample_type
        else:
            self.up_type = self.transpose_type

        self.encoder_input_sizes = [
            int(self.base_channels * pow(self.channel_multiplier, i - 1)) for i in range(self.num_encoders)
        ]
        self.encoder_output_sizes = [
            int(self.base_channels * pow(self.channel_multiplier, i)) for i in range(self.num_encoders)
        ]

        self.max_num_channels = self.encoder_output_sizes[-1]

    def skip_fn(self, x, y, mode="sum"):
        assert y.shape[2:] <= x.shape[2:]
        if x.shape[2:] > y.shape[2:]:
            print("Warning: skipping row/col in skip_fn() due to odd dimensions throughout the architecture.")
            x = x[:, :, : y.shape[2], : y.shape[3]]  # skip last row/col if necessary

        if mode == "sum":
            assert x.shape[1] == y.shape[1]
            x = x + y
        elif mode == "concat":
            x = torch.cat([x, y], dim=1)
        return x

    def get_axonal_delays(self):
        self.delays = 0

    def build_encoders(self):
        encoders = nn.ModuleList()
        for i, (input_size, output_size) in enumerate(zip(self.encoder_input_sizes, self.encoder_output_sizes)):
            if i == 0:
                input_size = self.num_bins
            encoders.append(
                self.ff_type(
                    input_size,
                    output_size,
                    kernel_size=self.kernel_size,
                    stride=self.encoder_stride,
                    activation=self.ff_act,
                    norm=self.norm,
                )
            )
        return encoders

    def build_recurrent_encoders(self):
        encoders = nn.ModuleList()
        for i, (input_size, output_size) in enumerate(zip(self.encoder_input_sizes, self.encoder_output_sizes)):
            if i == 0:
                input_size = self.num_bins
            encoders.append(
                self.rec_type(
                    input_size,
                    output_size,
                    kernel_size=self.kernel_size,
                    stride=self.encoder_stride,
                    recurrent_block_type=self.recurrent_block_type,
                    activation_ff=self.ff_act,
                    activation_rec=self.rec_act,
                    norm=self.norm,
                )
            )
        return encoders

    def build_resblocks(self):
        resblocks = nn.ModuleList()
        for i in range(self.num_residual_blocks):
            resblocks.append(
                self.res_type(
                    self.max_num_channels,
                    self.max_num_channels,
                    activation=self.ff_act,
                    norm=self.norm,
                )
            )
        return resblocks

    def build_decoders(self):
        decoder_input_sizes = reversed(self.encoder_output_sizes)
        decoder_output_sizes = reversed(self.encoder_input_sizes)
        decoders = nn.ModuleList()
        for input_size, output_size in zip(decoder_input_sizes, decoder_output_sizes):
            decoders.append(
                self.up_type(
                    input_size if self.skip_type == "sum" else 2 * input_size,
                    output_size,
                    kernel_size=self.kernel_size,
                    activation=self.ff_act,
                    norm=self.norm,
                )
            )
        return decoders

    def build_multires_prediction_decoders(self):
        decoder_input_sizes = reversed(self.encoder_output_sizes)
        decoder_output_sizes = reversed(self.encoder_input_sizes)
        decoders = nn.ModuleList()
        for i, (input_size, output_size) in enumerate(zip(decoder_input_sizes, decoder_output_sizes)):
            input_size = 2 * input_size if self.skip_type == "concat" else input_size
            prediction_channels = 0 if i == 0 else self.num_output_channels
            decoders.append(
                self.up_type(
                    input_size + prediction_channels,
                    output_size,
                    kernel_size=self.kernel_size,
                    activation=self.ff_act,
                    norm=self.norm,
                )
            )
        return decoders

    def build_prediction_layer(self):
        return self.ff_type(
            2 * self.base_channels if self.skip_type == "concat" else self.base_channels,
            self.num_output_channels,
            kernel_size=1,
            activation=self.final_activation,
            norm=self.norm,
            w_scale=self.final_w_scale,
            bias=self.final_bias,
        )

    def build_multires_prediction_layer(self):
        preds = nn.ModuleList()
        decoder_output_sizes = reversed(self.encoder_input_sizes)
        for output_size in decoder_output_sizes:
            preds.append(
                self.ff_type(
                    output_size,
                    self.num_output_channels,
                    1,
                    activation=self.final_activation,
                    norm=self.norm,
                    w_scale=self.final_w_scale,
                    bias=self.final_bias,
                )
            )
        return preds


class MultiResUNetRecurrent(BaseUNet):
    """
    Recurrent UNet architecture where every encoder is followed by a recurrent convolutional block.
    Symmetric, skip connections on every encoding layer (concat/sum).
    Predictions at each decoding layer.
    Predictions are added as skip connection (concat) to the input of the subsequent layer.
    """

    rec_type = RecurrentConvLayer

    def __init__(self, kwargs):
        super().__init__(**kwargs)

        self.encoders = self.build_recurrent_encoders()
        self.resblocks = self.build_resblocks()
        self.decoders = self.build_multires_prediction_decoders()
        self.preds = self.build_multires_prediction_layer()
        self.num_states = self.num_encoders
        self.states = [None] * self.num_states

    def forward(self, x):
        """
        :param x: N x num_input_channels x H x W
        :return: [N x num_output_channels x H x W for i in range(self.num_encoders)]
        """

        # encoder
        blocks = []
        for i, encoder in enumerate(self.encoders):
            x, self.states[i] = encoder(x, self.states[i])
            blocks.append(x)

        # residual blocks
        for resblock in self.resblocks:
            x, _ = resblock(x)

        # decoder and multires predictions
        predictions = []
        for i, (decoder, pred) in enumerate(zip(self.decoders, self.preds)):
            x = self.skip_fn(x, blocks[self.num_encoders - i - 1], mode=self.skip_type)
            if i > 0:
                x = self.skip_fn(predictions[-1], x, mode="concat")
            x = decoder(x)
            predictions.append(pred(x))

        return predictions

models/arch.py
tudelft/taming_event_flow
"""
Adapted from UZH-RPG https://github.com/uzh-rpg/rpg_e2vid
"""

from abc import abstractmethod

import numpy as np
import torch


class BaseModel(torch.nn.Module):
    """
    Base class for all models
    """

    @abstractmethod
    def forward(self, *inputs):
        """
        Forward pass logic

        :return: Model output
        """
        raise NotImplementedError

    def __str__(self):
        """
        Model prints with number of trainable parameters
        """
        model_parameters = filter(lambda p: p.requires_grad, self.parameters())
        params = sum([np.prod(p.size()) for p in model_parameters])
        return super().__str__() + "\nTrainable parameters: {}".format(params)

models/base.py
tudelft/taming_event_flow
from .arch import *
from .base import BaseModel
from .model_util import copy_states, ImagePadder


class RecEVFlowNet(BaseModel):
    """
    Recurrent version of the EV-FlowNet model, as described in the paper "Self-Supervised Learning of
    Event-based Optical Flow with Spiking Neural Networks", Hagenaars and Paredes-Vallés et al., NeurIPS 2021.
    """

    net_type = MultiResUNetRecurrent
    recurrent_block_type = "convgru"
    activations = ["relu", None]

    def __init__(self, kwargs, num_bins=2, key="flow", min_size=16):
        super().__init__()
        self.image_padder = ImagePadder(min_size=min_size)

        self.key = key
        arch_kwargs = {
            "num_bins": num_bins,
            "base_channels": 64,
            "num_encoders": 4,
            "num_residual_blocks": 2,
            "num_output_channels": 2,
            "skip_type": "sum",
            "norm": None,
            "use_upsample_conv": True,
            "kernel_size": 3,
            "encoder_stride": 2,
            "channel_multiplier": 2,
            "final_activation": "tanh",
            "activations": self.activations,
            "recurrent_block_type": self.recurrent_block_type,
        }
        arch_kwargs.update(kwargs)  # udpate params with config
        arch_kwargs.pop("name", None)
        self.arch = self.net_type(arch_kwargs)
        self.num_encoders = arch_kwargs["num_encoders"]

    @property
    def states(self):
        return copy_states(self.arch.states)

    @states.setter
    def states(self, states):
        self.arch.states = states

    def detach_states(self):
        detached_states = []
        for state in self.arch.states:
            if type(state) is tuple:
                tmp = []
                for hidden in state:
                    tmp.append(hidden.detach())
                detached_states.append(tuple(tmp))
            else:
                detached_states.append(state.detach())
        self.arch.states = detached_states

    def reset_states(self):
        self.arch.states = [None] * self.arch.num_states

    def forward(self, x):

        # image padding
        x = self.image_padder.pad(x).contiguous()

        # forward pass
        multires_flow = self.arch.forward(x)

        # upsample flow estimates to the original input resolution
        flow_list = []
        for i, flow in enumerate(multires_flow):
            scaling_h = x.shape[2] / flow.shape[2]
            scaling_w = x.shape[3] / flow.shape[3]
            scaling_flow = 2 ** (self.num_encoders - i - 1)
            upflow = scaling_flow * torch.nn.functional.interpolate(
                flow, scale_factor=(scaling_h, scaling_w), mode="bilinear", align_corners=False
            )
            upflow = self.image_padder.unpad(upflow)
            flow_list.append(upflow)

        return {self.key: flow_list}

models/model.py
tudelft/taming_event_flow
import copy

import torch


def recursive_clone(tensor):
    """
    Assumes tensor is a torch.tensor with 'clone()' method, possibly
    inside nested iterable.
    E.g., tensor = [(pytorch_tensor, pytorch_tensor), ...]
    """
    if hasattr(tensor, "clone"):
        return tensor.clone()
    try:
        return type(tensor)(recursive_clone(t) for t in tensor)
    except TypeError:
        print("{} is not iterable and has no clone() method.".format(tensor))


def copy_states(states):
    """
    Simple deepcopy if list of Nones, else clone.
    """
    if states[0] is None:
        return copy.deepcopy(states)
    return recursive_clone(states)


class ImagePadder(object):
    """
    From E-RAFT: https://github.com/uzh-rpg/E-RAFT
    """

    # =================================================================== #
    # In some networks, the image gets downsized. This is a problem, if   #
    # the to-be-downsized image has odd dimensions ([15x20]->[7.5x10]).   #
    # To prevent this, the input image of the network needs to be a       #
    # multiple of a minimum size (min_size)                               #
    # The ImagePadder makes sure, that the input image is of such a size, #
    # and if not, it pads the image accordingly.                          #
    # =================================================================== #

    def __init__(self, min_size=64):
        # --------------------------------------------------------------- #
        # The min_size additionally ensures, that the smallest image      #
        # does not get too small                                          #
        # --------------------------------------------------------------- #
        self.min_size = min_size
        self.pad_height = None
        self.pad_width = None

    def pad(self, image):
        # --------------------------------------------------------------- #
        # If necessary, this function pads the image on the left & top    #
        # --------------------------------------------------------------- #
        height, width = image.shape[-2:]
        if self.pad_width is None:
            self.pad_height = (self.min_size - height % self.min_size) % self.min_size
            self.pad_width = (self.min_size - width % self.min_size) % self.min_size
        else:
            pad_height = (self.min_size - height % self.min_size) % self.min_size
            pad_width = (self.min_size - width % self.min_size) % self.min_size
            if pad_height != self.pad_height or pad_width != self.pad_width:
                raise
        return torch.nn.ZeroPad2d((self.pad_width, 0, self.pad_height, 0))(image)

    def unpad(self, image):
        # --------------------------------------------------------------- #
        # Removes the padded rows & columns                               #
        # --------------------------------------------------------------- #
        return image[..., self.pad_height :, self.pad_width :]

models/model_util.py
tudelft/taming_event_flow
import math

import torch
import torch.nn as nn
import torch.nn.functional as f


class ConvLayer(nn.Module):
    """
    Convolutional layer.
    Default: bias, ReLU, no downsampling, no batch norm.
    """

    def __init__(
        self,
        in_channels,
        out_channels,
        kernel_size,
        stride=1,
        activation="relu",
        norm=None,
        BN_momentum=0.1,
        w_scale=None,
        padding=None,
        bias=None,
    ):
        super(ConvLayer, self).__init__()

        if padding is None:
            padding = kernel_size // 2
        if bias is None:
            bias = False if norm == "BN" else True
        if w_scale is None:
            w_scale = math.sqrt(1 / in_channels)

        self.conv2d = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, bias=bias)
        nn.init.uniform_(self.conv2d.weight, -w_scale, w_scale)
        if bias:
            nn.init.zeros_(self.conv2d.bias)

        if activation is not None:
            if hasattr(torch, activation):
                self.activation = getattr(torch, activation)
        else:
            self.activation = None

        self.norm = norm
        if norm == "BN":
            self.norm_layer = nn.BatchNorm2d(out_channels, momentum=BN_momentum)
        elif norm == "IN":
            self.norm_layer = nn.InstanceNorm2d(out_channels, track_running_stats=True)

    def forward(self, x):
        out = self.conv2d(x)

        if self.norm in ["BN", "IN"]:
            out = self.norm_layer(out)

        if self.activation is not None:
            out = self.activation(out)

        return out


class RecurrentConvLayer(nn.Module):
    """
    Layer comprised of a convolution followed by a recurrent convolutional block.
    Default: bias, ReLU, no downsampling, no batch norm, ConvGRU.
    """

    def __init__(
        self,
        in_channels,
        out_channels,
        kernel_size=3,
        stride=1,
        recurrent_block_type="convgru",
        activation_ff="relu",
        activation_rec=None,
        norm=None,
        BN_momentum=0.1,
    ):
        super(RecurrentConvLayer, self).__init__()

        assert recurrent_block_type in ["convgru"]
        self.recurrent_block_type = recurrent_block_type
        if recurrent_block_type == "convgru":
            RecurrentBlock = ConvGRU
        else:
            raise NotImplementedError

        self.conv = ConvLayer(
            in_channels,
            out_channels,
            kernel_size,
            stride,
            activation_ff,
            norm,
            BN_momentum=BN_momentum,
        )
        self.recurrent_block = RecurrentBlock(
            input_size=out_channels, hidden_size=out_channels, kernel_size=3, activation=activation_rec
        )

    def forward(self, x, prev_state):
        x = self.conv(x)
        x, state = self.recurrent_block(x, prev_state)
        return x, state


class ConvGRU(nn.Module):
    """
    Convolutional GRU cell.
    Adapted from https://github.com/jacobkimmel/pytorch_convgru/blob/master/convgru.py
    """

    def __init__(self, input_size, hidden_size, kernel_size, activation=None):
        super().__init__()
        padding = kernel_size // 2
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.reset_gate = nn.Conv2d(input_size + hidden_size, hidden_size, kernel_size, padding=padding)
        self.update_gate = nn.Conv2d(input_size + hidden_size, hidden_size, kernel_size, padding=padding)
        self.out_gate = nn.Conv2d(input_size + hidden_size, hidden_size, kernel_size, padding=padding)
        assert activation is None, "ConvGRU activation cannot be set (just for compatibility)"

        nn.init.orthogonal_(self.reset_gate.weight)
        nn.init.orthogonal_(self.update_gate.weight)
        nn.init.orthogonal_(self.out_gate.weight)
        nn.init.constant_(self.reset_gate.bias, 0.0)
        nn.init.constant_(self.update_gate.bias, 0.0)
        nn.init.constant_(self.out_gate.bias, 0.0)

    def forward(self, input_, prev_state):

        # get batch and spatial sizes
        batch_size = input_.data.size()[0]
        spatial_size = input_.data.size()[2:]

        # generate empty prev_state, if None is provided
        if prev_state is None:
            state_size = [batch_size, self.hidden_size] + list(spatial_size)
            prev_state = torch.zeros(state_size, dtype=input_.dtype, device=input_.device)

        # data size is [batch, channel, height, width]
        stacked_inputs = torch.cat([input_, prev_state], dim=1)
        update = torch.sigmoid(self.update_gate(stacked_inputs))
        reset = torch.sigmoid(self.reset_gate(stacked_inputs))
        out_inputs = torch.tanh(self.out_gate(torch.cat([input_, prev_state * reset], dim=1)))
        new_state = prev_state * (1 - update) + out_inputs * update

        return new_state, new_state


class ResidualBlock(nn.Module):
    """
    Residual block as in "Deep residual learning for image recognition", He et al. 2016.
    Default: bias, ReLU, no downsampling, no batch norm.
    """

    def __init__(
        self,
        in_channels,
        out_channels,
        kernel_size=3,
        stride=1,
        activation="relu",
        downsample=None,
        norm=None,
        BN_momentum=0.1,
    ):
        super(ResidualBlock, self).__init__()
        bias = False if norm == "BN" else True
        self.conv1 = nn.Conv2d(
            in_channels,
            out_channels,
            kernel_size=kernel_size,
            stride=stride,
            padding=kernel_size // 2,
            bias=bias,
        )

        if activation is not None:
            if hasattr(torch, activation):
                self.activation = getattr(torch, activation)
        else:
            self.activation = None

        self.norm = norm
        if norm == "BN":
            self.bn1 = nn.BatchNorm2d(out_channels, momentum=BN_momentum)
            self.bn2 = nn.BatchNorm2d(out_channels, momentum=BN_momentum)
        elif norm == "IN":
            self.bn1 = nn.InstanceNorm2d(out_channels, track_running_stats=True)
            self.bn2 = nn.InstanceNorm2d(out_channels, track_running_stats=True)

        self.conv2 = nn.Conv2d(
            out_channels,
            out_channels,
            kernel_size=kernel_size,
            stride=stride,
            padding=kernel_size // 2,
            bias=bias,
        )
        self.downsample = downsample

    def forward(self, x):
        residual = x
        out1 = self.conv1(x)
        if self.norm in ["BN", "IN"]:
            out1 = self.bn1(out1)

        if self.activation is not None:
            out1 = self.activation(out1)

        out2 = self.conv2(out1)
        if self.norm in ["BN", "IN"]:
            out2 = self.bn2(out2)

        if self.downsample:
            residual = self.downsample(x)

        out2 += residual
        if self.activation is not None:
            out2 = self.activation(out2)

        return out2, out1


class UpsampleConvLayer(nn.Module):
    """
    Upsampling layer (bilinear interpolation + Conv2d) to increase spatial resolution (x2) in a decoder.
    Default: bias, ReLU, no downsampling, no batch norm.
    """

    def __init__(
        self,
        in_channels,
        out_channels,
        kernel_size,
        stride=1,
        activation="relu",
        norm=None,
    ):
        super(UpsampleConvLayer, self).__init__()

        bias = False if norm == "BN" else True
        padding = kernel_size // 2
        self.conv2d = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, bias=bias)

        if activation is not None:
            if hasattr(torch, activation):
                self.activation = getattr(torch, activation)
        else:
            self.activation = None

        self.norm = norm
        if norm == "BN":
            self.norm_layer = nn.BatchNorm2d(out_channels)
        elif norm == "IN":
            self.norm_layer = nn.InstanceNorm2d(out_channels, track_running_stats=True)

    def forward(self, x):
        x_upsampled = f.interpolate(x, scale_factor=2, mode="bilinear", align_corners=False)
        out = self.conv2d(x_upsampled)

        if self.norm in ["BN", "IN"]:
            out = self.norm_layer(out)

        if self.activation is not None:
            out = self.activation(out)

        return out


class TransposedConvLayer(nn.Module):
    """
    Transposed convolutional layer to increase spatial resolution (x2) in a decoder.
    Default: bias, ReLU, no downsampling, no batch norm.
    """

    def __init__(
        self,
        in_channels,
        out_channels,
        kernel_size,
        activation="relu",
        norm=None,
    ):
        super(TransposedConvLayer, self).__init__()

        bias = False if norm == "BN" else True
        padding = kernel_size // 2
        self.transposed_conv2d = nn.ConvTranspose2d(
            in_channels,
            out_channels,
            kernel_size,
            stride=2,
            padding=padding,
            output_padding=1,
            bias=bias,
        )

        if activation is not None:
            if hasattr(torch, activation):
                self.activation = getattr(torch, activation)
        else:
            self.activation = None

        self.norm = norm
        if norm == "BN":
            self.norm_layer = nn.BatchNorm2d(out_channels)
        elif norm == "IN":
            self.norm_layer = nn.InstanceNorm2d(out_channels, track_running_stats=True)

    def forward(self, x):
        out = self.transposed_conv2d(x)

        if self.norm in ["BN", "IN"]:
            out = self.norm_layer(out)

        if self.activation is not None:
            out = self.activation(out)

        return out

models/submodules.py
tudelft/taming_event_flow
import argparse
import os
import numpy as np


def retrieve_eval(args):
    eval_id = args.eval_id
    if args.eval_id < 0:
        eval_id = 0
        for file in os.listdir(args.path + args.runid + "/"):
            if file == ".DS_Store":
                continue
            tmp = int(file.split(".")[0].split("_")[-1])
            eval_id = tmp + 1 if tmp + 1 > eval_id else eval_id
        eval_id -= 1
    path_from = args.path + args.runid + "/" + "eval_" + str(eval_id) + "/"
    print("Preparing submission for eval_{0}".format(eval_id))

    return path_from


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("runid")
    parser.add_argument("--path", default="dsec_submissions/")
    parser.add_argument("--eval_id", default=-1, type=int)
    args = parser.parse_args()

    # retrieve last eval run unless specified
    path_from = retrieve_eval(args)

    # retrieve folders in directory
    entry = "/flow_bw/"
    folders = os.listdir(path_from)
    for folder in folders:
        if folder in [".DS_Store", "submission"]:
            continue

        # retrieve files in folder with png extension
        files = os.listdir(path_from + folder + entry)
        indices = []
        for file in files:
            indices.append(int(file.split(".")[0]))
        indices.sort()

        # fixing pred-gt alignment
        flags = np.load(args.path + folder + "_flag.npy")
        flags = np.roll(flags, -1)

        # select gt maps to be submitted
        flow_timestamp = np.genfromtxt(args.path + folder + ".txt", skip_header=1, delimiter=",")
        flow_filenames = flow_timestamp[:, -1]

        selected_indices = []
        for i in range(len(indices)):
            if flags[i] == 1:
                selected_indices.append(indices[i])

        # create new folder
        if not os.path.exists(path_from + "submission/"):
            os.makedirs(path_from + "submission/")
        if not os.path.exists(path_from + "submission/" + folder + "/"):
            os.makedirs(path_from + "submission/" + folder + "/")

        # copy files to new folder with the right name
        for i in range(len(selected_indices)):
            filename = path_from + "submission/" + folder + "/" + str(int(flow_filenames[i])).zfill(6) + ".png"
            os.system("cp " + path_from + folder + entry + str(selected_indices[i]).zfill(9) + ".png " + filename)

        print(folder)

prepare_dsec_submission.py
tudelft/taming_event_flow
import argparse

import mlflow
import torch
from torch.optim import *
from torch.utils.tensorboard import SummaryWriter

from configs.parser import YAMLParser
from dataloader.h5 import H5Loader
from loss.flow import *
from models.model import *
from utils.utils import load_model, save_diff, save_model
from utils.visualization import Visualization


def train(args, config_parser):
    """
    Main function of the training pipeline for event-based optical flow estimation.
    :param args: arguments of the script
    :param config_parser: YAMLParser object with config data
    """

    mlflow.set_tracking_uri(args.path_mlflow)

    # configs
    config = config_parser.config
    mlflow.set_experiment(config["experiment"])
    run = mlflow.start_run()
    runid = run.to_dictionary()["info"]["run_id"]
    mlflow.log_params(config)
    mlflow.log_param("prev_runid", args.prev_runid)
    config = config_parser.combine_entries(config)
    print("MLflow dir:", mlflow.active_run().info.artifact_uri[:-9])

    # log git diff
    save_diff("train_diff.txt")
    tb_writer = SummaryWriter(log_dir=args.path_mlflow + "mlruns/0/" + runid + "/")

    # initialize settings
    device = config_parser.device
    kwargs = config_parser.loader_kwargs
    config["loader"]["device"] = device

    # visualization tool
    if config["vis"]["enabled"]:
        vis = Visualization(config)

    # data loader
    data = H5Loader(config, shuffle=True, path_cache=args.path_cache)
    dataloader = torch.utils.data.DataLoader(
        data,
        drop_last=True,
        batch_size=config["loader"]["batch_size"],
        collate_fn=data.custom_collate,
        worker_init_fn=config_parser.worker_init_fn,
        **kwargs,
    )

    # model initialization and settings
    num_bins = 2 if config["data"]["voxel"] is None else config["data"]["voxel"]
    model = eval(config["model"]["name"])(config["model"].copy(), num_bins, key="flow")
    model = model.to(device)
    model, epoch = load_model(args.prev_runid, model, device, curr_run=run, tb_writer=tb_writer)
    model.train()

    # loss functions
    loss_function = eval(config["loss"]["warping"])(config, device)

    # optimizers
    optimizer = eval(config["optimizer"]["name"])(model.parameters(), lr=config["optimizer"]["lr"])
    optimizer.zero_grad()

    # simulation variables
    train_loss = 0
    best_loss = 1.0e6
    end_train = False
    data.epoch = epoch

    # dataloader loop
    while True:
        for inputs in dataloader:

            if data.new_seq:
                data.new_seq = False
                loss_function.reset()
                model.reset_states()
                optimizer.zero_grad()

            if data.seq_num >= len(data.files):
                tb_writer.add_scalar("loss", train_loss / data.samples, data.epoch)
                mlflow.log_metric("loss", train_loss / data.samples, step=data.epoch)
                with torch.no_grad():
                    if train_loss / data.samples < best_loss:
                        save_model(model)
                        best_loss = train_loss / data.samples

                data.epoch += 1
                data.samples = 0
                train_loss = 0
                data.seq_num = data.seq_num % len(data.files)
                if data.epoch == config["loader"]["n_epochs"]:
                    end_train = True
                    break

            # forward pass (flow in px/input_time)
            x = model(inputs["net_input"].to(device))
            for i in range(len(x["flow"])):
                x["flow"][i] = x["flow"][i] * config["loss"]["flow_scaling"]

            # event-flow association
            loss_function.update(
                x["flow"],
                inputs["event_list"].to(device),
                inputs["event_list_pol_mask"].to(device),
                inputs["d_event_list"].to(device),
                inputs["d_event_list_pol_mask"].to(device),
            )

            # loss computation
            if loss_function.num_passes >= config["data"]["passes_loss"]:
                data.samples += config["loader"]["batch_size"]

                loss = loss_function()
                train_loss += loss.item()
                loss.backward()

                if config["loss"]["clip_grad"] is not None:
                    torch.nn.utils.clip_grad.clip_grad_norm_(model.parameters(), config["loss"]["clip_grad"])

                optimizer.step()
                optimizer.zero_grad()

                if config["vis"]["enabled"] and config["loader"]["batch_size"] == 1:
                    vis.data["flow"] = x["flow"][-1].clone()

                model.detach_states()
                loss_function.reset()

                with torch.no_grad():
                    if config["vis"]["enabled"] and config["loader"]["batch_size"] == 1:
                        vis.step(inputs)

                if config["vis"]["verbose"]:
                    print(
                        "Train Epoch: {:04d} [{:03d}/{:03d} ({:03d}%)] Loss: {:.6f}".format(
                            data.epoch,
                            data.seq_num,
                            len(data.files),
                            int(100 * data.seq_num / len(data.files)),
                            train_loss / data.samples,
                        ),
                        end="\r",
                    )

        if end_train:
            break

    mlflow.end_run()


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--config",
        default="configs/train_flow.yml",
        help="training configuration",
    )
    parser.add_argument(
        "--path_mlflow",
        default="",
        help="location of the mlflow ui",
    )
    parser.add_argument(
        "--path_cache",
        default="",
        help="location of the cache version of the formatted dataset",
    )
    parser.add_argument(
        "--prev_runid",
        default="",
        help="pre-trained model to use as starting point",
    )
    args = parser.parse_args()

    # launch training
    train(args, YAMLParser(args.config))

train_flow.py
tudelft/taming_event_flow
import torch
import torch.nn.functional as f


def event_propagation(events_ts, events_idx, flow, tref):
    """
    Warp the input events according to the provided optical flow map.
    :param events: [batch_size x N x 4] input events (ts, y, x, p)
    :param flow: [batch_size x N x 2] optical flows (y, x)
    :param tref: reference time toward which events are warped
    :return warped event indices
    """

    return events_idx + (tref - events_ts) * flow


def get_event_flow(flow_map_x, flow_map_y, event_loc):
    """
    Sample optical flow maps using event indices
    :param flow_map_x: [batch_size x H x W] horizontal optical flow map
    :param flow_map_y: [batch_size x H x W] vertical optical flow map
    :param event_loc: [batch_size x N x 2] event locations
    :return event_flow: [batch_size x N x 2] per-event optical flow (y, x)
    """

    _, h, w = flow_map_x.shape

    # flow vector per input event
    event_idx = event_loc.clone()
    event_idx[..., 0] = 2 * event_idx[..., 0] / (h - 1) - 1
    event_idx[..., 1] = 2 * event_idx[..., 1] / (w - 1) - 1
    event_idx = torch.roll(event_idx, 1, dims=-1).unsqueeze(2)  # needs to be (x, y) and not (y, x)

    event_flow_x = f.grid_sample(flow_map_x.unsqueeze(1), event_idx, mode="bilinear", align_corners=True)
    event_flow_y = f.grid_sample(flow_map_y.unsqueeze(1), event_idx, mode="bilinear", align_corners=True)
    event_flow_x = event_flow_x.squeeze(1)
    event_flow_y = event_flow_y.squeeze(1)
    event_flow = torch.cat([event_flow_y, event_flow_x], dim=2)

    return event_flow


def purge_unfeasible(event_loc, event_pol_mask, res):
    """
    Purge events that are warped outside the image space.
    :param event_loc: [batch_size x N x 2] warped event location
    :param event_pol_mask: [batch_size x N x 2] polarity mask of warped events
    :return event_loc: [batch_size x N x 2] masked warped event location
    :return event_pol_mask: [batch_size x N x 2] masked polarity mask of warped events
    """

    mask = (
        (event_loc[:, :, 0:1] >= 0)
        * (event_loc[:, :, 0:1] <= res[0] - 1.0)
        * (event_loc[:, :, 1:2] >= 0)
        * (event_loc[:, :, 1:2] <= res[1] - 1.0)
    )
    event_loc = event_loc * mask
    event_pol_mask = event_pol_mask * mask
    return event_loc, event_pol_mask


def get_interpolation(warped_events, res, round_idx=False, zeros=None):
    """
    Warp the input events according to the provided optical flow map and compute the bilinar interpolation
    (or rounding) weights to distribute the events to the closes (integer) locations in the image space.
    :param events: [batch_size x N x 4] input events (ts, y, x, p)
    :param flow: [batch_size x N x 2] optical flows (y, x)
    :param tref: reference time toward which events are warped
    :param res: resolution of the image space
    :param round_idx: whether or not to round the event locations instead of doing bilinear interp. (default = False)
    :return interpolated event indices
    :return interpolation weights
    """

    if round_idx:

        # no bilinear interpolation
        idx = torch.round(warped_events)
        weights = torch.ones(idx.shape, device=warped_events.device)

    else:

        # get scattering indices
        top_y = torch.floor(warped_events[:, :, 0:1])
        bot_y = torch.floor(warped_events[:, :, 0:1] + 1)
        left_x = torch.floor(warped_events[:, :, 1:2])
        right_x = torch.floor(warped_events[:, :, 1:2] + 1)

        top_left = torch.cat([top_y, left_x], dim=2)
        top_right = torch.cat([top_y, right_x], dim=2)
        bottom_left = torch.cat([bot_y, left_x], dim=2)
        bottom_right = torch.cat([bot_y, right_x], dim=2)
        idx = torch.cat([top_left, top_right, bottom_left, bottom_right], dim=1)

        # get scattering interpolation weights
        warped_events = torch.cat([warped_events for i in range(4)], dim=1)
        if zeros is None:
            zeros = torch.zeros(warped_events.shape, device=warped_events.device)
        weights = torch.max(zeros, 1 - torch.abs(warped_events - idx))

    # purge unfeasible indices
    mask = (idx[:, :, 0:1] >= 0) * (idx[:, :, 0:1] < res[0]) * (idx[:, :, 1:2] >= 0) * (idx[:, :, 1:2] < res[1])
    idx *= mask

    # make unfeasible weights zero
    weights = torch.prod(weights, dim=-1, keepdim=True) * mask  # bilinear interpolation

    # prepare indices
    idx[:, :, 0] *= res[1]  # torch.view is row-major
    idx = torch.sum(idx, dim=2, keepdim=True)

    return idx, weights


def interpolate(idx, weights, res, polarity_mask=None, zeros=None):
    """
    Create an image-like representation of the warped events.
    :param idx: [batch_size x N x 1] warped event locations
    :param weights: [batch_size x N x 1] interpolation weights for the warped events
    :param res: resolution of the image space
    :param polarity_mask: [batch_size x N x 1] polarity mask for the warped events (default = None)
    :return image of warped events
    """

    if polarity_mask is not None:
        weights = weights * polarity_mask

    if zeros is None:
        iwe = torch.zeros((idx.shape[0], res[0] * res[1], 1), device=idx.device)
    else:
        iwe = zeros.clone()

    iwe = iwe.scatter_add_(1, idx.long(), weights)
    iwe = iwe.view((idx.shape[0], 1, res[0], res[1]))
    return iwe


def deblur_events(flow, event_list, res, round_idx=True, polarity_mask=None, round_flow=True):
    """
    Deblur the input events given an optical flow map.
    Event timestamp needs to be normalized between 0 and 1.
    :param flow: [batch_size x 2 x H x W] optical flow map
    :param events: [batch_size x N x 4] input events (ts, y, x, p)
    :param res: resolution of the image space
    :param round_idx: whether or not to round the event locations instead of doing bilinear interp.
    :param polarity_mask: [batch_size x N x 1] polarity mask for the warped events
    :param round_flow: whether or not to associate events with the closest flow vector
    :return iwe: [batch_size x 1 x H x W] image of warped events
    """

    # flow vector per input event
    flow_idx = event_list[:, :, 1:3].clone()
    mask_unfeasible = (
        (flow_idx[:, :, 0:1] >= 0)
        * (flow_idx[:, :, 0:1] < res[0])
        * (flow_idx[:, :, 1:2] >= 0)
        * (flow_idx[:, :, 1:2] < res[1])
    )
    flow_idx *= mask_unfeasible

    if not round_flow:

        top_y = torch.floor(flow_idx[:, :, 0:1])
        bot_y = torch.floor(flow_idx[:, :, 0:1] + 1)
        left_x = torch.floor(flow_idx[:, :, 1:2])
        right_x = torch.floor(flow_idx[:, :, 1:2] + 1)

        top_left = torch.cat([top_y, left_x], dim=2)
        top_right = torch.cat([top_y, right_x], dim=2)
        bottom_left = torch.cat([bot_y, left_x], dim=2)
        bottom_right = torch.cat([bot_y, right_x], dim=2)
        idx = torch.cat([top_left, top_right, bottom_left, bottom_right], dim=1)

        og_idx = torch.cat([flow_idx for i in range(4)], dim=1)
        zeros = torch.zeros(idx.shape, device=idx.device)
        interp_weights = torch.max(zeros, 1 - torch.abs(og_idx - idx))

        mask_y = (idx[:, :, 0:1] >= 0) * (idx[:, :, 0:1] < res[0])
        mask_x = (idx[:, :, 1:2] >= 0) * (idx[:, :, 1:2] < res[1])
        mask = mask_y * mask_x
        flow_idx = idx * mask
        interp_weights = torch.prod(interp_weights, dim=-1, keepdim=True) * mask  # bilinear interpolation

    flow_idx[:, :, 0] *= res[1]  # torch.view is row-major
    flow_idx = torch.sum(flow_idx, dim=2)

    # get flow for every event in the list
    flow = flow.view(flow.shape[0], 2, -1)
    event_flowy = torch.gather(flow[:, 1, :], 1, flow_idx.long())  # vertical component
    event_flowx = torch.gather(flow[:, 0, :], 1, flow_idx.long())  # horizontal component
    event_flowy = event_flowy.view(event_flowy.shape[0], event_flowy.shape[1], 1)
    event_flowx = event_flowx.view(event_flowx.shape[0], event_flowx.shape[1], 1)

    # bilinear interpolation of the optical flow vectors
    if not round_flow:
        N = event_list.shape[1]  # number of events
        event_flowy = (
            interp_weights[:, 0 * N : 1 * N, :] * event_flowy[:, 0 * N : 1 * N, :]
            + interp_weights[:, 1 * N : 2 * N, :] * event_flowy[:, 1 * N : 2 * N, :]
            + interp_weights[:, 2 * N : 3 * N, :] * event_flowy[:, 2 * N : 3 * N, :]
            + interp_weights[:, 3 * N : 4 * N, :] * event_flowy[:, 3 * N : 4 * N, :]
        )
        event_flowx = (
            interp_weights[:, 0 * N : 1 * N, :] * event_flowx[:, 0 * N : 1 * N, :]
            + interp_weights[:, 1 * N : 2 * N, :] * event_flowx[:, 1 * N : 2 * N, :]
            + interp_weights[:, 2 * N : 3 * N, :] * event_flowx[:, 2 * N : 3 * N, :]
            + interp_weights[:, 3 * N : 4 * N, :] * event_flowx[:, 3 * N : 4 * N, :]
        )

    event_flow = torch.cat([event_flowy, event_flowx], dim=2)

    # interpolate forward
    fw_events = event_propagation(event_list[:, :, 0:1], event_list[:, :, 1:3], event_flow, 1)
    fw_idx, fw_weights = get_interpolation(fw_events, res, round_idx=round_idx)
    if not round_idx and polarity_mask is not None:
        polarity_mask = torch.cat([polarity_mask for i in range(4)], dim=1)
        mask_unfeasible = torch.cat([mask_unfeasible for i in range(4)], dim=1)
    fw_weights *= mask_unfeasible

    # image of (forward) warped events
    iwe = interpolate(fw_idx, fw_weights, res, polarity_mask=polarity_mask)

    return iwe


def compute_pol_iwe(flow, event_list, res, pol_mask, round_idx=True, round_flow=True):
    """
    Create a per-polarity image of warped events given an optical flow map.
    :param flow: [batch_size x 2 x H x W] optical flow map
    :param event_list: [batch_size x N x 4] input events (ts, y, x, p)
    :param res: resolution of the image space
    :param pol_mask: [batch_size x N x 2] polarity mask
    :param round_idx: whether or not to round the event locations instead of doing bilinear interp.
    :param round_flow: whether or not to associate events with the closest flow vector
    :return iwe: [batch_size x 2 x H x W] image of warped events
    """

    iwe_pos = deblur_events(
        flow,
        event_list,
        res,
        round_idx=round_idx,
        polarity_mask=pol_mask[:, :, 0:1],
        round_flow=round_flow,
    )
    iwe_neg = deblur_events(
        flow,
        event_list,
        res,
        round_idx=round_idx,
        polarity_mask=pol_mask[:, :, 1:2],
        round_flow=round_flow,
    )
    iwe = torch.cat([iwe_pos, iwe_neg], dim=1)

    return iwe

utils/iwe.py
tudelft/taming_event_flow
import os
import yaml

import torch
import mlflow


def log_config(path_results, runid, config):
    """
    Log configuration file to MlFlow run.
    """

    eval_id = 0
    for file in os.listdir(path_results):
        if file.endswith(".yml"):
            tmp = int(file.split(".")[0].split("_")[-1])
            eval_id = tmp + 1 if tmp + 1 > eval_id else eval_id
    yaml_filename = path_results + "eval_" + str(eval_id) + ".yml"
    with open(yaml_filename, "w") as outfile:
        yaml.dump(config, outfile, default_flow_style=False)

    mlflow.start_run(runid)
    mlflow.log_artifact(yaml_filename)
    mlflow.end_run()

    return eval_id


def log_results(runid, results, path, eval_id):
    """
    Log validation results as artifacts to MlFlow run.
    """

    yaml_filename = path + "metrics_" + str(eval_id) + ".yml"
    with open(yaml_filename, "w") as outfile:
        yaml.dump(results, outfile, default_flow_style=False)

    mlflow.start_run(runid)
    mlflow.log_artifact(yaml_filename)
    mlflow.end_run()

utils/mlflow.py
tudelft/taming_event_flow
import os

import mlflow
import numpy as np
import pandas as pd
import torch


def load_model(prev_runid, model, device, curr_run=None, tb_writer=None):
    try:
        run = mlflow.get_run(prev_runid)
    except:
        return model, 0

    model_dir = run.info.artifact_uri + "/model/data/model.pth"
    if model_dir[:7] == "file://":
        model_dir = model_dir[7:]

    starting_epoch = 0
    if os.path.isfile(model_dir):
        model_loaded = torch.load(model_dir, map_location=device).state_dict()

        # check for input-dependent layers
        for key in model_loaded.keys():
            if key.split(".")[1] == "pooling" and key.split(".")[-1] in ["weight", "weight_f"]:
                model.encoder_unet.pooling = model.encoder_unet.build_pooling(model_loaded[key].shape).to(device)
                model.encoder_unet.get_axonal_delays()

        new_params = model.state_dict()
        new_params.update(model_loaded)
        model.load_state_dict(new_params)

        loss_file = run.info.artifact_uri[:-9] + "metrics/loss"
        if os.path.isfile(run.info.artifact_uri[:-9] + "metrics/loss"):
            loss = np.genfromtxt(loss_file)
            if curr_run is not None:
                if not os.path.exists(curr_run.info.artifact_uri[:-9] + "metrics/"):
                    os.makedirs(curr_run.info.artifact_uri[:-9] + "metrics/")
                for i in range(loss.shape[0]):
                    mlflow.log_metric("loss", loss[i, 1], step=int(loss[i, 2]))
                    if tb_writer is not None:
                        tb_writer.add_scalar("loss", loss[i, 1], int(loss[i, 2]))
            starting_epoch = int(loss[-1][-1])

        print("Model restored from " + prev_runid + "\n")
    else:
        print("No model found at " + prev_runid + "\n")

    return model, starting_epoch


def create_model_dir(path_results, runid):
    path_results += runid + "/"
    if not os.path.exists(path_results):
        os.makedirs(path_results)
    print("Results stored at " + path_results + "\n")
    return path_results


def save_model(model):
    mlflow.pytorch.log_model(model, "model", conda_env={"dependencies": []})


def save_state_dict(runid, state_dict):
    mlflow.start_run(runid)
    mlflow.pytorch.log_state_dict(state_dict, "state_dict")
    mlflow.end_run()


def load_state_dict(runid, dir="state_dict/", filename="state_dict.pth"):
    run = mlflow.get_run(runid)
    model_dir = run.info.artifact_uri + "/" + dir
    if model_dir[:7] == "file://":
        model_dir = model_dir[7:]

    model_dict = None
    if os.path.isfile(model_dir + filename):
        model_dict = torch.load(model_dir + filename, map_location=torch.device("cpu"))
        print("Model restored from " + runid)
    else:
        print("No model found at " + runid)

    return model_dict


def save_csv(data, fname):
    # create file if not there
    path = mlflow.get_artifact_uri(artifact_path=fname)
    if path[:7] == "file://":  # to_csv() doesn't work with 'file://'
        path = path[7:]
    if not os.path.isfile(path):
        mlflow.log_text("", fname)
        pd.DataFrame(data).to_csv(path)
    # else append
    else:
        pd.DataFrame(data).to_csv(path, mode="a", header=False)


def save_diff(fname="git_diff.txt"):
    # .txt to allow showing in mlflow
    path = mlflow.get_artifact_uri(artifact_path=fname)
    if path[:7] == "file://":
        path = path[7:]
    mlflow.log_text("", fname)
    os.system(f"git diff > {path}")


def binary_search_array(array, x, left=None, right=None, side="left"):
    left = 0 if left is None else left
    right = len(array) - 1 if right is None else right
    mid = left + (right - left) // 2

    if left > right:
        return left if side == "left" else right

    if array[mid] == x:
        return mid

    if x < array[mid]:
        return binary_search_array(array, x, left=left, right=mid - 1, side=side)

    return binary_search_array(array, x, left=mid + 1, right=right, side=side)


def initialize_quant_results(results, filename, metrics):
    if filename not in results.keys():
        results[filename] = {}
    for metric in metrics:
        if metric not in results[filename].keys():
            results[filename][metric] = {}
            results[filename][metric]["metric"] = 0
            results[filename][metric]["it"] = 0

    return results

utils/utils.py
tudelft/taming_event_flow
import os

import cv2
import matplotlib
import numpy as np


class Visualization:
    """
    Utility class for the visualization and storage of rendered image-like representation of multiple elements of the pipeline.
    """

    def __init__(self, kwargs, eval_id=-1, path_results=None):
        self.img_idx = 0
        self.px = kwargs["vis"]["px"]
        self.res = kwargs["loader"]["resolution"]
        self.color_scheme = "green_red"  # gray / blue_red / green_red
        self.show_rendered = kwargs["vis"]["enabled"]
        self.store_rendered = kwargs["vis"]["store"]

        if eval_id >= 0 and path_results is not None:
            self.store_dir = path_results + "results/"
            self.store_dir = self.store_dir + "eval_" + str(eval_id) + "/"
            if not os.path.exists(self.store_dir):
                os.makedirs(self.store_dir)
            self.store_file = None

        self.data = {}
        self.keys = [
            "events",
            "events_window",
            "events_dynamic",
            "flow",
            "flow_window",
            "flow_dynamic",
            "flow_bw",
            "iwe",
            "iwe_fw_window",
            "iwe_bw_window",
            "iwe_fw_dynamic",
            "iwe_bw_dynamic",
            "flow_gt",
            "depth_gt",
        ]
        self.title = [
            "Input events",
            "Input events - Eval window",
            "Input events - Dynamic window",
            "Estimated flow",
            "Estimated flow - Eval window",
            "Estimated flow - Dynamic window",
            "Estimated flow - Backward",
            "IWE",
            "Forward IWE - Eval window",
            "Backward IWE - Eval window",
            "Forward IWE - Dynamic window",
            "Backward IWE - Dynamic window",
            "Ground truth flow",
            "Ground truth depth",
        ]

        self.reset_image_ph()

    def step(self, inputs, sequence=None, ts=None, show=None):
        """
        Main function of the visualization workflow.
        :param inputs: input data (output of the dataloader)
        """

        # render images
        self.render(inputs, show)

        # live display
        if self.show_rendered:
            self.update(show)

        # store rendered images
        if self.store_rendered and sequence is not None:
            self.store(sequence, ts, show)

        # reset image placeholders
        self.reset_image_ph()

    def reset_image_ph(self):
        """
        Initialize/Reset image placeholders.
        """
        for key in self.keys:
            self.data[key] = None

    def render(self, inputs, show=None):
        """
        Rendering tool.
        :param inputs: input data (output of the dataloader)
        """

        self.data["events"] = inputs["event_cnt"] if "event_cnt" in inputs.keys() else None
        if self.data["events"] is None:
            self.data["events"] = inputs["net_input"] if "net_input" in inputs.keys() else None

        self.data["flow_gt"] = inputs["gtflow"] if "gtflow" in inputs.keys() else None
        self.data["depth_gt"] = inputs["gtdepth"] if "gtdepth" in inputs.keys() else None

        # optical flow error
        if self.data["flow_bw"] is not None and self.data["flow_gt"] is not None:
            self.data["error_flow"] = (
                (self.data["flow_bw"].cpu() - self.data["flow_gt"]).pow(2).sum(1).sqrt().unsqueeze(1)
            )
            gtflow_mask = (self.data["flow_gt"][:, 0:1, :, :] == 0.0) * (self.data["flow_gt"][:, 1:2, :, :] == 0.0)
            self.data["error_flow"] *= ~gtflow_mask
            if "error_flow" not in self.keys:
                self.keys.append("error_flow")
                self.title.append("AEE (capped at 30px)")

        for key in self.keys:
            if show is not None:
                if key not in show:
                    continue

            if self.data[key] is not None:
                self.data[key] = self.data[key].detach()

            # input events
            if key.split("_")[0] == "events" and self.data[key] is not None:
                self.data[key] = (
                    self.data[key]
                    .cpu()
                    .numpy()
                    .transpose(0, 2, 3, 1)
                    .reshape((self.data[key].shape[2], self.data[key].shape[3], 2))
                )
                self.data[key] = self.events_to_image(self.data[key])

            # optical flow
            elif key.split("_")[0] == "flow" and self.data[key] is not None:
                self.data[key] = (
                    self.data[key]
                    .cpu()
                    .numpy()
                    .transpose(0, 2, 3, 1)
                    .reshape((self.data[key].shape[2], self.data[key].shape[3], 2))
                )
                if key != "flow_bw":
                    self.data[key] = self.flow_to_image(self.data[key])
                else:
                    self.data[key] = self.data[key] * 128 + 2**15
                    self.data[key] = self.data[key].astype(np.uint16)
                    self.data[key] = np.pad(self.data[key], ((0, 0), (0, 0), (0, 1)), constant_values=0)
                    self.data[key] = np.flip(self.data[key], axis=-1)

            # optical flow error
            elif key == "error_flow" and self.data[key] is not None:
                self.data[key] = (
                    self.data[key]
                    .cpu()
                    .numpy()
                    .transpose(0, 2, 3, 1)
                    .reshape((self.data[key].shape[2], self.data[key].shape[3], 1))
                )
                self.data[key] = self.minmax_norm(self.data[key], max=30, min=0)
                self.data[key] *= 255
                self.data[key] = self.data[key].astype(np.uint8)
                self.data[key] = cv2.applyColorMap(self.data[key], cv2.COLORMAP_VIRIDIS)

            # image of warped events
            elif key.split("_")[0] == "iwe" and self.data[key] is not None:
                self.data[key] = (
                    self.data[key]
                    .cpu()
                    .numpy()
                    .transpose(0, 2, 3, 1)
                    .reshape((self.data[key].shape[2], self.data[key].shape[3], 2))
                )
                self.data[key] = self.events_to_image(self.data[key])

    def update(self, show=None):
        """
        Live visualization of the previously-rendered images.
        """

        for i, key in enumerate(self.keys):
            if show is not None:
                if key not in show:
                    continue

            if key not in ["flow_bw"] and self.data[key] is not None:
                cv2.namedWindow(self.title[i], cv2.WINDOW_NORMAL)
                cv2.resizeWindow(self.title[i], int(self.px), int(self.px))
                cv2.imshow(self.title[i], self.data[key])

        cv2.waitKey(1)

    def store(self, sequence, ts=None, show=None):
        """
        Store previously-rendered images.
        :param sequence: name of the sequence
        :param ts: timestamp of the images to be stored
        """

        # check if new sequence
        path_to = self.store_dir + sequence + "/"
        if not os.path.exists(path_to):
            os.makedirs(path_to)
            for key in self.keys:
                os.makedirs(path_to + key + "/")
            if self.store_file is not None:
                self.store_file.close()
            self.store_file = open(path_to + "timestamps.txt", "w")
            self.img_idx = 0

        # store images
        for key in self.keys:
            if show is not None:
                if key not in show:
                    continue

            if not os.path.exists(path_to + key + "/"):
                os.makedirs(path_to + key + "/")
            if self.data[key] is not None:
                filename = path_to + key + "/%09d.png" % self.img_idx
                cv2.imwrite(filename, self.data[key])

        # store timestamps
        if ts is not None:
            self.store_file.write(str(ts) + "\n")
            self.store_file.flush()

        self.img_idx += 1
        cv2.waitKey(1)

    @staticmethod
    def flow_to_image(flow):
        """
        Use the optical flow color scheme from the supplementary materials of the paper 'Back to Event
        Basics: Self-Supervised Image Reconstruction for Event Cameras via Photometric Constancy',
        Paredes-Valles et al., CVPR'21.
        :param flow: [H x W x 2] optical flow map (horizontal, vertical in dim=2)
        :return: [H x W x 3] color-encoded optical flow in BGR format
        """
        mag = np.linalg.norm(flow, axis=2)
        min_mag = np.min(mag)
        mag_range = np.max(mag) - min_mag

        ang = np.arctan2(flow[:, :, 1], flow[:, :, 0]) + np.pi
        ang *= 1.0 / np.pi / 2.0

        hsv = np.zeros([flow.shape[0], flow.shape[1], 3])
        hsv[:, :, 0] = ang
        hsv[:, :, 1] = 1.0
        hsv[:, :, 2] = mag - min_mag
        if mag_range != 0.0:
            hsv[:, :, 2] /= mag_range

        flow_rgb = matplotlib.colors.hsv_to_rgb(hsv)
        flow_rgb = (255 * flow_rgb).astype(np.uint8)
        return cv2.cvtColor(flow_rgb, cv2.COLOR_RGB2BGR)

    @staticmethod
    def events_to_image(event_cnt, color_scheme="green_red"):
        """
        Format events into an image.
        :param event_cnt: [H x W x 2] event count map
        :param color_scheme: gray / blue_red / green_red
        """
        pos = event_cnt[:, :, 0]
        neg = event_cnt[:, :, 1]
        pos_max = np.percentile(pos, 99)
        pos_min = np.percentile(pos, 1)
        neg_max = np.percentile(neg, 99)
        neg_min = np.percentile(neg, 1)
        max = pos_max if pos_max > neg_max else neg_max

        if pos_min != max:
            pos = (pos - pos_min) / (max - pos_min)
        if neg_min != max:
            neg = (neg - neg_min) / (max - neg_min)

        pos = np.clip(pos, 0, 1)
        neg = np.clip(neg, 0, 1)

        event_image = np.ones((event_cnt.shape[0], event_cnt.shape[1]))
        if color_scheme == "gray":
            event_image *= 0.5
            pos *= 0.5
            neg *= -0.5
            event_image += pos + neg

        elif color_scheme == "green_red":
            event_image = np.repeat(event_image[:, :, np.newaxis], 3, axis=2)
            event_image *= 0
            mask_pos = pos > 0
            mask_neg = neg > 0
            mask_not_pos = pos == 0
            mask_not_neg = neg == 0

            event_image[:, :, 0][mask_pos] = 0
            event_image[:, :, 1][mask_pos] = pos[mask_pos]
            event_image[:, :, 2][mask_pos * mask_not_neg] = 0
            event_image[:, :, 2][mask_neg] = neg[mask_neg]
            event_image[:, :, 0][mask_neg] = 0
            event_image[:, :, 1][mask_neg * mask_not_pos] = 0

        elif color_scheme == "rpg":
            event_image = np.repeat(event_image[:, :, np.newaxis], 3, axis=2)
            mask_pos = pos > 0
            mask_neg = neg > 0

            event_image[:, :, 0][mask_neg] = 1
            event_image[:, :, 1][mask_neg] = 0
            event_image[:, :, 2][mask_neg] = 0
            event_image[:, :, 0][mask_pos] = 0
            event_image[:, :, 1][mask_pos] = 0
            event_image[:, :, 2][mask_pos] = 1

        elif color_scheme == "prophesee":
            event_image = np.repeat(event_image[:, :, np.newaxis], 3, axis=2)
            mask_pos = pos > 0
            mask_neg = neg > 0

            event_image[:, :, 0][mask_neg] = 0.24313725490196078
            event_image[:, :, 1][mask_neg] = 0.11764705882352941
            event_image[:, :, 2][mask_neg] = 0.047058823529411764
            event_image[:, :, 0][mask_pos] = 0.6352941176470588
            event_image[:, :, 1][mask_pos] = 0.4235294117647059
            event_image[:, :, 2][mask_pos] = 0.23529411764705882

        else:
            print("Visualization error: Unknown color scheme for event images.")
            raise AttributeError

        event_image = (255 * event_image).astype(np.uint8)
        return event_image

    @staticmethod
    def minmax_norm(x, max=None, min=None):
        """
        Robust min-max normalization.
        :param x: [H x W x 1]
        :return x: [H x W x 1] normalized x
        """

        if max is None:
            max = np.percentile(x, 99)
        if min is None:
            min = np.percentile(x, 1)

        den = max - min
        if den != 0:
            x = (x - min) / den
        return np.clip(x, 0, 1)

utils/visualization.py
tudelft/taming_event_flow