使用albumentations进行数据增强保存图像

最新推荐文章于 2024-07-02 09:25:27 发布

H A I

最新推荐文章于 2024-07-02 09:25:27 发布

阅读量357

点赞数 5

文章标签：人工智能 python 深度学习

本文链接：https://blog.csdn.net/weixin_44199723/article/details/138836659

版权

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

# 读取图片
image = cv2.imread("test_2_0000_0000.png")
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

# 定义数据增强的transform
transform = A.Compose([
    # A.RandomRotate90(),  # 随机旋转90度
    # A.HorizontalFlip(p=0.5),  # 水平翻转
    # A.VerticalFlip(p=0.5),  # 垂直翻转
    # A.RandomResizedCrop(256, 256, scale=(0.5, 1.0)),  # 随机缩放
    A.ColorJitter(p=0.5),  # 色彩调整
    ToTensorV2()  # 转换为PyTorch张量
])

# 对图片进行数据增强
augmented = transform(image=image)

# 获取增强后的图像
augmented_image = augmented["image"].permute(1, 2, 0).numpy()  # 转换为NumPy数组

# 将增强后的图像保存为文件
cv2.imwrite("augmented_image.jpg", cv2.cvtColor(augmented_image, cv2.COLOR_RGB2BGR))