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
  25 DNNReturnType proc_from_dnn_to_frame(AVFrame *frame, DNNData *output, void *log_ctx)
  26 {
  27     struct SwsContext *sws_ctx;
  28     int bytewidth = av_image_get_linesize(frame->format, frame->width, 0);
  29     if (output->dt != DNN_FLOAT) {
  30         av_log(log_ctx, AV_LOG_ERROR, "do not support data type rather than DNN_FLOAT\n");
  31         return DNN_ERROR;
  32     }
  33
  34     switch (frame->format) {
  35     case AV_PIX_FMT_RGB24:
  36     case AV_PIX_FMT_BGR24:
  37         sws_ctx = sws_getContext(frame->width * 3,
  38                                  frame->height,
  39                                  AV_PIX_FMT_GRAYF32,
  40                                  frame->width * 3,
  41                                  frame->height,
  42                                  AV_PIX_FMT_GRAY8,
  43                                  0, NULL, NULL, NULL);
  44         if (!sws_ctx) {
  45             av_log(log_ctx, AV_LOG_ERROR, "Impossible to create scale context for the conversion "
  46                 "fmt:%s s:%dx%d -> fmt:%s s:%dx%d\n",
  47                 av_get_pix_fmt_name(AV_PIX_FMT_GRAYF32), frame->width * 3, frame->height,
  48                 av_get_pix_fmt_name(AV_PIX_FMT_GRAY8),   frame->width * 3, frame->height);
  49             return DNN_ERROR;
  50         }
  51         sws_scale(sws_ctx, (const uint8_t *[4]){(const uint8_t *)output->data, 0, 0, 0},
  52                            (const int[4]){frame->width * 3 * sizeof(float), 0, 0, 0}, 0, frame->height,
  53                            (uint8_t * const*)frame->data, frame->linesize);
  54         sws_freeContext(sws_ctx);
  55         return DNN_SUCCESS;
  56     case AV_PIX_FMT_GRAYF32:
  57         av_image_copy_plane(frame->data[0], frame->linesize[0],
  58                             output->data, bytewidth,
  59                             bytewidth, frame->height);
  60         return DNN_SUCCESS;
  61     case AV_PIX_FMT_YUV420P:
  62     case AV_PIX_FMT_YUV422P:
  63     case AV_PIX_FMT_YUV444P:
  64     case AV_PIX_FMT_YUV410P:
  65     case AV_PIX_FMT_YUV411P:
  66     case AV_PIX_FMT_GRAY8:
  67         sws_ctx = sws_getContext(frame->width,
  68                                  frame->height,
  69                                  AV_PIX_FMT_GRAYF32,
  70                                  frame->width,
  71                                  frame->height,
  72                                  AV_PIX_FMT_GRAY8,
  73                                  0, NULL, NULL, NULL);
  74         if (!sws_ctx) {
  75             av_log(log_ctx, AV_LOG_ERROR, "Impossible to create scale context for the conversion "
  76                 "fmt:%s s:%dx%d -> fmt:%s s:%dx%d\n",
  77                 av_get_pix_fmt_name(AV_PIX_FMT_GRAYF32), frame->width, frame->height,
  78                 av_get_pix_fmt_name(AV_PIX_FMT_GRAY8),   frame->width, frame->height);
  79             return DNN_ERROR;
  80         }
  81         sws_scale(sws_ctx, (const uint8_t *[4]){(const uint8_t *)output->data, 0, 0, 0},
  82                            (const int[4]){frame->width * sizeof(float), 0, 0, 0}, 0, frame->height,
  83                            (uint8_t * const*)frame->data, frame->linesize);
  84         sws_freeContext(sws_ctx);
  85         return DNN_SUCCESS;
  86     default:
  87         av_log(log_ctx, AV_LOG_ERROR, "do not support frame format %s\n",
  88                av_get_pix_fmt_name(frame->format));
  89         return DNN_ERROR;
  90     }
  91
  92     return DNN_SUCCESS;
  93 }
  94
  95 DNNReturnType proc_from_frame_to_dnn(AVFrame *frame, DNNData *input, void *log_ctx)
  96 {
  97     struct SwsContext *sws_ctx;
  98     int bytewidth = av_image_get_linesize(frame->format, frame->width, 0);
  99     if (input->dt != DNN_FLOAT) {
 100         av_log(log_ctx, AV_LOG_ERROR, "do not support data type rather than DNN_FLOAT\n");
 101         return DNN_ERROR;
 102     }
 103
 104     switch (frame->format) {
 105     case AV_PIX_FMT_RGB24:
 106     case AV_PIX_FMT_BGR24:
 107         sws_ctx = sws_getContext(frame->width * 3,
 108                                  frame->height,
 109                                  AV_PIX_FMT_GRAY8,
 110                                  frame->width * 3,
 111                                  frame->height,
 112                                  AV_PIX_FMT_GRAYF32,
 113                                  0, NULL, NULL, NULL);
 114         if (!sws_ctx) {
 115             av_log(log_ctx, AV_LOG_ERROR, "Impossible to create scale context for the conversion "
 116                 "fmt:%s s:%dx%d -> fmt:%s s:%dx%d\n",
 117                 av_get_pix_fmt_name(AV_PIX_FMT_GRAY8),  frame->width * 3, frame->height,
 118                 av_get_pix_fmt_name(AV_PIX_FMT_GRAYF32),frame->width * 3, frame->height);
 119             return DNN_ERROR;
 120         }
 121         sws_scale(sws_ctx, (const uint8_t **)frame->data,
 122                            frame->linesize, 0, frame->height,
 123                            (uint8_t * const*)(&input->data),
 124                            (const int [4]){frame->width * 3 * sizeof(float), 0, 0, 0});
 125         sws_freeContext(sws_ctx);
 126         break;
 127     case AV_PIX_FMT_GRAYF32:
 128         av_image_copy_plane(input->data, bytewidth,
 129                             frame->data[0], frame->linesize[0],
 130                             bytewidth, frame->height);
 131         break;
 132     case AV_PIX_FMT_YUV420P:
 133     case AV_PIX_FMT_YUV422P:
 134     case AV_PIX_FMT_YUV444P:
 135     case AV_PIX_FMT_YUV410P:
 136     case AV_PIX_FMT_YUV411P:
 137     case AV_PIX_FMT_GRAY8:
 138         sws_ctx = sws_getContext(frame->width,
 139                                  frame->height,
 140                                  AV_PIX_FMT_GRAY8,
 141                                  frame->width,
 142                                  frame->height,
 143                                  AV_PIX_FMT_GRAYF32,
 144                                  0, NULL, NULL, NULL);
 145         if (!sws_ctx) {
 146             av_log(log_ctx, AV_LOG_ERROR, "Impossible to create scale context for the conversion "
 147                 "fmt:%s s:%dx%d -> fmt:%s s:%dx%d\n",
 148                 av_get_pix_fmt_name(AV_PIX_FMT_GRAY8),  frame->width, frame->height,
 149                 av_get_pix_fmt_name(AV_PIX_FMT_GRAYF32),frame->width, frame->height);
 150             return DNN_ERROR;
 151         }
 152         sws_scale(sws_ctx, (const uint8_t **)frame->data,
 153                            frame->linesize, 0, frame->height,
 154                            (uint8_t * const*)(&input->data),
 155                            (const int [4]){frame->width * sizeof(float), 0, 0, 0});
 156         sws_freeContext(sws_ctx);
 157         break;
 158     default:
 159         av_log(log_ctx, AV_LOG_ERROR, "do not support frame format %s\n",
 160                av_get_pix_fmt_name(frame->format));
 161         return DNN_ERROR;
 162     }
 163
 164     return DNN_SUCCESS;
 165 }