Spaces:
Sleeping
Sleeping
File size: 8,658 Bytes
7e2a2a5 |
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 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 |
import utils.augmentation as transforms
import sys
from PIL import Image
import numpy as np
import random
import cv2
import os
class Transforms():
def __init__(self, config, input_grayscale_flag=False, output_grayscale_flag=False, method=Image.BICUBIC, convert=True):
self.config = config
self.input_grayscale_flag = input_grayscale_flag
self.output_grayscale_flag = output_grayscale_flag
self.method = method
self.convert = convert
self.transform_list = []
def create_transforms_from_list(self, preprocess_list):
if self.input_grayscale_flag:
if self.output_grayscale_flag:
self.transform_list.append(transforms.Grayscale())
else:
self.transform_list.append(transforms.Grayscale(1, 3))
elif self.output_grayscale_flag:
self.transform_list.append(transforms.Grayscale(3, 1))
if 'resize' in preprocess_list:
if self.config['dataset']['load_size'] < 10000:
osize = [self.config['dataset']['load_size'], self.config['dataset']['load_size']]
else:
osize = [self.config['dataset']['load_size'] // 10000, self.config['dataset']['load_size'] % 10000]
self.transform_list.append(transforms.Resize(osize, self.method))
elif 'scale_width' in preprocess_list:
self.transform_list.append(transforms.ScaleWidth(self.config['dataset']['load_size'], self.method))
if 'crop' in preprocess_list:
if 'crop_pos' in self.config['dataset']:
self.transform_list.append(transforms.Crop(self.config['dataset']['crop_pos'], self.config['dataset']['crop_size']))
else:
self.transform_list.append(transforms.RandomCrop(self.config['dataset']['crop_size']))
if 'add_lighting' in preprocess_list:
self.transform_list.append(transforms.ColorJitter())
if 'random_affine' in preprocess_list:
self.transform_list.append(transforms.RandomAffine(20, translate=(0.2, 0.2), scale=(0.2, 0.2)))
if 'random_rotate' in preprocess_list:
self.transform_list.append(transforms.RandomRotation(20))
if 'random_blur' in preprocess_list:
self.transform_list.append(transforms.RandomBlur(0.2))
if 'add_gauss_noise' in preprocess_list:
self.transform_list.append(transforms.NoiseTransform("gauss"))
if 'add_s&p_noise' in preprocess_list:
self.transform_list.append(transforms.NoiseTransform("s&p"))
if 'add_poisson_noise' in preprocess_list:
self.transform_list.append(transforms.NoiseTransform("poisson"))
if 'add_speckle_noise' in preprocess_list:
self.transform_list.append(transforms.NoiseTransform("speckle"))
if 'add_band_noise' in preprocess_list:
self.transform_list.append(transforms.NoiseTransform("band"))
if preprocess_list == 'none':
self.transform_list.append(transforms.MakePower2(base=4, method=self.method))
if not self.config['dataset']['no_flip']:
if 'flip' in self.config['dataset']:
self.transform_list.append(transforms.Flip(self.config['dataset']['flip']))
else:
self.transform_list.append(transforms.RandomHorizontalFlip())
if self.convert:
self.transform_list += [transforms.ToTensor()]
if self.input_grayscale_flag:
if self.output_grayscale_flag:
self.transform_list += [transforms.Normalize((0.5,), (0.5,))]
else:
self.transform_list += [transforms.Normalize((0.5,), (0.5,), (0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]
elif self.output_grayscale_flag:
self.transform_list += [transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5), (0.5,), (0.5,))]
else:
self.transform_list += [transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]
def get_transforms(self):
return self.transform_list
def compose_transforms(self):
return transforms.JointCompose(self.transform_list)
def check_create_shuffled_order(data_list, order):
# returns the order used to shuffle all paired data
if order is None: # Does not perform shuffling. Return normal order.
order = np.arange(0, len(data_list)).tolist()
else:
if not isinstance(order, list): # order is -1, which means has not been created.
order = np.arange(0, len(data_list)).tolist() # create the shuffle order.
random.shuffle(order)
# otherwise shuffle order already exists and we do nothing.
return order
def check_equal_length(list1, list2, data):
if len(list1) != len(list2):
print("different length in paired data types. Please double check your data.")
print("length of current data type: ", len(list1))
print("----------------current lengths for all data types-------------------")
for k, v in data.items():
print("%s: %d" % (k, len(v)))
sys.exit()
def check_img_loaded(path):
img = cv2.imread(path)
if img is None or img.size == 0:
print("image loading failed for " + path + '. Please double check.')
return False
return True
def check_numpy_loaded(path):
try:
arr = np.load(path)
except Exception as e:
print("numpy loading failed for " + path + '. Please double check.')
return False
return True
# custom, paired, numpy_paired, unpaired, numpy_unpaired, landmark
def check_old_config_val_possible(old_style_config):
for data_type in old_style_config['dataset']['data_type']:
if data_type == 'custom':
if old_style_config['dataset']['custom_val_data'] == {}:
return False
elif data_type == 'paired' or data_type == 'numpy_paired':
keyword = ''.join(data_type.split('_'))
filelist_not_exist = old_style_config['dataset']['paired_val_filelist'] == ''
filefolders_not_exist = old_style_config['dataset']['paired_valA_folder'] == '' or \
old_style_config['dataset']['paired_valB_folder'] == ''
dataroot_contains_no_val_folders = not os.path.exists(
os.path.join(old_style_config['dataset']['dataroot'], 'val' + keyword + 'A')) \
or not os.path.exists(
os.path.join(old_style_config['dataset']['dataroot'], 'val' + keyword + 'B'))
if filelist_not_exist and filefolders_not_exist and dataroot_contains_no_val_folders:
return False
elif data_type == 'unpaired' or data_type == 'numpy_unpaired':
keyword = ''.join(data_type.split('_'))
filelist_not_exist = old_style_config['dataset']['unpaired_valA_filelist'] == '' or \
old_style_config['dataset']['unpaired_valB_filelist'] == ''
filefolders_not_exist = old_style_config['dataset']['unpaired_valA_folder'] == '' or \
old_style_config['dataset']['unpaired_valB_folder'] == ''
dataroot_contains_no_val_folders = not os.path.exists(
os.path.join(old_style_config['dataset']['dataroot'], 'val' + keyword + 'A')) \
or not os.path.exists(
os.path.join(old_style_config['dataset']['dataroot'], 'val' + keyword + 'B'))
if filelist_not_exist and filefolders_not_exist and dataroot_contains_no_val_folders:
return False
elif data_type == 'landmark':
filelist_not_exist = old_style_config['dataset']['paired_val_filelist'] == ''
filefolders_not_exist = old_style_config['dataset']['paired_valA_folder'] == '' or \
old_style_config['dataset']['paired_valB_folder'] == '' or \
not os.path.exists(old_style_config['dataset']['paired_valA_lmk_folder']) or \
not os.path.exists(old_style_config['dataset']['paired_valB_lmk_folder'])
dataroot_contains_no_val_folders = not os.path.exists(
os.path.join(old_style_config['dataset']['dataroot'], 'valpairedA_lmk')) \
or not os.path.exists(
os.path.join(old_style_config['dataset']['dataroot'], 'valpairedB_lmk'))
if filelist_not_exist and filefolders_not_exist and dataroot_contains_no_val_folders:
return False
return True
|