Coverage for src/evutils/vis/plot3d.py: 33%

1"""Matplotlib 3D visualization utilities for event data. 

2 

3This module provides functions to plot event streams, 3D histograms, 

4and time surfaces using matplotlib's 3D plotting capabilities. 

5 

6``cv2`` and ``matplotlib`` are imported lazily, inside the functions that 

7use them, so ``import evutils.vis`` (and importing this module) stays cheap 

8and does not drag in OpenCV / the full matplotlib stack. 

9""" 

10from __future__ import annotations 

11 

12from typing import TYPE_CHECKING, Optional, Union 

13 

14import numpy as np 

15 

16if TYPE_CHECKING: 

17 from matplotlib.colors import Colormap 

18 from matplotlib.figure import Figure 

19 from mpl_toolkits.mplot3d import Axes3D 

20 

21def plot_3d(events: np.ndarray, 

22 width: int =1280, 

23 height: int = 720, 

24 colormap: Union[str, Colormap] ='Spectral', 

25 time_scale: float = 1000.0, 

26 time_unit: str = 'ms', 

27 fig: Optional[Figure] = None, 

28 ax: Optional[Axes3D] = None) -> tuple[Optional[Figure], Optional[Axes3D]]: 

29 """Plot a 3D scatter plot of events. 

30  

31 Parameters 

32 ---------- 

33 events : np.ndarray 

34 Array of events with fields 't', 'x', 'y', and 'p'. 

35 width : int, optional 

36 Width of the event frame, by default 1280. 

37 height : int, optional 

38 Height of the event frame, by default 720. 

39 colormap : Union[str, Colormap], optional 

40 Colormap to use for the event polarities, by default 'Spectral'. 

41 fig : Optional[plt.Figure], optional 

42 Matplotlib figure to plot on, by default None (a new figure will be created). 

43 ax : Optional[Axes3D], optional 

44 Matplotlib 3D axis to plot on, by default None (a new axis will be created). 

45 

46 Returns 

47 ------- 

48 tuple 

49 A tuple containing the figure and axis objects. 

50 

51 Raises 

52 ------ 

53 ValueError 

54 If the colormap is not found or invalid. 

55  

56 """ 

57 import matplotlib.pyplot as plt 

58 from mpl_toolkits.mplot3d import Axes3D # noqa: F401 (registers the '3d' projection) 

59 

60 ts = events['t'] / time_scale 

61 

62 # Create color map: red for polarity=0, blue for polarity=1 

63 # Resolve colormap 

64 if isinstance(colormap, str): 

65 colormap = plt.get_cmap(colormap) 

66 if colormap is None: 

67 raise ValueError(f"Colormap not found or invalid.") 

68 

69 colors = colormap(events['p'].astype(np.float32)) # Normalize polarity to [0, 1] for colormap 

70 

71 # Create figure/axis if not provided 

72 if ax is None: 

73 fig = plt.figure(figsize=(10, 7)) if fig is None else fig 

74 ax = fig.add_subplot(111, projection='3d') 

75 

76 if len(events) == 0: 

77 return fig, ax 

78 

79 # Plot as points 

80 ax.scatter(events['x'], events['y'], ts, c=colors, s=1, alpha=0.2) 

81 

82 # Labels 

83 ax.set_xlabel('x') 

84 ax.set_ylabel('y') 

85 ax.set_zlabel(f'time ({time_unit})') 

86 ax.set_title('Event Stream 3D Plot') 

87 ax.set_xlim([0, width]) 

88 ax.set_ylim([height, 0]) 

89 ax.set_zlim([ts.min(), ts.max()]) 

90 ax.set_box_aspect([width, height, max(width, height)]) # Aspect ratio 

91 

92 return fig, ax 

93 

94def plot_3d_histogram(histogram: np.ndarray, 

95 down_sample: int = 4, 

96 fig: Optional[Figure] = None, 

97 ax: Optional[Axes3D] = None) -> tuple[Optional[Figure], Optional[Axes3D]]: 

98 """Plot a 3D histogram of events. 

99 

100 Parameters 

101 ---------- 

102 histogram : np.ndarray 

103 3D histogram array with shape (depth, height, width). 

104 down_sample : int, optional 

105 Downsampling factor for the histogram, by default 4. 

106 fig : Optional[plt.Figure], optional 

107 Matplotlib figure to plot on, by default None (a new figure will be created). 

108 ax : Optional[Axes3D], optional 

109 Matplotlib 3D axis to plot on, by default None (a new axis will be created). 

110 

111 Returns 

112 ------- 

113 tuple 

114 A tuple containing the figure and axis objects. 

115 

116 """ 

117 import cv2 

118 import matplotlib.pyplot as plt 

119 from mpl_toolkits.mplot3d import Axes3D # noqa: F401 (registers the '3d' projection) 

120 

121 # Create figure/axis if not provided 

122 if ax is None: 

123 fig = plt.figure(figsize=(10, 7)) if fig is None else fig 

124 ax = fig.add_subplot(111, projection='3d') 

125 # Create a meshgrid for the histogram 

126 

127 assert histogram.ndim == 4, "Histogram must be a 4D array." 

128 n_bins, height, width, ch = histogram.shape 

129 

130 height_downsampled = height // down_sample 

131 width_downsampled = width // down_sample 

132 x, y = np.meshgrid(np.linspace(0, width - 1, width_downsampled), 

133 np.linspace(0, height - 1, height_downsampled)) 

134 

135 

136 for bin in range(n_bins): 

137 print(f"Plotting bin {bin+1}/{n_bins}") 

138 image = histogram[bin] 

139 

140 # Downsample the image for better visualization 

141 image = cv2.resize(image, (width_downsampled, height_downsampled), interpolation=cv2.INTER_LINEAR) 

142 

143 # Conver to RGBA 

144 image = cv2.cvtColor(image, cv2.COLOR_BGR2RGBA) 

145 

146 # Set the A channel to 0 where the pixel is black 

147 image[:, :, 3] = np.where(np.all(image[:, :, :3] == 0, axis=-1), 0, 255) 

148 

149 z = np.full(x.shape, bin) 

150 

151 # Plot the surface 

152 ax.plot_surface(x, y, z, facecolors=image/255, rstride=1, cstride=1, shade=False) 

153 

154 # Create a meshgrid for the histogram 

155 

156 ax.set_xlabel('x') 

157 ax.set_ylabel('y') 

158 ax.set_zlabel('bins') 

159 ax.set_title('Voxel Histogram') 

160 ax.set_xlim([0, width]) 

161 ax.set_ylim([height, 0]) 

162 ax.set_zlim([0, n_bins]) 

163 ax.set_box_aspect([width, height, max(width, height)]) # Aspect ratio 

164 

165 return fig, ax 

166 

167def plot_3d_timesurface(events: np.ndarray, 

168 width: int =1280, 

169 height: int = 720, 

170 tau: int = 10_000, 

171 fig: Optional[Figure] = None, 

172 ax: Optional[Axes3D] = None) -> tuple[Optional[Figure], Optional[Axes3D]]: 

173 """Plot a 3D time surface of events. 

174 

175 Parameters 

176 ---------- 

177 events : np.ndarray 

178 Array of events with fields 't', 'x', 'y', and 'p'. 

179 width : int, optional 

180 Width of the event frame, by default 1280. 

181 height : int, optional 

182 Height of the event frame, by default 720. 

183 tau : int, optional 

184 Time constant for the time surface exponential decay in microseconds, by default 10_000. 

185 fig : Optional[plt.Figure], optional 

186 Matplotlib figure to plot on, by default None (a new figure will be created). 

187 ax : Optional[Axes3D], optional 

188 Matplotlib 3D axis to plot on, by default None (a new axis will be created). 

189 

190 Returns 

191 ------- 

192 tuple 

193 A tuple containing the figure and axis objects. 

194 

195 """ 

196 import matplotlib.pyplot as plt 

197 from mpl_toolkits.mplot3d import Axes3D # noqa: F401 (registers the '3d' projection) 

198 

199 ts = events['t'] # Convert timestamps to seconds 

200 

201 # Create figure/axis if not provided 

202 if ax is None: 

203 fig = plt.figure(figsize=(10, 7)) if fig is None else fig 

204 ax = fig.add_subplot(111, projection='3d') 

205 

206 if len(events) == 0: 

207 return fig, ax 

208 

209 ts_ref = ts[-1] # Reference timestamp for normalization 

210 

211 p_sign = (events['p'].astype(np.float32) * 2 - 1) 

212 

213 colors = np.exp(-(ts_ref - ts) / tau) * p_sign # Normalize polarity 

214 

215 ts = ts / 1_000 # Convert timestamps to milliseconds 

216 ax.scatter(events['x'], events['y'], ts, c=colors, s=1, alpha=0.2) 

217 

218 # Labels 

219 ax.set_xlabel('x') 

220 ax.set_ylabel('y') 

221 ax.set_zlabel('time (ms)') 

222 ax.set_title('Event Stream 3D Plot') 

223 ax.set_xlim([0, width]) 

224 ax.set_ylim([height, 0]) 

225 ax.set_zlim([ts.min(), ts.max()]) 

226 ax.set_box_aspect([width, height, max(width, height)]) # Aspect ratio 

227 

228 return fig, ax