Floating point events + image reconstruction additional arg

python/conversion/h5writer.py

@@ -1,23 +1,27 @@
 from pathlib import Path
 import weakref
+from typing import Union
 
 import h5py
 import numpy as np
 
-from data.format import Events
+from data.format import Events, EventsFloat
+
 
 class H5Writer:
-    def __init__(self, outfile: Path):
+    def __init__(self, outfile: Path, floating_point_coords=False):
         assert not outfile.exists(), str(outfile)
         self.h5f = h5py.File(str(outfile), 'w')
         self._finalizer = weakref.finalize(self, self.close_callback, self.h5f)
 
+        coord_dtype = "u2" if not floating_point_coords else "f4"
+
         # create hdf5 datasets
         shape = (2**16,)
         maxshape = (None,)
         compression = 'lzf'
-        self.h5f.create_dataset('x', shape=shape, dtype='u2', chunks=shape, maxshape=maxshape, compression=compression)
-        self.h5f.create_dataset('y', shape=shape, dtype='u2', chunks=shape, maxshape=maxshape, compression=compression)
+        self.h5f.create_dataset('x', shape=shape, dtype=coord_dtype, chunks=shape, maxshape=maxshape, compression=compression)
+        self.h5f.create_dataset('y', shape=shape, dtype=coord_dtype, chunks=shape, maxshape=maxshape, compression=compression)
         self.h5f.create_dataset('p', shape=shape, dtype='u1', chunks=shape, maxshape=maxshape, compression=compression)
         self.h5f.create_dataset('t', shape=shape, dtype='i8', chunks=shape, maxshape=maxshape, compression=compression)
         self.row_idx = 0
@@ -26,7 +30,7 @@ def __init__(self, outfile: Path):
     def close_callback(h5f: h5py.File):
         h5f.close()
 
-    def add_data(self, events: Events):
+    def add_data(self, events: Union[Events, EventsFloat]):
         current_size = events.size
         new_size = self.row_idx + current_size
         self.h5f['x'].resize(new_size, axis=0)

python/data/format.py

@@ -28,6 +28,31 @@ def __post_init__(self):
         if self.size > 0:
             assert np.max(self.p) <= 1
 
+
+@dataclass(frozen=True)
+class EventsFloat:
+    x: np.ndarray
+    y: np.ndarray
+    p: np.ndarray
+    t: np.ndarray
+
+    size: int = field(init=False)
+
+    def __post_init__(self):
+        assert self.x.dtype == np.float32
+        assert self.y.dtype == np.float32
+        assert self.p.dtype == np.uint8
+        assert self.t.dtype == np.int64
+
+        assert self.x.shape == self.y.shape == self.p.shape == self.t.shape
+        assert self.x.ndim == 1
+
+        # Without the frozen option, we could just do: self.size = self.x.size
+        super().__setattr__('size', self.x.size)
+
+        if self.size > 0:
+            assert np.max(self.p) <= 1
+
 @dataclass(frozen=True)
 class EventsForReconstruction:
     events: Events

python/offline_reconstruction.py

@@ -37,7 +37,8 @@ def download_checkpoint(path_to_model):
     parser.add_argument('--freq_hz', '-fhz', type=int, default=0, help='Frequency for saving the reconstructed images from events')
     parser.add_argument('--timestamps_file', '-tsf', help='Path to txt file containing image reconstruction timestamps')
     parser.add_argument('--upsample_rate', '-u', type=int, default=1, help='Multiplies the number of reconstructions, which effectively lowers the time window of events for E2VID. These intermediate reconstructions will not be saved to disk.')
-    parser.add_argument('--verbose', '-v',  action='store_true', default=False, help='Verbose output')
+    parser.add_argument('--verbose', '-v',  action='store_true', help='Verbose output')
+    parser.add_argument('--index_by_order', '-i',  action='store_true', help='Index reconstrutions with 0,1,2,3...')
 
     set_inference_options(parser)
 
@@ -79,14 +80,15 @@ def download_checkpoint(path_to_model):
     print('Will write images to: {}'.format(os.path.join(args.output_folder, args.dataset_name)))
     grid = VoxelGrid(model.num_bins, args.width, args.height, upsample_rate=args.upsample_rate)
     pbar = tqdm.tqdm(total=len(data_provider))
-    for events in data_provider:
+    for j, events in enumerate(data_provider):
         if events.events.size > 0:
             sliced_events = grid.event_slicer(events.events, events.t_reconstruction)
             for i in range(len(sliced_events)):
                 event_grid, _ = grid.events_to_voxel_grid(sliced_events[i])
                 event_tensor = torch.from_numpy(event_grid)
                 if i== len(sliced_events) - 1:
-                    image_reconstructor.update_reconstruction(event_tensor, int(events.t_reconstruction)*1000, save=True)
+                    index = j if args.index_by_order else int(events.t_reconstruction)*1000
+                    image_reconstructor.update_reconstruction(event_tensor, index, save=True)
                     pbar.update(1)
                 else:
                     image_reconstructor.update_reconstruction(event_tensor)