framereader: refactor seeking (#22926)

* refactor seeking

* ckeck packets.empty()

Co-authored-by: Willem Melching <willem.melching@gmail.com>

selfdrive/ui/replay/framereader.cc

@@ -38,8 +38,8 @@ FrameReader::FrameReader(bool local_cache, int chunk_size, int retries) : FileRe
 }
 
 FrameReader::~FrameReader() {
-  for (auto &f : frames_) {
-    av_packet_unref(&f.pkt);
+  for (AVPacket *pkt : packets) {
+    av_packet_free(&pkt);
   }
 
   if (decoder_ctx) avcodec_free_context(&decoder_ctx);
@@ -106,18 +106,20 @@ bool FrameReader::load(const std::string &url, bool no_cuda, std::atomic<bool> *
   ret = avcodec_open2(decoder_ctx, decoder, NULL);
   if (ret < 0) return false;
 
-  frames_.reserve(60 * 20);  // 20fps, one minute
+  packets.reserve(60 * 20);  // 20fps, one minute
   while (!(abort && *abort)) {
-    Frame &frame = frames_.emplace_back();
-    ret = av_read_frame(input_ctx, &frame.pkt);
+    AVPacket *pkt = av_packet_alloc();
+    ret = av_read_frame(input_ctx, pkt);
     if (ret < 0) {
-      frames_.pop_back();
+      av_packet_free(&pkt);
       valid_ = (ret == AVERROR_EOF);
       break;
     }
+    packets.push_back(pkt);
     // some stream seems to contian no keyframes
-    key_frames_count_ += frame.pkt.flags & AV_PKT_FLAG_KEY;
+    key_frames_count_ += pkt->flags & AV_PKT_FLAG_KEY;
   }
+  valid_ = valid_ && !packets.empty();
   return valid_;
 }
 
@@ -151,35 +153,29 @@ bool FrameReader::initHardwareDecoder(AVHWDeviceType hw_device_type) {
 
 bool FrameReader::get(int idx, uint8_t *rgb, uint8_t *yuv) {
   assert(rgb || yuv);
-  if (!valid_ || idx < 0 || idx >= frames_.size()) {
+  if (!valid_ || idx < 0 || idx >= packets.size()) {
     return false;
   }
   return decode(idx, rgb, yuv);
 }
 
 bool FrameReader::decode(int idx, uint8_t *rgb, uint8_t *yuv) {
-  auto get_keyframe = [=](int idx) {
-    for (int i = idx; i >= 0 && key_frames_count_ > 1; --i) {
-      if (frames_[i].pkt.flags & AV_PKT_FLAG_KEY) return i;
-    }
-    return idx;
-  };
-
   int from_idx = idx;
-  if (idx > 0 && !frames_[idx].decoded && !frames_[idx - 1].decoded) {
-    // find the previous keyframe
-    from_idx = get_keyframe(idx);
+  if (idx != prev_idx + 1 && key_frames_count_ > 1) {
+    // seeking to the nearest key frame
+    for (int i = idx; i >= 0; --i) {
+      if (packets[i]->flags & AV_PKT_FLAG_KEY) {
+        from_idx = i;
+        break;
+      }
+    }
   }
+  prev_idx = idx;
 
   for (int i = from_idx; i <= idx; ++i) {
-    Frame &frame = frames_[i];
-    if ((!frame.decoded || i == idx) && !frame.failed) {
-      AVFrame *f = decodeFrame(&frame.pkt);
-      frame.decoded = f != nullptr;
-      frame.failed = !frame.decoded;
-      if (frame.decoded && i == idx) {
-        return copyBuffers(f, rgb, yuv);
-      }
+    AVFrame *f = decodeFrame(packets[i]);
+    if (f && i == idx) {
+      return copyBuffers(f, rgb, yuv);
     }
   }
   return false;

selfdrive/ui/replay/framereader.h

@@ -23,7 +23,7 @@ public:
   bool get(int idx, uint8_t *rgb, uint8_t *yuv);
   int getRGBSize() const { return width * height * 3; }
   int getYUVSize() const { return width * height * 3 / 2; }
-  size_t getFrameCount() const { return frames_.size(); }
+  size_t getFrameCount() const { return packets.size(); }
   bool valid() const { return valid_; }
 
   int width = 0, height = 0;
@@ -34,12 +34,7 @@ private:
   AVFrame * decodeFrame(AVPacket *pkt);
   bool copyBuffers(AVFrame *f, uint8_t *rgb, uint8_t *yuv);
 
-  struct Frame {
-    AVPacket pkt = {};
-    int decoded = false;
-    bool failed = false;
-  };
-  std::vector<Frame> frames_;
+  std::vector<AVPacket*> packets;
   std::unique_ptr<AVFrame, AVFrameDeleter>av_frame_, hw_frame;
   AVFormatContext *input_ctx = nullptr;
   AVCodecContext *decoder_ctx = nullptr;
@@ -50,5 +45,6 @@ private:
   AVPixelFormat hw_pix_fmt = AV_PIX_FMT_NONE;
   AVBufferRef *hw_device_ctx = nullptr;
   std::vector<uint8_t> nv12toyuv_buffer;
+  int prev_idx = -1;
   inline static std::atomic<bool> has_cuda_device = true;
 };