图像分割：将分割结果可视化（批量处理）

语义分割的输出结果往往没有叠加在原图像上，因此很难直观看出训练结果的好坏，因此考虑将分割结果叠加在原图上，这样就能直观看出分割的细节。

代码：

# demo.py
import os
import cv2
import numpy as np


def overlay_and_save(original_folder, segmentation_folder, output_folder):
    # Gets the list in the original image folder
    original_files = os.listdir(original_folder)

    # check output folder
    output_path = os.path.join(segmentation_folder, output_folder)
    os.makedirs(output_path, exist_ok=True)

    for original_file in original_files:
        original_path = os.path.join(original_folder, original_file)
        segmentation_path = os.path.join(segmentation_folder, os.path.splitext(original_file)[0] + ".png")

        original_img = cv2.imread(original_path)
        segmentation_result = cv2.imread(segmentation_path)

        # check
        if original_img is None or segmentation_result is None:
            print(f"Failed to read {original_file}. Skipping...")
            continue

        overlay_result = plot_segmentation(original_img, segmentation_result)

        output_file_path = os.path.join(output_path, f"{original_file}_overlay")
        cv2.imwrite(output_file_path, overlay_result)
    print("Done!")


def plot_segmentation(img, segmentation_result, colors=None, alpha=0.5) -> np.ndarray:
    """
    Visualize segmentation results.

    Parameters:
    - img: numpy.ndarray
        Original image with shape `(H, W, 3)`.
    - segmentation_result: numpy.ndarray
        Segmentation result with shape `(H, W, C)`, where `C` is the number of classes.
    - colors: numpy.ndarray
        Colors for each class, shape `(C, 3)`.
        If None, generate random colors for each class.
    - alpha: float, optional, default 0.5
        Transparency of the plotted segmentation result.

    Returns:
    - numpy.ndarray
        Image plotted with segmentation colors, shape `(H, W, 3)`.
    """
    if colors is None:
        # colors = np.random.random((segmentation_result.shape[0], 3)) * 255
        colors = np.array([
            [0, 0, 0],  # _background_
            [0, 255, 255],  # lane
        ])
    else:
        if colors.shape != (2, 3):
            # if colors.shape[0] < segmentation_result.shape[-1]:
            raise RuntimeError(
                f"colors count: {colors.shape[0]} is less than classes count: {segmentation_result.shape[-1]}"
            )
    # print(segmentation_result.shape)

    for class_idx in range(2):
        # for class_idx in range(segmentation_result.shape[-1]):
        class_mask = segmentation_result[:, :, class_idx]
        color = colors[class_idx]

        img = np.where(np.stack([class_mask] * 3, axis=-1), img * (1 - alpha) + color * alpha, img)

    return img.astype(np.uint8)


original_folder = "your_original_image_path/"
segmentation_folder = "your_segmentation_image_path/"
output_folder = "Plot"
overlay_and_save(original_folder, segmentation_folder, output_folder)

参考：

 https://butui.me/posts/plot-mask-on-image/