File size: 1,870 Bytes
28e6f98 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 | # -*- encoding: utf-8 -*-
#Time :2022/02/24 18:14:15
#Author :Hao Chen
#FileName :trans_lib.py
#Version :2.0
import cv2
import torch
import numpy as np
import albumentations as A
def get_albu_transforms(type="train", img_size = (192, 192)):
if type == 'train':
compose = [
# A.VerticalFlip(p=0.5),
# A.HorizontalFlip(p=0.5),
A.ShiftScaleRotate(shift_limit=0.2, scale_limit=(-0.2, 0.2),
rotate_limit=5, p=0.5),
A.OneOf([
A.GridDistortion(num_steps=1, distort_limit=0.3, p=1.0),
A.ElasticTransform(alpha=2, sigma=5, p=1.0)
], p=0.5),
A.Resize(img_size[0], img_size[1])]
else:
compose = [A.Resize(img_size[0], img_size[1])]
return A.Compose(compose, p=1.0, additional_targets={'image2': 'image',
'image3': 'image',
'image4': 'image',
'image5': 'image',
'image6': 'image',
"mask2": "mask"})
# Beta function
def gamma_concern(img, gamma):
mean = torch.mean(img)
img = (img - mean) * gamma
img = img + mean
img = torch.clip(img, 0, 1)
return img
def gamma_power(img, gamma, direction=0):
if direction == 1:
img = 1 - img
img = torch.pow(img, gamma)
img = img / torch.max(img)
if direction == 1:
img = 1 - img
return img
def gamma_exp(img, gamma, direction=0):
if direction == 1:
img = 1 - img
img = torch.exp(img * gamma)
img = img / torch.max(img)
if direction == 1:
img = 1 - img
return img
|