Unified · Split
| @@ -32,6 +32,7 @@ · 31 unmodified lines | ||
| 32 | 32 | #include "dnn_interface.h" |
| 33 | 33 | #include "formats.h" |
| 34 | 34 | #include "internal.h" |
| 35 | +#include "libswscale/swscale.h" | |
| 35 | 36 | |
| 36 | 37 | typedef struct DnnProcessingContext { |
| 37 | 38 | const AVClass *class; |
| @@ -47,6 +48,9 @@ typedef struct DnnProcessingContext { · 9 unmodified lines | ||
| 47 | 48 | // input & output of the model at execution time |
| 48 | 49 | DNNData input; |
| 49 | 50 | DNNData output; |
| 51 | + | |
| 52 | + struct SwsContext *sws_gray8_to_grayf32; | |
| 53 | + struct SwsContext *sws_grayf32_to_gray8; | |
| 50 | 54 | } DnnProcessingContext; |
| 51 | 55 | |
| 52 | 56 | #define OFFSET(x) offsetof(DnnProcessingContext, x) |
| @@ -211,6 +215,45 @@ static int config_input(AVFilterLink *inlink) · 158 unmodified lines | ||
| 211 | 215 | return 0; |
| 212 | 216 | } |
| 213 | 217 | |
| 218 | +static int prepare_sws_context(AVFilterLink *outlink) | |
| 219 | +{ | |
| 220 | + AVFilterContext *context = outlink->src; | |
| 221 | + DnnProcessingContext *ctx = context->priv; | |
| 222 | + AVFilterLink *inlink = context->inputs[0]; | |
| 223 | + enum AVPixelFormat fmt = inlink->format; | |
| 224 | + DNNDataType input_dt = ctx->input.dt; | |
| 225 | + DNNDataType output_dt = ctx->output.dt; | |
| 226 | + | |
| 227 | + switch (fmt) { | |
| 228 | + case AV_PIX_FMT_RGB24: | |
| 229 | + case AV_PIX_FMT_BGR24: | |
| 230 | + if (input_dt == DNN_FLOAT) { | |
| 231 | + ctx->sws_gray8_to_grayf32 = sws_getContext(inlink->w * 3, | |
| 232 | + inlink->h, | |
| 233 | + AV_PIX_FMT_GRAY8, | |
| 234 | + inlink->w * 3, | |
| 235 | + inlink->h, | |
| 236 | + AV_PIX_FMT_GRAYF32, | |
| 237 | + 0, NULL, NULL, NULL); | |
| 238 | + } | |
| 239 | + if (output_dt == DNN_FLOAT) { | |
| 240 | + ctx->sws_grayf32_to_gray8 = sws_getContext(outlink->w * 3, | |
| 241 | + outlink->h, | |
| 242 | + AV_PIX_FMT_GRAYF32, | |
| 243 | + outlink->w * 3, | |
| 244 | + outlink->h, | |
| 245 | + AV_PIX_FMT_GRAY8, | |
| 246 | + 0, NULL, NULL, NULL); | |
| 247 | + } | |
| 248 | + return 0; | |
| 249 | + default: | |
| 250 | + //do nothing | |
| 251 | + break; | |
| 252 | + } | |
| 253 | + | |
| 254 | + return 0; | |
| 255 | +} | |
| 256 | + | |
| 214 | 257 | static int config_output(AVFilterLink *outlink) |
| 215 | 258 | { |
| 216 | 259 | AVFilterContext *context = outlink->src; |
| @@ -227,25 +270,23 @@ static int config_output(AVFilterLink *outlink) · 10 unmodified lines | ||
| 227 | 270 | outlink->w = ctx->output.width; |
| 228 | 271 | outlink->h = ctx->output.height; |
| 229 | 272 | |
| 273 | + prepare_sws_context(outlink); | |
| 274 | + | |
| 230 | 275 | return 0; |
| 231 | 276 | } |
| 232 | 277 | |
| 233 | −static int copy_from_frame_to_dnn(DNNData *dnn_input, const AVFrame *frame) | |
| 278 | +static int copy_from_frame_to_dnn(DnnProcessingContext *ctx, const AVFrame *frame) | |
| 234 | 279 | { |
| 235 | 280 | int bytewidth = av_image_get_linesize(frame->format, frame->width, 0); |
| 281 | + DNNData *dnn_input = &ctx->input; | |
| 236 | 282 | |
| 237 | 283 | switch (frame->format) { |
| 238 | 284 | case AV_PIX_FMT_RGB24: |
| 239 | 285 | case AV_PIX_FMT_BGR24: |
| 240 | 286 | if (dnn_input->dt == DNN_FLOAT) { |
| 241 | − float *dnn_input_data = dnn_input->data; | |
| 242 | − for (int i = 0; i < frame->height; i++) { | |
| 243 | − for(int j = 0; j < frame->width * 3; j++) { | |
| 244 | − int k = i * frame->linesize[0] + j; | |
| 245 | − int t = i * frame->width * 3 + j; | |
| 246 | − dnn_input_data[t] = frame->data[0][k] / 255.0f; | |
| 247 | − } | |
| 248 | − } | |
| 287 | + sws_scale(ctx->sws_gray8_to_grayf32, (const uint8_t **)frame->data, frame->linesize, | |
| 288 | + 0, frame->height, (uint8_t * const*)(&dnn_input->data), | |
| 289 | + (const int [4]){frame->linesize[0] * sizeof(float), 0, 0, 0}); | |
| 249 | 290 | } else { |
| 250 | 291 | av_assert0(dnn_input->dt == DNN_UINT8); |
| 251 | 292 | av_image_copy_plane(dnn_input->data, bytewidth, |
| @@ -266,22 +307,19 @@ static int copy_from_frame_to_dnn(DNNData *dnn_input, const AVFrame *frame) · 14 unmodified lines | ||
| 266 | 307 | return 0; |
| 267 | 308 | } |
| 268 | 309 | |
| 269 | −static int copy_from_dnn_to_frame(AVFrame *frame, const DNNData *dnn_output) | |
| 310 | +static int copy_from_dnn_to_frame(DnnProcessingContext *ctx, AVFrame *frame) | |
| 270 | 311 | { |
| 271 | 312 | int bytewidth = av_image_get_linesize(frame->format, frame->width, 0); |
| 313 | + DNNData *dnn_output = &ctx->output; | |
| 272 | 314 | |
| 273 | 315 | switch (frame->format) { |
| 274 | 316 | case AV_PIX_FMT_RGB24: |
| 275 | 317 | case AV_PIX_FMT_BGR24: |
| 276 | 318 | if (dnn_output->dt == DNN_FLOAT) { |
| 277 | − float *dnn_output_data = dnn_output->data; | |
| 278 | − for (int i = 0; i < frame->height; i++) { | |
| 279 | − for(int j = 0; j < frame->width * 3; j++) { | |
| 280 | − int k = i * frame->linesize[0] + j; | |
| 281 | − int t = i * frame->width * 3 + j; | |
| 282 | − frame->data[0][k] = av_clip_uintp2((int)(dnn_output_data[t] * 255.0f), 8); | |
| 283 | − } | |
| 284 | − } | |
| 319 | + sws_scale(ctx->sws_grayf32_to_gray8, (const uint8_t *[4]){(const uint8_t *)dnn_output->data, 0, 0, 0}, | |
| 320 | + (const int[4]){frame->linesize[0] * sizeof(float), 0, 0, 0}, | |
| 321 | + 0, frame->height, (uint8_t * const*)frame->data, frame->linesize); | |
| 322 | + | |
| 285 | 323 | } else { |
| 286 | 324 | av_assert0(dnn_output->dt == DNN_UINT8); |
| 287 | 325 | av_image_copy_plane(frame->data[0], frame->linesize[0], |
| @@ -318,7 +356,7 @@ static int filter_frame(AVFilterLink *inlink, AVFrame *in) · 30 unmodified lines | ||
| 318 | 356 | DNNReturnType dnn_result; |
| 319 | 357 | AVFrame *out; |
| 320 | 358 | |
| 321 | − copy_from_frame_to_dnn(&ctx->input, in); | |
| 359 | + copy_from_frame_to_dnn(ctx, in); | |
| 322 | 360 | |
| 323 | 361 | dnn_result = (ctx->dnn_module->execute_model)(ctx->model, &ctx->output, 1); |
| 324 | 362 | if (dnn_result != DNN_SUCCESS){ |
| @@ -334,7 +372,7 @@ static int filter_frame(AVFilterLink *inlink, AVFrame *in) · 9 unmodified lines | ||
| 334 | 372 | } |
| 335 | 373 | |
| 336 | 374 | av_frame_copy_props(out, in); |
| 337 | − copy_from_dnn_to_frame(out, &ctx->output); | |
| 375 | + copy_from_dnn_to_frame(ctx, out); | |
| 338 | 376 | av_frame_free(&in); |
| 339 | 377 | return ff_filter_frame(outlink, out); |
| 340 | 378 | } |
| @@ -343,6 +381,9 @@ static av_cold void uninit(AVFilterContext *ctx) · 2 unmodified lines | ||
| 343 | 381 | { |
| 344 | 382 | DnnProcessingContext *context = ctx->priv; |
| 345 | 383 | |
| 384 | + sws_freeContext(context->sws_gray8_to_grayf32); | |
| 385 | + sws_freeContext(context->sws_grayf32_to_gray8); | |
| 386 | + | |
| 346 | 387 | if (context->dnn_module) |
| 347 | 388 | (context->dnn_module->free_model)(&context->model); |
| 348 | 389 | |