Floating point events + image reconstruction additional arg

python/conversion/h5writer.py

@@ -1,32 +1,36 @@
 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.h5f.create_dataset('t', shape=shape, dtype='u4', chunks=shape, maxshape=maxshape, compression=compression)
         self.row_idx = 0
 
     @staticmethod
     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

@@ -17,7 +17,32 @@ def __post_init__(self):
         assert self.x.dtype == np.uint16
         assert self.y.dtype == np.uint16
         assert self.p.dtype == np.uint8
-        assert self.t.dtype == np.int64
+        assert self.t.dtype == np.uint32
+
+        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 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.uint32
 
         assert self.x.shape == self.y.shape == self.p.shape == self.t.shape
         assert self.x.ndim == 1
@@ -54,7 +79,7 @@ def __post_init__(self):
                 x=np.array([0, 1], dtype='uint16'),
                 y=np.array([1, 2], dtype='uint16'),
                 p=np.array([0, 0], dtype='uint8'),
-                t=np.array([0, 5], dtype='int64')),
+                t=np.array([0, 5], dtype='uint32')),
             width=2,
             height=3,
             t_reconstruction=6)

python/data/provider.py

@@ -137,7 +137,7 @@ def get_events_until(self, time: int) -> Optional[Events]:
                         np.array([], dtype=np.uint16),
                         np.array([], dtype=np.uint16),
                         np.array([], dtype=np.uint8),
-                        np.array([], dtype=np.int64))
+                        np.array([], dtype=np.uint32))
             else:
                 retrieved_events = Events(
                         self.local_buffer.x[indices],

python/offline_reconstruction.py

@@ -38,6 +38,7 @@ def download_checkpoint(path_to_model):
     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('--index_by_order', '-i',  action='store_true', default=False, 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)