训练一个yolov8模型后,用官方的预测脚本一般是:
results = model.predict(img_path, save=True, save_dir=output_folder)
运行此代码会直接在run里面生成一个文件夹,保存预测图像。如果要获取分割后的mask点,或mask的轮廓点,可以这样:
1、先定义颜色列表,多个mask时可以用不同颜色表示
# 定义颜色列表color_palette = [[0, 0, 255], # 红色[0, 255, 0], # 绿色[255, 0, 0], # 蓝色[0, 255, 255], # 黄色[255, 0, 255] # 紫色
]# 随机生成颜色
def generate_random_color():return [random.randint(0, 255) for _ in range(3)]
2、将mask叠加到原图上
image_org = cv2.imread(img_path)# 进行预测
results = model.predict(img_path, save=True, save_dir=output_folder)
show_mask_or_contour = True
if results:print(f"Predicted for {img_path}:")for r in results:if not r.masks:continueif show_mask_or_contour:# 获取每个mask的轮廓for i, polygon in enumerate(r.masks.xy):# 将浮点数坐标点转换为整数类型polygon = [(int(point[0]), int(point[1])) for point in polygon]# 选择或生成颜色if i < len(color_palette):color = color_palette[i]else:color = generate_random_color() # 如果预定义颜色用尽,生成新的随机颜色for x, y in polygon:image_org[y, x] = color else:# 获取每个mask所有点for idx, mask in enumerate(r.masks):if mask.data.dim() == 3 and (mask.data.size(0) == 1 or mask.data.size(2) == 1):maskdata= mask.data.squeeze() # 这会移除大小为1的维度# 缩放掩码到目标图像的尺寸resized_mask = cv2.resize(maskdata.cpu().numpy(), (mask.orig_shape[1], mask.orig_shape[0]),interpolation=cv2.INTER_NEAREST)# 获取掩码中非零点的坐标y_indices, x_indices = np.where(resized_mask > 0)# 选择或生成颜色if idx < len(color_palette):color = color_palette[idx]else:color = generate_random_color() # 如果预定义颜色用尽,生成新的随机颜色ratio = 0.3 设置了透明度for y, x in zip(y_indices, x_indices):# 提取原始颜色original_color = image_org[y, x]# 计算新颜色new_color = [int(c * ratio + original_color[i] * (1 - ratio)) for i, c in enumerate(color)]# 更新图像image_org[y, x] = new_color
牙结石分割结果:
发现轮廓点并不是连续点,如有同学有不一样的结果,谢谢告知