Coverage for src/evutils/processing/_masking.py: 93%

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

2 

3import numpy as np 

4 

5 

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

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

8 

9 Parameters 

10 ---------- 

11 events : np.ndarray 

12 Array of events to be masked. 

13 mask : np.ndarray 

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

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

16  

17 Returns 

18 ------- 

19 np.ndarray 

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

21 

22 Examples 

23 -------- 

24 >>> import numpy as np 

25 >>> from evutils.processing import mask_events 

26 >>> events = np.array( 

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

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

29 ... ) 

30 >>> mask = np.array([ 

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

32 ... [0, 0, 0], 

33 ... [0, 0, 1] 

34 ... ]) 

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

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

37 [(0, 0), (2, 2)] 

38 

39 """ 

40 # Check if mask is a 2D array 

41 if mask.ndim != 2: 

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

43 

44 if len(events) == 0: 

45 return events 

46 

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

48 if events['x'].max() < 0 or events['y'].max() < 0: 

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

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

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

52 

53 x = events['x'] 

54 y = events['y'] 

55 valid_events = mask[y, x] > 0 

56 

57 return events[valid_events]