class FESetsGenerator:
def __init__(self, dst_shape, focal_len=350):
self._generator = FishEyeGenerator(focal_len, dst_shape)
self._F_RAND_FLAG = False
self._F_RANGE = [200, 400]
self._EXT_RAND_FLAG = False
self._EXT_PARAM_RANGE = [5, 5, 10, 0.3, 0.3, 0.4]
self._generator.set_ext_param_range(self._EXT_PARAM_RANGE)
def set_ext_param_range(self, ext_param):
for i in range(6):
self._EXT_PARAM_RANGE[i] = ext_param[i]
self._generator.set_ext_param_range(self._EXT_PARAM_RANGE)
def rand_ext_params(self):
self._EXT_RAND_FLAG = True
def set_ext_params(self, ext_params):
self._generator.set_ext_params(ext_params)
self._EXT_RAND_FLAG = False
def set_f(self, focal_len):
self._generator.set_f(focal_len)
self._F_RAND_FLAG = False
def rand_f(self, f_range=[200, 400]):
self._F_RANGE = f_range
self._F_RAND_FLAG = True
def generate(self, src_dir, src_annot_dir, dst_dir, dst_annot_dir, prefix):
image_list = sorted(
[image for image in os.listdir(src_dir) if image.endswith(".png")])
count = 0
for image in image_list:
src_image = cv2.imread(src_dir + image)
src_annot_image = cv2.imread(src_annot_dir + image, 0)
if self._F_RAND_FLAG:
self._generator.rand_f(self._F_RANGE)
if self._EXT_RAND_FLAG:
self._generator.rand_ext_params()
result1 = self._generator.transFromColor(src_image)
cv2.imwrite(dst_dir + prefix + image, result1)
print("Image", count, "Done!")
result2 = self._generator.transFromGray(src_annot_image)
cv2.imwrite(dst_annot_dir + prefix + image, result2)
print("Image annot", count, "Done!")
count += 1
print("ALL Done!")
class MyTransform(object):
def __init__(self, focal_len, shape=None):
self._transformer = FishEyeGenerator(focal_len, shape)
self._F_RAND_FLAG = False
self._F_RANGE = [200, 400]
self._EXT_RAND_FLAG = False
self._EXT_PARAM_RANGE = [0, 0, 0, 0, 0, 0]
self._transformer.set_ext_param_range(self._EXT_PARAM_RANGE)
self._RAND_CROP = False
self._rand_crop_rate = 0.8
self._NORMAL_SCALE = False
self._scale_range = [0.5, 2]
self._FISH_SCALE = False
self._fish_scale_range = [0.5, 2]
self._NORMAL_TRANSLATE = False
self._trans_range = [-20, 20]
def set_crop(self, rand=True, rate=0.8):
self._RAND_CROP = rand
self._rand_crop_rate = rate
def set_bkg(self, bkg_label=20, bkg_color=[0, 0, 0]):
self._transformer.set_bkg(bkg_label, bkg_color)
def set_ext_param_range(self, ext_param):
self._EXT_PARAM_RANGE = list(ext_param)
self._transformer.set_ext_param_range(self._EXT_PARAM_RANGE)
def rand_ext_params(self):
self._EXT_RAND_FLAG = True
def set_ext_params(self, ext_params):
self._transformer.set_ext_params(ext_params)
self._EXT_RAND_FLAG = False
def set_f(self, focal_len):
self._transformer.set_f(focal_len)
self._F_RAND_FLAG = False
def rand_f(self, f_range=[200, 400]):
self._F_RANGE = f_range
self._F_RAND_FLAG = True
def _rand_crop(self, image, annot):
rows, cols, channels = image.shape
new_rows = math.floor(rows * self._rand_crop_rate)
new_cols = math.floor(cols * self._rand_crop_rate)
row_start = math.floor((rows - new_rows) * random.random())
col_start = math.floor((cols - new_cols) * random.random())
crop_image = image[row_start:row_start + new_rows,
col_start:col_start + new_cols]
crop_annot = annot[row_start:row_start + new_rows,
col_start:col_start + new_cols]
return crop_image, crop_annot
def _fish_scale(self, image, annot):
borderValue = 20
rate = (random.random() *
(self._fish_scale_range[1] - self._fish_scale_range[0]) +
self._fish_scale_range[0])
if rate == 1:
return image, annot
rows, cols = annot.shape
image = cv2.resize(image, None, fx=rate, fy=rate)
annot = cv2.resize(annot,
None,
fx=rate,
fy=rate,
interpolation=cv2.INTER_NEAREST)
if rate < 1:
dst_image = np.ones((rows, cols, 3), dtype=np.uint8) * 0
dst_annot = np.ones((rows, cols), dtype=np.uint8) * borderValue
row_start = rows // 2 - annot.shape[0] // 2
col_start = cols // 2 - annot.shape[1] // 2
dst_image[row_start:row_start + annot.shape[0],
col_start:col_start + annot.shape[1], ] = image
dst_annot[row_start:row_start + annot.shape[0],
col_start:col_start + annot.shape[1], ] = annot
return dst_image, dst_annot
if rate > 1:
row_start = image.shape[0] // 2 - rows // 2
col_start = image.shape[1] // 2 - cols // 2
crop_image = image[row_start:row_start + rows,
col_start:col_start + cols]
crop_annot = annot[row_start:row_start + rows,
col_start:col_start + cols]
return crop_image, crop_annot
def __call__(self, image, annot):
if self._RAND_CROP:
image, annot = self._rand_crop(image, annot)
if self._NORMAL_SCALE:
scale_rate = (random.random() *
(self._scale_range[1] - self._scale_range[0]) +
self._scale_range[0])
image = cv2.resize(image, None, fx=scale_rate, fy=scale_rate)
annot = cv2.resize(
annot,
None,
fx=scale_rate,
fy=scale_rate,
interpolation=cv2.INTER_NEAREST,
)
if self._F_RAND_FLAG:
self._transformer.rand_f(self._F_RANGE)
if self._EXT_RAND_FLAG:
self._transformer.rand_ext_params()
dst_image = self._transformer.transFromColor(image)
dst_annot = self._transformer.transFromGray(annot, reuse=True)
if self._NORMAL_TRANSLATE:
x_shift = (random.random() *
(self._trans_range[1] - self._trans_range[0]) +
self._trans_range[0])
y_shift = (random.random() *
(self._trans_range[1] - self._trans_range[0]) +
self._trans_range[0])
M = np.array([[1, 0, x_shift], [0, 1, y_shift]], dtype=np.float32)
sz = (dst_annot.shape[1], dst_annot.shape[0])
dst_image = cv2.warpAffine(dst_image, M, sz)
dst_annot = cv2.warpAffine(dst_annot,
M,
sz,
flags=cv2.INTER_NEAREST,
borderValue=20)
if self._FISH_SCALE:
dst_image, dst_annot = self._fish_scale(dst_image, dst_annot)
dst_image = Image.fromarray(dst_image)
dst_annot = Image.fromarray(dst_annot)
brightness, contrast, hue, saturation = 0.1, 0.1, 0.1, 0.1
dst_image = ColorJitter(brightness, contrast, saturation,
hue)(dst_image)
if random.random() < 0.5:
dst_image = dst_image.transpose(Image.FLIP_LEFT_RIGHT)
dst_annot = dst_annot.transpose(Image.FLIP_LEFT_RIGHT)
dst_annot = np.asarray(dst_annot)
dst_annot = torch.from_numpy(dst_annot)
dst_image = ToTensor()(dst_image)
# image = Normalize(mean=[.485, .456, .406], std=[.229, .224, .225])(image)
return dst_image, dst_annot