From 7d10fa28148a846f20b9c44f8556683290d57a34 Mon Sep 17 00:00:00 2001
From: k4yt3x <i@k4yt3x.com>
Date: Tue, 21 Jan 2025 00:00:00 +0000
Subject: [PATCH] fix(encoder): re-add PTS calculation after processing for
 interpolators

Signed-off-by: k4yt3x <i@k4yt3x.com>
---
 include/libvideo2x/encoder.h | 2 +-
 src/encoder.cpp              | 5 ++++-
 src/libvideo2x.cpp           | 2 +-
 3 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/include/libvideo2x/encoder.h b/include/libvideo2x/encoder.h
index c485c1d..33fc086 100644
--- a/include/libvideo2x/encoder.h
+++ b/include/libvideo2x/encoder.h
@@ -63,7 +63,7 @@ class Encoder {
         int in_vstream_idx
     );
 
-    int write_frame(AVFrame* frame);
+    int write_frame(AVFrame* frame, int64_t frame_idx);
     int flush();
 
     AVCodecContext* get_encoder_context() const;
diff --git a/src/encoder.cpp b/src/encoder.cpp
index 48a90b3..dfb6e31 100644
--- a/src/encoder.cpp
+++ b/src/encoder.cpp
@@ -267,13 +267,16 @@ int Encoder::init(
 }
 
 [[gnu::target_clones("arch=x86-64-v4", "arch=x86-64-v3", "default")]]
-int Encoder::write_frame(AVFrame* frame) {
+int Encoder::write_frame(AVFrame* frame, int64_t frame_idx) {
     AVFrame* converted_frame = nullptr;
     int ret;
 
     // Let the encoder decide the frame type
     frame->pict_type = AV_PICTURE_TYPE_NONE;
 
+    // Calculate this frame's presentation timestamp (PTS)
+    frame->pts = av_rescale_q(frame_idx, av_inv_q(enc_ctx_->framerate), enc_ctx_->time_base);
+
     // Convert the frame to the encoder's pixel format if needed
     if (frame->format != enc_ctx_->pix_fmt) {
         converted_frame = conversions::convert_avframe_pix_fmt(frame, enc_ctx_->pix_fmt);
diff --git a/src/libvideo2x.cpp b/src/libvideo2x.cpp
index fbcc185..ee66983 100644
--- a/src/libvideo2x.cpp
+++ b/src/libvideo2x.cpp
@@ -314,7 +314,7 @@ int VideoProcessor::write_frame(AVFrame* frame, encoder::Encoder& encoder) {
     int ret = 0;
 
     if (!benchmark_) {
-        ret = encoder.write_frame(frame);
+        ret = encoder.write_frame(frame, frame_idx_);
         if (ret < 0) {
             av_strerror(ret, errbuf, sizeof(errbuf));
             logger()->critical("Error encoding/writing frame: {}", errbuf);