har0ke/FFmpeg
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

avfilter/vf_dnn_processing.c: add planar yuv format support

Browse Source 
Only the Y channel is handled by dnn, the UV channels are copied
without changes.

The command to use srcnn.pb (see vf_sr) looks like:
./ffmpeg -i 480p.jpg -vf format=yuv420p,scale=w=iw*2:h=ih*2,dnn_processing=dnn_backend=tensorflow:model=srcnn.pb:input=x:output=y -y srcnn.jpg

Signed-off-by: Guo, Yejun <yejun.guo@intel.com>
Reviewed-by: Pedro Arthur <bygrandao@gmail.com>
This commit is contained in:
Guo, Yejun 2020-02-21 14:20:48 +08:00
parent d86a8c056b
commit bd50453894
2 changed files with 81 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
doc/filters.texi
View File

@ -9196,6 +9196,8 @@ Set the output name of the dnn network.
@end table @end table
@subsection Examples
@itemize @itemize
@item @item
Halve the red channle of the frame with format rgb24: Halve the red channle of the frame with format rgb24:
@ -9209,6 +9211,12 @@ Halve the pixel value of the frame with format gray32f:
ffmpeg -i input.jpg -vf format=grayf32,dnn_processing=model=halve_gray_float.model:input=dnn_in:output=dnn_out:dnn_backend=native -y out.native.png ffmpeg -i input.jpg -vf format=grayf32,dnn_processing=model=halve_gray_float.model:input=dnn_in:output=dnn_out:dnn_backend=native -y out.native.png
@end example @end example
@item
Handle the Y channel with srcnn.pb (see @ref{sr} filter) for frame with yuv420p (planar YUV formats supported):
@example
./ffmpeg -i 480p.jpg -vf format=yuv420p,scale=w=iw*2:h=ih*2,dnn_processing=dnn_backend=tensorflow:model=srcnn.pb:input=x:output=y -y srcnn.jpg
@end example
@end itemize @end itemize
@section drawbox @section drawbox
@ -17311,6 +17319,7 @@ Set quality level. The value @code{max} can be used to set the maximum level,
currently @code{6}. currently @code{6}.
@end table @end table
@anchor{sr}
@section sr @section sr
Scale the input by applying one of the super-resolution methods based on Scale the input by applying one of the super-resolution methods based on

72
libavfilter/vf_dnn_processing.c
View File

@ -110,6 +110,8 @@ static int query_formats(AVFilterContext *context)
static const enum AVPixelFormat pix_fmts[] = { static const enum AVPixelFormat pix_fmts[] = {
AV_PIX_FMT_RGB24, AV_PIX_FMT_BGR24, AV_PIX_FMT_RGB24, AV_PIX_FMT_BGR24,
AV_PIX_FMT_GRAY8, AV_PIX_FMT_GRAYF32, AV_PIX_FMT_GRAY8, AV_PIX_FMT_GRAYF32,
AV_PIX_FMT_YUV420P, AV_PIX_FMT_YUV422P,
AV_PIX_FMT_YUV444P, AV_PIX_FMT_YUV410P, AV_PIX_FMT_YUV411P,
AV_PIX_FMT_NONE AV_PIX_FMT_NONE
}; };
AVFilterFormats *fmts_list = ff_make_format_list(pix_fmts); AVFilterFormats *fmts_list = ff_make_format_list(pix_fmts);
@ -163,6 +165,11 @@ static int check_modelinput_inlink(const DNNData *model_input, const AVFilterLin
} }
return 0; return 0;
case AV_PIX_FMT_GRAYF32: case AV_PIX_FMT_GRAYF32:
case AV_PIX_FMT_YUV420P:
case AV_PIX_FMT_YUV422P:
case AV_PIX_FMT_YUV444P:
case AV_PIX_FMT_YUV410P:
case AV_PIX_FMT_YUV411P:
if (model_input->channels != 1) { if (model_input->channels != 1) {
LOG_FORMAT_CHANNEL_MISMATCH(); LOG_FORMAT_CHANNEL_MISMATCH();
return AVERROR(EIO); return AVERROR(EIO);
@ -246,6 +253,28 @@ static int prepare_sws_context(AVFilterLink *outlink)
0, NULL, NULL, NULL); 0, NULL, NULL, NULL);
} }
return 0; return 0;
case AV_PIX_FMT_YUV420P:
case AV_PIX_FMT_YUV422P:
case AV_PIX_FMT_YUV444P:
case AV_PIX_FMT_YUV410P:
case AV_PIX_FMT_YUV411P:
av_assert0(input_dt == DNN_FLOAT);
av_assert0(output_dt == DNN_FLOAT);
ctx->sws_gray8_to_grayf32 = sws_getContext(inlink->w,
inlink->h,
AV_PIX_FMT_GRAY8,
inlink->w,
inlink->h,
AV_PIX_FMT_GRAYF32,
0, NULL, NULL, NULL);
ctx->sws_grayf32_to_gray8 = sws_getContext(outlink->w,
outlink->h,
AV_PIX_FMT_GRAYF32,
outlink->w,
outlink->h,
AV_PIX_FMT_GRAY8,
0, NULL, NULL, NULL);
return 0;
default: default:
//do nothing //do nothing
break; break;
@ -300,6 +329,15 @@ static int copy_from_frame_to_dnn(DnnProcessingContext *ctx, const AVFrame *fram
frame->data[0], frame->linesize[0], frame->data[0], frame->linesize[0],
bytewidth, frame->height); bytewidth, frame->height);
return 0; return 0;
case AV_PIX_FMT_YUV420P:
case AV_PIX_FMT_YUV422P:
case AV_PIX_FMT_YUV444P:
case AV_PIX_FMT_YUV410P:
case AV_PIX_FMT_YUV411P:
sws_scale(ctx->sws_gray8_to_grayf32, (const uint8_t **)frame->data, frame->linesize,
0, frame->height, (uint8_t * const*)(&dnn_input->data),
(const int [4]){frame->width * sizeof(float), 0, 0, 0});
return 0;
default: default:
return AVERROR(EIO); return AVERROR(EIO);
} }
@ -341,6 +379,15 @@ static int copy_from_dnn_to_frame(DnnProcessingContext *ctx, AVFrame *frame)
dnn_output->data, bytewidth, dnn_output->data, bytewidth,
bytewidth, frame->height); bytewidth, frame->height);
return 0; return 0;
case AV_PIX_FMT_YUV420P:
case AV_PIX_FMT_YUV422P:
case AV_PIX_FMT_YUV444P:
case AV_PIX_FMT_YUV410P:
case AV_PIX_FMT_YUV411P:
sws_scale(ctx->sws_grayf32_to_gray8, (const uint8_t *[4]){(const uint8_t *)dnn_output->data, 0, 0, 0},
(const int[4]){frame->width * sizeof(float), 0, 0, 0},
0, frame->height, (uint8_t * const*)frame->data, frame->linesize);
return 0;
default: default:
return AVERROR(EIO); return AVERROR(EIO);
} }
@ -348,6 +395,27 @@ static int copy_from_dnn_to_frame(DnnProcessingContext *ctx, AVFrame *frame)
return 0; return 0;
} }
static av_always_inline int isPlanarYUV(enum AVPixelFormat pix_fmt)
{
const AVPixFmtDescriptor *desc = av_pix_fmt_desc_get(pix_fmt);
av_assert0(desc);
return !(desc->flags & AV_PIX_FMT_FLAG_RGB) && desc->nb_components == 3;
}
static int copy_uv_planes(DnnProcessingContext *ctx, AVFrame *out, const AVFrame *in)
{
const AVPixFmtDescriptor *desc = av_pix_fmt_desc_get(in->format);
int uv_height = AV_CEIL_RSHIFT(in->height, desc->log2_chroma_h);
for (int i = 1; i < 3; ++i) {
int bytewidth = av_image_get_linesize(in->format, in->width, i);
av_image_copy_plane(out->data[i], out->linesize[i],
in->data[i], in->linesize[i],
bytewidth, uv_height);
}
return 0;
}
static int filter_frame(AVFilterLink *inlink, AVFrame *in) static int filter_frame(AVFilterLink *inlink, AVFrame *in)
{ {
AVFilterContext *context = inlink->dst; AVFilterContext *context = inlink->dst;
@ -373,6 +441,10 @@ static int filter_frame(AVFilterLink *inlink, AVFrame *in)
av_frame_copy_props(out, in); av_frame_copy_props(out, in);
copy_from_dnn_to_frame(ctx, out); copy_from_dnn_to_frame(ctx, out);
if (isPlanarYUV(in->format))
copy_uv_planes(ctx, out, in);
av_frame_free(&in); av_frame_free(&in);
return ff_filter_frame(outlink, out); return ff_filter_frame(outlink, out);
} }