fix models and update images

README.md

@@ -32,8 +32,8 @@ 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|
+|AX650|AX650_RTIGEV| 139.803 ms |
+|AX637|AX637_RTIGEV| 385.72 ms |
 
 
 ## How to use
@@ -56,20 +56,22 @@ pip install axengine-0.1.3-py3-none-any.whl
 
 Maybe None.
 
-#### Inference with AX650 Host, such as M4N-Dock(爱芯派Pro)
+#### Inference with AX650 Host (such as M4N-Dock, a.k.a. 爱芯派 Pro) or AX637 Host
 
-Input image:
+Input image (take 1 pair of images as examples):
 
-![](samples/1_image.png)
+![Left image](./demo-imgs/Adirondack/im0.png)
+![Right image](./demo-imgs/Adirondack/im1.png)
 
-run 
-```
-python3 infer_ax.py
+run:
+
+```sh
+# the default target_chip is AX650, the default input is all images in `demo-img` folder
+python3 infer.py 
 ```
 
 Output image:  
-ONNX result  
-![](ONNX-Disparity-Map-rt.png)
 
-AXmodel result  
-![](AXModel-Disparity-Map-rt.png)
+AXmodel result (take 2 images as examples):
+![AX650 output image example](./output-imgs/AX650/AXModel_Disparity_Map-rt-Adirondack.png)
+![AX637 output image example](./output-imgs/AX637/AXModel_Disparity_Map-rt-2.png)

infer_ax650.py → infer.py

@@ -12,34 +12,35 @@ 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):
+def visualize_disparity(disparity_map, title, name="test"):
     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")
+    plt.savefig(f"{title}-rt-{name}.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))
+    if args.target_chip == "AX637":
+        print("\033[91mWarning: AX637 uses quant_axmodel, which can not be run by onnxruntime, \
+so we use AX650's onnx model for comparison\033[0m")
+    ort_session = ort.InferenceSession(f"./models/AX650.onnx")
+    ax_session = axe.InferenceSession(f"./models/{args.target_chip}_RTIGEV.axmodel")
     for (imfile1, imfile2) in tqdm(list(zip(left_images, right_images))):
         image1 = load_image(imfile1)
         image2 = load_image(imfile2)
+        img_name = Path(imfile1).parent.name
 
-        # 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)}
@@ -54,20 +55,19 @@ def demo(args):
         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")
-
-
+        # print("disp_onnx",disp_onnx)
+        # print("disp_ax",disp_ax)
+        visualize_disparity(disp_onnx, title="ONNX_Disparity_Map", name=img_name)
+        visualize_disparity(disp_ax, title="AXModel_Disparity_Map", name=img_name)
 
 
 if __name__ == '__main__':
     parser = argparse.ArgumentParser()
     parser.add_argument('-l', '--left_imgs', help="path to all first (left) frames",
-                        default="demo-imgs/im0.png")
+                        default="./demo-imgs/*/im0.png")
+    parser.add_argument('-t', '--target_chip', help="target chip for inference",
+                        default="AX650", choices=["AX637", "AX650"])
     parser.add_argument('-r', '--right_imgs', help="path to all second (right) frames",
-                        default="demo-imgs/im1.png")
-
+                        default="./demo-imgs/*/im1.png")
     args = parser.parse_args()
     demo(args)

infer_ax637.py

@@ -1,127 +0,0 @@
-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_p1.axmodel

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

models-ax637/rt_p1.onnx

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

models-ax637/rt_iterfn.axmodel → models/AX637_RTIGEV.axmodel

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

models-ax637/rt_iterfn.onnx → models/AX637_quant_axmodel.onnx

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

models/AX650.onnx

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

models/AX650_RTIGEV.axmodel

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

models/rt_sceneflow.axmodel

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

models/rt_sceneflow.onnx

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