SegmentAnything官网demo使用vue+python实现

一、效果&准备工作

1.效果

没啥好说的，低质量复刻SAM官网 https://segment-anything.com/
Github：https://github.com/Gwen0x4c3/segment-anything-vue

需要提一点：所有生成embedding和mask的操作都是python后端做的，计算mask不是onnxruntime-web实现的，前端只负责了把rle编码的mask解码后画到canvas上，会有几十毫秒的网络传输延迟。我不会react和typescript，官网F12里的源代码太难懂了，生成的svg总是与期望的不一样

主页

鼠标移动动态分割（Hover）

throttle了一下，修改代码里的throttle delay，反应更快些，我觉得没必要已经够了，设置的150ms

点选前景背景（Click）

蓝色前景，红色背景，对应clickType分别为1和0

分割（Cut out object）

同官网，分割出该区域需要的最小矩形框部分

分割所有（Everything）

随便做了下，实在做不出官网的效果，可能模型也有问题 ，我用的vit_b，懒得试了，这功能对我来说没卵用

2.准备工作

安装依赖

前端使用了Vue3+ElementPlus（https://element-plus.org/zh-CN/#/zh-CN）+axios+lz-string，npm安装一下。

后端是fastapi（https://fastapi.tiangolo.com/），FastAPI 依赖 Python 3.8 及更高版本。

安装 FastAPI

pip install fastapi

另外我们还需要一个 ASGI 服务器，生产环境可以使用 Uvicorn 或者 Hypercorn：

pip install "uvicorn[standard]"

要用的js文件

@/util/request.js

import axios from "axios";
import { ElMessage } from "element-plus";

axios.interceptors.request.use(
    config => {
        return config;
    },
    error => {
        return Promise.reject(error);
    }
);

axios.interceptors.response.use(
    response => {
        if (response.data.success != null && !response.data.success) {
            return Promise.reject(response.data)
        }
        return response.data;
    },
    error => {
        console.log('error: ', error)
        ElMessage.error(' ');
        return Promise.reject(error);
    }
);

export default axios;

然后在main.js中绑定

import axios from './util/request.js'
axios.defaults.baseURL = 'http://localhost:9000'
axios.defaults.headers.post['Content-Type'] = 'application/x-www-form-urlencoded';
app.config.globalProperties.$http = axios

@/util/throttle.js

function throttle(func, delay) {
    let timer = null; // 定时器变量

    return function() {
        const context = this; // 保存this指向
        const args = arguments; // 保存参数列表

        if (!timer) {
            timer = setTimeout(() => {
                func.apply(context, args); // 调用原始函数并传入上下文和参数
                clearTimeout(timer); // 清除计时器
                timer = null; // 重置计时器为null
            }, delay);
        }
    };
}
export default throttle

@/util/mask_utils.js

/**
 * Parses RLE from compressed string
 * @param {Array<number>} input
 * @returns array of integers
 */
export const rleFrString = (input) => {
    let result = [];
    let charIndex = 0;
    while (charIndex < input.length) {
        let value = 0,
            k = 0,
            more = 1;
        while (more) {
            let c = input.charCodeAt(charIndex) - 48;
            value |= (c & 0x1f) << (5 * k);
            more = c & 0x20;
            charIndex++;
            k++;
            if (!more && c & 0x10) value |= -1 << (5 * k);
        }
        if (result.length > 2) value += result[result.length - 2];
        result.push(value);
    }
    return result;
};

/**
 * Parse RLE to mask array
 * @param rows
 * @param cols
 * @param counts
 * @returns {Uint8Array}
 */
export const decodeRleCounts = ([rows, cols], counts) => {
    let arr = new Uint8Array(rows * cols)
    let i = 0
    let flag = 0
    for (let k of counts) {
        while (k-- > 0) {
            arr[i++] = flag
        }
        flag = (flag + 1) % 2
    }
    return arr
};

/**
 * Parse Everything mode counts array to mask array
 * @param rows
 * @param cols
 * @param counts
 * @returns {Uint8Array}
 */
export const decodeEverythingMask = ([rows, cols], counts) => {
    let arr = new Uint8Array(rows * cols)
    let k = 0;
    for (let i = 0; i < counts.length; i += 2) {
        for (let j = 0; j < counts[i]; j++) {
            arr[k++] = counts[i + 1]
        }
    }
    return arr;
};

/**
 * Get globally unique color in the mask
 * @param category
 * @param colorMap
 * @returns {*}
 */
export const getUniqueColor = (category, colorMap) => {
    // 该种类没有颜色
    if (!colorMap.hasOwnProperty(category)) {
        // 生成唯一的颜色
        while (true) {
            const color = {
                r: Math.floor(Math.random() * 256),
                g: Math.floor(Math.random() * 256),
                b: Math.floor(Math.random() * 256)
            }
            // 检查颜色映射中是否已存在相同的颜色
            const existingColors = Object.values(colorMap);
            const isDuplicateColor = existingColors.some((existingColor) => {
                return color.r === existingColor.r && color.g === existingColor.g && color.b === existingColor.b;
            });
            // 如果不存在相同颜色，结束循环
            if (!isDuplicateColor) {
                colorMap[category] = color;
                break
            }
        }
        console.log("生成唯一颜色", category, colorMap[category])
        return colorMap[category]
    } else {
        return colorMap[category]
    }
}

/**
 * Cut out specific area of image uncovered by mask
 * @param w image's natural width
 * @param h image's natural height
 * @param image source image
 * @param canvas mask canvas
 * @param callback function to solve the image blob
 */
export const cutOutImage = ({w, h}, image, canvas, callback) => {
    const resultCanvas = document.createElement('canvas'),
        resultCtx = resultCanvas.getContext('2d', {willReadFrequently: true}),
        originalCtx = canvas.getContext('2d', {willReadFrequently: true});
    resultCanvas.width = w;
    resultCanvas.height = h;
    resultCtx.drawImage(image, 0, 0, w, h)
    const maskDataArray = originalCtx.getImageData(0, 0, w, h).data;
    const imageData = resultCtx.getImageData(0, 0, w, h);
    const imageDataArray = imageData.data
    // 将mask的部分去掉
    for (let i = 0; i < maskDataArray.length; i += 4) {
        const alpha = maskDataArray[i + 3];
        if (alpha !== 0) { // 不等于0，是mask区域
            imageDataArray[i + 3] = 0;
        }
    }
    // 计算被分割出来的部分的矩形框
    let minX = w;
    let minY = h;
    let maxX = 0;
    let maxY = 0;
    for (let y = 0; y < h; y++) {
        for (let x = 0; x < w; x++) {
            const alpha = imageDataArray[(y * w + x) * 4 + 3];
            if (alpha !== 0) {
                minX = Math.min(minX, x);
                minY = Math.min(minY, y);
                maxX = Math.max(maxX, x);
                maxY = Math.max(maxY, y);
            }
        }
    }
    const width = maxX - minX + 1;
    const height = maxY - minY + 1;
    const startX = minX;
    const startY = minY;
    resultCtx.putImageData(imageData, 0, 0)
    // 创建一个新的canvas来存储特定区域的图像
    const croppedCanvas = document.createElement("canvas");
    const croppedContext = croppedCanvas.getContext("2d");
    croppedCanvas.width = width;
    croppedCanvas.height = height;
    // 将特定区域绘制到新canvas上
    croppedContext.drawImage(resultCanvas, startX, startY, width, height, 0, 0, width, height);
    croppedCanvas.toBlob(blob => {
        if (callback) {
            callback(blob)
        }
    }, "image/png");
}

/**
 * Cut out specific area of image covered by target color mask
 * PS: 我写的这代码有问题，比较color的时候tmd明明mask canvas中有这个颜色，
 * 就是说不存在这颜色，所以不用这个函数，改成下面的了
 * @param w image's natural width
 * @param h image's natural height
 * @param image source image
 * @param canvas mask canvas
 * @param color target color
 * @param callback function to solve the image blob
 */
export const cutOutImageWithMaskColor = ({w, h}, image, canvas, color, callback) => {
    const resultCanvas = document.createElement('canvas'),
        resultCtx = resultCanvas.getContext('2d', {willReadFrequently: true}),
        originalCtx = canvas.getContext('2d', {willReadFrequently: true});
    resultCanvas.width = w;
    resultCanvas.height = h;
    resultCtx.drawImage(image, 0, 0, w, h)
    const maskDataArray = originalCtx.getImageData(0, 0, w, h).data;
    const imageData = resultCtx.getImageData(0, 0, w, h);
    const imageDataArray = imageData.data

    let find = false

    // 比较mask的color和目标color
    for (let i = 0; i < maskDataArray.length; i += 4) {
        const r = maskDataArray[i],
            g = maskDataArray[i + 1],
            b = maskDataArray[i + 2];
        if (r != color.r || g != color.g || b != color.b) { // 颜色与目标颜色不相同，是mask区域
            // 设置alpha为0
            imageDataArray[i + 3] = 0;
        } else {
            find = true
        }
    }
    // 计算被分割出来的部分的矩形框
    let minX = w;
    let minY = h;
    let maxX = 0;
    let maxY = 0;
    for (let y = 0; y < h; y++) {
        for (let x = 0; x < w; x++) {
            const alpha = imageDataArray[(y * w + x) * 4 + 3];
            if (alpha !== 0) {
                minX = Math.min(minX, x);
                minY = Math.min(minY, y);
                maxX = Math.max(maxX, x);
                maxY = Math.max(maxY, y);
            }
        }
    }
    const width = maxX - minX + 1;
    const height = maxY - minY + 1;
    const startX = minX;
    const startY = minY;
    // console.log(`矩形宽度：${width}`);
    // console.log(`矩形高度：${height}`);
    // console.log(`起点坐标：(${startX}, ${startY})`);
    resultCtx.putImageData(imageData, 0, 0)
    // 创建一个新的canvas来存储特定区域的图像
    const croppedCanvas = document.createElement("canvas");
    const croppedContext = croppedCanvas.getContext("2d");
    croppedCanvas.width = width;
    croppedCanvas.height = height;
    // 将特定区域绘制到新canvas上
    croppedContext.drawImage(resultCanvas, startX, startY, width, height, 0, 0, width, height);
    croppedCanvas.toBlob(blob => {
        if (callback) {
            callback(blob)
        }
    }, "image/png");
}

/**
 * Cut out specific area whose category is target category
 * @param w image's natural width
 * @param h image's natural height
 * @param image source image
 * @param arr original mask array that stores all pixel's category
 * @param category target category
 * @param callback function to solve the image blob
 */
export const cutOutImageWithCategory = ({w, h}, image, arr, category, callback) => {
    const resultCanvas = document.createElement('canvas'),
        resultCtx = resultCanvas.getContext('2d', {willReadFrequently: true});
    resultCanvas.width = w;
    resultCanvas.height = h;
    resultCtx.drawImage(image, 0, 0, w, h)
    const imageData = resultCtx.getImageData(0, 0, w, h);
    const imageDataArray = imageData.data
    // 比较mask的类别和目标类别
    let i = 0
    for(let y = 0; y < h; y++){
        for(let x = 0; x < w; x++){
            if (category != arr[i++]) { // 类别不相同，是mask区域
                // 设置alpha为0
                imageDataArray[3 + (w * y + x) * 4] = 0;
            }
        }
    }
    // 计算被分割出来的部分的矩形框
    let minX = w;
    let minY = h;
    let maxX = 0;
    let maxY = 0;
    for (let y = 0; y < h; y++) {
        for (let x = 0; x < w; x++) {
            const alpha = imageDataArray[(y * w + x) * 4 + 3];
            if (alpha !== 0) {
                minX = Math.min(minX, x);
                minY = Math.min(minY, y);
                maxX = Math.max(maxX, x);
                maxY = Math.max(maxY, y);
            }
        }
    }
    const width = maxX - minX + 1;
    const height = maxY - minY + 1;
    const startX = minX;
    const startY = minY;
    resultCtx.putImageData(imageData, 0, 0)
    // 创建一个新的canvas来存储特定区域的图像
    const croppedCanvas = document.createElement("canvas");
    const croppedContext = croppedCanvas.getContext("2d");
    croppedCanvas.width = width;
    croppedCanvas.height = height;
    // 将特定区域绘制到新canvas上
    croppedContext.drawImage(resultCanvas, startX, startY, width, height, 0, 0, width, height);
    croppedCanvas.toBlob(blob => {
        if (callback) {
            callback(blob)
        }
    }, "image/png");
}

二、后端代码

1.SAM下载

首先从github上下载SAM的代码https://github.com/facebookresearch/segment-anything

然后下载模型文件，保存到项目根目录/checkpoints中，

• default or vit_h: ViT-H SAM model.
• vit_l: ViT-L SAM model.
• vit_b: ViT-B SAM model.

2.后端代码

在项目根目录下创建main.py

main.py

import os
import time

from PIL import Image
import numpy as np
import io
import base64
from segment_anything import SamPredictor, SamAutomaticMaskGenerator, sam_model_registry
from pycocotools import mask as mask_utils
import lzstring


def init():
    # your model path
    checkpoint = "checkpoints/sam_vit_b_01ec64.pth"
    model_type = "vit_b"
    sam = sam_model_registry[model_type](checkpoint=checkpoint)
    sam.to(device='cuda')
    predictor = SamPredictor(sam)
    mask_generator = SamAutomaticMaskGenerator(sam)
    return predictor, mask_generator


predictor, mask_generator = init()

from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware

app = FastAPI()
app.add_middleware(
    CORSMiddleware,
    allow_origins="*",
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

last_image = ""
last_logit = None


@app.post("/segment")
def process_image(body: dict):
    global last_image, last_logit
    print("start processing image", time.time())
    path = body["path"]
    is_first_segment = False
    # 看上次分割的图片是不是该图片
    if path != last_image:  # 不是该图片，重新生成图像embedding
        pil_image = Image.open(path)
        np_image = np.array(pil_image)
        predictor.set_image(np_image)
        last_image = path
        is_first_segment = True
        print("第一次识别该图片，获取embedding")
    # 获取mask
    clicks = body["clicks"]
    input_points = []
    input_labels = []
    for click in clicks:
        input_points.append([click["x"], click["y"]])
        input_labels.append(click["clickType"])
    print("input_points:{}, input_labels:{}".format(input_points, input_labels))
    input_points = np.array(input_points)
    input_labels = np.array(input_labels)
    masks, scores, logits = predictor.predict(
        point_coords=input_points,
        point_labels=input_labels,
        mask_input=last_logit[None, :, :] if not is_first_segment else None,
        multimask_output=is_first_segment  # 第一次产生3个结果，选择最优的
    )
    # 设置mask_input，为下一次做准备
    best = np.argmax(scores)
    last_logit = logits[best, :, :]
    masks = masks[best, :, :]
    # print(mask_utils.encode(np.asfortranarray(masks))["counts"])
    # numpy_array = np.frombuffer(mask_utils.encode(np.asfortranarray(masks))["counts"], dtype=np.uint8)
    # print("Uint8Array([" + ", ".join(map(str, numpy_array)) + "])")
    source_mask = mask_utils.encode(np.asfortranarray(masks))["counts"].decode("utf-8")
    # print(source_mask)
    lzs = lzstring.LZString()
    encoded = lzs.compressToEncodedURIComponent(source_mask)
    print("process finished", time.time())
    return {"shape": masks.shape, "mask": encoded}


@app.get("/everything")
def segment_everything(path: str):
    start_time = time.time()
    print("start segment_everything", start_time)
    pil_image = Image.open(path)
    np_image = np.array(pil_image)
    masks = mask_generator.generate(np_image)
    sorted_anns = sorted(masks, key=(lambda x: x['area']), reverse=True)
    img = np.zeros((sorted_anns[0]['segmentation'].shape[0], sorted_anns[0]['segmentation'].shape[1]), dtype=np.uint8)
    for idx, ann in enumerate(sorted_anns, 0):
        img[ann['segmentation']] = idx
    #看一下mask是什么样
    #plt.figure(figsize=(10,10))
	#plt.imshow(img) 
	#plt.show()
    # 压缩数组
    result = my_compress(img)
    end_time = time.time()
    print("finished segment_everything", end_time)
    print("time cost", end_time - start_time)
    return {"shape": img.shape, "mask": result}


@app.get('/automatic_masks')
def automatic_masks(path: str):
    pil_image = Image.open(path)
    np_image = np.array(pil_image)
    mask = mask_generator.generate(np_image)
    sorted_anns = sorted(mask, key=(lambda x: x['area']), reverse=True)
    lzs = lzstring.LZString()
    res = []
    for ann in sorted_anns:
        m = ann['segmentation']
        source_mask = mask_utils.encode(m)['counts'].decode("utf-8")
        encoded = lzs.compressToEncodedURIComponent(source_mask)
        r = {
            "encodedMask": encoded,
            "point_coord": ann['point_coords'][0],
        }
        res.append(r)
    return res


# 就是将连续的数字统计个数，然后把[个数，数字]放到result中，类似rle算法
# 比如[[1,1,1,2,3,2,2,4,4],[3,3,4...]]
# result是[3,1,  1,2,  1,3,  2,2,  2,4,  2,3,...]
def my_compress(img):
    result = []
    last_pixel = img[0][0]
    count = 0
    for line in img:
        for pixel in line:
            if pixel == last_pixel:
                count += 1
            else:
                result.append(count)
                result.append(int(last_pixel))
                last_pixel = pixel
                count = 1
    result.append(count)
    result.append(int(last_pixel))
    return result

3.原神启动

在cmd或者pycharm终端，cd到项目根目录下，输入uvicorn main:app --port 8006，启动服务器

三、前端代码

1.页面代码

template

<template>
  <div class="segment-container">
    <ElScrollbar class="tool-box">
      <div class="image-section">
        <div class="title">
          <div style="padding-left:15px">
            <el-icon><Picture /></el-icon><span style="font-size: 18px;font-weight: 550;">展示图像</span>
            <el-icon class="header-icon"></el-icon>
          </div>
        </div>
        <ElScrollbar height="350px">
          <div v-if="cutOuts.length === 0">
            <p>未进行抠图</p>
            <p>左键设置区域为前景</p>
            <p>右键设置区域为背景</p>
          </div>
          <img v-for="src in cutOuts" :src="src" alt="加载中"
               @click="openInNewTab(src)"/>
        </ElScrollbar>
      </div>
      <div class="options-section">
        <span class="option" @click="reset">重置</span>
        <span :class="'option'+(clicks.length===0?' disabled':'')" @click="undo">撤销</span>
        <span :class="'option'+(clickHistory.length===0?' disabled':'')" @click="redo">恢复</span>
      </div>
      <button :class="'segmentation-button'+(lock||clicks.length===0?' disabled':'')"
              @click="cutImage">分割</button>
      <button :class="'segmentation-button'+(lock||isEverything?' disabled':'')"
              @click="segmentEverything">分割所有</button>
    </ElScrollbar>
    <div class="segment-box">
      <div class="segment-wrapper" :style="{'left': left + 'px'}">
        <img v-show="path" id="segment-image" :src="url" :style="{width:w, height:h}" alt="加载失败" crossorigin="anonymous"
             @mousedown="handleMouseDown" @mouseenter="canvasVisible = true"
             @mouseout="() => {if (!this.clicks.length&&!this.isEverything) this.canvasVisible = false}"/>
        <canvas v-show="path && canvasVisible" id="segment-canvas" :width="originalSize.w" :height="originalSize.h"></canvas>
        <div id="point-box" :style="{width:w, height:h}"></div>
      </div>

    </div>
  </div>
</template>

script

<script>
import throttle from "@/util/throttle";
import LZString from "lz-string";
import {
  rleFrString,
  decodeRleCounts,
  decodeEverythingMask,
  getUniqueColor,
  cutOutImage,
  cutOutImageWithMaskColor, cutOutImageWithCategory
} from "@/util/mask_utils";
import {ElCollapse, ElCollapseItem, ElScrollbar} from "element-plus";
import {Picture} from '@element-plus/icons-vue'
export default {
  name: "segment",
  components: {
    ElCollapse, ElCollapseItem, ElScrollbar, Picture
  },
  data() {
    return {
      image: null,
      clicks: [],
      clickHistory: [],
      originalSize: {w: 0, h: 0},
      w: 0,
      h: 0,
      left: 0,
      scale: 1,
      url: null, // url用来设置成img的src展示
      path: null, // path是该图片在文件系统中的绝对路径
      loading: false,
      lock: false,
      canvasVisible: true,
      // cutOuts: ['http://localhost:9000/p/2024/01/19/112ce48bd76e47c7900863a3a0147853.jpg', 'http://localhost:9000/p/2024/01/19/112ce48bd76e47c7900863a3a0147853.jpg'],
      cutOuts: [],
      isEverything: false
    }
  },
  mounted() {
    this.init()
  },
  methods: {
    async init() {
      this.loading = true
      // 从路由获取id
      let id = this.$route.params.id
      if (!id) {
        this.$message.error('未选择图片')
        return
      }
      this.id = id
      // 获取图片信息
      try {
        const { path, url } = await this.getPathAndUrl()
        this.loadImage(path, url)
      } catch (e) {
        console.error(e)
        this.$message.error(e)
      }
    },
    async getPathAndUrl() {
      let res = await this.$http.get("/photo/path/" + this.id)
      console.log(res)
      return res.data
    },
    loadImage(path, url) {
      let image = new Image();
      image.src = this.$photo_base + url;
      image.onload = () => {
        let w = image.width, h = image.height
        let nw, nh
        let body = document.querySelector('.segment-box')
        let mw = body.clientWidth, mh = body.clientHeight
        let ratio = w / h
        if (ratio * mh > mw) {
          nw = mw
          nh = mw / ratio
        } else {
          nh = mh
          nw = ratio * mh
        }
        this.originalSize = {w, h}
        nw = parseInt(nw)
        nh = parseInt(nh)
        this.w = nw + 'px'
        this.h = nh + 'px'
        this.left = (mw - nw) / 2
        this.scale = nw / w
        this.url = this.$photo_base + url
        this.path = path
        console.log((this.scale > 1 ? '放大' : '缩小') + w + ' --> ' + nw)
        const img = document.getElementById('segment-image')
        img.addEventListener('contextmenu', e => e.preventDefault())
        img.addEventListener('mousemove', throttle(this.handleMouseMove, 150))
        const canvas = document.getElementById('segment-canvas')
        canvas.style.transform = `scale(${this.scale})`
      }
    },
    getClick(e) {
      let click = {
        x: e.offsetX,
        y: e.offsetY,
      }
      const imageScale = this.scale
      click.x /= imageScale;
      click.y /= imageScale;
      if(e.which === 3){ // 右键
        click.clickType = 0
      } else if(e.which === 1 || e.which === 0) { // 左键
        click.clickType = 1
      }
      return click
    },
    handleMouseMove(e) {
      if (this.isEverything) { // 分割所有模式，返回
        return;
      }
      if (this.clicks.length !== 0) { // 选择了点
        return;
      }
      if (this.lock) {
        return;
      }
      this.lock = true;
      let click = this.getClick(e);
      requestIdleCallback(() => {
        this.getMask([click])
      })
    },
    handleMouseDown(e) {
      e.preventDefault();
      e.stopPropagation();
      if (e.button === 1) {
        return;
      }
      // 如果是“分割所有”模式，返回
      if (this.isEverything) {
        return;
      }
      if (this.lock) {
        return;
      }
      this.lock = true
      let click = this.getClick(e);
      this.placePoint(e.offsetX, e.offsetY, click.clickType)
      this.clicks.push(click);
      requestIdleCallback(() => {
        this.getMask()
      })
    },
    placePoint(x, y, clickType) {
      let box = document.getElementById('point-box')
      let point = document.createElement('div')
      point.className = 'segment-point' + (clickType ? '' : ' negative')
      point.style = `position: absolute;
                      width: 10px;
                      height: 10px;
                      border-radius: 50%;
                      background-color: ${clickType?'#409EFF':'#F56C6C '};
                      left: ${x-5}px;
                      top: ${y-5}px`
      // 点的id是在clicks数组中的下标索引
      point.id = 'point-' + this.clicks.length
      box.appendChild(point)
    },
    removePoint(i) {
      const selector = 'point-' + i
      let point = document.getElementById(selector)
      if (point != null) {
        point.remove()
      }
    },
    getMask(clicks) {
      // 如果clicks为空，则是mouse move产生的click
      if (clicks == null) {
        clicks = this.clicks
      }
      const data = {
        path: this.path,
        clicks: clicks
      }
      console.log(data)
      this.$http.post('http://localhost:8006/segment', data, {
        headers: {
          "Content-Type": "application/json"
        }
      }).then(res => {
        const shape = res.shape
        const maskenc = LZString.decompressFromEncodedURIComponent(res.mask);
        const decoded = rleFrString(maskenc)
        this.drawCanvas(shape, decodeRleCounts(shape, decoded))
        this.lock = false
      }).catch(err => {
        console.error(err)
        this.$message.error("生成失败")
        this.lock = false
      })
    },
    segmentEverything() {
      if (this.isEverything) { // 上一次刚点过了
        return;
      }
      if (this.lock) {
        return;
      }
      this.lock = true
      this.reset()
      this.isEverything = true
      this.canvasVisible = true
      this.$http.get("http://localhost:8006/everything?path=" + this.path).then(res => {
        const shape = res.shape
        const counts = res.mask
        this.drawEverythingCanvas(shape, decodeEverythingMask(shape, counts))
      }).catch(err => {
        console.error(err)
        this.$message.error("生成失败")
      })
    },
    drawCanvas(shape, arr) {
      let height = shape[0],
          width = shape[1]
      console.log("height: ", height, " width: ", width)
      let canvas = document.getElementById('segment-canvas'),
          canvasCtx = canvas.getContext("2d"),
          imgData = canvasCtx.getImageData(0, 0, width, height),
          pixelData = imgData.data
      let i = 0
      for(let x = 0; x < width; x++){
        for(let y = 0; y < height; y++){
          if (arr[i++] === 0) { // 如果是0，是背景，遮住
            pixelData[0 + (width * y + x) * 4] = 40;
            pixelData[1 + (width * y + x) * 4] = 40;
            pixelData[2 + (width * y + x) * 4] = 40;
            pixelData[3 + (width * y + x) * 4] = 190;
          } else {
            pixelData[3 + (width * y + x) * 4] = 0;
          }
        }
      }
      canvasCtx.putImageData(imgData, 0, 0)
    },
    drawEverythingCanvas(shape, arr) {
      const height = shape[0],
          width = shape[1]
      console.log("height: ", height, " width: ", width)
      let canvas = document.getElementById('segment-canvas'),
          canvasCtx = canvas.getContext("2d"),
          imgData = canvasCtx.getImageData(0, 0, width, height),
          pixelData = imgData.data;
      const colorMap = {}
      let i = 0
      for(let y = 0; y < height; y++){
        for(let x = 0; x < width; x++){
          const category = arr[i++]
          const color = getUniqueColor(category, colorMap)
          pixelData[0 + (width * y + x) * 4] = color.r;
          pixelData[1 + (width * y + x) * 4] = color.g;
          pixelData[2 + (width * y + x) * 4] = color.b;
          pixelData[3 + (width * y + x) * 4] = 150;
        }
      }
      // 显示在图片上
      canvasCtx.putImageData(imgData, 0, 0)
      // 开始分割每一个mask的图片
      const image = document.getElementById('segment-image')
      Object.keys(colorMap).forEach(category => {
        cutOutImageWithCategory(this.originalSize, image, arr, category, blob => {
          const url = URL.createObjectURL(blob);
          this.cutOuts = [url, ...this.cutOuts]
        })
      })
    },
    reset() {
      for (let i = 0; i < this.clicks.length; i++) {
        this.removePoint(i)
      }
      this.clicks = []
      this.clickHistory = []
      this.isEverything = false
      this.clearCanvas()
    },
    undo() {
      if (this.clicks.length === 0)
        return
      const idx = this.clicks.length - 1
      const click = this.clicks[idx]
      this.clickHistory.push(click)
      this.clicks.splice(idx, 1)
      this.removePoint(idx)
      if (this.clicks.length) {
        this.getMask()
      } else {
        this.clearCanvas()
      }
    },
    redo() {
      if (this.clickHistory.length === 0)
        return
      const idx = this.clickHistory.length - 1
      const click = this.clickHistory[idx]
      console.log(this.clicks, this.clickHistory, click)
      this.placePoint(click.x * this.scale, click.y * this.scale, click.clickType)
      this.clicks.push(click)
      this.clickHistory.splice(idx, 1)
      this.getMask()
    },
    clearCanvas() {
      let canvas = document.getElementById('segment-canvas')
      canvas.getContext('2d').clearRect(0, 0, canvas.width, canvas.height)
    },
    cutImage() {
      if (this.lock || this.clicks.length === 0) {
        return;
      }
      const canvas = document.getElementById('segment-canvas'),
          image = document.getElementById('segment-image')
      const {w, h} = this.originalSize
      cutOutImage(this.originalSize, image, canvas, blob => {
        const url = URL.createObjectURL(blob);
        this.cutOuts = [url, ...this.cutOuts]
        // 不需要之后用下面的清除文件
        // URL.revokeObjectURL(url);
      })
    },
    openInNewTab(src) {
      window.open(src, '_blank')
    }
  }
}
</script>

style

<style scoped lang="scss">
.segment-container {
  position: relative;
}

.tool-box {
  position: absolute;
  left: 20px;
  top: 20px;
  width: 200px;
  height: 600px;
  border-radius: 20px;
  //background: pink;
  overflow: auto;
  box-shadow: 0 0 5px rgb(150, 150, 150);
  box-sizing: border-box;
  padding: 10px;

  .image-section {
    height: fit-content;
    width: 100%;
    .title {
      height: 48px;
      line-height: 48px;
      border-bottom: 1px solid lightgray;
      margin-bottom: 15px;
    }
  }

  .image-section img {
    max-width: 85%;
    max-height: 140px;
    margin: 10px auto;
    padding: 10px;
    box-sizing: border-box;
    object-fit: contain;
    display: block;
    transition: .3s;
    cursor: pointer;
  }
  .image-section img:hover {
    background: rgba(0, 30, 160, 0.3);
  }

  .image-section p {
    text-align: center;
  }

  .options-section {
    margin-top: 5px;
    display: flex;
    justify-content: space-between;
    align-items: center;
    padding: 10px;
    box-sizing: border-box;
    border: 3px solid lightgray;
    border-radius: 20px;
  }
  .options-section:hover {
    border: 3px solid #59ACFF;
  }

  .option {
    font-size: 15px;
    padding: 5px 10px;
    cursor: pointer;
  }
  .option:hover {
    color: #59ACFF;
  }
  .option.disabled {
    color: gray;
    cursor: not-allowed;
  }

  .segmentation-button {
    margin-top: 5px;
    width: 100%;
    height: 40px;
    background-color: white;
    color: rgb(40, 40, 40);
    font-size: 17px;
    cursor: pointer;
    border: 3px solid lightgray;
    border-radius: 20px;
  }
  .segmentation-button:hover {
    border: 3px solid #59ACFF;
  }
  .segmentation-button.disabled {
    color: lightgray;
    cursor: not-allowed;
  }
}

.segment-box {
  position: relative;
  margin-left: calc(220px);
  width: calc(100% - 220px);
  height: calc(100vh - 80px);
  //background: #42b983;
  .segment-wrapper {
    position: absolute;
    left: 0;
    top: 0;
  }
  #segment-canvas {
    position: absolute;
    left: 0;
    top: 0;
    pointer-events: none;
    transform-origin: left top;
    z-index: 1;
  }
  #point-box {
    position: absolute;
    left: 0;
    top: 0;
    z-index: 2;
    pointer-events: none;
  }
  .segment-point {
    position: absolute;
    width: 10px;
    height: 10px;
    border-radius: 50%;
    background-color: #409EFF;
  }
  .segment-point.negative {
    background-color: #F56C6C;
  }
}
</style>

2.代码说明

• 本项目没做上传图片分割，就是简单的选择本地图片分割，data中url是img的src，path是绝对路径用来传给python后端进行分割，我是从我项目的系统获取的，请自行修改代码成你的图片路径，如src: “/assets/test.jpg”, path:“D:/project/segment/assets/test.jpg”
• 由于pycocotools的rle encode是从上到下进行统计连续的0和1，为了方便，我在【@/util/mask_utils.js:decodeRleCounts】解码Click点选产生的mask时将（H，W）的矩阵转成了（W，H）顺序存储的Uint8array；而在Everything分割所有时，我没有使用pycocotools的encode，而是main.py中的my_compress函数编码的，是从左到右进行压缩，因此矩阵解码后仍然是（H，W）的矩阵，所以在drawCanvas和drawEverythingCanvas中的二层循环xy的顺序不一样，我实在懒得改了，就这样就可以了。

关于上面所提rle，可以在项目根目录/notebooks/predictor_example.ipynb中产生mask的位置添加代码自行观察他编码的rle，他只支持矩阵元素为0或1，result的第一个位置是0的个数，不管矩阵是不是0开头。

• [0,0,1,1,0,1,0]，rle counts是[2（两个0）, 2（两个1）, 1（一个0）, 1（一个1）, 1（一个0）]；

• [1,1,1,1,1,0]，rle counts是[0（零个0）,5（五个1）,1（一个0）]

def decode_rle(rle_string): # 这是将pycocotools的counts编码的字符串转成counts数组，而非转成原矩阵
    result = []
    char_index = 0
    
    while char_index < len(rle_string):
        value = 0
        k = 0
        more = 1
        
        while more:
            c = ord(rle_string[char_index]) - 48
            value |= (c & 0x1f) << (5 * k)
            more = c & 0x20
            char_index += 1
            k += 1
            if not more and c & 0x10:
                value |= -1 << (5 * k)
        
        if len(result) > 2:
            value += result[-2]
        result.append(value)
    return result

from pycocotools import mask as mask_utils
import numpy as np
mask = np.array([[1,1,0,1,1,0],[1,1,1,1,1,1],[0,1,1,1,0,0],[1,1,1,1,1,1]])
mask = np.asfortranarray(mask, dtype=np.uint8)
print("原mask：\n{}".format(mask))
res = mask_utils.encode(mask)
print("encode：{}".format(res))
print("rle counts：{}".format(decode_rle(res["counts"].decode("utf-8"))))
# 转置后好看
print("转置：{}".format(mask.transpose()))
# flatten后更好看
print("flatten：{}".format(mask.transpose().flatten()))
#numpy_array = np.frombuffer(res["counts"], dtype=np.uint8)
# 打印numpy数组作为uint8array的格式
#print("Uint8Array([" + ", ".join(map(str, numpy_array)) + "])")

输出：