add ax637 models

.gitattributes

@@ -39,3 +39,9 @@ demo-imgs/im0.png filter=lfs diff=lfs merge=lfs -text
 demo-imgs/im1.png filter=lfs diff=lfs merge=lfs -text
 AXModel-Disparity-Map-rt.png filter=lfs diff=lfs merge=lfs -text
 ONNX-Disparity-Map-rt.png filter=lfs diff=lfs merge=lfs -text
+models-ax637/rt_p1.onnx filter=lfs diff=lfs merge=lfs -text
+models-ax637/rt_up.axmodel filter=lfs diff=lfs merge=lfs -text
+models-ax637/rt_up.onnx filter=lfs diff=lfs merge=lfs -text
+models-ax637/rt_iterfn.axmodel filter=lfs diff=lfs merge=lfs -text
+models-ax637/rt_iterfn.onnx filter=lfs diff=lfs merge=lfs -text
+models-ax637/rt_p1.axmodel filter=lfs diff=lfs merge=lfs -text

README.md

@@ -32,7 +32,7 @@ For those who are interested in model conversion, you can try to export axmodel
 |Chips|Models |Time|
 |--|--|--|
 |AX650|rt_sceneflow|142.0 ms |
-
+|AX637|rt_sceneflow(分成三段 rt_p1 + rt_iterfn + rt_up)|123.0 + 29.3*18 + 9.1 = 659.5 ms|
 
 
 ## How to use

infer_ax637.py

@@ -0,0 +1,127 @@
+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")
+
+
+class Session:
+    def __init__(self, path_p1, path_iterfn, path_up, iters=18):
+        self.iters = iters
+        if os.path.splitext(path_p1)[1] == ".onnx":
+            self.session_p1 = ort.InferenceSession(path_p1)
+        elif os.path.splitext(path_p1)[1] == ".axmodel":
+            self.session_p1 = axe.InferenceSession(path_p1)
+        else:
+            raise NotImplementedError
+        
+        if os.path.splitext(path_iterfn)[1] == ".onnx":
+            self.session_iterfn = ort.InferenceSession(path_iterfn)
+        elif os.path.splitext(path_iterfn)[1] == ".axmodel":
+            self.session_iterfn = axe.InferenceSession(path_iterfn)
+        else:
+            raise NotImplementedError
+        
+        if os.path.splitext(path_up)[1] == ".onnx":
+            self.session_up = ort.InferenceSession(path_up)
+        elif os.path.splitext(path_up)[1] == ".axmodel":
+            self.session_up = axe.InferenceSession(path_up)
+        else:
+            raise NotImplementedError
+        
+
+
+
+    def infer(self, img_left, img_right):
+        inputs_p1 = {"left": img_left, "right": img_right}
+        outputs_p1 = self.session_p1.run(["disp",  "net", "stem_2x", "c0", "c1", "c2"], inputs_p1)
+
+        disp = outputs_p1[0]
+        net = outputs_p1[1]
+        stem_2x = outputs_p1[2]
+        c0 = outputs_p1[3]
+        c1 = outputs_p1[4]
+        c2 = outputs_p1[5]
+        for itr in range(self.iters):
+            # disp, mask_feat_4, net = self.iter_fn(disp, net, context)
+            inputs_iter = {"disp_in":disp, "net_in":net,  "c0":c0, "c1":c1, "c2":c2 }
+            outputs_iter = self.session_iterfn.run(["disp", "mask_feat_4", "net"], inputs_iter)
+            disp = outputs_iter[0]
+            net = outputs_iter[2]
+
+        mask_feat_4 = outputs_iter[1]
+        inputs_up = {"disp":disp, "mask_feat_4":mask_feat_4, "stem_2x":stem_2x}
+        outputs = self.session_up.run(None, inputs_up)[0]
+        outputs = outputs.squeeze()
+        return outputs
+
+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 = Session("models-ax637/rt_p1.onnx", "models-ax637/rt_iterfn.onnx", "models-ax637/rt_up.onnx", args.valid_iters)
+        ax_session = Session("models-ax637/rt_p1.axmodel", 
+                             "models-ax637/rt_iterfn.axmodel", 
+                             "models-ax637/rt_up.axmodel",
+                             args.valid_iters)
+
+        input_l_np = image1.cpu().numpy()
+        input_r_np = image2.cpu().numpy()
+        # ax_inputs_p1 = {"left": input_l_np.transpose(0,2,3,1).astype(np.uint8), "right": input_r_np.transpose(0,2,3,1).astype(np.uint8)}
+
+        disp_ax = ax_session.infer(input_l_np.transpose(0,2,3,1).astype(np.uint8), 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_p1 = {"left": input_l_np, "right": input_r_np}
+        
+        disp_onnx = ort_session.infer(input_l_np, input_r_np)
+
+        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")
+    parser.add_argument('--valid_iters', type=int, default=18, help='number of flow-field updates during forward pass')
+
+    args = parser.parse_args()
+    demo(args)

models-ax637/rt_iterfn.axmodel

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

models-ax637/rt_iterfn.onnx

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

models-ax637/rt_p1.axmodel

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

models-ax637/rt_p1.onnx

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

models-ax637/rt_up.axmodel

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

models-ax637/rt_up.onnx

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