def transformers(img_size):
return A.Compose(
[
## A.RandomResizedCrop(img_size, img_size),
A.CenterCrop(img_size, img_size, p=1.0),
A.HorizontalFlip(p=1.0),
A.ShiftScaleRotate(p=1.0),
A.ColorJitter(p=1.0),
A.Cutout(p=1.0),
],
bbox_params=A.BboxParams('pascal_voc', label_fields=['category_ids']),
keypoint_params=A.KeypointParams('xy'))
def get_video_load_transform(self):
"""Load videos at double the resoluton as HRNet downsamples.
If we aren't generating labels, however, we then need to downsample x2."""
video_height, video_width = self.label_generation_dim
return A.Compose([
A.PadIfNeeded(min_height=video_height,
min_width=video_width,
border_mode=cv2.BORDER_CONSTANT,
mask_value=0),
A.CenterCrop(height=video_height, width=video_width)],
)
def get_transforms(size: int, scope: str = 'geometric', crop='random'):
augs = {
'strong':
albu.Compose([
albu.HorizontalFlip(),
albu.ShiftScaleRotate(shift_limit=0.0,
scale_limit=0.2,
rotate_limit=20,
p=.4),
albu.ElasticTransform(),
albu.OpticalDistortion(),
albu.OneOf([
albu.CLAHE(clip_limit=2),
albu.IAASharpen(),
albu.IAAEmboss(),
albu.RandomBrightnessContrast(),
albu.RandomGamma()
],
p=0.5),
albu.OneOf([
albu.RGBShift(),
albu.HueSaturationValue(),
], p=0.5),
]),
'weak':
albu.Compose([
albu.HorizontalFlip(),
]),
'geometric':
albu.OneOf([
albu.HorizontalFlip(always_apply=True),
albu.ShiftScaleRotate(always_apply=True),
albu.Transpose(always_apply=True),
albu.OpticalDistortion(always_apply=True),
albu.ElasticTransform(always_apply=True),
])
}
aug_fn = augs[scope]
crop_fn = {
'random': albu.RandomCrop(size, size, always_apply=True),
'center': albu.CenterCrop(size, size, always_apply=True)
}[crop]
pad = albu.PadIfNeeded(size, size)
pipeline = albu.Compose([aug_fn, crop_fn, pad],
additional_targets={'target': 'image'})
def process(a, b):
r = pipeline(image=a, target=b)
return r['image'], r['target']
return process
def __init__(self,
size: int,
extra_size: int = 32,
transforms: Union[List[Callable], None] = None,
normalize: bool = True,
to_tensor: bool = True) -> None:
if transforms is None:
transforms = [
A.Resize(size + extra_size, size + extra_size),
A.CenterCrop(size, size)
]
super().__init__(transforms, normalize, to_tensor)
def crop(imsize, mode):
x, y = imsize
if mode == 'train':
cropper = A.RandomCrop(height=x, width=y, always_apply=True, p=1)
else:
cropper = A.CenterCrop(height=x, width=y, always_apply=True, p=1)
return A.Compose([cropper],
p=1,
additional_targets={
'image{}'.format(_): 'image'
for _ in range(1, 201)
})
def get_aug(name='default', input_shape=[48, 48, 3]):
if name == 'default':
augmentations = A.Compose([
A.RandomBrightnessContrast(p=0.4),
A.RandomGamma(p=0.4),
A.HueSaturationValue(hue_shift_limit=20,
sat_shift_limit=30, val_shift_limit=30, p=0.4),
A.CLAHE(p=0.4),
A.Blur(blur_limit=1, p=0.3),
A.GaussNoise(var_limit=(50, 80), p=0.3)
], p=1)
elif name == 'plates':
augmentations = A.Compose([
A.RandomBrightnessContrast(p=0.4),
A.RandomGamma(p=0.4),
A.HueSaturationValue(hue_shift_limit=20,
sat_shift_limit=30,
val_shift_limit=30,
p=0.4),
A.CLAHE(p=0.4),
A.HorizontalFlip(p=0.5),
A.VerticalFlip(p=0.5),
A.Blur(blur_limit=1, p=0.3),
A.GaussNoise(var_limit=(50, 80), p=0.3),
A.RandomCrop(p=0.8, height=2*input_shape[1]/3, width=2*input_shape[0]/3)
], p=1)
elif name == 'deepfake':
augmentations = A.Compose([
A.HorizontalFlip(p=0.5),
], p=1)
elif name == 'plates2':
augmentations = A.Compose([
A.CLAHE(clip_limit=(1,4),p=0.3),
A.HorizontalFlip(p=0.5),
A.VerticalFlip(p=0.5),
A.RandomBrightness(limit=0.2, p=0.3),
A.RandomContrast(limit=0.2, p=0.3),
# A.Rotate(limit=360, p=0.9),
A.RandomRotate90(p=0.3),
A.HueSaturationValue(hue_shift_limit=(-50,50),
sat_shift_limit=(-15,15),
val_shift_limit=(-15,15),
p=0.5),
# A.Blur(blur_limit=(5,7), p=0.3),
A.GaussNoise(var_limit=(10, 50), p=0.3),
A.CenterCrop(p=1, height=2*input_shape[1]//3, width=2*input_shape[0]//3),
A.Resize(p=1, height=input_shape[1], width=input_shape[0])
], p=1)
else:
augmentations = None
return augmentations
def __init__(
self,
data_csv,
data_root,
segmentation_root,
size=None,
random_crop=False,
interpolation="bicubic",
):
self.n_labels = 182
self.data_csv = data_csv
self.data_root = data_root
self.segmentation_root = segmentation_root
with open(self.data_csv, "r") as f:
self.image_paths = f.read().splitlines()
self._length = len(self.image_paths)
self.labels = {
"relative_file_path_": [l for l in self.image_paths],
"file_path_": [
os.path.join(self.data_root + 'images', l)
for l in self.image_paths
],
"segmentation_path_": [
os.path.join(self.data_root + "segmentations",
l.replace("image", "label"))
for l in self.image_paths
],
}
size = None if size is not None and size <= 0 else size
self.size = size
if self.size is not None:
self.interpolation = interpolation
self.interpolation = {
"nearest": cv2.INTER_NEAREST,
"bilinear": cv2.INTER_LINEAR,
"bicubic": cv2.INTER_CUBIC,
"area": cv2.INTER_AREA,
"lanczos": cv2.INTER_LANCZOS4
}[self.interpolation]
self.image_rescaler = albumentations.SmallestMaxSize(
max_size=self.size, interpolation=self.interpolation)
self.segmentation_rescaler = albumentations.SmallestMaxSize(
max_size=self.size, interpolation=cv2.INTER_NEAREST)
self.center_crop = not random_crop
if self.center_crop:
self.cropper = albumentations.CenterCrop(height=self.size,
width=self.size)
else:
self.cropper = albumentations.RandomCrop(height=self.size,
width=self.size)
self.preprocessor = self.cropper
def transform_val(self, sample):
compose = A.Compose([
A.PadIfNeeded(self.base_size[0], self.base_size[1], p=1),
A.CenterCrop(self.crop_size[0], self.crop_size[1], p=1),
A.Normalize(mean=self.mean, std=self.std, p=1),
],
additional_targets={
'image': 'image',
'label': 'mask'
})
return self.toTorchTensor(compose(**sample))
def _get_valid_transforms(image_side_length: int,
mean: Tuple[float, float,
float] = (0.548, 0.504, 0.479),
std: Tuple[float, float,
float] = (0.237, 0.247, 0.246)):
val_transforms = A.Compose([
SmallestMaxSize(max_size=image_side_length, always_apply=True),
A.CenterCrop(image_side_length, image_side_length, always_apply=True),
A.Normalize(mean=mean, std=std, max_pixel_value=255.0, p=1.0),
ToTensorV2(p=1.0),
])
return val_transforms
def __init__(self,
used_img_size,
final_img_size,
transform_params,
custom_additional_targets=None):
if custom_additional_targets is None:
custom_additional_targets = {
"image2": "image",
"image3": "image",
"image4": "image"
}
self._custom_additional_targets = custom_additional_targets
self._ratio = max(
float(final_img_size[0]) / used_img_size[0],
float(final_img_size[1]) / used_img_size[1])
self._final_img_size = final_img_size
self._scale_compose = [
albumentations.Resize(height=int(used_img_size[0] * self._ratio),
width=int(used_img_size[1] * self._ratio),
always_apply=True),
albumentations.CenterCrop(height=self._final_img_size[0],
width=self._final_img_size[1],
always_apply=True,
p=1)
]
self._normalize_transform = albumentations.Normalize()
self._normalize_no_transform = albumentations.Normalize(mean=(0, 0, 0),
std=(1, 1, 1))
self._train_compose = self._scale_compose
if "flip" in transform_params and transform_params["flip"]:
flip_compose = [albumentations.HorizontalFlip()]
self._train_compose = flip_compose + self._train_compose
if "filters" in transform_params and transform_params["filters"]:
random_compose = [
albumentations.RandomBrightnessContrast(brightness_limit=(-0.2,
0.2),
contrast_limit=(-0.2,
0.2),
p=0.5),
albumentations.RandomGamma(gamma_limit=(90, 110), p=0.5),
albumentations.ChannelShuffle(p=0.5),
]
self._train_compose = random_compose + self._train_compose
if "normalize" in transform_params and transform_params["normalize"]:
self._train_compose.append(albumentations.Normalize())
else:
self._train_compose.append(
albumentations.Normalize(mean=(0, 0, 0), std=(1, 1, 1)))
def __init__(self, config=None, size=None, random_crop=False, interpolation="bicubic", crop_size=None):
self.split = self.get_split()
self.n_labels = 151 # unknown + 150
self.data_csv = {"train": "data/ade20k_train.txt",
"validation": "data/ade20k_test.txt"}[self.split]
self.data_root = "data/ade20k_root"
with open(os.path.join(self.data_root, "sceneCategories.txt"), "r") as f:
self.scene_categories = f.read().splitlines()
self.scene_categories = dict(line.split() for line in self.scene_categories)
with open(self.data_csv, "r") as f:
self.image_paths = f.read().splitlines()
self._length = len(self.image_paths)
self.labels = {
"relative_file_path_": [l for l in self.image_paths],
"file_path_": [os.path.join(self.data_root, "images", l)
for l in self.image_paths],
"relative_segmentation_path_": [l.replace(".jpg", ".png")
for l in self.image_paths],
"segmentation_path_": [os.path.join(self.data_root, "annotations",
l.replace(".jpg", ".png"))
for l in self.image_paths],
"scene_category": [self.scene_categories[l.split("/")[1].replace(".jpg", "")]
for l in self.image_paths],
}
size = None if size is not None and size<=0 else size
self.size = size
if crop_size is None:
self.crop_size = size if size is not None else None
else:
self.crop_size = crop_size
if self.size is not None:
self.interpolation = interpolation
self.interpolation = {
"nearest": cv2.INTER_NEAREST,
"bilinear": cv2.INTER_LINEAR,
"bicubic": cv2.INTER_CUBIC,
"area": cv2.INTER_AREA,
"lanczos": cv2.INTER_LANCZOS4}[self.interpolation]
self.image_rescaler = albumentations.SmallestMaxSize(max_size=self.size,
interpolation=self.interpolation)
self.segmentation_rescaler = albumentations.SmallestMaxSize(max_size=self.size,
interpolation=cv2.INTER_NEAREST)
if crop_size is not None:
self.center_crop = not random_crop
if self.center_crop:
self.cropper = albumentations.CenterCrop(height=self.crop_size, width=self.crop_size)
else:
self.cropper = albumentations.RandomCrop(height=self.crop_size, width=self.crop_size)
self.preprocessor = self.cropper
def composeAugmentation(self):
if (source == "train" :)
self.augment = albumentations.Compose(
[
albumentations.HorizontalFlip(),
albumentations.Rotate(10, always_apply=True),
albumentations.SmallestMaxSize(self.sizeX+8, always_apply=True),
albumentations.CenterCrop(self.sizeX, self.sizeY, always_apply=True),
albumentations.GridDistortion(always_apply=False),
albumentations.IAAAffine(rotate=2, shear=5, always_apply=False),
#albumentations.OpticalDistortion(),
albumentations.ElasticTransform(alpha=128, sigma=64, always_apply=True, alpha_affine=0),
albumentations.RandomBrightnessContrast(0.1, 0.1, always_apply=True),
]
)
if (source == "val"):
self.augment = albumentations.Compose(
[
albumentations.SmallestMaxSize(self.sizeX+8, always_apply=True),
albumentations.CenterCrop(self.sizeX, self.sizeY, always_apply=True),
]
)
def mask_annotations(image):
mask_classifier = cached_classifier()
transform = A.Compose([
A.SmallestMaxSize(max_size=256, interpolation=cv2.INTER_CUBIC),
A.CenterCrop(height=224, width=224),
A.Normalize(),
])
crop_transformed = transform(image=image)['image']
model_input = torch.from_numpy(np.transpose(crop_transformed, (2, 0, 1)))
prediction = [mask_classifier(model_input.unsqueeze(0))[0].item()]
return prediction
def hard_augmentations(image_size: Tuple[int, int], rot_angle=30):
pad_h, pad_w = padding_for_rotation(image_size, rot_angle)
crop_height = int(image_size[0] + pad_h * 2)
crop_width = int(image_size[1] + pad_w * 2)
crop_transform = A.Compose([
A.RandomSizedCrop((int(crop_height * 0.75), int(crop_height * 1.25)), crop_height, crop_width),
A.ShiftScaleRotate(shift_limit=0, scale_limit=0, rotate_limit=rot_angle, border_mode=cv2.BORDER_CONSTANT),
A.CenterCrop(image_size[0], image_size[1]),
])
return A.Compose([
# spatial transform
A.PadIfNeeded(int(crop_height * 1.25), int(crop_height * 1.25)),
A.OneOf([
crop_transform,
A.RandomSizedCrop((image_size[0], int(image_size[0] * 1.25)), image_size[0], image_size[1], p=0.25),
A.RandomSizedCrop((image_size[0], int(image_size[0] * 1.5)), image_size[0], image_size[1], p=0.25),
A.RandomSizedCrop((image_size[0], int(image_size[0] * 2)), image_size[0], image_size[1], p=0.25),
A.Resize(image_size[0], image_size[1], p=0.75)
], p=1.0),
# add occasion blur/sharpening
A.OneOf([
A.GaussianBlur(),
A.MotionBlur(),
A.IAASharpen(),
A.JpegCompression(quality_lower=75, p=0.25),
]),
# D4 augmentations
A.Compose([
A.HorizontalFlip(),
]),
# spatial-preserving augmentations
A.OneOf([
A.Cutout(),
A.GaussNoise(),
]),
A.OneOf([
A.RandomBrightnessContrast(),
A.CLAHE(),
A.HueSaturationValue(),
A.RGBShift(),
A.RandomGamma()
]),
A.Normalize(),
ToTensor()
])
def transform_image(image):
# We don't need augmentations in testing and validation.
# All we need here is to resize the PIL image and transform it into Tensor.
test_tfm = A.Compose([
A.Resize(height=224, width=224),
A.CenterCrop(height=180, width=180),
A.Normalize(mean=[0.9369], std=[0.2202], max_pixel_value=1),
ToTensorV2(),
#transforms.Normalize([0.9252], [0.2451])
])
#cv2.imwrite(str(esun_timestamp)+'.jpg', image)
#sample = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
sample = image.astype(np.float32) / 255
return test_tfm(image=sample)["image"].unsqueeze(0)
def get_trm(cfg, is_train=True):
normalize_transform = T.Normalize(mean=cfg.INPUT.PIXEL_MEAN,
std=cfg.INPUT.PIXEL_STD)
if is_train:
main_transform = A.Compose(
[
A.Resize(cfg.INPUT.RESIZE_TRAIN[0], cfg.INPUT.RESIZE_TRAIN[1]),
A.RandomCrop(cfg.INPUT.SIZE_TRAIN[0], cfg.INPUT.SIZE_TRAIN[1]),
A.HorizontalFlip(p=cfg.INPUT.PROB),
# A.RandomRotate90(p=1),
# A.ShiftScaleRotate(p=0.5, shift_limit=0.1, scale_limit=.1, rotate_limit=10),
# A.OneOf([A.RandomContrast(limit=0.2, p=0.1),
# A.RandomGamma(gamma_limit=(90, 110), p=0.5),
# A.RandomBrightness(limit=0.2, p=0.2)], p=0.5),
# A.CoarseDropout()
],
p=1)
if cfg.INPUT.USE_AUTOAUG:
image_transform_list = [
Augmentation(autoaug_policy()),
T.ToTensor(),
normalize_transform,
RandomErasing(probability=cfg.INPUT.RANDOM_ERASE.RE_PROB,
sh=cfg.INPUT.RANDOM_ERASE.RE_MAX_RATIO,
mean=cfg.INPUT.PIXEL_MEAN),
]
else:
image_transform_list = [
T.ToTensor(),
normalize_transform,
RandomErasing(probability=cfg.INPUT.RANDOM_ERASE.RE_PROB,
sh=cfg.INPUT.RANDOM_ERASE.RE_MAX_RATIO,
mean=cfg.INPUT.PIXEL_MEAN),
]
image_transform = T.Compose(image_transform_list)
else:
main_transform = A.Compose\
([
A.Resize(cfg.INPUT.RESIZE_TEST[0], cfg.INPUT.RESIZE_TEST[1]),
A.CenterCrop(cfg.INPUT.SIZE_TEST[0], cfg.INPUT.SIZE_TEST[1])
])
image_transform = T.Compose([
T.ToTensor(),
normalize_transform,
])
return main_transform, image_transform
def get_transforms(crop_size: int, size: int, validate=False):
randomCrop = albu.Compose([albu.RandomCrop(crop_size, crop_size, always_apply=True), albu.Resize(size, size),
albu.VerticalFlip(), albu.RandomRotate90(always_apply=True)],
additional_targets={'target': 'image'})
centerCrop = albu.Compose([albu.CenterCrop(crop_size, crop_size, always_apply=True), albu.Resize(size, size),
albu.VerticalFlip(), albu.RandomRotate90(always_apply=True)],
additional_targets={'target': 'image'})
pipeline = randomCrop if not validate else randomCrop
def process(a, b, c):
r = pipeline(image=a, target=b, mask=c)
return r['image'], r['target'], r['mask']
return process
def __init__(self, data_dir, classes=('background', 'road')):
# all list contains all items:
# (rgb_image, rgb_mask, rgb_visual, vis_image, vis_mask, vis_visual, nir_image, nir_mask, nir_visual)
f_list = np.genfromtxt(os.path.join(data_dir, 'all_list.txt'),
dtype='str')
# generate lists
self.rgb_image_list = [
os.path.join(data_dir, file) for file in f_list[:, 0]
]
self.rgb_label_list = [
os.path.join(data_dir, file) for file in f_list[:, 1]
]
self.vis_image_list = [
os.path.join(data_dir, file) for file in f_list[:, 3]
]
self.vis_label_list = [
os.path.join(data_dir, file) for file in f_list[:, 4]
]
self.nir_image_list = [
os.path.join(data_dir, file) for file in f_list[:, 6]
]
self.nir_label_list = [
os.path.join(data_dir, file) for file in f_list[:, 7]
]
# crop to right size, applied to rgb image only
self.crop_rgb = albu.CenterCrop(704, 1280) # crop from 736*1280
self.crop_vis = albu.CenterCrop(256, 480) # crop from 272*512
self.crop_nir = albu.CenterCrop(192, 384) # crop from 217*409
# parse the mask values
self.classes = [self.CLASSES.index(cls.lower()) for cls in classes]
# load correction matrix for hsi data correction
self.vis_matrix = np.transpose(
np.loadtxt(os.path.join(data_dir, 'vis_correction.txt'),
dtype=np.float32))
self.nir_matrix = np.transpose(
np.loadtxt(os.path.join(data_dir, 'nir_correction.txt'),
dtype=np.float32))
def __init__(self, directory: str, use_augmentation: bool, image_height: int = 480, image_width: int = 480):
self.directory = pathlib.Path(directory)
self.use_augmentation = use_augmentation
assert self.directory.exists()
assert self.directory.is_dir()
self.labelme_paths = []
self.categories = collections.defaultdict(list)
for labelme_path in self.directory.rglob('*.json'):
with open(labelme_path, 'r') as labelme_file:
labelme_json = json.load(labelme_file)
required_keys = ['version', 'flags', 'shapes', 'imagePath', 'imageData', 'imageHeight', 'imageWidth']
assert all(key in labelme_json for key in required_keys), (required_keys, labelme_json.keys())
self.labelme_paths += [labelme_path]
for shape in labelme_json['shapes']:
label = shape['label']
self.categories[label] += [labelme_path]
for category, paths in self.categories.items():
for path in paths:
logging.debug(f'{category} - {path}')
self.categories = sorted(list(self.categories.keys()))
logging.info(f'loaded {len(self)} annotations from {self.directory}')
logging.info(f'use augmentation: {self.use_augmentation}')
logging.info(f'categories: {self.categories}')
aug_transforms = [ToTensor()]
if self.use_augmentation:
aug_transforms = [
alb.HueSaturationValue(always_apply=True),
alb.RandomBrightnessContrast(always_apply=True),
alb.HorizontalFlip(),
alb.RandomGamma(always_apply=True),
alb.ShiftScaleRotate(shift_limit=0.2, scale_limit=0.2, always_apply=True),
alb.PadIfNeeded(min_height=image_height, min_width=image_width, always_apply=True),
alb.RandomCrop(image_height, image_width, always_apply=True),
] + aug_transforms
else:
aug_transforms = [
alb.PadIfNeeded(min_height=image_height, min_width=image_width, always_apply=True),
alb.CenterCrop(image_height, image_width, always_apply=True),
] + aug_transforms
self.transforms = alb.Compose(transforms=aug_transforms)
def common_aug(mode, params):
'''
:param mode: 'train', 'test', 'inference'
:param params:
'''
#aug_params = params.get('augm_params', dict())
augs_list = []
assert mode in {'train', 'debug', 'inference'}
if mode == 'train':
augs_list.append(
albumentations.PadIfNeeded(min_height=params.data.net_hw[0],
min_width=params.data.net_hw[1],
border_mode=cv2.BORDER_REPLICATE,
always_apply=True))
augs_list.append(
albumentations.RandomCrop(height=params.data.net_hw[0],
width=params.data.net_hw[1],
always_apply=True))
if params.augmentation.rotate_limit:
augs_list.append(
T.Rotate(limit=params.augmentation.rotate_limit,
border_mode=cv2.BORDER_CONSTANT,
always_apply=True))
# augs_list.append(T.OpticalDistortion(border_mode=cv2.BORDER_CONSTANT)) - can't handle boundboxes
elif mode == 'debug':
augs_list.append(
albumentations.CenterCrop(height=params.data.net_hw[0],
width=params.data.net_hw[1],
always_apply=True))
if mode != 'inference':
if params.augmentation.get('blur_limit', 4):
augs_list.append(
T.Blur(blur_limit=params.augmentation.get('blur_limit', 4)))
if params.augmentation.get('RandomBrightnessContrast', True):
augs_list.append(T.RandomBrightnessContrast())
#augs_list.append(T.MotionBlur())
if params.augmentation.get('JpegCompression', True):
augs_list.append(
T.JpegCompression(quality_lower=30, quality_upper=100))
#augs_list.append(T.VerticalFlip())
if params.augmentation.get('HorizontalFlip', True):
augs_list.append(T.HorizontalFlip())
return albumentations.ReplayCompose(augs_list,
p=1.,
bbox_params={
'format': 'albumentations',
'min_visibility': 0.5
})
def __init__(self, args, mode='train'):
self.args = args
root_dir = args.dataset
batch_size = args.batch_size
augment = args.augment
img_size = args.img_size
self.resize_size = (int(img_size / 0.875), int(img_size / 0.875))
self.crop_size = (img_size, img_size)
self.root_dir = os.path.join(root_dir, mode)
self.img_path_list = []
self.label_list = []
self.class_names = []
for class_index, class_name in enumerate(os.listdir(self.root_dir)):
class_dir = os.path.join(self.root_dir, class_name)
for img_name in os.listdir(class_dir):
self.img_path_list.append(os.path.join(class_dir, img_name))
self.label_list.append(class_index)
self.class_names.append(class_name)
if mode == 'train':
pad_len = len(self.img_path_list) % batch_size
img_path_list_len = len(self.img_path_list)
for _ in range(pad_len):
rand_index = np.random.randint(0, img_path_list_len)
self.img_path_list.append(self.img_path_list[rand_index])
self.label_list.append(self.label_list[rand_index])
self.data_index = np.arange(0, len(self.label_list))
np.random.shuffle(self.data_index)
else:
self.data_index = np.arange(0, len(self.label_list))
self.img_path_list = np.array(self.img_path_list)
self.label_list = np.array(self.label_list)
self.augment = augment
self.mode = mode
self.batch_size = batch_size
self.eppch_index = 0
self.num_class = len(self.class_names)
self.random_crop_transform = A.Compose([
A.RandomCrop(width=self.crop_size[0], height=self.crop_size[1]),
])
self.center_crop_transform = A.Compose([
A.CenterCrop(width=self.crop_size[0], height=self.crop_size[1]),
])
self.auto_augment = auto_augment.AutoAugment()
self.rand_augment = rand_augment.RandAugment(magnitude=10.)
def validate_augment(image, pixel_avg, stdev):
"""
Simply normalize image values and do a crop.
"""
height, width = image.shape[0], image.shape[1]
crop_width = min(224, width)
crop_height = min(224, height)
augmenter = albumentations.Compose([
albumentations.CenterCrop(crop_height, crop_width)])
output_image = ((augmenter(image=image)["image"] - pixel_avg) / 255)/stdev
return AlexNetSequence.pad_image(output_image)
def _val_sync_transform(self, sample):
compose = A.Compose([
A.PadIfNeeded(self.base_size, self.base_size, p=1),
A.CenterCrop(self.crop_size, self.crop_size, p=1),
A.Normalize(),
],
additional_targets={
'image': 'image',
'label': 'mask'
})
sample = compose(**sample)
sample["image"] = self._transform(sample["image"])
sample["label"] = self._mask_transform(sample["label"])
return sample
def _preprocess_img(self, img, mask, augment, size):
img_pil = Image.open(img.numpy()).resize((556, 556))
mask_pil = Image.open(mask.numpy()).resize((556, 556))
img, mask = self._pil_to_np(img_pil, mask_pil)
if augment:
img, mask = self._augment_img(img, mask, size[0], size[1])
else:
if self._crop is None:
self._crop = A.Compose(
[A.CenterCrop(height=size[0], width=size[1])])
cropped = self._crop(image=img, mask=mask)
img = cropped["image"]
mask = cropped["mask"]
return np.stack((img, ) * 3, axis=-1), mask[..., None]
def get_train_transforms_no_distortion(cfg):
return A.Compose([
A.CenterCrop(cfg["img_size"], cfg["img_size"], p=1.0),
A.HueSaturationValue(),
A.RandomBrightnessContrast(brightness_limit=(-0.3, 0.3),
contrast_limit=(-0.1, 0.1),
p=0.5),
A.Normalize(
mean=[0.56019358, 0.52410121, 0.501457],
std=[0.23318603, 0.24300033, 0.24567522],
max_pixel_value=255.0,
p=1.0,
),
ToTensorV2(p=1.0),
])
def __init__(self,
base_dir=None,
filename=None,
n_class=3,
target_size=224,
affine_augmenter=None,
image_augmenter=None,
debug=False):
self.base_dir = base_dir
self.base_dir = Path(base_dir)
self.n_class = n_class
self.debug = debug
self.img_paths = []
self.bbox = []
self.labels = []
self.euler_binned = []
with open(self.base_dir / filename, 'r') as f:
for i, line in enumerate(f.readlines()):
ls = line.strip().split(' ')
bbox = [float(ls[i]) for i in range(1, 5)]
pose = [float(ls[i]) for i in range(5, 8)]
pose[1] = -pose[1]
if True and (abs(pose[0]) > 99 or abs(pose[1]) > 99
or abs(pose[2]) > 99):
continue
img_path = self.base_dir / ls[0]
if not (os.path.exists(img_path)):
print(f"[INFO] File not exists: {img_path}")
continue
self.labels.append(np.array(pose))
self.bbox.append(bbox)
self.img_paths.append(ls[0])
self.labels_sort_idx = np.argsort(
-np.mean(np.abs(self.labels), axis=1))
# self.resizer = albu.Compose([albu.Resize(target_size[0], target_size[1], p=1.)])
self.resizer = albu.Compose([
albu.SmallestMaxSize(target_size, p=1.),
albu.CenterCrop(target_size, target_size, p=1.)
])
self.affine_augmenter = affine_augmenter
self.image_augmenter = image_augmenter
def make_aug(
img_size: int,
crop_size: Optional[int],
starategy: str,
mode: str,
severity: int = 3,
width: int = 3,
):
assert mode in ["train", "val", "test"]
assert starategy in ["mix", "rand"]
if crop_size is None:
crop_size = img_size
if starategy == "mix":
rand_aug = RandomAugMix(severity=severity,
width=width,
alpha=1.0,
p=1.0)
else:
rand_aug = RandomAug(severity=severity, n=width, p=0.75)
if mode == "train":
return A.Compose([
A.RandomResizedCrop(
crop_size,
crop_size,
scale=(0.9, 1),
p=1,
interpolation=cv2.INTER_LANCZOS4,
),
rand_aug,
A.Cutout(
max_h_size=int(crop_size * 0.075),
max_w_size=int(crop_size * 0.075),
num_holes=10,
p=0.5,
),
A.Normalize(mean=IMAGENET_DEFAULT_MEAN, std=IMAGENET_DEFAULT_STD),
ToTensorV2(),
], )
if mode in ["val", "test"]:
return A.Compose([
A.Resize(img_size, img_size, interpolation=cv2.INTER_LANCZOS4),
A.CenterCrop(
crop_size, crop_size, p=1 if crop_size != img_size else 0),
A.Normalize(mean=IMAGENET_DEFAULT_MEAN, std=IMAGENET_DEFAULT_STD),
ToTensorV2(),
], )
def __init__(self, base_dir=None, split='train', affine_augmenter=None, image_augmenter=None,
target_size=224, filename=None, use_bined=False, n_class=4, debug=False):
self.base_dir = base_dir
self.base_dir = Path(base_dir)
self.split = split
self.use_bined = use_bined
self.n_class = n_class
self.debug = debug
self.img_paths = []
self.bbox = []
self.labels = []
self.euler_binned = []
with open(self.base_dir / filename) as f:
for i, line in enumerate(f.readlines()):
ls = line.strip()
mat_path = self.base_dir / ls.replace('.jpg', '.mat')
bbox, pose = get_pt_ypr_from_mat(mat_path, pt3d=True)
if True and (abs(pose[0])>99 or abs(pose[1])>99 or abs(pose[2])>99):
continue
if use_bined:
yaw_pitch_bins = np.array([-60, -40, -20, 20, 40, 60])
roll_bins = np.array(range(-81, 82, 9))
self.euler_binned.append([np.digitize(pose[0], yaw_pitch_bins),np.digitize(pose[1], yaw_pitch_bins),np.digitize(pose[2], roll_bins)])
self.labels.append(np.array(pose))
self.bbox.append(bbox)
self.img_paths.append(ls)
self.labels_sort_idx = np.argsort(-np.mean(np.abs(self.labels), axis=1))
if 'train' in self.split:
self.resizer = albu.Compose([albu.SmallestMaxSize(target_size, p=1.),
albu.RandomScale(scale_limit=(-0.2, 0.2), p=0.1),
albu.PadIfNeeded(min_height=target_size, min_width=target_size, value=0, p=1),
albu.RandomCrop(target_size, target_size, p=1.)])
else:
# self.resizer = albu.Compose([albu.Resize(target_size[0], target_size[1], p=1.)])
self.resizer = albu.Compose([albu.SmallestMaxSize(target_size, p=1.),
albu.CenterCrop(target_size, target_size, p=1.)])
self.affine_augmenter = affine_augmenter
self.image_augmenter = image_augmenter
def __init__(self, resize, mean, std, **args):
self.transform = A.Compose([
A.CenterCrop(350, 300, p = 1.0),
A.Resize(256, 256, p=1.0),
A.Compose([
A.HorizontalFlip(p = 0.5),
A.ShiftScaleRotate(p = 0.5),
A.IAAAffine(scale = [0.8, 1.2], rotate = [-10, 10], shear = 10, p = 0.5),
A.OneOf([
A.RandomBrightness(limit = 0.2, p = 0.5),
A.RandomBrightnessContrast(brightness_limit = 0.2, contrast_limit = 0.1, p = 0.5),
]),
], p = 0.5),
A.Normalize(mean = mean, std = std, max_pixel_value = 255.0, p = 1.0),
ToTensorV2(p = 1.0),
])
def _sync_transform(self, sample):
compose = A.Compose([
A.PadIfNeeded(self.base_size, self.base_size, p=1),
# A.RandomSizedCrop((self.crop_size - 10, self.crop_size + 10), self.crop_size, self.crop_size, p=1),
A.Resize(self.crop_size + 20, self.crop_size + 20),
A.CenterCrop(self.crop_size, self.crop_size, p=1),
A.HorizontalFlip(),
A.Blur(p=0.3),
A.GaussNoise(p=0.3),
A.Normalize(),
], additional_targets={'image': 'image', 'label': 'mask'})
sample = compose(**sample)
sample["image"] = self._transform(sample["image"])
sample["label"] = self._mask_transform(sample["label"])
return sample
