带有不同大小蒙版的Albumentations

pip install albumentations

import cv2
import numpy as np
import albumentations as A
from albumentations.pytorch import ToTensorV2

# 加载图像
image = cv2.imread("path/to/your/image.jpg")
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

# 定义不同大小的蒙版
mask1 = np.zeros((256, 256), dtype=np.uint8)
cv2.circle(mask1, (128, 128), 64, (255, 255, 255), -1)

mask2 = np.zeros((256, 256), dtype=np.uint8)
cv2.rectangle(mask2, (64, 64), (192, 192), (255, 255, 255), -1)

# 定义 Albumentations 变换
transform = A.Compose([
    A.RandomBrightnessContrast(p=0.5),
    A.ShiftScaleRotate(shift_limit=0.1, scale_limit=0.2, rotate_limit=20, p=0.5),
    A.CoarseDropout(max_holes=8, max_height=32, max_width=32, p=0.5),
    A.CoarseDropout(max_holes=4, max_height=64, max_width=64, p=0.5),
    A.ToTensorV2()
])

# 应用 Albumentations 变换
augmented_image1 = transform(image=image, mask=mask1)['image']
augmented_image2 = transform(image=image, mask=mask2)['image']

# 显示结果
import matplotlib.pyplot as plt

plt.figure(figsize=(10, 5))
plt.subplot(1, 3, 1)
plt.imshow(image)
plt.title("Original Image")

plt.subplot(1, 3, 2)
plt.imshow(augmented_image1)
plt.title("Augmented Image with Small Mask")

plt.subplot(1, 3, 3)
plt.imshow(augmented_image2)
plt.title("Augmented Image with Large Mask")

plt.show()