Coverage for src/evutils/io/_dat.py: 78%

1"""Prophesee DAT (.dat) CD-event decoder/encoder. 

2 

3DAT is a fixed-record format: an ASCII ``% ...`` header, two bytes (event type 

4``0x0C`` = CD, event size ``0x08``), then 8-byte little-endian events 

5(``uint32`` timestamp + ``uint32`` data with ``x[0:13]``, ``y[14:27]``, 

6``p[28]``). Decoding is done by the native ``DAT_parse_chunk_soa`` (which also 

7tracks the 32-bit timestamp overflow); encoding is vectorised numpy. 

8""" 

9from __future__ import annotations 

10 

11import io 

12from datetime import datetime 

13 

14from typing import Any 

15 

16import numpy as np 

17 

18from ..types import EventArray, TriggerArray 

19from .common import EventDecoder, EventEncoder 

20from ._native_core import ( 

21 EVUTILS_PARSE_ERROR, 

22 EventSoABuffers, 

23 TriggerSoABuffers, 

24 decode_all_soa, 

25 events_view, 

26 triggers_view, 

27 parse_step, 

28) 

29from ._native_dat import ( 

30 DatInput, 

31 DatParser, 

32) 

33from ._source import ByteSource 

34 

35_EMPTY_EVENTS = EventArray.empty() 

36 

37DAT_EVENT_TYPE_CD = 0x0C 

38DAT_EVENT_SIZE = 0x08 

39 

40 

41class EventDecoder_Dat(EventDecoder): 

42 """Decode Prophesee DAT files into ``EventArray`` chunks. 

43  

44 Parameters 

45 ---------- 

46 source 

47 Byte source to read from. 

48 chunk_size 

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

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

51 

52 References 

53 ---------- 

54 [1] Prophesee DAT file format 

55 https://docs.prophesee.ai/stable/data/file_formats/dat.html 

56 

57 """ 

58 

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

60 super().__init__(source, chunk_size) 

61 self._header: dict[str, Any] = {} 

62 self._buf: Any = None 

63 self._payload_off: int = 0 

64 self._words: Any = None # uint32 view of the payload (2 words / event) 

65 self._offset: int = 0 # current uint32-word offset 

66 self._parser: Any = None 

67 self._events: Any = None 

68 self._triggers: Any = None 

69 

70 # ------------------------------------------------------------------ # 

71 def _parse_header(self, buf: Any) -> int: 

72 """Scan the leading ``%`` ASCII header; return the byte offset of the 

73 first non-header byte (the event-type byte). 

74 """ 

75 mv = memoryview(buf) 

76 n = len(mv) 

77 off = 0 

78 # Header lines start with "% "; the binary section (event-type byte, 

79 # then records) begins at the first line without that prefix. 

80 while off + 1 < n and mv[off] == 0x25 and mv[off + 1] == 0x20: 

81 window = bytes(mv[off:off + 8192]) 

82 rel = window.find(b"\n") 

83 if rel < 0: 

84 break 

85 self._consume_header_line(window[:rel]) 

86 off += rel + 1 

87 return off 

88 

89 def _consume_header_line(self, line: bytes) -> None: 

90 try: 

91 parts = line.decode("ascii", "ignore").strip().split() 

92 except ValueError: 

93 return 

94 # e.g. ["%", "Width", "1280"] 

95 if len(parts) >= 3: 

96 key = parts[1].lower() 

97 try: 

98 if key == "width": 

99 self._width = int(parts[2]) 

100 elif key == "height": 

101 self._height = int(parts[2]) 

102 except ValueError: 

103 pass 

104 

105 def init(self) -> None: 

106 """Initialize the DAT reader. 

107 

108 Returns 

109 ------- 

110 None 

111 

112 """ 

113 if self._is_initialized: 

114 return 

115 

116 if self._source.mappable(): 

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

118 else: 

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

120 

121 off = self._parse_header(self._buf) 

122 # Two-byte binary header: event type, event size. 

123 if off + 2 <= len(self._buf): 

124 event_size = self._buf[off + 1] 

125 off += 2 

126 if event_size not in (0, DAT_EVENT_SIZE): 

127 raise NotImplementedError( 

128 f"DAT event size {event_size} not supported (only 8-byte CD)" 

129 ) 

130 self._payload_off = off 

131 

132 n_events = (len(self._buf) - off) // 8 

133 if n_events > 0: 

134 self._words = np.frombuffer( 

135 self._buf, dtype=np.uint32, count=n_events * 2, offset=off 

136 ) 

137 else: 

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

139 

140 self._offset = 0 

141 self._parser = DatParser() 

142 self._input_cls = DatInput 

143 self._word_dtype = np.uint32 

144 cap = int(self._chunk_size) 

145 self._events = EventSoABuffers(cap) 

146 self._triggers = TriggerSoABuffers(1) # DAT CD files have no triggers 

147 self._is_initialized = True 

148 

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

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

151 See :meth:`EventDecoder_EVT.parse_step`. 

152 """ 

153 if not self._is_initialized: 

154 self.init() 

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

156 self._eof = True 

157 return 0 

158 appended, self._offset = parse_step( 

159 self._words, self._offset, DatInput, self._parser, events, triggers, 

160 ) 

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 | tuple[EventArray, TriggerArray]': 

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 if self.read_external_triggers: 

173 from ..types import TriggerArray 

174 return _EMPTY_EVENTS, TriggerArray.empty() 

175 return _EMPTY_EVENTS 

176 

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

178 ev.reset() 

179 tr.reset() 

180 appended = 0 

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

182 appended = self.parse_step(ev, tr) 

183 

184 n = ev.size 

185 if n == 0: 

186 if self.read_external_triggers: 

187 from ..types import TriggerArray 

188 return _EMPTY_EVENTS, TriggerArray.empty() 

189 return _EMPTY_EVENTS 

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

191 if self.read_external_triggers: 

192 return events_view(ev), triggers_view(tr) 

193 return events_view(ev) 

194 

195 def read_all(self) -> 'EventArray | tuple[EventArray, TriggerArray]': 

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

197 if not self._is_initialized: 

198 self.init() 

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

200 self._eof = True 

201 return _EMPTY_EVENTS 

202 # Exactly one event per two uint32 words. 

203 out, self._offset = decode_all_soa( 

204 self._words, self._offset, DatInput, self._parser, 

205 est_events_per_word=0.5, 

206 ) 

207 self._eof = True 

208 return out 

209 

210 def reset(self) -> None: 

211 """Reset the DAT reader to the beginning. 

212 

213 Returns 

214 ------- 

215 None 

216 

217 """ 

218 self._offset = 0 

219 self._eof = False 

220 if self._parser is not None: 

221 self._parser.reset() 

222 

223 def tell(self) -> int: 

224 """Get the current byte offset. 

225 

226 Returns 

227 ------- 

228 int 

229 Current byte offset. 

230 

231 """ 

232 return self._payload_off + self._offset * 4 

233 

234 def close(self) -> None: 

235 """Close the DAT reader. 

236 

237 Returns 

238 ------- 

239 None 

240 

241 """ 

242 self._words = None 

243 self._buf = None 

244 

245 

246class EventEncoder_Dat(EventEncoder): 

247 """Encode events into a Prophesee DAT (.dat) CD file. 

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 Recording timestamp (defaults to now). 

257 version : int 

258 DAT version (defaults to 2). 

259 

260 References 

261 ---------- 

262 [1] Prophesee DAT file format 

263 https://docs.prophesee.ai/stable/data/file_formats/dat.html 

264 

265 """ 

266 

267 def __init__(self, writable: io.BufferedWriter, width: int = 1280, height: int = 720, 

268 dt: datetime | None = None, version: int = 2): 

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

270 self._version = version 

271 

272 def init(self) -> None: 

273 """Initialize the DAT writer. 

274 

275 Returns 

276 ------- 

277 None 

278 

279 """ 

280 if self._is_initialized: 

281 return 

282 header = ( 

283 "% Data file containing CD events.\n" 

284 f"% Version {self._version}\n" 

285 f"% Date {self._dt.strftime('%Y-%m-%d %H:%M:%S')}\n" 

286 f"% Height {self._height}\n" 

287 f"% Width {self._width}\n" 

288 ) 

289 self._fd.write(header.encode("ascii")) 

290 self._fd.write(bytes([DAT_EVENT_TYPE_CD, DAT_EVENT_SIZE])) 

291 self._is_initialized = True 

292 

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

294 """Write events to the DAT file. 

295 

296 Parameters 

297 ---------- 

298 events : np.ndarray or EventArray 

299 Array of events to write. 

300 

301 Returns 

302 ------- 

303 int 

304 Number of written events. 

305 

306 """ 

307 if not self._is_initialized: 

308 self.init() 

309 

310 if isinstance(events, EventArray): 

311 t, x, y, p = events.t, events.x, events.y, events.p 

312 else: 

313 t, x, y, p = events["t"], events["x"], events["y"], events["p"] 

314 

315 n = len(t) 

316 out = np.empty(n * 2, dtype=np.uint32) 

317 out[0::2] = (np.asarray(t).astype(np.uint64) & np.uint64(0xFFFFFFFF)).astype(np.uint32) 

318 out[1::2] = ( 

319 (x.astype(np.uint32) & np.uint32(0x3FFF)) 

320 | ((y.astype(np.uint32) & np.uint32(0x3FFF)) << np.uint32(14)) 

321 | ((p.astype(np.uint32) & np.uint32(0x1)) << np.uint32(28)) 

322 ) 

323 out.tofile(self._fd) 

324 self._n_written_events += n 

325 return n