Unified · Split
| @@ -32,6 +32,7 @@ · 31 unmodified lines | |||
| 32 | #include "dnn_interface.h" | 32 | #include "dnn_interface.h" |
| 33 | #include "formats.h" | 33 | #include "formats.h" |
| 34 | #include "internal.h" | 34 | #include "internal.h" |
| 35 | #include "libswscale/swscale.h" | ||
| 35 | 36 | ||
| 36 | typedef struct DnnProcessingContext { | 37 | typedef struct DnnProcessingContext { |
| 37 | const AVClass *class; | 38 | const AVClass *class; |
| @@ -47,6 +48,9 @@ typedef struct DnnProcessingContext { · 9 unmodified lines | |||
| 47 | // input & output of the model at execution time | 48 | // input & output of the model at execution time |
| 48 | DNNData input; | 49 | DNNData input; |
| 49 | DNNData output; | 50 | DNNData output; |
| 51 | |||
| 52 | struct SwsContext *sws_gray8_to_grayf32; | ||
| 53 | struct SwsContext *sws_grayf32_to_gray8; | ||
| 50 | } DnnProcessingContext; | 54 | } DnnProcessingContext; |
| 51 | 55 | ||
| 52 | #define OFFSET(x) offsetof(DnnProcessingContext, x) | 56 | #define OFFSET(x) offsetof(DnnProcessingContext, x) |
| @@ -211,6 +215,45 @@ static int config_input(AVFilterLink *inlink) · 158 unmodified lines | |||
| 211 | return 0; | 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 | static int config_output(AVFilterLink *outlink) | 257 | static int config_output(AVFilterLink *outlink) |
| 215 | { | 258 | { |
| 216 | AVFilterContext *context = outlink->src; | 259 | AVFilterContext *context = outlink->src; |
| @@ -227,25 +270,23 @@ static int config_output(AVFilterLink *outlink) · 10 unmodified lines | |||
| 227 | outlink->w = ctx->output.width; | 270 | outlink->w = ctx->output.width; |
| 228 | outlink->h = ctx->output.height; | 271 | outlink->h = ctx->output.height; |
| 229 | 272 | ||
| 273 | prepare_sws_context(outlink); | ||
| 274 | |||
| 230 | return 0; | 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 | int bytewidth = av_image_get_linesize(frame->format, frame->width, 0); | 280 | int bytewidth = av_image_get_linesize(frame->format, frame->width, 0); |
| 281 | DNNData *dnn_input = &ctx->input; | ||
| 236 | 282 | ||
| 237 | switch (frame->format) { | 283 | switch (frame->format) { |
| 238 | case AV_PIX_FMT_RGB24: | 284 | case AV_PIX_FMT_RGB24: |
| 239 | case AV_PIX_FMT_BGR24: | 285 | case AV_PIX_FMT_BGR24: |
| 240 | if (dnn_input->dt == DNN_FLOAT) { | 286 | if (dnn_input->dt == DNN_FLOAT) { |
| 241 | float *dnn_input_data = dnn_input->data; | 287 | sws_scale(ctx->sws_gray8_to_grayf32, (const uint8_t **)frame->data, frame->linesize, |
| 242 | for (int i = 0; i < frame->height; i++) { | 288 | 0, frame->height, (uint8_t * const*)(&dnn_input->data), |
| 243 | for(int j = 0; j < frame->width * 3; j++) { | 289 | (const int [4]){frame->linesize[0] * sizeof(float), 0, 0, 0}); |
| 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 | } | ||
| 249 | } else { | 290 | } else { |
| 250 | av_assert0(dnn_input->dt == DNN_UINT8); | 291 | av_assert0(dnn_input->dt == DNN_UINT8); |
| 251 | av_image_copy_plane(dnn_input->data, bytewidth, | 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 | return 0; | 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 | int bytewidth = av_image_get_linesize(frame->format, frame->width, 0); | 312 | int bytewidth = av_image_get_linesize(frame->format, frame->width, 0); |
| 313 | DNNData *dnn_output = &ctx->output; | ||
| 272 | 314 | ||
| 273 | switch (frame->format) { | 315 | switch (frame->format) { |
| 274 | case AV_PIX_FMT_RGB24: | 316 | case AV_PIX_FMT_RGB24: |
| 275 | case AV_PIX_FMT_BGR24: | 317 | case AV_PIX_FMT_BGR24: |
| 276 | if (dnn_output->dt == DNN_FLOAT) { | 318 | if (dnn_output->dt == DNN_FLOAT) { |
| 277 | float *dnn_output_data = dnn_output->data; | 319 | sws_scale(ctx->sws_grayf32_to_gray8, (const uint8_t *[4]){(const uint8_t *)dnn_output->data, 0, 0, 0}, |
| 278 | for (int i = 0; i < frame->height; i++) { | 320 | (const int[4]){frame->linesize[0] * sizeof(float), 0, 0, 0}, |
| 279 | for(int j = 0; j < frame->width * 3; j++) { | 321 | 0, frame->height, (uint8_t * const*)frame->data, frame->linesize); |
| 280 | int k = i * frame->linesize[0] + j; | 322 | |
| 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 | } | ||
| 285 | } else { | 323 | } else { |
| 286 | av_assert0(dnn_output->dt == DNN_UINT8); | 324 | av_assert0(dnn_output->dt == DNN_UINT8); |
| 287 | av_image_copy_plane(frame->data[0], frame->linesize[0], | 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 | DNNReturnType dnn_result; | 356 | DNNReturnType dnn_result; |
| 319 | AVFrame *out; | 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 | dnn_result = (ctx->dnn_module->execute_model)(ctx->model, &ctx->output, 1); | 361 | dnn_result = (ctx->dnn_module->execute_model)(ctx->model, &ctx->output, 1); |
| 324 | if (dnn_result != DNN_SUCCESS){ | 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 | av_frame_copy_props(out, in); | 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 | av_frame_free(&in); | 376 | av_frame_free(&in); |
| 339 | return ff_filter_frame(outlink, out); | 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 | DnnProcessingContext *context = ctx->priv; | 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 | if (context->dnn_module) | 387 | if (context->dnn_module) |
| 347 | (context->dnn_module->free_model)(&context->model); | 388 | (context->dnn_module->free_model)(&context->model); |
| 348 | 389 | ||