【语义分割】将VOC格式分割掩码覆盖到原图生成标注可视化

修改文件路径，并且根据掩码像素值修改BGR色盘。

单图：

import cv2
import numpy as np

# Load the original image and the mask image
original_image_path = 'VOCdevkit/VOC/JPEGImages/1.jpg'
mask_image_path = 'VOCdevkit/VOC/SegmentationClass/1.png'
original_image = cv2.imread(original_image_path)
mask_image = cv2.imread(mask_image_path, cv2.IMREAD_UNCHANGED)

# Resize mask to match the original image if necessary
if mask_image.shape[:2] != original_image.shape[:2]:
    mask_image = cv2.resize(mask_image, (original_image.shape[1], original_image.shape[0]), interpolation=cv2.INTER_NEAREST)

# If the mask is not single channel, convert it to grayscale
if len(mask_image.shape) > 2:
    mask_image = cv2.cvtColor(mask_image, cv2.COLOR_BGR2GRAY)

# Get the unique values from the mask image which represent different categories
unique_values = np.unique(mask_image)

# Define colors for the unique categories we found in the mask
# We'll use a different color for each category value.
# The colors are in BGR format as OpenCV uses BGR by default.
category_colors = {
    0: [0, 0, 0],       # Black for background
    132: [153, 76, 0],  
    135: [0, 0, 153],  
    167: [0, 76, 153],  
    179: [204, 0, 102] 
}

# Create an empty colored mask with the same size as the original mask
colored_mask = np.zeros_like(original_image, dtype=np.uint8)

# Apply colors to the mask
for value in unique_values:
    # We check if the value is in our category colors, if not we continue to the next value
    if value in category_colors:
        colored_mask[mask_image == value] = category_colors[value]

# Blend the colored mask with the original image
overlayed_image = cv2.addWeighted(original_image, 1, colored_mask, 0.7, 0)

# Save the result
output_path = 'data/ground-truth/image.png'
cv2.imwrite(output_path, overlayed_image)
print(f"Image with colored mask overlay saved to {output_path}")

批量处理：

import cv2
import numpy as np
import os

# Set the directories for the images and masks
original_images_dir = 'VOCdevkit/VOC/JPEGImages'
mask_images_dir = 'VOCdevkit/VOC/SegmentationClass'
output_images_dir = 'VOCdevkit/VOC/visual'

# Ensure the output directory exists
if not os.path.exists(output_images_dir):
    os.makedirs(output_images_dir)

# Define the colors for the categories found in the mask images
category_colors = {
    0: [0, 0, 0],
    132: [153, 76, 0],
    135: [0, 0, 153],
    181: [0, 76, 153],
    167: [0, 76, 153],
    179: [204, 0, 102]
}

# Process images
for image_filename in os.listdir(original_images_dir):
    # Construct the full path to the original image
    original_image_path = os.path.join(original_images_dir, image_filename)

    # Check if the file is a JPG image
    if not image_filename.lower().endswith('.jpg'):
        continue  # Skip non-JPG files

    # Construct the corresponding mask image path (assuming same filename but with .png extension)
    mask_image_filename = os.path.splitext(image_filename)[0] + '.png'
    mask_image_path = os.path.join(mask_images_dir, mask_image_filename)

    # Check if the corresponding mask image exists
    if not os.path.exists(mask_image_path):
        print(f"No corresponding mask found for {image_filename}, skipping...")
        continue

    # Read the original image and mask image
    original_image = cv2.imread(original_image_path)
    mask_image = cv2.imread(mask_image_path, cv2.IMREAD_UNCHANGED)

    # Resize the mask to match the original image if necessary
    if mask_image.shape[:2] != original_image.shape[:2]:
        mask_image = cv2.resize(mask_image, (original_image.shape[1], original_image.shape[0]), interpolation=cv2.INTER_NEAREST)

    # If the mask is not single channel, convert it to grayscale
    if len(mask_image.shape) > 2:
        mask_image = cv2.cvtColor(mask_image, cv2.COLOR_BGR2GRAY)

    # Apply colors to the mask
    colored_mask = np.zeros_like(original_image, dtype=np.uint8)
    for value in np.unique(mask_image):
        if value in category_colors:
            colored_mask[mask_image == value] = category_colors[value]

    # Blend the colored mask with the original image
    overlayed_image = cv2.addWeighted(original_image, 1, colored_mask, 0.5, 0)

    # Save the result to the output directory
    output_image_path = os.path.join(output_images_dir, image_filename)
    cv2.imwrite(output_image_path, overlayed_image)
    print(f"Saved overlayed image for {image_filename} to {output_image_path}")