From 8c3b475812a0ec43f0ffd69eb82224e1a618fd34 Mon Sep 17 00:00:00 2001
From: Ricardo Constantino <wiiaboo@gmail.com>
Date: Wed, 9 Nov 2016 05:38:55 +0000
Subject: [PATCH] obs-ffmpeg: Fix assumption about plane height with i444

Affected both FFmpeg and nvenc encoders.
Fixes issue 667.
---
 plugins/obs-ffmpeg/obs-ffmpeg-nvenc.c  | 9 ++++++---
 plugins/obs-ffmpeg/obs-ffmpeg-output.c | 9 ++++++---
 2 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/plugins/obs-ffmpeg/obs-ffmpeg-nvenc.c b/plugins/obs-ffmpeg/obs-ffmpeg-nvenc.c
index 270383ff3..e63fe77b5 100644
--- a/plugins/obs-ffmpeg/obs-ffmpeg-nvenc.c
+++ b/plugins/obs-ffmpeg/obs-ffmpeg-nvenc.c
@@ -23,6 +23,7 @@
 #include <obs-avc.h>
 
 #include <libavutil/opt.h>
+#include <libavutil/pixdesc.h>
 #include <libavformat/avformat.h>
 
 #include "obs-ffmpeg-formats.h"
@@ -298,8 +299,10 @@ fail:
 }
 
 static inline void copy_data(AVPicture *pic, const struct encoder_frame *frame,
-		int height)
+		int height, enum AVPixelFormat format)
 {
+	int h_chroma_shift, v_chroma_shift;
+	av_pix_fmt_get_chroma_sub_sample(format, &h_chroma_shift, &v_chroma_shift);
 	for (int plane = 0; plane < MAX_AV_PLANES; plane++) {
 		if (!frame->data[plane])
 			continue;
@@ -308,7 +311,7 @@ static inline void copy_data(AVPicture *pic, const struct encoder_frame *frame,
 		int pic_rowsize   = pic->linesize[plane];
 		int bytes = frame_rowsize < pic_rowsize ?
 			frame_rowsize : pic_rowsize;
-		int plane_height = plane == 0 ? height : height/2;
+		int plane_height = height >> (plane ? v_chroma_shift : 0);
 
 		for (int y = 0; y < plane_height; y++) {
 			int pos_frame = y * frame_rowsize;
@@ -331,7 +334,7 @@ static bool nvenc_encode(void *data, struct encoder_frame *frame,
 
 	av_init_packet(&av_pkt);
 
-	copy_data(&enc->dst_picture, frame, enc->height);
+	copy_data(&enc->dst_picture, frame, enc->height, enc->context->pix_fmt);
 
 	enc->vframe->pts = frame->pts;
 	ret = avcodec_encode_video2(enc->context, &av_pkt, enc->vframe,
diff --git a/plugins/obs-ffmpeg/obs-ffmpeg-output.c b/plugins/obs-ffmpeg/obs-ffmpeg-output.c
index f08bd7c80..e23adf120 100644
--- a/plugins/obs-ffmpeg/obs-ffmpeg-output.c
+++ b/plugins/obs-ffmpeg/obs-ffmpeg-output.c
@@ -23,6 +23,7 @@
 #include <util/platform.h>
 
 #include <libavutil/opt.h>
+#include <libavutil/pixdesc.h>
 #include <libavformat/avformat.h>
 #include <libswscale/swscale.h>
 
@@ -620,8 +621,10 @@ static void ffmpeg_output_destroy(void *data)
 }
 
 static inline void copy_data(AVPicture *pic, const struct video_data *frame,
-		int height)
+		int height, enum AVPixelFormat format)
 {
+	int h_chroma_shift, v_chroma_shift;
+	av_pix_fmt_get_chroma_sub_sample(format, &h_chroma_shift, &v_chroma_shift);
 	for (int plane = 0; plane < MAX_AV_PLANES; plane++) {
 		if (!frame->data[plane])
 			continue;
@@ -630,7 +633,7 @@ static inline void copy_data(AVPicture *pic, const struct video_data *frame,
 		int pic_rowsize   = pic->linesize[plane];
 		int bytes = frame_rowsize < pic_rowsize ?
 			frame_rowsize : pic_rowsize;
-		int plane_height = plane == 0 ? height : height/2;
+		int plane_height = height >> (plane ? v_chroma_shift : 0);
 
 		for (int y = 0; y < plane_height; y++) {
 			int pos_frame = y * frame_rowsize;
@@ -669,7 +672,7 @@ static void receive_video(void *param, struct video_data *frame)
 				0, data->config.height, data->dst_picture.data,
 				data->dst_picture.linesize);
 	else
-		copy_data(&data->dst_picture, frame, context->height);
+		copy_data(&data->dst_picture, frame, context->height, context->pix_fmt);
 
 	if (data->output->flags & AVFMT_RAWPICTURE) {
 		packet.flags        |= AV_PKT_FLAG_KEY;