训练一个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
牙结石分割结果:
发现轮廓点并不是连续点,如有同学有不一样的结果,谢谢告知