Add appropriate pre- to input image and postprocessing to validate function

convert_to_mixed.py

@@ -1,8 +1,5 @@
-import sys
 import numpy as np
-from PIL import Image
 import onnx
-import onnxruntime as ort
 from onnxconverter_common import auto_mixed_precision_model_path
 import argparse
 from rtmo_gpu import RTMO_GPU, draw_skeleton
@@ -24,28 +21,49 @@ def detect_model_input_size(model_path):
             return tuple(dims[2:4])  # Return (height, width)
     raise ValueError("Input node 'input' not found in the model")
 
-def load_and_preprocess_image(image_path, size=(224, 224)):
-    image = Image.open(image_path).convert('RGB')
-    image = image.resize(size)
-    image = np.array(image).astype(np.float32) / 255.
-    image = np.transpose(image, (2, 0, 1))
-    image = np.expand_dims(image, axis=0)
+def load_and_preprocess_image(image_path, preprocesss=None):
+
+    image = cv2.imread(image_path)
+
+    if preprocesss is not None:
+        image  = preprocesss(image)
+
     return image
 
-def validate_pose(res1, res2):
-    rtmo = RTMO_GPU(args.source_model_path)
-    poses1 = rtmo.postprocess(res1)
-    poses2 = rtmo.postprocess(res2)
+def compare_result(res1, res2):
+    keypoints1, scores1 = res1
+    keypoints2, scores2 = res2
+
+    from termcolor import colored
+
+    for d1, d2 in zip(keypoints1, keypoints2):
+        for i, (j1, j2) in enumerate(zip(d1, d2)):
+            x1, y1 = j1
+            x2, y2 = j2
+            print(f"Joint-{i}: X: {colored(x1,'green')} VS {colored(x2, 'blue')} Y: {colored(y1, 'green')} VS {colored(y2, 'blue')}")    
+
+    for d1, d2 in zip(scores1, scores2):
+        for i, (s1, s2) in enumerate(zip(d1, d2)):
+            print(f"Joint-{i}: S: {colored(s1,'green')} VS {colored(s2, 'blue')}")    
+
+
+def validate_pose(res1, res2, postprocess=None):
 
-    for k1, k2 in zip(poses1, poses2):
-        if not np.allclose(k1, k2, rtol=args.rtol, atol=args.atol):
+    if postprocess is not None:
+        res1 = postprocess(res1)
+        res2 = postprocess(res2)
+
+    compare_result(res1, res2)
+
+    for r1, r2 in zip(res1, res2):
+        if not np.allclose(r1, r2, rtol=args.rtol, atol=args.atol):
             return False
     return True
 
 def infer_on_image(onnx_model, model_input_size, test_image_path):
     body = RTMO_GPU(onnx_model=onnx_model, 
         model_input_size=model_input_size, 
-        is_yolo_nas_pose=False)
+        is_yolo_nas_pose=args.yolo_nas_pose)
 
     frame = cv2.imread(test_image_path)
     img_show = frame.copy()
@@ -63,12 +81,29 @@ def infer_on_image(onnx_model, model_input_size, test_image_path):
 
 def main(args):
     model_input_size = detect_model_input_size(args.source_model_path)
-    input_feed = {'input':load_and_preprocess_image(args.test_image_path, size=model_input_size)}
+
+    body = RTMO_GPU(onnx_model=args.source_model_path,
+                    model_input_size=model_input_size,
+                    is_yolo_nas_pose=args.yolo_nas_pose)
+    
+    def preprocess(image, body, is_yolo_nas_pose):
+
+        img, _ = body.preprocess(image)
+
+        # build input to (1, 3, H, W)
+        img = img.transpose(2, 0, 1)
+        img = np.ascontiguousarray(img, dtype=np.float32 if not is_yolo_nas_pose else np.uint8)
+        img = img[None, :, :, :]
+        return img
+    
+    image = load_and_preprocess_image(args.test_image_path, lambda img: preprocess(img, body, args.yolo_nas_pose))
+
+    input_feed = {'input': image}
 
     auto_mixed_precision_model_path.auto_convert_mixed_precision_model_path(source_model_path=args.source_model_path, 
                                     input_feed=input_feed, 
                                     target_model_path=args.target_model_path,
-                                    customized_validate_func=validate_pose, 
+                                    customized_validate_func=lambda res1,res2:validate_pose(res1, res2, body.postprocess), 
                                     rtol=args.rtol, atol=args.atol,
                                     provider=PROVIDERS, 
                                     keep_io_types=True,
@@ -83,6 +118,7 @@ if __name__ == "__main__":
     parser.add_argument("test_image_path", type=str, help="Path to a test image for validating the model conversion.")
     parser.add_argument('--rtol', type=float, default=0.01, help=' the relative tolerance to do validation')
     parser.add_argument('--atol', type=float, default=0.001, help=' the absolute tolerance to do validation')
+    parser.add_argument('--yolo_nas_pose', action='store_true', help='Use YOLO NAS Pose (flat format only) instead of RTMO Model')
 
     args = parser.parse_args()