From b44d4134d96499f1e5fb408de39d767a5e90de9a Mon Sep 17 00:00:00 2001
From: Daniel <daniel@mail.de>
Date: Mon, 2 Dec 2024 13:45:26 +0100
Subject: [PATCH] Move fp16 casting into onnx models.

---
 extras/mmdeploy/README.md                     |  2 +-
 .../{add_norm_step.py => add_extra_steps.py}  | 30 ++++++++++++++-----
 2 files changed, 23 insertions(+), 9 deletions(-)
 rename extras/mmdeploy/{add_norm_step.py => add_extra_steps.py} (78%)

diff --git a/extras/mmdeploy/README.md b/extras/mmdeploy/README.md
index 85944bf..9c00b37 100644
--- a/extras/mmdeploy/README.md
+++ b/extras/mmdeploy/README.md
@@ -43,5 +43,5 @@ mv /mmdeploy/work_dir/end2end.onnx /RapidPoseTriangulation/extras/mmdeploy/expor
 ```
 
 ```bash
-python3 /RapidPoseTriangulation/extras/mmdeploy/add_norm_step.py
+python3 /RapidPoseTriangulation/extras/mmdeploy/add_extra_steps.py
 ```
diff --git a/extras/mmdeploy/add_norm_step.py b/extras/mmdeploy/add_extra_steps.py
similarity index 78%
rename from extras/mmdeploy/add_norm_step.py
rename to extras/mmdeploy/add_extra_steps.py
index 581a4a7..cc44fdb 100644
--- a/extras/mmdeploy/add_norm_step.py
+++ b/extras/mmdeploy/add_extra_steps.py
@@ -1,6 +1,6 @@
 import numpy as np
 import onnx
-from onnx import helper, numpy_helper
+from onnx import helper, numpy_helper, TensorProto
 
 # ==================================================================================================
 
@@ -44,22 +44,35 @@ def add_steps_to_onnx(model_path):
 
     # Define layer names, assuming the first input is the image tensor
     input_name = graph.input[0].name
-    mean_added_output = "mean_added_output"
-    normalized_output = "normalized_output"
+
+    # Set input type to always be float32
+    graph.input[0].type.tensor_type.elem_type = TensorProto.FLOAT
+
+    # Create to cast the float32 if needed
+    cast_type = 10 if use_fp16 else 1
+    casted_output = "casted_output"
+    cast_node = helper.make_node(
+        "Cast",
+        inputs=[input_name],
+        outputs=[casted_output],
+        to=cast_type,
+    )
 
     # Node to add mean
+    mean_added_output = "mean_added_output"
     mean_add_node = helper.make_node(
         "Add",
-        inputs=[input_name, "norm_mean"],
+        inputs=[casted_output, "norm_mean"],
         outputs=[mean_added_output],
         name="Mean_Addition",
     )
 
     # Node to multiply by std
+    std_mult_output = "std_mult_output"
     std_mul_node = helper.make_node(
         "Mul",
         inputs=[mean_added_output, "norm_std"],
-        outputs=[normalized_output],
+        outputs=[std_mult_output],
         name="Std_Multiplication",
     )
 
@@ -67,11 +80,12 @@ def add_steps_to_onnx(model_path):
     for node in graph.node:
         for idx, input_name_in_node in enumerate(node.input):
             if input_name_in_node == input_name:
-                node.input[idx] = normalized_output
+                node.input[idx] = std_mult_output
 
     # Add the new nodes to the graph
-    graph.node.insert(0, mean_add_node)
-    graph.node.insert(1, std_mul_node)
+    graph.node.insert(0, cast_node)
+    graph.node.insert(1, mean_add_node)
+    graph.node.insert(2, std_mul_node)
 
     path = model_path.replace(".onnx", "_with-norm.onnx")
     onnx.save(model, path)