赞
踩
- import os
- import cv2
- import random
- import torch
- import numpy as np
- import pandas as pd
- from torch.utils.data import Dataset
- from imgaug import augmenters as iaa
- from utils.process_labels import encode_labels, decode_labels, decode_color_labels
imgaug是一个图像数据扩增工具
- def crop_resize_data(image, label=None, image_size=(1024, 384), offset=690):
- roi_image = image[offset:, :]
- if label is not None:
- roi_label = label[offset:, :]
- train_image = cv2.resize(roi_image, image_size, interpolation=cv2.INTER_LINEAR) #image做的是线性插值
- train_label = cv2.resize(roi_label, image_size, interpolation=cv2.INTER_NEAREST) #label做的是最邻近插值
- return train_image, train_label
- else:
- train_image = cv2.resize(roi_image, image_size, interpolation=cv2.INTER_LINEAR)
- return train_image
图像可能会有部分是没用的,比如说,我们想实现车道线分割任务,但是,图像中天空占比很大,这时候需要将天空剪裁掉来节省我们的显存资源。
- class ImageAug(object):
- image, mask = sample
-
- if np.random.uniform(0,1) > 0.5:
- seq = iaa.Sequential([iaa.OneOf([
- iaa.AdditiveGaussianNoise(scale=(0, 0.2 * 255)), # 加高斯白噪声
- iaa.Sharpen(alpha=(alpha=(0.1,0.3), lightness=(0.7,1.3)), # 锐化
- iaa.GaussianBlur(sigma=(o,1.0))])]) # 高斯模糊
-
- image = seq.augment_image(image)
- return image, mask
- class DeformAug(object):
- def __call__(self, sample):
- image, mask = sample
- seq = iaa.Sequential([iaa.CropAndPad(percent=(-0.05,0.1))])
- seg_to = seq.to_deterministic()
- image = seg_to.augment_image(image)
- mask = seg_to.augment_image(mask)
- return image, mask
- class ScaleAug(object):
- def __cal__(self, sample):
- image, mask = sample
- scale = random.uniform(0.7,1.5) # 按照均匀分布在0.7~1.5之间取值
- h,w,_ = image.shape
-
- aug_image = image.copy() # 复制一下,不要修改原数据
- aug_mask = mask.copy()
-
- aug_image = cv2.resize(aug_image, (int(scale * w), int(scale * h)))
- aug_mask = cv2.resize(aug_mask, (int(scale * w, int(scale * h)))
- if (scale < 1.0):
- new_h, new_w, _ = aug_image.shape
- pre_h_pad = int((h - new_h) / 2)
- pre_w_pad = int((w - new_w) / 2)
- pad_list = [[pre_h_pad, h - new_h - pre_h_pad], [pre_w_pad, w - new_w - pre_w_pad], [0, 0]]
- aug_image = np.pad(aug_image, pad_list, mode="constant")
- aug_mask = np.pad(aug_mask, pad_list[:2], mode="constant")
- if (scale > 1.0):
- new_h, new_w, _ = aug_image.shape
- pre_h_crop = int ((new_h - h) / 2)
- pre_w_crop = int ((new_w - w) / 2)
- post_h_crop = h + pre_h_crop
- post_w_crop = w + pre_w_crop
- aug_image = aug_image[pre_h_crop:post_h_crop, pre_w_crop:post_w_crop]
- aug_mask = aug_mask[pre_h_crop:post_h_crop, pre_w_crop:post_w_crop]
- return aug_image, aug_mask

- class ToTensor(object):
- def __call__(self, sample):
- image, mask = sample
- image = np.transpose(image, (2,0,1))
- image = image.astype(np.float32)
- mask = mask.astype(np.long)
- return {'image':torch.from_numpy(image.copy()),
- 'mask':torch.from_numpy(mask.copy())}
Copyright © 2003-2013 www.wpsshop.cn 版权所有,并保留所有权利。