diff mbox series

[FFmpeg-devel,1/4] dnn: add DCO_RGB color order to enum DNNColorOrder

Message ID 20210430030711.30216-1-ting.fu@intel.com
State Accepted
Headers show
Series [FFmpeg-devel,1/4] dnn: add DCO_RGB color order to enum DNNColorOrder | expand

Checks

Context Check Description
andriy/x86_make success Make finished
andriy/x86_make_fate success Make fate finished
andriy/PPC64_make success Make finished
andriy/PPC64_make_fate success Make fate finished

Commit Message

Fu, Ting April 30, 2021, 3:07 a.m. UTC 
Adding DCO_RGB color order to DNNColorOrder, since tensorflow model
needs this kind of color oder as input.

Signed-off-by: Ting Fu <ting.fu@intel.com>
---
 libavfilter/dnn/dnn_backend_tf.c |  1 +
 libavfilter/dnn/dnn_io_proc.c    | 14 +++++++++++---
 libavfilter/dnn_interface.h      |  1 +
 3 files changed, 13 insertions(+), 3 deletions(-)
diff mbox series

Patch

diff --git a/libavfilter/dnn/dnn_backend_tf.c b/libavfilter/dnn/dnn_backend_tf.c
index 076dd3d6a9..f70e3d4659 100644
--- a/libavfilter/dnn/dnn_backend_tf.c
+++ b/libavfilter/dnn/dnn_backend_tf.c
@@ -143,6 +143,7 @@  static DNNReturnType get_input_tf(void *model, DNNData *input, const char *input
 
     tf_output.index = 0;
     input->dt = TF_OperationOutputType(tf_output);
+    input->order = DCO_RGB;
 
     status = TF_NewStatus();
     TF_GraphGetTensorShape(tf_model->graph, tf_output, dims, 4, status);
diff --git a/libavfilter/dnn/dnn_io_proc.c b/libavfilter/dnn/dnn_io_proc.c
index e104cc5064..5f6ce36b96 100644
--- a/libavfilter/dnn/dnn_io_proc.c
+++ b/libavfilter/dnn/dnn_io_proc.c
@@ -167,11 +167,19 @@  static DNNReturnType proc_from_frame_to_dnn_frameprocessing(AVFrame *frame, DNND
 
 static enum AVPixelFormat get_pixel_format(DNNData *data)
 {
-    if (data->dt == DNN_UINT8 && data->order == DCO_BGR) {
-        return AV_PIX_FMT_BGR24;
+    if (data->dt == DNN_UINT8) {
+        switch (data->order) {
+        case DCO_BGR:
+            return AV_PIX_FMT_BGR24;
+        case DCO_RGB:
+            return AV_PIX_FMT_RGB24;
+        default:
+            av_assert0(!"unsupported data pixel format.\n");
+            return AV_PIX_FMT_BGR24;
+        }
     }
 
-    av_assert0(!"not supported yet.\n");
+    av_assert0(!"unsupported data type.\n");
     return AV_PIX_FMT_BGR24;
 }
 
diff --git a/libavfilter/dnn_interface.h b/libavfilter/dnn_interface.h
index ae5a488341..92c3b0fc6e 100644
--- a/libavfilter/dnn_interface.h
+++ b/libavfilter/dnn_interface.h
@@ -39,6 +39,7 @@  typedef enum {DNN_FLOAT = 1, DNN_UINT8 = 4} DNNDataType;
 typedef enum {
     DCO_NONE,
     DCO_BGR,
+    DCO_RGB,
 } DNNColorOrder;
 
 typedef enum {