YOLOV8图像分割预测后输出mask图

训练一个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:
           continue

		if 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

牙结石分割结果：