encoder, ffmpeg/tools: Fix and improve initialization behavior

Correctly sets all color settings for the context and frames, which should result in better playback in players that support these. Unfortunately it does not fix the bug that VLC and MPC-HC incorrectly assume that the ProRes encoded content is in Partial range, however most editing software does correctly detect it.
Xaymar · Sep 27, 2019 · 2ebf90f · 2ebf90f
1 parent a32f8dd
commit 2ebf90f
Show file tree

Hide file tree

Showing 3 changed files with 59 additions and 20 deletions.
diff --git a/source/encoder.cpp b/source/encoder.cpp
@@ -632,25 +632,23 @@ obsffmpeg::encoder::encoder(obs_data_t* settings, obs_encoder_t* encoder, bool i
 			}
 		}
 
-		_context->width                   = voi->width;
-		_context->height                  = voi->height;
-		_context->colorspace              = ffmpeg::tools::obs_videocolorspace_to_avcolorspace(voi->colorspace);
-		_context->color_range             = ffmpeg::tools::obs_videorangetype_to_avcolorrange(voi->range);
+		_context->width  = voi->width;
+		_context->height = voi->height;
+		ffmpeg::tools::setup_obs_color(voi->colorspace, voi->range, _context);
+
 		_context->pix_fmt                 = _pixfmt_target;
 		_context->field_order             = AV_FIELD_PROGRESSIVE;
-		_context->time_base.num           = voi->fps_den;
-		_context->time_base.den           = voi->fps_num;
 		_context->ticks_per_frame         = 1;
 		_context->sample_aspect_ratio.num = _context->sample_aspect_ratio.den = 1;
+		_context->framerate.num = _context->time_base.den = voi->fps_num;
+		_context->framerate.den = _context->time_base.num = voi->fps_den;
 
 		_swscale.set_source_size(_context->width, _context->height);
-		_swscale.set_source_color(_context->color_range, _context->colorspace);
-		_swscale.set_source_full_range(voi->range == VIDEO_RANGE_FULL);
+		_swscale.set_source_color(_context->color_range == AVCOL_RANGE_JPEG, _context->colorspace);
 		_swscale.set_source_format(_pixfmt_source);
 
 		_swscale.set_target_size(_context->width, _context->height);
-		_swscale.set_target_color(_context->color_range, _context->colorspace);
-		_swscale.set_target_full_range(voi->range == VIDEO_RANGE_FULL);
+		_swscale.set_target_color(_context->color_range == AVCOL_RANGE_JPEG, _context->colorspace);
 		_swscale.set_target_format(_pixfmt_target);
 
 		// Create Scaler
@@ -679,8 +677,9 @@ obsffmpeg::encoder::encoder(obs_data_t* settings, obs_encoder_t* encoder, bool i
 		          ffmpeg::tools::get_pixel_format_name(_swscale.get_target_format()),
 		          ffmpeg::tools::get_color_space_name(_swscale.get_target_colorspace()),
 		          _swscale.is_target_full_range() ? "Full" : "Partial");
-		PLOG_INFO("[%s]   Framerate: %ld/%ld (%f", id, _context->time_base.num, _context->time_base.den,
-		          _context->time_base.num / _context->time_base.den);
+		PLOG_INFO("[%s]   Framerate: %ld/%ld (%f FPS)", id, _context->time_base.den, _context->time_base.num,
+		          static_cast<double_t>(_context->time_base.den)
+		              / static_cast<double_t>(_context->time_base.num));
 	}
 
 	// Update settings
@@ -866,11 +865,13 @@ bool obsffmpeg::encoder::video_encode(encoder_frame* frame, encoder_packet* pack
 		ScopeProfiler profile("convert");
 #endif
 
-		vframe->height      = _context->height;
-		vframe->format      = _context->pix_fmt;
-		vframe->color_range = _context->color_range;
-		vframe->colorspace  = _context->colorspace;
-		vframe->pts         = frame->pts;
+		vframe->height          = _context->height;
+		vframe->format          = _context->pix_fmt;
+		vframe->color_range     = _context->color_range;
+		vframe->colorspace      = _context->colorspace;
+		vframe->color_primaries = _context->color_primaries;
+		vframe->color_trc       = _context->color_trc;
+		vframe->pts             = frame->pts;
 
 		if ((_swscale.is_source_full_range() == _swscale.is_target_full_range())
 		    && (_swscale.get_source_colorspace() == _swscale.get_target_colorspace())
@@ -1016,6 +1017,9 @@ int obsffmpeg::encoder::receive_packet(bool* received_packet, struct encoder_pac
 			} else if (_codec->id == AV_CODEC_ID_HEVC) {
 				obsffmpeg::codecs::hevc::extract_header_sei(_current_packet.data, _current_packet.size,
 				                                            _extra_data, _sei_data);
+			} else if (_context->extradata != nullptr) {
+				_extra_data.resize(_context->extradata_size);
+				std::memcpy(_extra_data.data(), _context->extradata, _context->extradata_size);
 			}
 			_have_first_frame = true;
 		}

diff --git a/source/ffmpeg/tools.cpp b/source/ffmpeg/tools.cpp
@@ -220,10 +220,10 @@ AVColorRange ffmpeg::tools::obs_videorangetype_to_avcolorrange(video_range_type
 {
 	switch (v) {
 	case VIDEO_RANGE_DEFAULT:
-	case VIDEO_RANGE_FULL:
-		return AVCOL_RANGE_JPEG;
 	case VIDEO_RANGE_PARTIAL:
 		return AVCOL_RANGE_MPEG;
+	case VIDEO_RANGE_FULL:
+		return AVCOL_RANGE_JPEG;
 	}
 	throw std::invalid_argument("unknown range");
 }
@@ -378,3 +378,36 @@ AVPixelFormat ffmpeg::tools::get_best_compatible_format(const AVPixelFormat* lis
 
 	return best;
 }
+
+void ffmpeg::tools::setup_obs_color(video_colorspace colorspace, video_range_type range, AVCodecContext* context)
+{
+	std::map<video_colorspace, std::tuple<AVColorSpace, AVColorPrimaries, AVColorTransferCharacteristic>>
+	    colorspaces = {
+	        {VIDEO_CS_DEFAULT, {AVCOL_SPC_BT470BG, AVCOL_PRI_BT470BG, AVCOL_TRC_SMPTE170M}},
+	        {VIDEO_CS_601, {AVCOL_SPC_BT470BG, AVCOL_PRI_BT470BG, AVCOL_TRC_SMPTE170M}},
+	        {VIDEO_CS_709, {AVCOL_SPC_BT709, AVCOL_PRI_BT709, AVCOL_TRC_BT709}},
+	    };
+	std::map<video_range_type, AVColorRange> colorranges = {
+	    {VIDEO_RANGE_DEFAULT, AVCOL_RANGE_MPEG},
+	    {VIDEO_RANGE_PARTIAL, AVCOL_RANGE_MPEG},
+	    {VIDEO_RANGE_FULL, AVCOL_RANGE_JPEG},
+	};
+
+	{
+		auto found = colorspaces.find(colorspace);
+		if (found != colorspaces.end()) {
+			context->colorspace      = std::get<AVColorSpace>(found->second);
+			context->color_primaries = std::get<AVColorPrimaries>(found->second);
+			context->color_trc       = std::get<AVColorTransferCharacteristic>(found->second);
+		}
+	}
+	{
+		auto found = colorranges.find(range);
+		if (found != colorranges.end()) {
+			context->color_range = found->second;
+		}
+	}
+
+	// Downscaling should result in downscaling, not pixelation
+	context->chroma_sample_location = AVCHROMA_LOC_CENTER; 
+}
diff --git a/source/ffmpeg/tools.hpp b/source/ffmpeg/tools.hpp
@@ -28,8 +28,8 @@
 #include <vector>
 
 extern "C" {
-#include <libavutil/pixfmt.h>
 #include <libavcodec/avcodec.h>
+#include <libavutil/pixfmt.h>
 }
 
 namespace ffmpeg {
@@ -57,6 +57,8 @@ namespace ffmpeg {
 		std::vector<AVPixelFormat> get_software_formats(const AVPixelFormat* list);
 
 		AVPixelFormat get_best_compatible_format(const AVPixelFormat* list, AVPixelFormat source);
+
+		void setup_obs_color(video_colorspace colorspace, video_range_type range, AVCodecContext* context);
 	} // namespace tools
 } // namespace ffmpeg