Coverage for src/evutils/dense/_histogram.py: 93%

1"""Module for generating histogram-based representations from events.""" 

2 

3import numpy as np 

4from ..jit import lazy_njit_unwrapped_events 

5from ..types import EventArray 

6 

7@lazy_njit_unwrapped_events 

8def _histogram_jit(t, x, y, p, buffer, clip): 

9 height, width, _ = buffer.shape 

10 for i in range(len(t)): 

11 xi = x[i] 

12 yi = y[i] 

13 pi = p[i] 

14 if 0 <= xi < width and 0 <= yi < height: 

15 pi_mapped = 0 if pi == 1 else 2 

16 if buffer[yi, xi, pi_mapped] < clip: 

17 buffer[yi, xi, pi_mapped] += 1 

18 

19def histogram(events: 'np.ndarray | EventArray', width: int = 1280, height: int = 720, fill: bool = False, dtype: np.dtype | type = np.uint8) -> np.ndarray: 

20 """Generate a histogram frame from the events. 

21 

22 Parameters 

23 ---------- 

24 events : np.ndarray 

25 Array of events in the :class:`~evutils.types.Events` format 

26 width : int, optional 

27 Width of the frame, by default 1280 

28 height : int, optional 

29 Height of the frame, by default 720 

30 fill : bool, optional 

31 If True, the non-zero values are set to 255, by default False 

32 dtype : np.dtype, optional 

33 Data type of the output array, by default np.uint8 

34  

35 Returns 

36 ------- 

37 np.ndarray 

38 A numpy array with the histogram frame (height, width, 3) 

39 

40 Examples 

41 -------- 

42 >>> import numpy as np 

43 >>> from evutils.dense import histogram 

44 >>> events = np.array([(10, 20, 100, 1), (15, 25, 200, 0)], 

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

46 >>> frame = histogram(events, width=100, height=100) 

47 >>> frame.shape 

48 (100, 100, 3) 

49 """ 

50 buffer = np.zeros((height, width, 3), dtype=dtype) 

51 if len(events) == 0: 

52 return buffer 

53 

54 try: 

55 clip = np.iinfo(dtype).max 

56 except ValueError: 

57 clip = np.finfo(dtype).max 

58 

59 _histogram_jit(events, buffer, clip) 

60 

61 if fill: 

62 buffer[buffer > 0] = 255 

63 

64 return buffer 

65 

66@lazy_njit_unwrapped_events 

67def _wedge_histogram_jit(t, x, y, p, buffer): 

68 height, width, _ = buffer.shape 

69 for i in range(len(t)): 

70 xi = x[i] 

71 yi = y[i] 

72 pi = p[i] 

73 if 0 <= xi < width and 0 <= yi < height: 

74 pi_mapped = 0 if pi == 1 else 2 

75 buffer[yi, xi, pi_mapped] = 255 

76 

77def wedge_histogram(events: 'np.ndarray | EventArray', width: int = 1280, height: int = 720, tl: float = 30e6, dtype: np.dtype | type = np.uint8) -> np.ndarray: 

78 """Generate a wedge histogram frame from the events. 

79 

80 Parameters 

81 ---------- 

82 events : np.ndarray 

83 Array of events in the :class:`~evutils.types.Events` format 

84 width : int, optional 

85 Width of the frame, by default 1280 

86 height : int, optional 

87 Height of the frame, by default 720 

88 tl : float, optional 

89 Time limit for the frame in us, by default 30e6 

90 dtype : np.dtype, optional 

91 Data type of the output array, by default np.uint8 

92  

93 Returns 

94 ------- 

95 np.ndarray 

96 A numpy array with the wedge frame (height, width, 3) 

97 

98 Examples 

99 -------- 

100 >>> import numpy as np 

101 >>> from evutils.dense import wedge_histogram 

102 >>> events = np.array([(10, 20, 100, 1), (15, 25, 200, 0)], 

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

104 >>> frame = wedge_histogram(events, width=100, height=100) 

105 >>> frame.shape 

106 (100, 100, 3) 

107 """ 

108 buffer = np.zeros((height, width, 3), dtype=dtype) 

109 if len(events) == 0: 

110 return buffer 

111 

112 ts_norm = (events['t'] - events['t'][0])/tl 

113 sel = (height - events['y'])/height > ts_norm - 0.1 

114 ev_sel = events[sel] 

115 

116 if len(ev_sel) > 0: 

117 _wedge_histogram_jit(ev_sel, buffer) 

118 

119 return buffer