From cfda275192ef56e2e806548deb9c9a75c82aa213 Mon Sep 17 00:00:00 2001
From: "pengchao.hu" <pengchao.hu@sophgo.com>
Date: Thu, 2 Mar 2023 15:18:04 +0800
Subject: [PATCH] fix preprocess error

Change-Id: If68395d57e25b238c76556746d04b6f7ab18649a
---
 python/tools/model_deploy.py  | 14 +++++++++-----
 python/utils/preprocess.py    |  2 +-
 regression/script_test/run.sh |  9 +++++++++
 3 files changed, 19 insertions(+), 6 deletions(-)

diff --git a/python/tools/model_deploy.py b/python/tools/model_deploy.py
index 2c8fa05ab..40e9038c0 100755
--- a/python/tools/model_deploy.py
+++ b/python/tools/model_deploy.py
@@ -103,6 +103,11 @@ def _prepare_input_npz(self):
         num_inputs = len(self.test_input)
         self.do_validate = (0 < num_inputs)
         if not self.do_validate:
+            if self.customization_format == '':
+                ppa = preprocess()
+                input_op = self.module.inputs[0].op
+                ppa.load_config(input_op)
+                self.customization_format = getCustomFormat(ppa.pixel_format, ppa.channel_format)
             return
         self.inputs = {}
         #self.pre_inputs = {}
@@ -125,10 +130,9 @@ def _prepare_input_npz(self):
                         #use the first input_op's shape as input_shapes
                         input_shapes.append(Operation.shape(input_op))
                     ppa.load_config(input_op)
-                    if self.customization_format == None:
+                    if self.customization_format == '':
                         self.customization_format = getCustomFormat(ppa.pixel_format,
                                                                     ppa.channel_format)
-                        print(self.customization_format)
                     config = {
                         'input_shapes': input_shapes,
                         'resize_dims': ppa.resize_dims,
@@ -152,7 +156,7 @@ def _prepare_input_npz(self):
                     self.inputs[op.name] = data
         if self.aligned_input and not self.fuse_preprocess:
             ppa = preprocess()
-            if self.customization_format is None:
+            if self.customization_format == '':
                 #use the first input_op's pixel format
                 input_op = self.module.inputs[0].op
                 ppa.load_config(input_op)
@@ -247,7 +251,7 @@ def validate_model(self):
     # fuse preprocess
     parser.add_argument("--fuse_preprocess", action='store_true',
                         help="add tpu preprocesses (mean/scale/channel_swap) in the front of model")
-    parser.add_argument("--customization_format", default=None, type=str,
+    parser.add_argument("--customization_format", default='', type=str,
                         choices=supported_customization_format,
                         help="pixel format of input frame to the model")
     parser.add_argument("--aligned_input", action='store_true',
@@ -274,7 +278,7 @@ def validate_model(self):
     parser.add_argument("--debug", action='store_true', help='to keep all intermediate files for debug')
     # yapf: enable
     args = parser.parse_args()
-    if args.customization_format is not None and args.customization_format.startswith("YUV") >= 0:
+    if args.customization_format.startswith("YUV"):
         args.aligned_input = True
 
     tool = DeployTool(args)
diff --git a/python/utils/preprocess.py b/python/utils/preprocess.py
index 54edb8718..1901beaab 100644
--- a/python/utils/preprocess.py
+++ b/python/utils/preprocess.py
@@ -27,7 +27,7 @@ class YuvType(Enum):
     'YUV_NV21',
     'YUV_NV12',
     'RGBA_PLANAR',
-    None
+    ''
 ]
 
 customization_format_attributes = {
diff --git a/regression/script_test/run.sh b/regression/script_test/run.sh
index 7ac94c8e5..8dfa0f291 100755
--- a/regression/script_test/run.sh
+++ b/regression/script_test/run.sh
@@ -90,4 +90,13 @@ model_deploy.py \
   --test_reference mobilenet_v2_top_outputs.npz \
   --compare_all \
   --model mobilenet_v2_1684x_int8_fuse.bmodel
+
+# no test
+model_deploy.py \
+  --mlir mobilenet_v2.mlir \
+  --quantize INT8 \
+  --chip bm1684x \
+  --calibration_table mobilenet_v2_cali_table \
+  --fuse_preprocess \
+  --model mobilenet_v2_1684x_int8_fuse2.bmodel
 popd