Labelme加载AI(Segment-Anything)模型进行图像标注

Make_magic 2024-07-10 12:01:13 阅读 66

labelme是使用python写的基于QT的跨平台图像标注工具,可用来标注分类、检测、分割、关键点等常见的视觉任务,支持VOC格式和COCO等的导出,代码简单易读,是非常利用上手的良心工具。

在这里插入图片描述

第一步:

下载源码进行安装。

<code>git clone https://github.com/wkentaro/labelme.git

cd labelme

pip install -e .

第二步:

找到源码所在路径进行修改。

(1)打开labelme/labelme/ai/init.py,源码如下:

MODELS = [

Model(

name="Segment-Anything (speed)",code>

encoder_weight=Weight(

url="https://github.com/wkentaro/labelme/releases/download/sam-20230416/sam_vit_b_01ec64.quantized.encoder.onnx", # NOQAcode>

md5="80fd8d0ab6c6ae8cb7b3bd5f368a752c",code>

),

decoder_weight=Weight(

url="https://github.com/wkentaro/labelme/releases/download/sam-20230416/sam_vit_b_01ec64.quantized.decoder.onnx", # NOQAcode>

md5="4253558be238c15fc265a7a876aaec82",code>

),

),

Model(

name="Segment-Anything (balanced)",code>

encoder_weight=Weight(

url="https://github.com/wkentaro/labelme/releases/download/sam-20230416/sam_vit_l_0b3195.quantized.encoder.onnx", # NOQAcode>

md5="080004dc9992724d360a49399d1ee24b",code>

),

decoder_weight=Weight(

url="https://github.com/wkentaro/labelme/releases/download/sam-20230416/sam_vit_l_0b3195.quantized.decoder.onnx", # NOQAcode>

md5="851b7faac91e8e23940ee1294231d5c7",code>

),

),

Model(

name="Segment-Anything (accuracy)",code>

encoder_weight=Weight(

url="https://github.com/wkentaro/labelme/releases/download/sam-20230416/sam_vit_h_4b8939.quantized.encoder.onnx", # NOQAcode>

md5="958b5710d25b198d765fb6b94798f49e",code>

),

decoder_weight=Weight(

url="https://github.com/wkentaro/labelme/releases/download/sam-20230416/sam_vit_h_4b8939.quantized.decoder.onnx", # NOQAcode>

md5="a997a408347aa081b17a3ffff9f42a80",code>

),

),

]

(2)在labelme/labelme/文件夹下自建一个文件夹model_file。

(3)依次输入以下几个网址下载onnx到model_file文件目录。

https://github.com/wkentaro/labelme/releases/download/sam-20230416/sam_vit_b_01ec64.quantized.encoder.onnx

https://github.com/wkentaro/labelme/releases/download/sam-20230416/sam_vit_b_01ec64.quantized.decoder.onnx

https://github.com/wkentaro/labelme/releases/download/sam-20230416/sam_vit_l_0b3195.quantized.encoder.onnx

https://github.com/wkentaro/labelme/releases/download/sam-20230416/sam_vit_l_0b3195.quantized.decoder.onnx

https://github.com/wkentaro/labelme/releases/download/sam-20230416/sam_vit_h_4b8939.quantized.encoder.onnx

https://github.com/wkentaro/labelme/releases/download/sam-20230416/sam_vit_h_4b8939.quantized.decoder.onnx

在这里插入图片描述

(4)修改labelme/labelme/ai/init.py,代码如下:

<code>import collections

from .models.segment_anything import SegmentAnythingModel # NOQA

Model = collections.namedtuple(

"Model", ["name", "encoder_weight", "decoder_weight"]

)

Weight = collections.namedtuple("Weight", ["url", "md5"])

# MODELS = [

# Model(

# name="Segment-Anything (speed)",code>

# encoder_weight=Weight(

# url="https://github.com/wkentaro/labelme/releases/download/sam-20230416/sam_vit_b_01ec64.quantized.encoder.onnx", # NOQAcode>

# md5="80fd8d0ab6c6ae8cb7b3bd5f368a752c",code>

# ),

# decoder_weight=Weight(

# url="https://github.com/wkentaro/labelme/releases/download/sam-20230416/sam_vit_b_01ec64.quantized.decoder.onnx", # NOQAcode>

# md5="4253558be238c15fc265a7a876aaec82",code>

# ),

# ),

# Model(

# name="Segment-Anything (balanced)",code>

# encoder_weight=Weight(

# url="https://github.com/wkentaro/labelme/releases/download/sam-20230416/sam_vit_l_0b3195.quantized.encoder.onnx", # NOQAcode>

# md5="080004dc9992724d360a49399d1ee24b",code>

# ),

# decoder_weight=Weight(

# url="https://github.com/wkentaro/labelme/releases/download/sam-20230416/sam_vit_l_0b3195.quantized.decoder.onnx", # NOQAcode>

# md5="851b7faac91e8e23940ee1294231d5c7",code>

# ),

# ),

# Model(

# name="Segment-Anything (accuracy)",code>

# encoder_weight=Weight(

# url="https://github.com/wkentaro/labelme/releases/download/sam-20230416/sam_vit_h_4b8939.quantized.encoder.onnx", # NOQAcode>

# md5="958b5710d25b198d765fb6b94798f49e",code>

# ),

# decoder_weight=Weight(

# url="https://github.com/wkentaro/labelme/releases/download/sam-20230416/sam_vit_h_4b8939.quantized.decoder.onnx", # NOQAcode>

# md5="a997a408347aa081b17a3ffff9f42a80",code>

# ),

# ),

# ]

MODELS = [

Model(

name="Segment-Anything (speed)",code>

encoder_weight=Weight(

url="E:\labelme\labelme\model_file\sam_vit_b_01ec64.quantized.encoder.onnx", # NOQAcode>

md5="80fd8d0ab6c6ae8cb7b3bd5f368a752c",code>

),

decoder_weight=Weight(

url="E:\labelme\labelme\model_file\sam_vit_b_01ec64.quantized.decoder.onnx", # NOQAcode>

md5="4253558be238c15fc265a7a876aaec82",code>

),

),

Model(

name="Segment-Anything (balanced)",code>

encoder_weight=Weight(

url="E:\labelme\labelme\model_file\sam_vit_l_0b3195.quantized.encoder.onnx", # NOQAcode>

md5="080004dc9992724d360a49399d1ee24b",code>

),

decoder_weight=Weight(

url="E:\labelme\labelme\model_file\sam_vit_l_0b3195.quantized.decoder.onnx", # NOQAcode>

md5="851b7faac91e8e23940ee1294231d5c7",code>

),

),

Model(

name="Segment-Anything (accuracy)",code>

encoder_weight=Weight(

url="E:\labelme\labelme\model_file\sam_vit_h_4b8939.quantized.encoder.onnx", # NOQAcode>

md5="958b5710d25b198d765fb6b94798f49e",code>

),

decoder_weight=Weight(

url="E:\labelme\labelme\model_file\sam_vit_h_4b8939.quantized.decoder.onnx", # NOQAcode>

md5="a997a408347aa081b17a3ffff9f42a80",code>

),

),

]

(5)修改labelme/labelme/widgets/canvas.py,代码如下:

def initializeAiModel(self, name):

if name not in [model.name for model in labelme.ai.MODELS]:

raise ValueError("Unsupported ai model: %s" % name)

model = [model for model in labelme.ai.MODELS if model.name == name][0]

if self._ai_model is not None and self._ai_model.name == model.name:

logger.debug("AI model is already initialized: %r" % model.name)

else:

logger.debug("Initializing AI model: %r" % model.name)

self._ai_model = labelme.ai.SegmentAnythingModel(

name=model.name,

# encoder_path=gdown.cached_download(

# url=model.encoder_weight.url,

# md5=model.encoder_weight.md5,

# ),

# decoder_path=gdown.cached_download(

# url=model.decoder_weight.url,

# md5=model.decoder_weight.md5,

# ),

encoder_path=model.encoder_weight.url,

decoder_path=model.decoder_weight.url,

)

self._ai_model.set_image(

image=labelme.utils.img_qt_to_arr(self.pixmap.toImage())

)

第三步:

启动labelme

cd labelme

labelme

在这里插入图片描述



声明

本文内容仅代表作者观点,或转载于其他网站,本站不以此文作为商业用途
如有涉及侵权,请联系本站进行删除
转载本站原创文章,请注明来源及作者。