SAM (General Image Segmentation Foundation Model)

1. Introduction

SAM is a promptable model proposed by Meta that segments anything. It was trained on more than 1 billion masks over 11 million images, achieving strong zero-shot generalization and breaking through the boundaries of segmentation. This example ports the model and algorithm of the SAM official open-source repository so that it can run inference tests on SOPHON BM1684X.

1. Features

  • Supports BM1684X (x86 PCIe, SoC, riscv PCIe).
  • The image compression (embedding) part supports FP16 1batch (BM1684X) model compilation and inference.
  • The image inference (mask_decoder) part supports FP32 1batch and FP16 1batch (BM1684X) model compilation and inference.
  • Supports OpenCV-based Python inference.
  • Supports model inference with single-point and box inputs, and outputs the highest-confidence mask or the top-three confidence masks.
  • Supports image testing.
  • Supports automatic mask generation without point/box input.

Note: This example runs image compression (embedding) and image inference (mask_decoder) as two separate bmodels; the last-layer resize of the image inference part is not compiled into the bmodel.

2. Project Directory

The author made many modifications to the demo project files; it is recommended to copy the author's files directly to the /data directory.

SAM
├─datasets ##weby以及python案例的图片保存
│      dog.jpg
│      groceries.jpg
│      truck.jpg
│
├─docs  ##帮助文档
│  │  boxShare_PC_Wifi.md
│  │  sam.md
│  │
│  └─image  ##文档中显示的图片
│          eth.png
│          ipv4.png
│          ping.png
│          regedit.png
│          result_0.jpg
│          result_auto.jpg
│          result_box_0.jpg
│          result_box_1.jpg
│          result_box_2.jpg
│          t2.png
│          t3.png
│          terminal.png
│          ui.png
│          uib.png
│          uip.png
│          wlan.png
│
├─models	##模型文件
│  └─BM1684X	##1684x的模型权重文件
│      ├─decode_bmodel
│      │      SAM-ViT-B_auto_multi_decoder_fp32_1b.bmodel
│      │      SAM-ViT-B_decoder_multi_mask_fp16_1b.bmodel
│      │      SAM-ViT-B_decoder_multi_mask_fp32_1b.bmodel
│      │      SAM-ViT-B_decoder_single_mask_fp16_1b.bmodel
│      │      SAM-ViT-B_decoder_single_mask_fp32_1b.bmodel
│      │
│      └─embedding_bmodel
│              SAM-ViT-B_embedding_fp16_1b.bmodel
│
├─python	##python脚本
│      amg.py
│      automatic_mask_generator.py
│      backend.py
│      predictor.py
│      sam_encoder.py
│      sam_model.py
│      sam_opencv.py
│      transforms.py
│
└─web_ui	web例程文件
    │  index.html
    │
    ├─components
    │      drawBox.png
    │      firstPage.png
    │      frontPage.png
    │      singlePoint.png
    │
    ├─css
    │      styles.css
    │
    ├─images
    │      dog.jpg
    │      groceries.jpg
    │      truck.jpg
    │
    └─scripts
            main.js

2. Running Steps

Check the network environment: because the subsequent interactive web page uses a fixed IP, here we use the method of sharing the PC network over Ethernet with the development board. For details, refer to the networking document.

1. Environment Preparation

Configure the Python Environment

Edit the .bashrc file to import the sophon Python environment:

sudo vim ~/.bashrc

Add the following line at the end of the file:

export PYTHONPATH=$PYTHONPATH:/opt/sophon/libsophon-current/lib:/opt/sophon/sophon-opencv-latest/opencv-python/

After :wq to save and exit, reload the terminal:

source ~/.bashrc

You can run echo $PYTHONPATH to verify the field.

The runtime environment also requires the following Python libraries:

pip3 install torch
##torchcision安装过慢，可指定清华源安装
pip3 install torchvision -i https://pypi.tuna.tsinghua.edu.cn/simple
pip3 install matplotlib
pip3 install flask flask-cors ##运行web交互案例所需，可选择性安装

After installation, you can check with pip show <package-name>.

2. Python Example Test

2.1 Parameter Description

The Python example mainly runs the sam_opencv.py file. The parameters are:

usage: sam_opencv.py [--input_image INPUT_PATH] [--input_point INPOINT_POINT]
                     [--embedding_bmodel EMBEDDING_BMODEL] [--bmodel BMODEL]
                     [--auto bool][--dev_id DEV_ID]

--input_image: 测试图片路径，需输入图片路径；
--input_point: 输入点的坐标，输入格式为x,y；或者输入框坐标，格式为x1,y1,x2,y2
--embedding_bmodel 用于图像压缩(embedding)的bmodel路径；
--decode_bmodel: 用于推理(mask_decode)的bmodel路径；
--dev_id: 用于推理的tpu设备id；
--auto: 是否启用自动分割，为bool，默认为0不开启，1为开启；

'''以下为automatic masks generator的可调参数，可控制采样点的密度以及去除低质量或重复mask的阈值'''
--points_per_side: 沿图像一侧采样的点数。总点数为points_per_side2^2。默认值为32；
--points_per_batch: 设置模型同时检测的点数。数字越大可能速度越快，但会使用更多GPU内存。默认值为64；
--pred_iou_thresh: [0,1]中的过滤阈值，模型的预测mask质量。默认值为0.88；
--stability_score_thresh: [0,1] 中的过滤阈值（截止值变化时掩模的稳定性）用于对模型的mask预测进行二值化。默认值为0.95；
--stability_score_offset: 计算稳定性分数时，偏移截止值的量。默认值为1.0；
--box_nms_thresh: 用于过滤重复mask的非极大值抑制框IoU截止。默认值为0.7；
--crop_nms_thresh: 用于非极大值抑制的框IoU截止，以过滤不同对象之间的重复mask。默认值为0.7；
--crop_overlap_ratio: 设置物体重叠的程度。在第一个裁剪层中，裁剪将重叠图像长度的这一部分。物体较多的后几层会缩小这种重叠。默认值为512 / 1500；
--crop_n_points_downscale_factor: 在层n中采样的每侧的点数按比例缩小"crop_n_points_downscale_factorn"^n。默认值为1；
--min_mask_region_area: 如果>0，将应用后处理来移除面积小于"min_mask_region_area"的mask来中断开连接的区域和孔。需要opencv。默认为0；
--output_mode: mask输出方式。可以是binary_mask、uncompressed_rle或coco_rle ，coco_rle需要pycocotools。对于大分辨率，binary_mask可能会消耗大量内存。默认为'binary_mask'；

2.2 Image Test

2.2.1 Point Input Test
cd /data/SAM
python3 python/sam_opencv.py --input_image datasets/truck.jpg --input_point 700,375 --embedding_bmodel models/BM1684X/embedding_bmodel/SAM-ViT-B_embedding_fp16_1b.bmodel --decode_bmodel models/BM1684X/decode_bmodel/SAM-ViT-B_decoder_single_mask_fp16_1b.bmodel  --dev_id 0

Results:

Terminal:

terminal

Images: located under results/ in the SAM directory.

r0

r1

r2

2.2.2 Box Input
python3 python/sam_opencv.py --input_image datasets/truck.jpg --input_point 100,300,1700,800 --embedding_bmodel models/BM1684X/embedding_bmodel/SAM-ViT-B_embedding_fp16_1b.bmodel --decode_bmodel models/BM1684X/decode_bmodel/SAM-ViT-B_decoder_multi_mask_fp16_1b.bmodel --dev_id 0

The effect and position are similar to the point input.

box_1

2.2.3 Automatic Segmentation

To use fully automatic mask generation without point and box input, set the input parameter auto to 1 and set --bmodel to the auto bmodel. Steps:

python3 python/sam_opencv.py --input_image datasets/dog.jpg --embedding_bmodel models/BM1684X/embedding_bmodel/SAM-ViT-B_embedding_fp16_1b.bmodel --decode_bmodel models/BM1684X/decode_bmodel/SAM-ViT-B_auto_multi_decoder_fp32_1b.bmodel --dev_id 0 --auto 1 --pred_iou_thresh 0.86

After running, the result images are saved under results/, and the inference time and other information are printed.

auto

t2

3. Web Example

The image files used for interaction are stored under the SAM/web_ui/images directory. The program automatically reads all *.jpg images in that directory and displays the image names in the front-end dropdown.

3.1 Start the Backend Program

The backend program is located in SAM/python/, and the script is named backend.py. This web_ui Python example does not need to be compiled and can be run directly.

3.1.1 Parameter Description
usage: backend.py [--embedding_bmodel EMBEDDING_BMODEL] [--bmodel BMODEL] [--dev_id DEV_ID]

--embedding_bmodel 用于图像压缩(embedding)的bmodel路径；
--bmodel: 用于推理(mask_decode)的bmodel路径；
--dev_id: 用于推理的tpu设备id；
3.1.2 Run Example
cd /data/SAM
python3 python/backend.py --embedding_bmodel models/BM1684X/embedding_bmodel/SAM-ViT-B_embedding_fp16_1b.bmodel --decode_bmodel models/BM1684X/decode_bmodel/SAM-ViT-B_decoder_single_mask_fp16_1b.bmodel --dev_id 0

When the following appears, the backend has started:

t3

3.2 Start the Frontend Service

The frontend program is in /data/SAM/web_ui and can be started with Python.

Keep the backend session window open and open a new session window for the frontend:

cd /data/SAM/web_ui/
python3 -m http.server 8080

Open the PC browser and enter 192.168.49.32:8080 in the address bar to enter the interactive interface. Click the Select image to load... dropdown to choose a preset image. Select Single Point to enter click mode, or Draw BOX to enter box-drawing mode.

ui

3.2.1 Click Mode

In click mode, wait for the image to load successfully, then click the region of interest. Wait 1–2 seconds for the page to draw the mask result.

uip

3.2.2 Box-Drawing Mode

In box-drawing mode, wait for the image to load successfully, then click and drag the mouse to select the region of interest. Wait 1–2 seconds for the page to draw the mask result.

uib

PS: You can check the running status of the backend and frontend in the original terminal; the frontend status can also be checked in the browser developer tools.