Fix RCV2 mean and scale values (#59)

luxonis · Dec 12, 2024 · 9d09fd0 · 9d09fd0
1 parent acb3c19
commit 9d09fd0
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/modelconverter/packages/rvc2/exporter.py b/modelconverter/packages/rvc2/exporter.py
@@ -120,15 +120,15 @@ def _export_openvino_ir(self) -> Path:
         scale_values_str = ""
         for name, inp in self.inputs.items():
             # Append mean values in a similar style
-            if inp.mean_values is not None:
+            if inp.is_color_input and inp.mean_values is not None:
                 if mean_values_str:
                     mean_values_str += ","
                 mean_values_str += (
                     f"{name}[{', '.join(str(v) for v in inp.mean_values)}]"
                 )
 
             # Append scale values in a similar style
-            if inp.scale_values is not None:
+            if inp.is_color_input and inp.scale_values is not None:
                 if scale_values_str:
                     scale_values_str += ","
                 scale_values_str += (

diff --git a/modelconverter/utils/config.py b/modelconverter/utils/config.py
@@ -141,6 +141,10 @@ class InputConfig(OutputConfig):
     def encoding_mismatch(self) -> bool:
         return self.encoding.from_ != self.encoding.to
 
+    @property
+    def is_color_input(self) -> bool:
+        return self.encoding.from_ in {Encoding.RGB, Encoding.BGR}
+
     @model_validator(mode="after")
     def _validate_grayscale_inputs(self) -> Self:
         if self.layout is None: