If we have waiting B-frames at the end of the encode, encode them as such.

[nageru] / h264encode.cpp

diff --git a/h264encode.cpp b/h264encode.cpp
index 641cf5039bb2db9214d15e9277120a43f3e38a61..d0da39d08c0925e79ab9c9cef418257c3b95bb2e 100644 (file)
--- a/h264encode.cpp
+++ b/h264encode.cpp
@@ -1610,8 +1610,8 @@ void H264Encoder::save_codeddata(storage_task task)
         vector<float> audio;
         {
              unique_lock<mutex> lock(frame_queue_mutex);
-             frame_queue_nonempty.wait(lock, [this]{ return copy_thread_should_quit || !pending_audio_frames.empty(); });
-             if (copy_thread_should_quit && pending_audio_frames.empty()) return;
+             frame_queue_nonempty.wait(lock, [this]{ return storage_thread_should_quit || !pending_audio_frames.empty(); });
+             if (storage_thread_should_quit && pending_audio_frames.empty()) return;
              auto it = pending_audio_frames.begin();
              if (it->first > task.pts) break;
              audio_pts = it->first;
@@ -1695,7 +1695,7 @@ void H264Encoder::storage_task_thread()
                        // wait until there's an encoded frame  
                        unique_lock<mutex> lock(storage_task_queue_mutex);
                        storage_task_queue_changed.wait(lock, [this]{ return storage_thread_should_quit || !storage_task_queue.empty(); });
-                       if (storage_thread_should_quit) return;
+                       if (storage_thread_should_quit && storage_task_queue.empty()) return;
                        current = move(storage_task_queue.front());
                        storage_task_queue.pop();
                }
@@ -1777,7 +1777,7 @@ H264Encoder::H264Encoder(QSurface *surface, int width, int height, HTTPD *httpd)
 
        storage_thread = thread(&H264Encoder::storage_task_thread, this);
 
-       copy_thread = thread([this]{
+       encode_thread = thread([this]{
                //SDL_GL_MakeCurrent(window, context);
                QOpenGLContext *context = create_context(this->surface);
                eglBindAPI(EGL_OPENGL_API);
@@ -1786,7 +1786,7 @@ H264Encoder::H264Encoder(QSurface *surface, int width, int height, HTTPD *httpd)
                                eglGetError());
                        exit(1);
                }
-               copy_thread_func();
+               encode_thread_func();
        });
 }
 
@@ -1794,13 +1794,14 @@ H264Encoder::~H264Encoder()
 {
        {
                unique_lock<mutex> lock(frame_queue_mutex);
-               copy_thread_should_quit = true;
+               encode_thread_should_quit = true;
                frame_queue_nonempty.notify_all();
        }
-       copy_thread.join();
+       encode_thread.join();
        {
                unique_lock<mutex> lock(storage_task_queue_mutex);
                storage_thread_should_quit = true;
+               frame_queue_nonempty.notify_all();
                storage_task_queue_changed.notify_all();
        }
        storage_thread.join();
@@ -1892,7 +1893,7 @@ void H264Encoder::end_frame(RefCountedGLsync fence, int64_t pts, const vector<Re
        frame_queue_nonempty.notify_all();
 }
 
-void H264Encoder::copy_thread_func()
+void H264Encoder::encode_thread_func()
 {
        int64_t last_dts = -1;
        int gop_start_display_frame_num = 0;
@@ -1911,9 +1912,15 @@ void H264Encoder::copy_thread_func()
                {
                        unique_lock<mutex> lock(frame_queue_mutex);
                        frame_queue_nonempty.wait(lock, [this, display_frame_num]{
-                               return copy_thread_should_quit || pending_video_frames.count(display_frame_num) != 0;
+                               return encode_thread_should_quit || pending_video_frames.count(display_frame_num) != 0;
                        });
-                       if (copy_thread_should_quit) {
+                       if (encode_thread_should_quit && pending_video_frames.count(display_frame_num) == 0) {
+                               // We have queued frames that were supposed to be B-frames,
+                               // but will be no P-frame to encode them against. Encode them all
+                               // as P-frames instead. Note that this happens under the mutex,
+                               // but nobody else uses it at this point, since we're shutting down,
+                               // so there's no contention.
+                               encode_remaining_frames_as_p(encoding_frame_num, gop_start_display_frame_num, last_dts);
                                return;
                        } else {
                                frame = move(pending_video_frames[display_frame_num]);
@@ -1935,6 +1942,23 @@ void H264Encoder::copy_thread_func()
        }
 }
 
+void H264Encoder::encode_remaining_frames_as_p(int encoding_frame_num, int gop_start_display_frame_num, int64_t last_dts)
+{
+       if (pending_video_frames.empty()) {
+               return;
+       }
+
+       for (auto &pending_frame : pending_video_frames) {
+               int display_frame_num = pending_frame.first;
+               assert(display_frame_num > 0);
+               PendingFrame frame = move(pending_frame.second);
+               int64_t dts = last_dts + (TIMEBASE / MAX_FPS);
+               printf("Finalizing encode: Encoding leftover frame %d as P-frame instead of B-frame.\n", display_frame_num);
+               encode_frame(frame, encoding_frame_num++, display_frame_num, gop_start_display_frame_num, FRAME_P, frame.pts, dts);
+               last_dts = dts;
+       }
+}
+
 void H264Encoder::encode_frame(H264Encoder::PendingFrame frame, int encoding_frame_num, int display_frame_num, int gop_start_display_frame_num,
                                int frame_type, int64_t pts, int64_t dts)
 {