def get_patch(self, input, size_must_mode=1):
scale = self.args.scale
if self.train:
input = utils.get_patch(input, patch_size=self.args.patch_size, scale=scale)
h, w, c = input.shape
new_h, new_w = h - h % size_must_mode, w - w % size_must_mode
input = input[:new_h, :new_w, :]
if not self.args.no_augment:
input = utils.data_augment(input)
else:
h, w, _ = input.shape
new_h, new_w = h - h % size_must_mode, w - w % size_must_mode
input = input[:new_h, :new_w, :]
return input
def get_patch(self, img1, img2, size_must_mode=1):
if self.train:
img1, img2 = utils.get_patch(img1,
img2,
patch_size=self.args.patch_size)
h, w, c = img1.shape
new_h, new_w = h - h % size_must_mode, w - w % size_must_mode
img1, img2 = img1[:new_h, :new_w, :], img2[:new_h, :new_w, :]
if not self.args.no_augment:
img1, img2 = utils.data_augment(img1, img2)
else:
# resize to accelerate the validation process
img1 = cv2.resize(img1, (960, 540))
img2 = cv2.resize(img2, (960, 540))
h, w, c = img1.shape
new_h, new_w = h - h % size_must_mode, w - w % size_must_mode
img1, img2 = img1[:new_h, :new_w, :], img2[:new_h, :new_w, :]
return img1, img2
def get_patch(self, input, gt, size_must_mode=1):
if self.train:
input, gt = utils.get_patch(input,
gt,
patch_size=self.args.patch_size)
h, w, c = input.shape
new_h, new_w = h - h % size_must_mode, w - w % size_must_mode
input, gt = input[:new_h, :new_w, :], gt[:new_h, :new_w, :]
if not self.args.no_augment:
input, gt = utils.data_augment(input, gt)
else:
h, w, c = input.shape
# resize to accelerate the validation process
input = cv2.resize(input, (960, 540))
gt = cv2.resize(gt, (960, 540))
new_h, new_w = h - h % size_must_mode, w - w % size_must_mode
input, gt = input[:new_h, :new_w, :], gt[:new_h, :new_w, :]
return input, gt
def __getitem__(self, index):
img_uri = self.img_files[index]
img_labels = self.labels[index]
# don't download, since it was already downloaded in the init
img_path = img_uri
img_name = ("_".join(map(str, img_path.split("_")[-5:])))
orig_img = PIL.Image.open(img_path).convert('RGB')
if orig_img is None:
raise Exception(
"Empty image: {img_path}".format(img_path=img_path))
if self.vis_batch and len(img_labels) > 0:
vis_orig_img = copy.deepcopy(orig_img)
labels = add_class_dimension_to_labels(img_labels)
labels = xyhw2xyxy_corner(labels, skip_class_dimension=True)
tmp_path = os.path.join(visualization_tmp_path,
img_name[:-4] + ".jpg")
visualize_and_save_to_local(vis_orig_img,
labels,
tmp_path,
box_color="green")
print(f'new image uploaded to {tmp_path}')
# First, handle image re-shaping
if self.ts:
scale = self.scales[index]
scaled_img = scale_image(orig_img, scale)
scaled_img_width, scaled_img_height = scaled_img.size
patch_width, patch_height = self.width, self.height
vert_pad, horiz_pad = pre_tile_padding(scaled_img_width,
scaled_img_height,
patch_width, patch_height)
padded_img = torchvision.transforms.functional.pad(
scaled_img,
padding=(horiz_pad, vert_pad, horiz_pad, vert_pad),
fill=(127, 127, 127),
padding_mode="constant")
padded_img_width, padded_img_height = padded_img.size
_, _, n_patches, _, _ = get_patch_spacings(padded_img_width,
padded_img_height,
patch_width,
patch_height)
patch_index = random.randint(0, n_patches - 1)
if self.debug_mode:
patch_index = 0
img, boundary = get_patch(padded_img, patch_width, patch_height,
patch_index)
else:
orig_img_width, orig_img_height = orig_img.size
vert_pad, horiz_pad, ratio = calculate_padding(
orig_img_height, orig_img_width, self.height, self.width)
img = torchvision.transforms.functional.pad(
orig_img,
padding=(horiz_pad, vert_pad, horiz_pad, vert_pad),
fill=(127, 127, 127),
padding_mode="constant")
img = torchvision.transforms.functional.resize(
img, (self.height, self.width))
# If no labels, no need to do augmentation (this should change in the future)
# so immediately return with the padded image and empty labels
if len(img_labels) == 0:
labels = torch.zeros((len(img_labels), 5))
img = torchvision.transforms.functional.to_tensor(img)
labels = F.pad(
labels,
pad=[0, 0, 0, self.num_targets_per_image - len(labels)],
mode="constant")
return img_uri, img, labels
# Next, handle label re-shaping
labels = add_class_dimension_to_labels(img_labels)
labels = xyhw2xyxy_corner(labels)
if self.ts:
labels = scale_labels(labels, self.scales[index])
labels = add_padding_on_each_side(labels, horiz_pad, vert_pad)
if self.vis_batch:
tmp_path = os.path.join(visualization_tmp_path,
img_name[:-4] + "_scaled.jpg")
visualize_and_save_to_local(padded_img,
labels,
tmp_path,
box_color="red")
labels_temp = filter_and_offset_labels(labels, boundary)
if self.vis_batch:
pre_vis_labels = copy.deepcopy(labels)
for i in range(n_patches):
vis_patch_img, boundary = get_patch(
padded_img, patch_width, patch_height, i)
labels = filter_and_offset_labels(pre_vis_labels, boundary)
tmp_path = os.path.join(visualization_tmp_path, img_name[:-4] + \
"_patch_{}.jpg".format(i))
visualize_and_save_to_local(vis_patch_img,
labels,
tmp_path,
box_color="blue")
if self.upload_dataset:
pre_vis_labels = copy.deepcopy(labels)
for i in range(n_patches):
vis_patch_img, boundary = get_patch(
padded_img, patch_width, patch_height, i)
labels = filter_and_offset_labels(pre_vis_labels, boundary)
tmp_path = os.path.join(visualization_tmp_path, img_name[:-4] + \
"_patch_{}.jpg".format(i))
upload_label_and_image_to_gcloud(vis_patch_img, labels,
tmp_path)
else:
labels = filter_and_offset_labels(labels, boundary)
else:
labels = add_padding_on_each_side(labels, horiz_pad, vert_pad)
labels = scale_labels(labels, ratio)
labels_temp = labels
if self.vis_batch:
tmp_path = os.path.join(visualization_tmp_path,
img_name[:-4] + "_pad_resized.jpg")
visualize_and_save_to_local(img,
labels,
tmp_path,
box_color="blue")
labels = labels_temp
if self.vis_batch and self.data_aug:
vis_aug_img = copy.deepcopy(img)
tmp_path = os.path.join(visualization_tmp_path,
img_name[:-4] + "_before_aug.jpg")
visualize_and_save_to_local(vis_aug_img,
labels,
tmp_path,
box_color="red")
if self.augment_hsv or self.data_aug:
if random.random() > 0.5:
img = self.jitter(img)
# no transformation on labels
# Augment image and labels
img_width, img_height = img.size
if self.augment_affine or self.data_aug:
if random.random() > 0:
angle = random.uniform(-10, 10)
translate = (random.uniform(-40,
40), random.uniform(-40,
40)) ## WORKS
scale = random.uniform(0.9, 1.1)
shear = random.uniform(-3, 3)
img = torchvision.transforms.functional.affine(img,
angle,
translate,
scale,
shear,
2,
fillcolor=(127,
127,
127))
labels = affine_labels(img_height, img_width, labels, -angle,
translate, scale, (-shear, 0))
if self.bw:
img = torchvision.transforms.functional.to_grayscale(
img, num_output_channels=1)
# random left-right flip
if self.lr_flip:
if random.random() > 0.5:
img = torchvision.transforms.functional.hflip(img)
# Is this correct?
# Not immediately obvious, when composed with the angle shift above
labels[:, 1] = img_width - labels[:, 1]
labels[:, 3] = img_width - labels[:, 3]
# GaussianBlur, needs further development
if self.blur:
if random.random() > 0.2:
arr = np.asarray(img)
angle = random.uniform(40, -40)
sigma = random.uniform(0, 3.00)
seq = iaa.Sequential([iaa.GaussianBlur(sigma=sigma)])
images_aug = seq.augment_images(arr)
img = PIL.Image.fromarray(np.uint8(images_aug), 'RGB')
#AdditiveGaussianNoise
if self.noise:
if random.random() > 0.3:
arr = np.asarray(img)
scale = random.uniform(0, 0.03 * 255)
seq = iaa.Sequential([
iaa.AdditiveGaussianNoise(loc=0,
scale=scale,
per_channel=0.5)
])
images_aug = seq.augment_images(arr)
img = PIL.Image.fromarray(np.uint8(images_aug), 'RGB')
#SigmoidContrast, need further development
if self.contrast:
if random.random() > 0.5:
arr = np.asarray(img)
cutoff = random.uniform(0.45, 0.75)
gain = random.randint(5, 10)
seq = iaa.Sequential(
[iaa.SigmoidContrast(gain=gain, cutoff=cutoff)])
images_aug = seq.augment_images(arr)
img = PIL.Image.fromarray(np.uint8(images_aug), 'RGB')
#Sharpen, need further development
if self.sharpen:
if random.random() > 0.3:
arr = np.asarray(img)
alpha = random.uniform(0, 0.5)
seq = iaa.Sharpen(alpha=alpha)
images_aug = seq.augment_images(arr)
img = PIL.Image.fromarray(np.uint8(images_aug), 'RGB')
if self.vis_batch and self.data_aug:
vis_post_aug_img = copy.deepcopy(img)
tmp_path = os.path.join(visualization_tmp_path,
img_name[:-4] + "_post_augmentation.jpg")
visualize_and_save_to_local(vis_post_aug_img,
labels,
tmp_path,
box_color="green")
if self.vis_batch:
self.vis_counter += 1
if self.vis_counter > (self.vis_batch - 1):
sys.exit('Finished visualizing enough images. Exiting!')
labels[:, 1:5] = xyxy2xywh(labels[:, 1:5])
labels[:, (1, 3)] /= self.width
labels[:, (2, 4)] /= self.height
img = torchvision.transforms.functional.to_tensor(img)
labels = F.pad(labels,
pad=[0, 0, 0, self.num_targets_per_image - len(labels)],
mode="constant")
if (labels < 0).sum() > 0:
raise Exception(f"labels for image {img_uri} have negative values")
return img_uri, img, labels