1. Anylabeling自动标记
使用以下代码
import numpy as np
import torch
import matplotlib.pyplot as plt
import cv2
import json
import sys
import os
from segment_anything import sam_model_registry, SamAutomaticMaskGenerator
def segment(imgPath):
# 通过opencv图取图像
print("开始读图像")
image = cv2.imread(imgPath)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
# 添加当前系统路径,添加模型文件路径
sys.path.append("..")
sam_checkpoint = "sam_vit_l_0b3195.pth"
model_type = "vit_l"
#设置运行推理的设备
device = "cuda"
# 创建sam模型推理对象
sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
sam.to(device=device)
mask_generator = SamAutomaticMaskGenerator(sam)
masks = mask_generator.generate(image)
# 给分割出来的物体上色,显示分割效果
show_anns(masks, imgPath)
torch.cuda.empty_cache()
def show_anns(anns, imgPath):
if len(anns) == 0:
return
# 对检测结果的字典对象进行排序
sorted_anns = sorted(anns, key=(lambda x: x['area']), reverse=True)
ax = plt.gca()
ax.set_autoscale_on(False)
img = np.ones((sorted_anns[0]['segmentation'].shape[0], sorted_anns[0]['segmentation'].shape[1], 4))
img[:,:,3] = 0
shapes = []
for ann in sorted_anns:
# 过滤面积比较小的物体
if ann['area'] >=10:
# 创建labelme格式
tempData = {
"label": "air", ###### 需要修改标签名 ###########
"points": [],
"group_id": None,
"shape_type": "polygon",
"flags": {
}
}
# 获取分割物体掩膜
m = ann['segmentation']
# 找出物体轮廓
objImg = np.zeros((m.shape[0], m.shape[1], 1), np.uint8)
objImg[m] = 255
contours, hierarchy = cv2.findContours(objImg, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
# 找出轮廓最大的
max_area = 0
maxIndex = 0
for i in range(0, len(contours)):
area = cv2.contourArea(contours[i])
if area >= max_area:
max_area = area
maxIndex = i
# 将每个物体轮廓点数限制在一定范围内
if len(contours[maxIndex