diff --git a/tonic/functional/drop_event.py b/tonic/functional/drop_event.py index d03887a..a1011a7 100644 --- a/tonic/functional/drop_event.py +++ b/tonic/functional/drop_event.py @@ -1,4 +1,4 @@ -from typing import Tuple, Union +from typing import Optional, Tuple, Union import numpy as np @@ -64,14 +64,14 @@ def drop_by_time_numpy( def drop_by_area_numpy( - events: np.ndarray, sensor_size: Tuple, area_ratio: Union[float, Tuple[float]] = 0.2 + events: np.ndarray, sensor_size: Optional[Tuple[int, int, int]] = None, area_ratio: Union[float, Tuple[float]] = 0.2 ): """Drops events located in a randomly chosen box area. The size of the box area is defined by a specified ratio of the sensor size. Args: events (np.ndarray): ndarray of shape [num_events, num_event_channels] - sensor_size (Tuple): size of the sensor that was used [W,H,P] + sensor_size (Optional[Tuple[int, int, int]]): size of the sensor that was used [W,H,P]. Defaults to None. area_ratio (Union[float, Tuple[float]], optional): Ratio of the sensor resolution that determines the size of the box area where events are dropped. - if a float, the value is used to calculate the size of the box area - if a tuple of 2 floats, the ratio is randomly chosen in [min, max) diff --git a/tonic/functional/spatial_jitter.py b/tonic/functional/spatial_jitter.py index f61dcaa..1d6b870 100644 --- a/tonic/functional/spatial_jitter.py +++ b/tonic/functional/spatial_jitter.py @@ -1,11 +1,11 @@ -from typing import List +from typing import Tuple import numpy as np def spatial_jitter_numpy( events: np.ndarray, - sensor_size: List[int], + sensor_size: Tuple[int, int, int], var_x: float = 1, var_y: float = 1, sigma_xy: float = 0, diff --git a/tonic/functional/to_frame.py b/tonic/functional/to_frame.py index af8dc23..5fb5e30 100644 --- a/tonic/functional/to_frame.py +++ b/tonic/functional/to_frame.py @@ -10,7 +10,7 @@ def to_frame_numpy( events, - sensor_size, + sensor_size = None, time_window=None, event_count=None, n_time_bins=None, @@ -23,7 +23,7 @@ def to_frame_numpy( Parameters: events: ndarray of shape [num_events, num_event_channels] - sensor_size: size of the sensor that was used [W,H,P] + sensor_size (None): size of the sensor that was used [W,H,P] time_window (None): window length in us. event_count (None): number of events per frame. n_time_bins (None): fixed number of frames, sliced along time axis. diff --git a/tonic/transforms.py b/tonic/transforms.py index 41302a0..dc1a590 100644 --- a/tonic/transforms.py +++ b/tonic/transforms.py @@ -61,7 +61,7 @@ def __call__(self, events: np.ndarray) -> np.ndarray: if type(self.size) == int: self.size = [self.size, self.size] offsets = (self.sensor_size[0] - self.size[0]) // 2, ( - self.sensor_size[1] - self.size[1] + self.sensor_size[1] - self.size[1] ) // 2 offset_idx = [max(offset, 0) for offset in offsets] cropped_events = events[ @@ -69,7 +69,7 @@ def __call__(self, events: np.ndarray) -> np.ndarray: & (events["x"] < (offset_idx[0] + self.size[0])) & (offset_idx[1] <= events["y"]) & (events["y"] < (offset_idx[1] + self.size[1])) - ] + ] cropped_events["x"] -= offsets[0] cropped_events["y"] -= offsets[1] return cropped_events @@ -189,7 +189,7 @@ class DropEventByArea: specified ratio of the sensor size. Args: - sensor_size (Tuple): size of the sensor that was used [W,H,P] + sensor_size (Optional[Tuple[int, int, int]]): size of the sensor that was used [W,H,P]. Defaults to None. area_ratio (Union[float, Tuple[float]], optional): Ratio of the sensor resolution that determines the size of the box area where events are dropped. - if a float, the value is used to calculate the size of the box area - if a tuple of 2 floats, the ratio is randomly chosen in [min, max) @@ -199,7 +199,7 @@ class DropEventByArea: >>> transform = tonic.transforms.DropEventByArea(sensor_size=(128,128,2), area_ratio=(0.1, 0.8)) """ - sensor_size: Tuple[int, int, int] + sensor_size: Optional[Tuple[int, int, int]] = None area_ratio: Union[float, Tuple[float, float]] = 0.2 def __call__(self, events): @@ -229,7 +229,7 @@ class DropPixel: def __call__(self, events): if len(events) == 0: - return events # return empty array + return events # return empty array if events.dtype.names is not None: # assert "x", "y", "p" in events.dtype.names @@ -788,10 +788,10 @@ class NumpyAsType: def __call__(self, events): source_is_structured_array = ( - hasattr(events.dtype, "names") and events.dtype.names != None + hasattr(events.dtype, "names") and events.dtype.names != None ) target_is_structured_array = ( - hasattr(self.dtype, "names") and self.dtype.names != None + hasattr(self.dtype, "names") and self.dtype.names != None ) if source_is_structured_array and not target_is_structured_array: return np.lib.recfunctions.structured_to_unstructured(events, self.dtype) @@ -897,15 +897,30 @@ class ToFrame: include_incomplete: bool = False def __call__(self, events): - # if events are empty, return a frame in the expected format if len(events) == 0: if self.time_window is not None or self.event_count is not None: - return np.zeros((1, self.sensor_size[2], self.sensor_size[0], self.sensor_size[1])) + return np.zeros( + (1, self.sensor_size[2], self.sensor_size[0], self.sensor_size[1]) + ) elif self.n_event_bins is not None: - return np.zeros((self.n_event_bins, self.sensor_size[2], self.sensor_size[0], self.sensor_size[1])) + return np.zeros( + ( + self.n_event_bins, + self.sensor_size[2], + self.sensor_size[0], + self.sensor_size[1], + ) + ) elif self.n_time_bins is not None: - return np.zeros((self.n_time_bins, self.sensor_size[2], self.sensor_size[0], self.sensor_size[1])) + return np.zeros( + ( + self.n_time_bins, + self.sensor_size[2], + self.sensor_size[0], + self.sensor_size[1], + ) + ) else: raise ValueError("No slicing method specified.") @@ -988,7 +1003,7 @@ class ToImage: smaller chunks that are then individually binned to frames. """ - sensor_size: Tuple[int, int, int] + sensor_size: Optional[Tuple[int, int, int]] def __call__(self, events): frames = functional.to_frame_numpy(