Coverage for src/evutils/dense/_frame.py: 100%

1"""Module for generating frame-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 _frame_gray_jit(t, x, y, p, buffer): 

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 if pi == 1: 

16 buffer[yi, xi] = 255 

17 else: 

18 buffer[yi, xi] = 0 

19 

20def frame_gray(events: 'np.ndarray | EventArray', width: int = 1280, height: int = 720, dtype: np.dtype | type = np.uint8) -> np.ndarray: 

21 """Generate a grayscale frame from the events. 

22 

23 Parameters 

24 ---------- 

25 events : np.ndarray 

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

27 width : int, optional 

28 Width of the frame, by default 1280 

29 height : int, optional 

30 Height of the frame, by default 720 

31 dtype : np.dtype, optional 

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

33  

34 Returns 

35 ------- 

36 np.ndarray 

37 A numpy array with the frame (height, width) 

38 

39 Examples 

40 -------- 

41 >>> import numpy as np 

42 >>> from evutils.dense import frame_gray 

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

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

45 >>> frame = frame_gray(events, width=100, height=100) 

46 >>> frame.shape 

47 (100, 100) 

48 """ 

49 buffer = np.full((height, width), 128, dtype=dtype) 

50 if len(events) > 0: 

51 _frame_gray_jit(events, buffer) 

52 return buffer 

53 

54@lazy_njit_unwrapped_events 

55def _frame_rgb_jit(t, x, y, p, buffer, pos_color, neg_color): 

56 height, width, _ = buffer.shape 

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

58 xi = x[i] 

59 yi = y[i] 

60 pi = p[i] 

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

62 if pi == 1: 

63 buffer[yi, xi, 0] = pos_color[0] 

64 buffer[yi, xi, 1] = pos_color[1] 

65 buffer[yi, xi, 2] = pos_color[2] 

66 else: 

67 buffer[yi, xi, 0] = neg_color[0] 

68 buffer[yi, xi, 1] = neg_color[1] 

69 buffer[yi, xi, 2] = neg_color[2] 

70 

71def frame_rgb(ev: 'np.ndarray | EventArray', width: int = 1280, height: int = 720, bg_color: tuple[int, int, int] | None = None, pos_color: tuple[int, int, int] | None = None, neg_color: tuple[int, int, int] | None = None) -> np.ndarray: 

72 """Generate an RGB frame from the events. 

73 

74 Parameters 

75 ---------- 

76 ev : np.ndarray 

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

78 width : int, optional 

79 Width of the frame, by default 1280 

80 height : int, optional 

81 Height of the frame, by default 720 

82 bg_color : np.ndarray, optional 

83 Background color, by default np.array((0, 0, 0)) 

84 pos_color : np.ndarray, optional 

85 Color for positive events, by default np.array((255, 0, 0)) 

86 neg_color : np.ndarray, optional 

87 Color for negative events, by default np.array((0, 0, 255)) 

88 

89 Returns 

90 ------- 

91 np.ndarray 

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

93 

94 Examples 

95 -------- 

96 >>> import numpy as np 

97 >>> from evutils.dense import frame_rgb 

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

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

100 >>> frame = frame_rgb(events, width=100, height=100) 

101 >>> frame.shape 

102 (100, 100, 3) 

103 """ 

104 if bg_color is None: bg_color = np.array([0, 0, 0], dtype=np.uint8) 

105 if pos_color is None: pos_color = np.array([255, 0, 0], dtype=np.uint8) 

106 if neg_color is None: neg_color = np.array([0, 0, 255], dtype=np.uint8) 

107 

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

109 buffer[:, :] = bg_color 

110 

111 if len(ev) > 0: 

112 _frame_rgb_jit(ev, buffer, pos_color, neg_color) 

113 

114 return buffer 

115 

116@lazy_njit_unwrapped_events 

117def _frame_diff_jit(t, x, y, p, buffer): 

118 height, width = buffer.shape 

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

120 xi = x[i] 

121 yi = y[i] 

122 pi = p[i] 

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

124 if pi == 1: 

125 buffer[yi, xi] += 1 

126 else: 

127 buffer[yi, xi] -= 1 

128 

129def frame_diff(ev: 'np.ndarray | EventArray', width: int = 1280, height: int = 720, dtype: np.dtype | type = np.int8) -> np.ndarray: 

130 """Generate a differential frame from the events. 

131 

132 Parameters 

133 ---------- 

134 ev : np.ndarray 

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

136 width : int, optional 

137 Width of the frame, by default 1280 

138 height : int, optional 

139 Height of the frame, by default 720 

140 dtype : np.dtype, optional 

141 Data type of the output array, by default np.int8 

142  

143 Returns 

144 ------- 

145 np.ndarray 

146 A numpy array with the frame (height, width) 

147 

148 Examples 

149 -------- 

150 >>> import numpy as np 

151 >>> from evutils.dense import frame_diff 

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

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

154 >>> frame = frame_diff(events, width=100, height=100) 

155 >>> frame.shape 

156 (100, 100) 

157 """ 

158 buffer = np.zeros((height, width), dtype=dtype) 

159 if len(ev) > 0: 

160 _frame_diff_jit(ev, buffer) 

161 return buffer 

162