Coverage for src/evutils/io/

1"""Prophesee AER CD-event decoder/encoder.

3AER is a raw 32-bit-per-event encoding with **no header and no timestamps**:

4``y[0:8]`` (9 bits), ``x[9:17]`` (9 bits), ``p[18]``. The 9-bit fields cap

5coordinates at 512 (e.g. GenX320). Decoding uses the native

6``AER_parse_chunk_soa``; encoding is vectorised numpy.

8Since the format carries no time information, the decoder's ``timestamps``

9parameter selects how the ``t`` column is generated:

11* ``"zero"`` (default) -- every event gets ``t = 0``;

12* ``"sequential"`` -- ``t = t_start + i * t_step``, generated in the native

13 parser and carried across chunks;

14* an array -- user-provided timestamps, assigned positionally (event ``i`` in

15 the stream gets ``timestamps[i]``).

16"""

17from __future__ import annotations

19import io

20from datetime import datetime

21from typing import Any

23import numpy as np

25from ..types import EventArray, TriggerArray

26from .common import EventDecoder, EventEncoder

27from ._native_core import (

28 EventSoABuffers,

29 TriggerSoABuffers,

30 decode_all_soa,

31 events_view,

32 parse_step,

33)

34from ._native_aer import (

35 AER_TS_SEQUENTIAL,

36 AER_TS_ZERO,

37 AerInput,

38 AerParser,

39)

40from ._source import ByteSource

42_EMPTY_EVENTS = EventArray.empty()

45class EventDecoder_AER(EventDecoder):

46 """Decode raw AER streams into ``EventArray`` chunks. Since AER is designed

47 for real-time streaming, it has no header and no timestamps; see the

48 ``timestamps`` parameter for how the ``t`` column is generated.

50 Parameters

51 ----------

52 source

53 Byte source to read from.

54 chunk_size

55 Maximum number of events produced per :meth:`read_chunk` call (the

56 native output-buffer capacity). Does not bound the file size.

57 timestamps : {"zero", "sequential"} or array_like, default "zero"

58 Timestamp generation mode: ``"zero"`` fills ``t = 0``,

59 ``"sequential"`` fills ``t = t_start + i * t_step``, and an integer

60 array assigns user-provided timestamps positionally (its length must

61 cover every decoded event).

62 t_start, t_step : int

63 Start value and per-event increment for ``"sequential"`` mode.

65 References

66 ----------

67 [1] Prophesee AER format: https://docs.prophesee.ai/stable/data/encoding_formats/aer.html

70 """

72 def __init__(self, source: ByteSource, chunk_size: int = 1_000_000,

73 timestamps: 'str | np.ndarray' = "zero",

74 t_start: int = 0, t_step: int = 1):

75 super().__init__(source, chunk_size)

76 self._buf: Any = None

77 self._words: Any = None # uint32 view of the payload (1 word / event)

78 self._offset: int = 0

79 self._parser: Any = None

80 self._events: Any = None

81 self._triggers: Any = None

83 self._custom_ts: np.ndarray | None = None

84 if isinstance(timestamps, str):

85 modes = {"zero": AER_TS_ZERO, "sequential": AER_TS_SEQUENTIAL}

86 if timestamps not in modes:

87 raise ValueError(

88 f"timestamps must be 'zero', 'sequential' or an array, got {timestamps!r}"

89 )

90 self._ts_mode = modes[timestamps]

91 else:

92 ts = np.ascontiguousarray(timestamps, dtype=np.int64)

93 if ts.ndim != 1:

94 raise ValueError("custom timestamps must be a 1-D array")

95 self._custom_ts = ts

96 self._ts_mode = AER_TS_ZERO # parser fills 0; overwritten below

97 self._t_start = int(t_start)

98 self._t_step = int(t_step)

99 self._n_decoded = 0 # stream position, for indexing the custom array

100

101 def init(self) -> None:

102 """Initialize the AER reader.

103

104 Returns

105 -------

106 None

107

108 """

109 if self._is_initialized:

110 return

111

112 if self._source.mappable():

113 self._buf = self._source.buffer()

114 else:

115 self._buf = memoryview(self._source.read(-1))

116

117 n_events = len(self._buf) // 4 # AER has no header, 4 bytes / event

118 if n_events > 0:

119 self._words = np.frombuffer(self._buf, dtype=np.uint32, count=n_events)

120 else:

121 self._words = np.empty(0, dtype=np.uint32)

122

123 if self._custom_ts is not None and len(self._custom_ts) < n_events:

124 raise ValueError(

125 f"custom timestamps array has {len(self._custom_ts)} entries, "

126 f"but the stream contains {n_events} events"

127 )

128

129 self._offset = 0

130 self._parser = AerParser(self._ts_mode, self._t_start, self._t_step)

131 self._input_cls = AerInput

132 self._word_dtype = np.uint32

133 cap = int(self._chunk_size)

134 self._events = EventSoABuffers(cap)

135 self._triggers = TriggerSoABuffers(1) # AER has no triggers

136 self._is_initialized = True

137

138 def _apply_custom_ts(self, t_out: np.ndarray) -> None:

139 """Overwrite ``t_out`` (the decoded slice's t column, int64/uint64 view)

140 with the next ``len(t_out)`` user-provided timestamps.

141 """

142 assert self._custom_ts is not None

143 n = len(t_out)

144 t_out.view(np.int64)[:] = self._custom_ts[self._n_decoded:self._n_decoded + n]

145

146 def parse_step(self, events: EventSoABuffers, triggers: TriggerSoABuffers) -> int:

147 """Run the parser once, appending into ``events``; advance the offset.

148 See :meth:`EventDecoder_EVT.parse_step`.

149 """

150 if not self._is_initialized:

151 self.init()

152 if self._words is None or self._offset >= len(self._words):

153 self._eof = True

154 return 0

155 appended, self._offset = parse_step(

156 self._words, self._offset, AerInput, self._parser, events, triggers,

157 )

158 if appended and self._custom_ts is not None:

159 self._apply_custom_ts(events.t[events.size - appended:events.size])

160 self._n_decoded += appended

161 if self._offset >= len(self._words):

162 self._eof = True

163 return appended

164

165 def read_chunk(self, delta_t_hint: int | None = None,

166 n_events_hint: int | None = None) -> EventArray:

167 if not self._is_initialized:

168 self.init()

169

170 if self._words is None or self._offset >= len(self._words):

171 self._eof = True

172 return _EMPTY_EVENTS

173

174 ev, tr = self._events, self._triggers

175 ev.reset()

176 tr.reset()

177 appended = 0

178 while appended == 0 and self._offset < len(self._words):

179 appended = self.parse_step(ev, tr)

180

181 n = ev.size

182 if n == 0:

183 return _EMPTY_EVENTS

184 # Zero-copy view (valid until the next read_chunk); see EVT decoder.

185 return events_view(ev)

186

187 def read_all(self) -> EventArray:

188 """Decode the whole remaining payload into one buffer (no per-chunk copy)."""

189 if not self._is_initialized:

190 self.init()

191 if self._words is None or self._offset >= len(self._words):

192 self._eof = True

193 return _EMPTY_EVENTS

194 start = self._n_decoded

195 # Exactly one event per uint32 word.

196 out, self._offset = decode_all_soa(

197 self._words, self._offset, AerInput, self._parser,

198 est_events_per_word=1.0,

199 )

200 self._n_decoded = start + len(out)

201 if self._custom_ts is not None and len(out) > 0:

202 out.t[:] = self._custom_ts[start:start + len(out)]

203 self._eof = True

204 return out

205

206 def reset(self) -> None:

207 """Reset the AER reader to the beginning.

208

209 Returns

210 -------

211 None

212

213 """

214 self._offset = 0

215 self._n_decoded = 0

216 self._eof = False

217 if self._parser is not None:

218 self._parser.reset()

219

220 def tell(self) -> int:

221 """Get the current byte offset.

222

223 Returns

224 -------

225 int

226 Current byte offset.

227

228 """

229 return self._offset * 4

230

231 def close(self) -> None:

232 """Close the AER reader.

233

234 Returns

235 -------

236 None

237

238 """

239 self._words = None

240 self._buf = None

241

242

243class EventEncoder_AER(EventEncoder):

244 """Encode events into a raw AER stream. Since AER is designed for real-time

245 streaming, it has no header and no timestamps.

246 Timestamps are dropped and coordinates are masked to 9 bits (values >= 512

247 are truncated), per the AER encoding.

248

249 Parameters

250 ----------

251 writable

252 Destination stream to write to.

253 width, height : int

254 Frame geometry written into the header.

255 dt : datetime, optional

256 No effect, since AER has no timestamps.

257

258 References

259 ----------

260 [1] Prophesee AER format: https://docs.prophesee.ai/stable/data/encoding_formats/aer.html

261

262

263 """

264

265 def __init__(self, writable: io.BufferedWriter, width: int = 512, height: int = 512, dt: datetime | None = None):

266 super().__init__(writable, width, height, dt)

267

268 def init(self) -> None:

269 """Initialize the AER writer.

270

271 Returns

272 -------

273 None

274

275 """

276 self._is_initialized = True # AER has no header

277

278 def write(self, events: 'np.ndarray | EventArray', triggers: 'np.ndarray | TriggerArray | None' = None) -> int:

279 """Write events to the AER file.

280

281 Parameters

282 ----------

283 events : np.ndarray or EventArray

284 Array of events to write.

285

286 Returns

287 -------

288 int

289 Number of written events.

290

291 """

292 if not self._is_initialized:

293 self.init()

294

295 if isinstance(events, EventArray):

296 x, y, p = events.x, events.y, events.p

297 else:

298 x, y, p = events["x"], events["y"], events["p"]

299

300 out = (

301 (y.astype(np.uint32) & np.uint32(0x1FF))

302 | ((x.astype(np.uint32) & np.uint32(0x1FF)) << np.uint32(9))

303 | ((p.astype(np.uint32) & np.uint32(0x1)) << np.uint32(18))

304 )

305 out.astype(np.uint32).tofile(self._fd)

306 self._n_written_events += len(out)

307 return len(out)

Coverage for src/evutils/io/_aer.py: 83%

126 statements