first commit

.gitattributes

@@ -33,3 +33,7 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+models/rt_sceneflow.axmodel filter=lfs diff=lfs merge=lfs -text
+models/rt_sceneflow.onnx filter=lfs diff=lfs merge=lfs -text
+demo-imgs/im0.png filter=lfs diff=lfs merge=lfs -text
+demo-imgs/im1.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,3 +1,5 @@
 ---
 license: mit
 ---
+
+RT IGEV from https://github.com/gangweiX/IGEV-plusplus

infer_ax.py

@@ -0,0 +1,73 @@
+import sys
+import argparse
+import glob
+import numpy as np
+import torch
+from tqdm import tqdm
+from pathlib import Path
+from PIL import Image
+from matplotlib import pyplot as plt
+import os
+import onnxruntime as ort
+import axengine as axe
+
+
+
+def load_image(imfile):
+    img = np.array(Image.open(imfile).resize((512,384))).astype(np.uint8)[..., :3]
+    img = torch.from_numpy(img).permute(2, 0, 1).float()
+    return img[None]
+
+
+def visualize_disparity(disparity_map, title):
+    plt.figure(figsize=(10, 6))
+    plt.imshow(disparity_map, cmap='jet')
+    plt.colorbar(label="Disparity")
+    plt.title(title)
+    plt.axis('off')
+    # plt.show()
+    plt.savefig(f"{title}-rt.png")
+
+
+def demo(args):
+    # PyTorch 和 ONNX 推理对比
+    left_images = sorted(glob.glob(args.left_imgs, recursive=True))
+    right_images = sorted(glob.glob(args.right_imgs, recursive=True))
+    for (imfile1, imfile2) in tqdm(list(zip(left_images, right_images))):
+        image1 = load_image(imfile1)
+        image2 = load_image(imfile2)
+
+        # ONNX 推理
+        ort_session = ort.InferenceSession("models/rt_sceneflow.onnx")
+        ax_session = axe.InferenceSession("models/rt_sceneflow.axmodel")
+        input_l_np = image1.cpu().numpy()
+        input_r_np = image2.cpu().numpy()
+        ax_inputs = {"left": input_l_np.transpose(0,2,3,1).astype(np.uint8), "right": input_r_np.transpose(0,2,3,1).astype(np.uint8)}
+
+        input_l_np = (2 * (input_l_np / 255.0) - 1.0)
+        input_r_np = (2 * (input_r_np / 255.0) - 1.0)
+        onnx_inputs = {"left": input_l_np, "right": input_r_np}
+        
+        onnx_outputs = ort_session.run(None, onnx_inputs)
+        disp_onnx = onnx_outputs[0].squeeze()
+
+        ax_outputs = ax_session.run(None, ax_inputs)
+        disp_ax = ax_outputs[0].squeeze()
+
+        print("disp_onnx",disp_onnx)
+        print("disp_ax",disp_ax)
+        visualize_disparity(disp_onnx, title="ONNX Disparity Map")
+        visualize_disparity(disp_ax, title="AXModel Disparity Map")
+
+
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('-l', '--left_imgs', help="path to all first (left) frames",
+                        default="demo-imgs/im0.png")
+    parser.add_argument('-r', '--right_imgs', help="path to all second (right) frames",
+                        default="demo-imgs/im1.png")
+
+    args = parser.parse_args()
+    demo(args)

models/rt_sceneflow.axmodel

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:183a426e7a36f02a1351aeb42b14320f2d350d6fae600fedae4ab2ab0a95ad1b
+size 112034124

models/rt_sceneflow.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e95dfb781f473647b147f0c2221359b39a8609687a58499d261ed7d6b7497a7b
+size 190012005