add lite inference

Author: Mark-ZhouWX

research/segment-anything/configs/export_lite.cfg

@@ -0,0 +1,2 @@
+[acl_init_options]
+ge.exec.precision_mode="allow_fp32_to_fp16"

research/segment-anything/export.py

@@ -0,0 +1,102 @@
+import argparse
+import os
+import mindspore as ms
+from mindspore import ops
+
+from segment_anything import sam_model_registry
+
+
+def main(args):
+    # Step0: prepare
+    os.makedirs(args.model_path_wo_ext, exist_ok=True)
+    model_path_wo_ext = os.path.join(args.model_path_wo_ext, 'sam_' + args.model_type)
+    mindir_path = os.path.join(model_path_wo_ext + '.mindir')
+    lite_path_wo_ext = os.path.join(model_path_wo_ext + f"_lite")
+    lite_path = os.path.join(model_path_wo_ext + f"_lite.mindir")
+    # model
+    model = sam_model_registry[args.model_type](checkpoint=args.checkpoint)
+
+    # Step 1: export mindir
+    if args.export_mindir:
+        # input
+        image = ops.ones(shape=(1, 3, 1024, 1024), dtype=ms.float32)  # b, 3, 1024, 1024
+        boxes = ops.ones(shape=(1, 1, 4), dtype=ms.float32)  # b, n, 4
+        inputs = (image, boxes)
+        model.set_inputs(*inputs)
+        print(f"start export mindir")
+        ms.export(model, *inputs, file_name=model_path_wo_ext, file_format="MINDIR")
+        print(f"finish export mindir")
+
+    print(f'mind ir path: {mindir_path}')
+    print(f'lite path wo_ext: {lite_path_wo_ext}')
+    print(f'lite path: {lite_path}')
+
+    # Step 2: convert lite
+    if args.convert_lite:
+        import mindspore_lite as mslite
+        optimize_dict = {"ascend": "ascend_oriented", "gpu": "gpu_oriented", "cpu": "general"}
+        converter = mslite.Converter()
+        converter.save_type = mslite.ModelType.MINDIR
+        converter.optimize = optimize_dict[args.device.lower()]
+
+        print(f"start convert lite")
+        converter.convert(
+            fmk_type=mslite.FmkType.MINDIR,
+            model_file=mindir_path,
+            output_file=lite_path_wo_ext,
+            config_file="./configs/export_lite.cfg",
+        )
+        print(converter)
+        print(f"finish convert lite")
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(
+        description=(
+            "Export online ckpt to offline mindir"
+        )
+    )
+
+    parser.add_argument(
+        "--model_path_wo_ext",
+        type=str,
+        default='./models/',
+        help=(
+            "Full path to the directory where the output model is saved, without file extension."
+        ),
+    )
+
+    parser.add_argument(
+        "--model-type",
+        type=str,
+        default='vit_b',
+        help="The type of model to load, in ['default', 'vit_h', 'vit_l', 'vit_b']",
+    )
+
+    parser.add_argument(
+        "--checkpoint",
+        type=str,
+        default='models/sam_vit_b-35e4849c.ckpt',
+        help="online checkpoint file that stores weight",
+    )
+
+    parser.add_argument("--device", type=str, default="Ascend", help="The device to run generation on.")
+
+    parser.add_argument(
+        "--export-mindir",
+        default=True,
+        help=(
+            "Button to enable export mindir."
+        ),
+    )
+
+    parser.add_argument(
+        "--convert-lite",
+        default=True,
+        help=(
+            "Button to enable convert lite."
+        ),
+    )
+    args = parser.parse_args()
+    print(args)
+    main(args)

research/segment-anything/inference_one_image.py renamed to research/segment-anything/inference.py

@@ -9,25 +9,11 @@
 from segment_anything.build_sam import sam_model_registry
 from segment_anything.dataset.transform import TransformPipeline, ImageNorm, ImageResizeAndPad
 import matplotlib.pyplot as plt
-import time
 
+from segment_anything.utils.utils import Timer
 from use_sam_with_promts import show_mask, show_box
 
 
-class Timer:
-    def __init__(self, name=''):
-        self.name = name
-        self.start = 0.0
-        self.end = 0.0
-
-    def __enter__(self):
-        self.start = time.time()
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        self.end = time.time()
-        print(f'{self.name} cost time {self.end - self.start:.3f}')
-
-
 def infer(args):
     ms.context.set_context(mode=args.mode, device_target=args.device)
 

research/segment-anything/lite_inference.py

@@ -0,0 +1,96 @@
+import argparse
+
+import cv2
+import mindspore_lite as mslite
+import numpy as np
+from matplotlib import pyplot as plt
+
+from segment_anything.dataset.transform import ImageResizeAndPad, ImageNorm, TransformPipeline
+from segment_anything.utils.utils import Timer
+from use_sam_with_promts import show_box, show_mask
+
+
+def set_context(device='Ascend', device_id=0):
+    context = mslite.Context()
+    context.target = [device.lower()]
+    if device.lower() == 'ascend':
+        context.ascend.device_id = device_id
+        context.ascend.precision_mode = "preferred_fp32"  # this line is important for keeping precision
+    elif device.lower() == 'gpu':
+        context.gpu.device_id = device_id
+    else:
+        raise NotImplementedError
+    return context
+
+
+def build_model(lite_mindir_path, context):
+    print(f'build model from: {lite_mindir_path}')
+    model = mslite.Model()
+    model.build_from_file(lite_mindir_path, mslite.ModelType.MINDIR, context)
+    return model
+
+
+def infer(args):
+    # Step0: prepare model
+    context = set_context(device=args.device, device_id=args.device_id)
+    model = build_model(args.model_path, context)
+
+    # Step1: data preparation
+    with Timer('preprocess'):
+        transform_list = [
+            ImageResizeAndPad(target_size=1024, apply_mask=False),
+            ImageNorm(),
+        ]
+        transform_pipeline = TransformPipeline(transform_list)
+
+        image_path = args.image_path
+        image_np = cv2.imread(image_path)
+        image_np = cv2.cvtColor(image_np, cv2.COLOR_BGR2RGB)
+        boxes_np = np.array([[425, 600, 700, 875]])
+
+        transformed = transform_pipeline(dict(image=image_np, boxes=boxes_np))
+        image, boxes, origin_hw = transformed['image'], transformed['boxes'], transformed['origin_hw']
+        # batch_size for speed test
+        # image = ms.Tensor(np.expand_dims(image, 0).repeat(8, axis=0))  # b, 3, 1023
+        # boxes = ms.Tensor(np.expand_dims(boxes, 0).repeat(8, axis=0))  # b, n, 4
+        image = np.expand_dims(image, 0)  # b, 3, 1023
+        boxes = np.expand_dims(boxes, 0)  # b, n, 4
+
+        inputs = model.get_inputs()
+        inputs[0].set_data_from_numpy(image.astype(np.float32))
+        inputs[1].set_data_from_numpy(boxes.astype(np.float32))
+
+
+    # Step2: inference
+    with Timer('model inference'):
+        mask_logits = model.predict(inputs)[0]   # (1, 1, 1024, 1024)
+
+    with Timer('Second time inference'):
+        mask_logits = model.predict(inputs)[0]   # (1, 1, 1024, 1024)
+
+    # Step3: post-process
+    with Timer('post-process'):
+        print(f'mask_logits', mask_logits)
+        mask_logits = mask_logits.get_data_to_numpy()[0, 0] > 0.0  # (1024, 1024)
+        mask_logits = mask_logits.astype(np.uint8)
+        final_mask = cv2.resize(mask_logits[:origin_hw[2], :origin_hw[3]], tuple((origin_hw[1], origin_hw[0])),
+                                interpolation=cv2.INTER_CUBIC)
+
+    # Step4: visualize
+    plt.imshow(image_np)
+    show_box(boxes_np[0], plt.gca())
+    show_mask(final_mask, plt.gca())
+    plt.savefig(args.image_path + '_lite_infer.jpg')
+    plt.show()
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description=("Runs inference on one image"))
+    parser.add_argument("--image_path", type=str, default='./images/truck.jpg', help="Path to an input image.")
+    parser.add_argument("--model-path", type=str, default='./models/sam_vit_b_lite.mindir',  help="mindir model path for lite inference")
+    parser.add_argument("--device", type=str, default="Ascend", help="The device to run generation on.")
+    parser.add_argument("--device_id", type=int, default=0, help="The device to run inference on.")
+
+    args = parser.parse_args()
+    print(args)
+    infer(args)

research/segment-anything/segment_anything/modeling/image_encoder.py

@@ -358,7 +358,7 @@ def add_decomposed_rel_pos(
     r_q = q.reshape(B, q_h, q_w, dim)
     dtype = r_q.dtype
     # rel_h = ops.einsum("bhwc,hkc->bhwk", r_q, Rh)
-    rel_h = ops.BatchMatMul(transpose_b=True)(r_q, ops.broadcast_to(ops.unsqueeze(Rh, 0).astype(dtype), (B, -1, -1, -1)))
+    rel_h = ops.BatchMatMul(transpose_b=True)(r_q, ops.unsqueeze(Rh, 0).astype(dtype).repeat(B, axis=0))
     # rel_w = ops.einsum("bhwc,wkc->bhwk", r_q, Rw)
     rel_w = ops.mul(ops.unsqueeze(r_q, -2), ops.unsqueeze(ops.unsqueeze(Rw, 0), 0).astype(dtype)).sum(axis=-1)
 

research/segment-anything/segment_anything/modeling/mask_decoder.py

@@ -114,7 +114,7 @@ def predict_masks(
         """Predicts masks. See 'forward' for more details."""
         # Concatenate output tokens
         output_tokens = ops.cat([self.iou_token.embedding_table, self.mask_tokens.embedding_table], axis=0)
-        output_tokens = output_tokens.unsqueeze(0).broadcast_to((sparse_prompt_embeddings.shape[0], -1, -1))
+        output_tokens = output_tokens.unsqueeze(0).repeat(sparse_prompt_embeddings.shape[0], axis=0)
         tokens = ops.cat((output_tokens, sparse_prompt_embeddings), axis=1)
 
         # Expand per-image data in batch direction to be per-mask

research/segment-anything/segment_anything/modeling/prompt_encoder.py

@@ -167,7 +167,8 @@ def construct(
             dense_embeddings = self._embed_masks(masks)
         else:
             dense_embeddings = self.no_mask_embed.embedding_table.reshape(1, -1, 1, 1).broadcast_to(
-                (bs, -1, self.image_embedding_size[0], self.image_embedding_size[1])
+                (bs, self.no_mask_embed.embedding_table.shape[1],
+                 self.image_embedding_size[0], self.image_embedding_size[1])
             )
 
         return sparse_embeddings, dense_embeddings

research/segment-anything/segment_anything/utils/utils.py

@@ -1,4 +1,5 @@
 import os
+import time
 from datetime import datetime
 
 import mindspore as ms
@@ -137,3 +138,17 @@ def set_directory_and_log(main_device, rank_id, rank_size, work_root, log_level,
                 hack_list = {'save_dir': save_dir, 'main_device': main_device}
                 cb.update(hack_list)
     return save_dir
+
+
+class Timer:
+    def __init__(self, name=''):
+        self.name = name
+        self.start = 0.0
+        self.end = 0.0
+
+    def __enter__(self):
+        self.start = time.time()
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.end = time.time()
+        print(f'{self.name} cost time {self.end - self.start:.3f}')