【YOLOv8分割】在 预测阶段 如何预测整个目录并输出二值mask图像

YOLOv8的原来的分割样式如图：

这里代码可以实现对整个目录进行预测（预测效果如上所示）：

import glob
from PIL import Image
from ultralytics import YOLO
import csv
import os
from os.path import join , basename
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn.functional as F
import cv2

# 模型路径
model = YOLO(r'F:\Deep_Learning\Model\YOLOv8_Seg\runs\segment\train\weights\best.pt')
# 图片路径
source = 'F:/CRACK500/val/images'
# 预测图片的保存目录
pred_dir = r'F:\Deep_Learning\Model\YOLOv8_Seg\Pre_Dir'


# 如果保存的话：
results = model(source=source,save=True, name='./Pre_Dir',show_labels=False,show_conf=False,boxes=False)

# 如果不保存的话：
# results = model(source=source,show_labels=False,show_conf=False,boxes=False)

for result in results:
    image_name = basename(result.path)  # 提取图片名称
    mask_name = f"{os.path.splitext(image_name)[0]}.png"  # 根据图片名称生成保存结果的名称
    pred_image_path = join(r'F:\Deep_Learning\Model\YOLOv8_Seg\Dataset\mask', mask_name)# 图片保存路径
    # 检测到裂缝时：
    if result.masks is not None and len(result.masks) > 0:
        masks_data = result.masks.data
        for index, mask in enumerate(masks_data):
            mask = mask.cpu().numpy() * 255

            # cv2.imwrite(f'./output_{index}.png', mask)


            cv2.imwrite(pred_image_path , mask)

其代码目录安排如下：

-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
但是上述图片不能放到论文上，需要转为二值图像，下图是转化后的预测图：

转换保存的目录存放位置：

转换代码：

import glob
from PIL import Image
from ultralytics import YOLO
import csv
import os
from os.path import join , basename
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn.functional as F
import cv2

# 模型路径
model = YOLO(r'F:\Deep_Learning\Model\YOLOv8_Seg\runs\segment\train\weights\best.pt')
# 图片路径(不能有中文的标点符号或中文！！！！)
source = 'F:/CRACK500/val/images'

# 如果保存的话：
#results = model(source=source,save=True, name='./Pre_Dir',show_labels=False,show_conf=False,boxes=False)

# 如果不保存的话：
results = model(source=source,show_labels=False,show_conf=False,boxes=False)

for result in results:
    image_name = basename(result.path)  # 提取图片名称

    mask_name = f"{os.path.splitext(image_name)[0]}.png"  # 根据图片名称生成保存结果的名称

    pred_image_path = join(r'F:\Deep_Learning\Model\YOLOv8_Seg\Dataset\mask', mask_name)
    # 检测到裂缝时：
    if result.masks is not None and len(result.masks) > 0:
        masks_data = result.masks.data
        for index, mask in enumerate(masks_data):
            mask = mask.cpu().numpy() * 255
            cv2.imwrite(pred_image_path , mask)
    # 检测不到裂缝时：
    else:
        width , height = 640 , 360
        black_image = np.zeros((height , width , 3) , dtype=np.uint8)
        # 保存全黑的图像为PNG文件
        cv2.imwrite(pred_image_path , black_image)