Coverage for src/evutils/filtering/_masking.py: 79%

1"""Module for applying spatial masks to event arrays.""" 

2 

3import numpy as np 

4 

5def mask_events(events: np.ndarray, mask: np.ndarray) -> np.ndarray: 

6 """Masks events based on a given mask. 

7 

8 Parameters 

9 ---------- 

10 events : np.ndarray 

11 Array of events to be masked. 

12 mask : np.ndarray 

13 A 2D mask array where the events will be checked against.  

14 The mask should have the same shape as the event frame size. 

15  

16 Returns 

17 ------- 

18 np.ndarray 

19 Array of events that fall within the valid regions of the mask. 

20 

21 Examples 

22 -------- 

23 >>> import numpy as np 

24 >>> from evutils.filtering import mask_events 

25 >>> events = np.array( 

26 ... [(0, 0, 100, 1), (1, 1, 200, 1), (2, 2, 300, 0)], 

27 ... dtype=[('x', 'u2'), ('y', 'u2'), ('t', 'i8'), ('p', 'i1')] 

28 ... ) 

29 >>> mask = np.array([ 

30 ... [1, 0, 0], 

31 ... [0, 0, 0], 

32 ... [0, 0, 1] 

33 ... ]) 

34 >>> masked_events = mask_events(events, mask) 

35 >>> masked_events[['x', 'y']].tolist() 

36 [(0, 0), (2, 2)] 

37 

38 """ 

39 # Check if mask is a 2D array 

40 if mask.ndim != 2: 

41 raise ValueError("Mask must be a 2D array") 

42 

43 if len(events) == 0: 

44 return events 

45 

46 if hasattr(events, 'dtype') and events.dtype.names and 'x' in events.dtype.names and 'y' in events.dtype.names: 

47 x = events['x'] 

48 y = events['y'] 

49 elif hasattr(events, 'x') and hasattr(events, 'y'): 

50 x = getattr(events, 'x') 

51 y = getattr(events, 'y') 

52 else: 

53 raise ValueError("events must be a structured array or object with 'x' and 'y' fields") 

54 

55 # Check if max x and y in events are within the mask dimensions 

56 if x.min() < 0 or y.min() < 0: 

57 raise ValueError("Events x and y coordinates must be non-negative") 

58 if x.max() >= mask.shape[1] or y.max() >= mask.shape[0]: 

59 raise ValueError("Events x and y coordinates must be within the mask dimensions") 

60 

61 valid_events = mask[y, x] > 0 

62 

63 return events[valid_events]