git.sesse.net Git - ffmpeg/blob - libavfilter/dnn/dnn_io_proc.c

   1 /*
   2  * Copyright (c) 2020
   3  *
   4  * This file is part of FFmpeg.
   5  *
   6  * FFmpeg is free software; you can redistribute it and/or
   7  * modify it under the terms of the GNU Lesser General Public
   8  * License as published by the Free Software Foundation; either
   9  * version 2.1 of the License, or (at your option) any later version.
  10  *
  11  * FFmpeg is distributed in the hope that it will be useful,
  12  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  13  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  14  * Lesser General Public License for more details.
  15  *
  16  * You should have received a copy of the GNU Lesser General Public
  17  * License along with FFmpeg; if not, write to the Free Software
  18  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
  19  */
  20
  21 #include "dnn_io_proc.h"
  22 #include "libavutil/imgutils.h"
  23 #include "libswscale/swscale.h"
  24 #include "libavutil/avassert.h"
  25
  26 DNNReturnType ff_proc_from_dnn_to_frame(AVFrame *frame, DNNData *output, void *log_ctx)
  27 {
  28     struct SwsContext *sws_ctx;
  29     int bytewidth = av_image_get_linesize(frame->format, frame->width, 0);
  30     if (output->dt != DNN_FLOAT) {
  31         avpriv_report_missing_feature(log_ctx, "data type rather than DNN_FLOAT");
  32         return DNN_ERROR;
  33     }
  34
  35     switch (frame->format) {
  36     case AV_PIX_FMT_RGB24:
  37     case AV_PIX_FMT_BGR24:
  38         sws_ctx = sws_getContext(frame->width * 3,
  39                                  frame->height,
  40                                  AV_PIX_FMT_GRAYF32,
  41                                  frame->width * 3,
  42                                  frame->height,
  43                                  AV_PIX_FMT_GRAY8,
  44                                  0, NULL, NULL, NULL);
  45         if (!sws_ctx) {
  46             av_log(log_ctx, AV_LOG_ERROR, "Impossible to create scale context for the conversion "
  47                 "fmt:%s s:%dx%d -> fmt:%s s:%dx%d\n",
  48                 av_get_pix_fmt_name(AV_PIX_FMT_GRAYF32), frame->width * 3, frame->height,
  49                 av_get_pix_fmt_name(AV_PIX_FMT_GRAY8),   frame->width * 3, frame->height);
  50             return DNN_ERROR;
  51         }
  52         sws_scale(sws_ctx, (const uint8_t *[4]){(const uint8_t *)output->data, 0, 0, 0},
  53                            (const int[4]){frame->width * 3 * sizeof(float), 0, 0, 0}, 0, frame->height,
  54                            (uint8_t * const*)frame->data, frame->linesize);
  55         sws_freeContext(sws_ctx);
  56         return DNN_SUCCESS;
  57     case AV_PIX_FMT_GRAYF32:
  58         av_image_copy_plane(frame->data[0], frame->linesize[0],
  59                             output->data, bytewidth,
  60                             bytewidth, frame->height);
  61         return DNN_SUCCESS;
  62     case AV_PIX_FMT_YUV420P:
  63     case AV_PIX_FMT_YUV422P:
  64     case AV_PIX_FMT_YUV444P:
  65     case AV_PIX_FMT_YUV410P:
  66     case AV_PIX_FMT_YUV411P:
  67     case AV_PIX_FMT_GRAY8:
  68     case AV_PIX_FMT_NV12:
  69         sws_ctx = sws_getContext(frame->width,
  70                                  frame->height,
  71                                  AV_PIX_FMT_GRAYF32,
  72                                  frame->width,
  73                                  frame->height,
  74                                  AV_PIX_FMT_GRAY8,
  75                                  0, NULL, NULL, NULL);
  76         if (!sws_ctx) {
  77             av_log(log_ctx, AV_LOG_ERROR, "Impossible to create scale context for the conversion "
  78                 "fmt:%s s:%dx%d -> fmt:%s s:%dx%d\n",
  79                 av_get_pix_fmt_name(AV_PIX_FMT_GRAYF32), frame->width, frame->height,
  80                 av_get_pix_fmt_name(AV_PIX_FMT_GRAY8),   frame->width, frame->height);
  81             return DNN_ERROR;
  82         }
  83         sws_scale(sws_ctx, (const uint8_t *[4]){(const uint8_t *)output->data, 0, 0, 0},
  84                            (const int[4]){frame->width * sizeof(float), 0, 0, 0}, 0, frame->height,
  85                            (uint8_t * const*)frame->data, frame->linesize);
  86         sws_freeContext(sws_ctx);
  87         return DNN_SUCCESS;
  88     default:
  89         avpriv_report_missing_feature(log_ctx, "%s", av_get_pix_fmt_name(frame->format));
  90         return DNN_ERROR;
  91     }
  92
  93     return DNN_SUCCESS;
  94 }
  95
  96 static DNNReturnType proc_from_frame_to_dnn_frameprocessing(AVFrame *frame, DNNData *input, void *log_ctx)
  97 {
  98     struct SwsContext *sws_ctx;
  99     int bytewidth = av_image_get_linesize(frame->format, frame->width, 0);
 100     if (input->dt != DNN_FLOAT) {
 101         avpriv_report_missing_feature(log_ctx, "data type rather than DNN_FLOAT");
 102         return DNN_ERROR;
 103     }
 104
 105     switch (frame->format) {
 106     case AV_PIX_FMT_RGB24:
 107     case AV_PIX_FMT_BGR24:
 108         sws_ctx = sws_getContext(frame->width * 3,
 109                                  frame->height,
 110                                  AV_PIX_FMT_GRAY8,
 111                                  frame->width * 3,
 112                                  frame->height,
 113                                  AV_PIX_FMT_GRAYF32,
 114                                  0, NULL, NULL, NULL);
 115         if (!sws_ctx) {
 116             av_log(log_ctx, AV_LOG_ERROR, "Impossible to create scale context for the conversion "
 117                 "fmt:%s s:%dx%d -> fmt:%s s:%dx%d\n",
 118                 av_get_pix_fmt_name(AV_PIX_FMT_GRAY8),  frame->width * 3, frame->height,
 119                 av_get_pix_fmt_name(AV_PIX_FMT_GRAYF32),frame->width * 3, frame->height);
 120             return DNN_ERROR;
 121         }
 122         sws_scale(sws_ctx, (const uint8_t **)frame->data,
 123                            frame->linesize, 0, frame->height,
 124                            (uint8_t * const*)(&input->data),
 125                            (const int [4]){frame->width * 3 * sizeof(float), 0, 0, 0});
 126         sws_freeContext(sws_ctx);
 127         break;
 128     case AV_PIX_FMT_GRAYF32:
 129         av_image_copy_plane(input->data, bytewidth,
 130                             frame->data[0], frame->linesize[0],
 131                             bytewidth, frame->height);
 132         break;
 133     case AV_PIX_FMT_YUV420P:
 134     case AV_PIX_FMT_YUV422P:
 135     case AV_PIX_FMT_YUV444P:
 136     case AV_PIX_FMT_YUV410P:
 137     case AV_PIX_FMT_YUV411P:
 138     case AV_PIX_FMT_GRAY8:
 139     case AV_PIX_FMT_NV12:
 140         sws_ctx = sws_getContext(frame->width,
 141                                  frame->height,
 142                                  AV_PIX_FMT_GRAY8,
 143                                  frame->width,
 144                                  frame->height,
 145                                  AV_PIX_FMT_GRAYF32,
 146                                  0, NULL, NULL, NULL);
 147         if (!sws_ctx) {
 148             av_log(log_ctx, AV_LOG_ERROR, "Impossible to create scale context for the conversion "
 149                 "fmt:%s s:%dx%d -> fmt:%s s:%dx%d\n",
 150                 av_get_pix_fmt_name(AV_PIX_FMT_GRAY8),  frame->width, frame->height,
 151                 av_get_pix_fmt_name(AV_PIX_FMT_GRAYF32),frame->width, frame->height);
 152             return DNN_ERROR;
 153         }
 154         sws_scale(sws_ctx, (const uint8_t **)frame->data,
 155                            frame->linesize, 0, frame->height,
 156                            (uint8_t * const*)(&input->data),
 157                            (const int [4]){frame->width * sizeof(float), 0, 0, 0});
 158         sws_freeContext(sws_ctx);
 159         break;
 160     default:
 161         avpriv_report_missing_feature(log_ctx, "%s", av_get_pix_fmt_name(frame->format));
 162         return DNN_ERROR;
 163     }
 164
 165     return DNN_SUCCESS;
 166 }
 167
 168 static enum AVPixelFormat get_pixel_format(DNNData *data)
 169 {
 170     if (data->dt == DNN_UINT8 && data->order == DCO_BGR) {
 171         return AV_PIX_FMT_BGR24;
 172     }
 173
 174     av_assert0(!"not supported yet.\n");
 175     return AV_PIX_FMT_BGR24;
 176 }
 177
 178 static DNNReturnType proc_from_frame_to_dnn_analytics(AVFrame *frame, DNNData *input, void *log_ctx)
 179 {
 180     struct SwsContext *sws_ctx;
 181     int linesizes[4];
 182     enum AVPixelFormat fmt = get_pixel_format(input);
 183     sws_ctx = sws_getContext(frame->width, frame->height, frame->format,
 184                              input->width, input->height, fmt,
 185                              SWS_FAST_BILINEAR, NULL, NULL, NULL);
 186     if (!sws_ctx) {
 187         av_log(log_ctx, AV_LOG_ERROR, "Impossible to create scale context for the conversion "
 188             "fmt:%s s:%dx%d -> fmt:%s s:%dx%d\n",
 189             av_get_pix_fmt_name(frame->format), frame->width, frame->height,
 190             av_get_pix_fmt_name(fmt), input->width, input->height);
 191         return DNN_ERROR;
 192     }
 193
 194     if (av_image_fill_linesizes(linesizes, fmt, input->width) < 0) {
 195         av_log(log_ctx, AV_LOG_ERROR, "unable to get linesizes with av_image_fill_linesizes");
 196         sws_freeContext(sws_ctx);
 197         return DNN_ERROR;
 198     }
 199
 200     sws_scale(sws_ctx, (const uint8_t *const *)frame->data, frame->linesize, 0, frame->height,
 201                        (uint8_t *const *)(&input->data), linesizes);
 202
 203     sws_freeContext(sws_ctx);
 204     return DNN_SUCCESS;
 205 }
 206
 207 DNNReturnType ff_proc_from_frame_to_dnn(AVFrame *frame, DNNData *input, DNNFunctionType func_type, void *log_ctx)
 208 {
 209     switch (func_type)
 210     {
 211     case DFT_PROCESS_FRAME:
 212         return proc_from_frame_to_dnn_frameprocessing(frame, input, log_ctx);
 213     case DFT_ANALYTICS_DETECT:
 214         return proc_from_frame_to_dnn_analytics(frame, input, log_ctx);
 215     default:
 216         avpriv_report_missing_feature(log_ctx, "model function type %d", func_type);
 217         return DNN_ERROR;
 218     }
 219 }