Validate only on pose result, add visual check

convert_to_mixed.py

@@ -5,6 +5,8 @@ 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
+import cv2
 
 PROVIDERS=[('TensorrtExecutionProvider', {'trt_fp16_enable':True,}), 'CUDAExecutionProvider', 'CPUExecutionProvider']
 
@@ -30,11 +32,34 @@ def load_and_preprocess_image(image_path, size=(224, 224)):
     image = np.expand_dims(image, axis=0)
     return image
 
-def infer(model_path, input_feed):
-    session = ort.InferenceSession(model_path, providers=PROVIDERS)
-    #input_name = session.get_inputs()[0].name
-    result = session.run(None, input_feed)
-    return result
+def validate_pose(res1, res2):
+    rtmo = RTMO_GPU(args.source_model_path)
+    poses1 = rtmo.postprocess(res1)
+    poses2 = rtmo.postprocess(res2)
+
+    for k1, k2 in zip(poses1, poses2):
+        if not np.allclose(k1, k2, 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)
+
+    frame = cv2.imread(test_image_path)
+    img_show = frame.copy()
+    keypoints, scores = body(img_show)
+
+    img_show = draw_skeleton(img_show,
+                            keypoints,
+                            scores,
+                            kpt_thr=0.3,
+                            line_width=2)
+    img_show = cv2.resize(img_show, (788, 525))
+    cv2.imshow(f'{args.target_model_path}', img_show)
+    cv2.waitKey(0)
+    cv2.destroyAllWindows()
 
 def main(args):
     model_input_size = detect_model_input_size(args.source_model_path)
@@ -43,28 +68,13 @@ def main(args):
     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=None, 
+                                    customized_validate_func=validate_pose, 
                                     rtol=args.rtol, atol=args.atol,
                                     provider=PROVIDERS, 
                                     keep_io_types=True,
                                     verbose=True)
 
-    original_result = infer(args.source_model_path, input_feed)
-    converted_result = infer(args.target_model_path, input_feed)
-    
-    is_close = np.allclose(original_result[0], converted_result[0], rtol=args.rtol, atol=args.atol)
-    print(f"Validation result: {'Success' if is_close else 'Failure'}")
-
-    print('Compare Joint Coordinates:')
-    original_keypoints, converted_keypoints = original_result[1][0, :, :, :2], converted_result[1][0, :, :, :2]
-    for n_det in range(original_keypoints.shape[0]):
-        print(f'\tDetection {n_det}:')
-        for n_joint in range(original_keypoints.shape[1]):
-            original_x, original_y = original_keypoints[n_det][n_joint]
-            converted_x, converted_y = converted_keypoints[n_det][n_joint]
-            print(f'\t\tJoint {n_joint}:')
-            print(f'\t\t\tX: {original_x:.2f}, {converted_x:.2f}')
-            print(f'\t\t\tY: {original_y:.2f}, {converted_y:.2f}')
+    infer_on_image(args.target_model_path, model_input_size, args.test_image_path)
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="Convert an ONNX model to mixed precision format.")