修改文件路径,并且根据掩码像素值修改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}")