Framereader: minor cleanup (#35577)

* No wrapping

* unused test

* another list

* mypy

* cleaner

* Revert "cleaner"

This reverts commit ccc1446b9d649d64b20175e22a66e135c44b21e5.

* mypy

selfdrive/test/process_replay/model_replay.py

@@ -190,7 +190,7 @@ def model_replay(lr, frs):
   print("----------------- Model Timing -----------------")
   print("------------------------------------------------")
   print(tabulate(rows, header, tablefmt="simple_grid", stralign="center", numalign="center", floatfmt=".4f"))
-  assert timings_ok
+  assert timings_ok or PC
 
   return msgs
 

selfdrive/test/process_replay/process_replay.py

@@ -297,7 +297,7 @@ class ProcessContainer:
             camera_state = getattr(m, m.which())
             camera_meta = meta_from_camera_state(m.which())
             assert frs is not None
-            img = frs[m.which()].get(camera_state.frameId)[0]
+            img = frs[m.which()].get(camera_state.frameId)
             self.vipc_server.send(camera_meta.stream, img.flatten().tobytes(),
                                   camera_state.frameId, camera_state.timestampSof, camera_state.timestampEof)
         self.msg_queue = []

selfdrive/ui/tests/test_ui/run.py

@@ -283,9 +283,9 @@ def create_screenshots():
       driver_img = frames[2]
   else:
     with open(frames_cache, 'wb') as f:
-      road_img = FrameReader(route.camera_paths()[segnum], pix_fmt="nv12").get(0)[0]
-      wide_road_img = FrameReader(route.ecamera_paths()[segnum], pix_fmt="nv12").get(0)[0]
-      driver_img = FrameReader(route.dcamera_paths()[segnum], pix_fmt="nv12").get(0)[0]
+      road_img = FrameReader(route.camera_paths()[segnum], pix_fmt="nv12").get(0)
+      wide_road_img = FrameReader(route.ecamera_paths()[segnum], pix_fmt="nv12").get(0)
+      driver_img = FrameReader(route.dcamera_paths()[segnum], pix_fmt="nv12").get(0)
       pickle.dump([road_img, wide_road_img, driver_img], f)
 
   STREAMS.append((VisionStreamType.VISION_STREAM_ROAD, cam.fcam, road_img.flatten().tobytes()))

tools/lib/framereader.py

@@ -159,9 +159,9 @@ class FrameReader:
     self.it: Iterator[tuple[int, np.ndarray]] | None = None
     self.fidx = -1
 
-  def get(self, fidx:int) -> list[np.ndarray]:
+  def get(self, fidx:int):
     if fidx in self._cache:  # If frame is cached, return it
-      return [self._cache[fidx]]
+      return self._cache[fidx]
     read_start = self.decoder.get_gop_start(fidx)
     if not self.it or fidx < self.fidx or read_start != self.decoder.get_gop_start(self.fidx):  # If the frame is in a different GOP, reset the iterator
       self.it = self.decoder.get_iterator(read_start)
@@ -169,4 +169,4 @@ class FrameReader:
     while self.fidx < fidx:
       self.fidx, frame = next(self.it)
       self._cache[self.fidx] = frame
-    return [self._cache[fidx]]  # TODO: return just frame
+    return self._cache[fidx]

tools/lib/tests/test_readers.py

@@ -1,63 +0,0 @@
-import pytest
-import requests
-import tempfile
-
-from collections import defaultdict
-import numpy as np
-from openpilot.tools.lib.framereader import FrameReader
-from openpilot.tools.lib.logreader import LogReader
-
-
-class TestReaders:
-  @pytest.mark.skip("skip for bandwidth reasons")
-  def test_logreader(self):
-    def _check_data(lr):
-      hist = defaultdict(int)
-      for l in lr:
-        hist[l.which()] += 1
-
-      assert hist['carControl'] == 6000
-      assert hist['logMessage'] == 6857
-
-    with tempfile.NamedTemporaryFile(suffix=".bz2") as fp:
-      r = requests.get("https://github.com/commaai/comma2k19/blob/master/Example_1/b0c9d2329ad1606b%7C2018-08-02--08-34-47/40/raw_log.bz2?raw=true", timeout=10)
-      fp.write(r.content)
-      fp.flush()
-
-      lr_file = LogReader(fp.name)
-      _check_data(lr_file)
-
-    lr_url = LogReader("https://github.com/commaai/comma2k19/blob/master/Example_1/b0c9d2329ad1606b%7C2018-08-02--08-34-47/40/raw_log.bz2?raw=true")
-    _check_data(lr_url)
-
-  @pytest.mark.skip("skip for bandwidth reasons")
-  def test_framereader(self):
-    def _check_data(f):
-      assert f.frame_count == 1200
-      assert f.w == 1164
-      assert f.h == 874
-
-      frame_first_30 = f.get(0, 30)
-      assert len(frame_first_30) == 30
-
-      print(frame_first_30[15])
-
-      print("frame_0")
-      frame_0 = f.get(0, 1)
-      frame_15 = f.get(15, 1)
-
-      print(frame_15[0])
-
-      assert np.all(frame_first_30[0] == frame_0[0])
-      assert np.all(frame_first_30[15] == frame_15[0])
-
-    with tempfile.NamedTemporaryFile(suffix=".hevc") as fp:
-      r = requests.get("https://github.com/commaai/comma2k19/blob/master/Example_1/b0c9d2329ad1606b%7C2018-08-02--08-34-47/40/video.hevc?raw=true", timeout=10)
-      fp.write(r.content)
-      fp.flush()
-
-      fr_file = FrameReader(fp.name)
-      _check_data(fr_file)
-
-    fr_url = FrameReader("https://github.com/commaai/comma2k19/blob/master/Example_1/b0c9d2329ad1606b%7C2018-08-02--08-34-47/40/video.hevc?raw=true")
-    _check_data(fr_url)

tools/replay/unlog_ci_segment.py

@@ -48,7 +48,7 @@ def replay(route, segment, loop):
     if w == 'roadCameraState':
       try:
         img = fr.get(frame_idx[msg.roadCameraState.frameId])
-        img = img[0][:, :, ::-1]  # Convert RGB to BGR, which is what the camera outputs
+        img = img[:, ::-1]  # Convert RGB to BGR, which is what the camera outputs
         msg.roadCameraState.image = img.flatten().tobytes()
       except (KeyError, ValueError):
         pass

tools/scripts/fetch_image_from_route.py

@@ -37,7 +37,7 @@ fr = FrameReader(segments[segment])
 if frame >= fr.frame_count:
   raise Exception("frame {frame} not found, got {fr.frame_count} frames")
 
-im = Image.fromarray(fr.get(frame)[0])
+im = Image.fromarray(fr.get(frame))
 fn = f"uxxx_{route.replace('|', '_')}_{segment}_{frame}.png"
 im.save(fn)
 print(f"saved {fn}")