语义分割数据集增强扩充python代码
时间: 2023-08-14 11:06:43 浏览: 99
以下是一个用于语义分割数据集增强扩充的Python代码示例:
```python
import cv2
import numpy as np
from albumentations import (
Compose,
VerticalFlip,
HorizontalFlip,
RandomRotate90,
ElasticTransform,
GridDistortion,
OpticalDistortion,
RandomBrightnessContrast,
)
def augment_data(image, mask):
# 定义数据增强的转换函数
augmentation = Compose([
VerticalFlip(p=0.5),
HorizontalFlip(p=0.5),
RandomRotate90(p=0.5),
ElasticTransform(p=0.3, alpha=120, sigma=120 * 0.05, alpha_affine=120 * 0.03),
GridDistortion(p=0.3),
OpticalDistortion(p=0.3, distort_limit=2, shift_limit=0.5),
RandomBrightnessContrast(p=0.2),
])
augmented = augmentation(image=image, mask=mask)
augmented_image = augmented['image']
augmented_mask = augmented['mask']
return augmented_image, augmented_mask
# 示例用法
image = cv2.imread('image.jpg')
mask = cv2.imread('mask.png', 0)
augmented_image, augmented_mask = augment_data(image, mask)
# 显示增强后的图像和掩码
cv2.imshow('Augmented Image', augmented_image)
cv2.imshow('Augmented Mask', augmented_mask)
cv2.waitKey(0)
cv2.destroyAllWindows()
```
在这个示例中,我们使用了`albumentations`库来实现数据增强。你可以根据需要自定义增强操作,并在`augment_data`函数中添加或删除转换函数。最后,你可以将增强后的图像和掩码用于训练语义分割模型。