[DEV] Adding latency model as a preprocessing of events in torch dataset

.vscode/launch.json

@@ -0,0 +1,61 @@
+{
+    // Use IntelliSense to learn about possible attributes.
+    // Hover to view descriptions of existing attributes.
+    // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387
+    "version": "0.2.0",
+    "configurations": [
+        {
+            "name": "eval",
+            "type": "python",
+            "request": "launch",
+            "program": "/mnt/data/latency-paper/EVREAL/eval.py",
+            "console": "integratedTerminal",
+            "justMyCode": false,
+            "purpose": [
+                "debug-in-terminal"
+            ],
+            "cwd": "/mnt/data/latency-paper/EVREAL",
+            "args": [
+                "-m",
+                "E2VID",
+                // "E2VID+",
+                // "ET-Net",
+                // "HyperE2VID",
+                "-c",
+                "std",
+                "-d",
+                "ECD",
+                "-qm",
+                "mse",
+                // "ssim",
+                // "lpips",
+            ],
+            "host": "127.0.0.1",
+            "pythonPath": "/root/cimaging-venv/bin/python"
+        },
+        {
+            "name": "debug_docker_predict",
+            "type": "python",
+            "request": "launch",
+            "program": "./pypkg/cimaging_ml/training/cimaging_main.py",
+            "console": "integratedTerminal",
+            "justMyCode": true,
+            "purpose": [
+                "debug-in-terminal"
+            ],
+            "cwd": "/root/metavision-computational-imaging/sdk/modules/cimaging_ml/python",
+            "args": [
+                "predict",
+                "--preset",
+                "user",
+                "--config",
+                "/root/metavision-computational-imaging/sdk/modules/cimaging_ml/python/pypkg/config/cimaging_ml_main/deblur/releases/latest/latest_ml_rggb.yaml",
+                "--experiment_name",
+                "debug_docker_predict",
+                "--no_clearml",
+            ],
+            "host": "127.0.0.1",
+            "pythonPath": "/root/cimaging-venv/bin/python"
+        },
+    ]
+}

config/dataset/ECD.json

@@ -1,5 +1,5 @@
 {
-    "root_path": "data/ECD",
+    "root_path": "/mnt/data/ECD",
     "sequences": 
     {
         "calibration":

config/eval/std.json

@@ -4,9 +4,10 @@
         "voxel_method": {
             "method": "between_frames"
         },
-        "keep_ratio": 1.0
+        "keep_ratio": 1.0,
+        "correct_latency": true
     },
-    "save_images": true,
+    "save_images": false,
     "histeq": "none",
     "eval_infer_all": false,
     "ts_tol_ms": 1.0,

config/eval/std_all.json

@@ -4,9 +4,10 @@
         "voxel_method": {
             "method": "between_frames"
         },
-        "keep_ratio": 1.0
+        "keep_ratio": 1.0,
+        "correct_latency": true
     },
-    "save_images": true,
+    "save_images": false,
     "histeq": "none",
     "eval_infer_all": true,
     "ts_tol_ms": 1.0,

dataset.py

@@ -7,14 +7,31 @@
 from torch.utils.data import Dataset
 
 # local modules
-from utils.event_utils import events_to_voxel_torch
+from utils.event_utils import events_to_polarity_fixed_bin_exposure_voxel_torch, events_to_voxel_torch
 from utils.util import read_json
 
+import sys
+
+sys.path.append('/root/metavision-computational-imaging/sdk/modules/cimaging_ml/python/pypkg')
+
+from cimaging_ml.applications.abstract.cimaging_pipeline_factory import instantiate_pipeline_from_ckpt
+from cimaging_ml.processings.event_processing import resample_histograms, HistogramPolarityFixedExposureProcessor
+
+LATENCY_CORRECTOR_CKPT = "/mnt/data/EventEnhancementPipeline-sparse-boxer-dog-52-0d71-last-epoch=199-step=25000"
+HISTO_BIN_SIZE_NS = 1e6     # =1ms
 
 class MemMapDataset(Dataset):
 
-    def __init__(self, data_path, sensor_resolution=None, num_bins=5,
-                 voxel_method=None, max_length=None, keep_ratio=1):
+    def __init__(
+    self,
+    data_path,
+    sensor_resolution=None,
+    num_bins=5,
+    voxel_method=None,
+    max_length=None,
+    keep_ratio=1,
+    correct_latency=False,
+):
         self.num_bins = num_bins
         self.data_path = data_path
         self.keep_ratio = keep_ratio
@@ -29,6 +46,18 @@ def __init__(self, data_path, sensor_resolution=None, num_bins=5,
 
         if max_length is not None:
             self.length = min(self.length, max_length + 1)
+        self.histo_processor = None
+        self.correct_latency = correct_latency
+        if self.correct_latency:
+            latency_corrector_model, _ = instantiate_pipeline_from_ckpt(LATENCY_CORRECTOR_CKPT)
+            self.latency_corrector_model = latency_corrector_model
+            self.histo_processor = HistogramPolarityFixedExposureProcessor(
+                bin_exposure=HISTO_BIN_SIZE_NS,
+                extra_events=0,
+                ignore_skew=False,
+                match_exposure=True,
+                histo_size_divisor=1
+            )
 
     def __getitem__(self, index):
 
@@ -212,7 +241,10 @@ def get_voxel_grid(self, xs, ys, ts, ps):
         create voxel grid merging positive and negative events (resulting in NUM_BINS x H x W tensor).
         """
         # generate voxel grid which has size self.num_bins x H x W
-        voxel_grid = events_to_voxel_torch(xs, ys, ts, ps, self.num_bins, sensor_size=self.sensor_resolution)
+        if not self.correct_latency:
+            voxel_grid = events_to_voxel_torch(xs, ys, ts, ps, self.num_bins, sensor_size=self.sensor_resolution)
+        else:
+            voxel_grid = self.correct_events_latency(xs, ys, ts, ps)
         return voxel_grid
 
     def get_frame(self, index):
@@ -292,3 +324,35 @@ def compute_frame_indices(self):
             frame_indices.append([start_idx, end_idx])
             start_idx = end_idx
         return frame_indices
+
+    def correct_events_latency(self, xs, ys, ts, ps):
+        """
+        Correct events latency
+        Parameters
+        ----------
+        events : events tensor
+        Returns
+        -------
+        events: events tensor with corrected latency
+        """
+        assert self.histo_processor is not None, "histo_processor is not defined"
+        # models from the literature assumes these thresholds
+        histo, histo_bin_size, histo_exposure_time = events_to_polarity_fixed_bin_exposure_voxel_torch(
+            self.histo_processor,
+            xs,
+            ys,
+            ts,
+            ps,
+        )
+        th_off = 1.0
+        th_on = 1.0
+
+        corrected_histo = self.latency_corrector_model(histo.unsqueeze(0))   # add batch dimension for inference
+        corrected_histo = -th_off * corrected_histo[:, 0] + th_on * corrected_histo[:, 1]
+        resampled_histo = resample_histograms(
+            corrected_histo,
+            self.num_bins,
+            histo_exposure_time / histo_bin_size,
+            simple_linear_interpolation=True,
+        ).detach().squeeze(0)
+        return resampled_histo

utils/event_utils.py

@@ -1,4 +1,8 @@
 import torch
+import numpy as np
+
+from cimaging_ml.processings.event_processing import compute_polarity_histograms
+from metavision_sdk_base import EventCD
 
 
 def events_to_image_torch(xs, ys, ps, device=None, sensor_size=(180, 240)):
@@ -57,3 +61,67 @@ def events_to_voxel_torch(xs, ys, ts, ps, num_bins, device=None, sensor_size=(18
         bins.append(vb)
     bins = torch.stack(bins)
     return bins
+
+def events_to_polarity_fixed_bin_exposure_voxel_torch(
+    histo_processor,
+    xs,
+    ys,
+    ts,
+    ps,
+    device=None,
+    sensor_size=(180, 240)
+):
+    """
+    Turn set of events to a voxel grid tensor, using temporal bilinear interpolation
+    Parameters
+    ----------
+    xs : list of event x coordinates (torch tensor)
+    ys : list of event y coordinates (torch tensor)
+    ts : list of event timestamps (torch tensor)
+    ps : list of event polarities (torch tensor)
+    device : device to put voxel grid. If left empty, same device as events
+    sensor_size : the size of the event sensor/output voxels
+    Returns
+    -------
+    voxel: voxel of the events between t0 and t1
+    """
+    if device is None:
+        device = xs.device
+    assert (len(xs) == len(ys) and len(ys) == len(ts) and len(ts) == len(ps))
+    dt = ts[-1] - ts[0]
+    dt *= 1e9
+
+    xs = xs.cpu().numpy().astype(dtype=np.int32)
+    ys = ys.cpu().numpy().astype(dtype=np.int32)
+    ps = ps.cpu().numpy().astype(dtype=np.int32)
+    ts = (ts.cpu().numpy() * 1e9).astype(dtype=np.int32)
+
+    # Create the structured array
+    events = np.empty(xs.shape, dtype=EventCD)
+    for i in range(len(events)):
+        events[i]["x"] = xs[i]
+        events[i]["y"] = ys[i]
+        events[i]["p"] = ps[i] if ps[i] == 1 else 0
+        events[i]["t"] = ts[i]
+
+    hist = histo_processor.process(
+        events,
+        sensor_size[0],
+        sensor_size[1],
+        row_exposure_time=dt,
+        exposure_map=None,
+        start_exposure_first_line=ts[0],
+        thoff=1.0,
+        thon=1.0,
+        sharp_ts_offset=None,
+    )
+    # hist = hist.clip(0, 100)
+
+    histo_bin_size, histo_exposure_time, n_histograms = histo_processor.compute_histo_bin_size(
+        dt,
+        None
+    )
+    # latency model expect even number of time bins so add zero padding if needed
+    if hist.shape[1] % 2 == 1:
+        hist = np.concatenate([hist, np.zeros_like(hist[:, :1])], axis=1)
+    return torch.tensor(hist, device=device, dtype=torch.float), histo_bin_size, histo_exposure_time