def _crop_img_2d(img_path, annt_path, out_dir, sub_array=[3, 2]):
get_sub_path = lambda sub_index, ori_path, out_dir: out_dir + '/' + os.path.splitext(
os.path.basename(ori_path))[0] + '-%d' % (
sub_index) + os.path.splitext(os.path.basename(ori_path))[1]
img_ori = Image.open(img_path)
img_size = img_ori.size
grid_array = ([0] + [
math.floor((i + 1) * img_size[0] / float(sub_array[0]))
for i in range(sub_array[0])
], [0] + [
math.floor((i + 1) * img_size[1] / float(sub_array[1]))
for i in range(sub_array[1])
])
annt_info, template_tree = ut.load_annt_file(annt_path)
sub_index = 0
for sub_index_x in range(sub_array[0]):
for sub_index_y in range(sub_array[1]):
sub_img_path = get_sub_path(sub_index, img_path, out_dir)
sub_annt_path = get_sub_path(sub_index, annt_path, out_dir)
x_min_img = grid_array[0][sub_index_x]
y_min_img = grid_array[1][sub_index_y]
x_max_img = grid_array[0][sub_index_x + 1] - 1
y_max_img = grid_array[1][sub_index_y + 1] - 1
img_box = np.array([x_min_img, y_min_img, x_max_img, y_max_img])
objs_info = _prune_bndbox(annt_info, img_box)
ut.save_annt_file(sub_annt_path, objs_info, template_tree)
sub_img = img_ori.crop(
(x_min_img, y_min_img, x_max_img + 1, y_max_img + 1))
sub_img.save(sub_img_path)
sub_index += 1
def _crop_img_slide_2d(img_path, annt_path, out_dir, sub_array=[3, 2]): get_sub_path = lambda sub_index, ori_path, out_dir:out_dir + '/' + os.path.splitext(os.path.basename(ori_path))[0] + '-%d'%(sub_index) + os.path.splitext(os.path.basename(ori_path))[1] img_ori = PIL.Image.open(img_path) img_size = img_ori.size sub_size = [math.floor(img_size[idx]/float(num_itm)) for idx, num_itm in enumrate(sub_array)] sub_array = [num_itm*2-1 for num_itm in sub_array] grid_array = ([math.floor(i*img_size[0]/float(sub_array[0]+1)) for i in range(sub_array[0])], [math.floor(i*img_size[1]/float(sub_array[1]+1)) for i in range(sub_array[1])]) annt_info, template_tree = ut.load_annt_file(annt_path) sub_index = 0 for sub_index_x in range(sub_array[0]): for sub_index_y in range(sub_array[1]): sub_img_path = get_sub_path(sub_index, img_path, out_dir) sub_annt_path = get_sub_path(sub_index, annt_path, out_dir) x_min_img = grid_array[0][sub_index_x] y_min_img = grid_array[1][sub_index_y] x_max_img = grid_array[0][sub_index_x] + sub_size[0] y_max_img = grid_array[1][sub_index_y] + sub_size[1] img_box = np.array([x_min_img, y_min_img, x_max_img, y_max_img]) objs_info = _prune_bndbox(annt_info, img_box) ut.save_annt_file(sub_annt_path, objs_info, template_tree) sub_img = img_ori.crop((x_min_img, y_min_img, x_max_img+1, y_max_img+1)) sub_img.save(sub_img_path) sub_index += 1
def _restore_img_slide_2d(img_path,
input_dir,
output_dir,
sub_array=[3, 2],
img_suffix='.jpg'):
get_sub_path = lambda sub_index, file_name, output_dir, file_suffix: output_dir + '/' + file_name + '-%d' % (
sub_index) + file_suffix
file_name = os.path.splitext(os.path.basename(img_path))[0]
img_ori = Image.open(img_path)
img_size = img_ori.size
sub_size = [
math.floor(img_size[idx] / float(num_itm))
for idx, num_itm in enumrate(sub_array)
]
sub_array = [num_itm * 2 - 1 for num_itm in sub_array]
grid_array = ([
math.floor(i * img_size[0] / float(sub_array[0] + 1))
for i in range(sub_array[0])
], [
math.floor(i * img_size[1] / float(sub_array[1] + 1))
for i in range(sub_array[1])
])
boxes = np.empty(shape=[0, 4], dtype=np.uint16)
gt_classes = []
ishards = np.empty(shape=[0], dtype=np.int32)
sub_index = 0
for sub_index_x in range(sub_array[0]):
for sub_index_y in range(sub_array[1]):
sub_annt_path = get_sub_path(sub_index, file_name, input_dir,
'.xml')
sub_index += 1
x_min_img = grid_array[0][sub_index_x]
y_min_img = grid_array[1][sub_index_y]
x_max_img = grid_array[0][sub_index_x] + sub_size[0]
y_max_img = grid_array[1][sub_index_y] + sub_size[1]
img_box = np.array([x_min_img, y_min_img, x_max_img, y_max_img])
objs_info, template_tree = ut.load_annt_file(sub_annt_path)
objs_info_restored = _restore_cord(objs_info, img_box)
if objs_info_restored == None:
continue
boxes = np.append(boxes, objs_info['boxes'], axis=0)
ishards = np.append(ishards, objs_info['gt_ishard'], axis=0)
gt_classes.extend(objs_info['gt_classes_name'])
objs_info = {
'boxes': boxes,
'gt_classes_name': gt_classes,
'gt_ishard': ishards
}
annt_path = output_dir + '/' + file_name + '.xml'
ut.save_annt_file(annt_path, objs_info)
def _aug_img_2d(img_path,
annt_path,
out_dir,
own_data_classes,
aug_factor=60,
trim_size=1024,
random_seed=False,
overlap_th=0.4,
ratio_th=2.5):
get_sub_path = lambda sub_index, ori_path, out_dir: os.path.join(
out_dir,
os.path.splitext(os.path.basename(ori_path))[0] + '-%d' %
(sub_index) + os.path.splitext(os.path.basename(ori_path))[1])
try:
im = imread(img_path)
annt_info, template_tree = ut.load_annt_file(annt_path)
except ValueError as e:
print(e)
return
except OSError as e:
print(e)
return
if len(annt_info['gt_classes_name']) == 0:
print('no bndboxes in this file.')
return
# random crop
if random_seed:
random.seed(random_seed)
gt_boxes_ori = torch.from_numpy(annt_info['boxes'].astype(np.float))
class_to_ind_dict = dict(
zip(own_data_classes, range(len(own_data_classes))))
ind_to_class_dict = dict(
zip(range(len(own_data_classes)), own_data_classes))
gt_classes_np = np.array(
[class_to_ind_dict[item] for item in annt_info['gt_classes_name']])
gt_classes_ori = torch.from_numpy(gt_classes_np)
gt_area = (gt_boxes_ori[:, 2] - gt_boxes_ori[:, 0]) * (gt_boxes_ori[:, 3] -
gt_boxes_ori[:, 1])
gt_ratio = (gt_boxes_ori[:, 2] -
gt_boxes_ori[:, 0]) / (gt_boxes_ori[:, 3] - gt_boxes_ori[:, 1])
gt_ratio[gt_ratio > ratio_th - 1e-5] = 0
gt_ratio[gt_ratio < 1.0 / ratio_th + 1e-5] = 0
if torch.sum(gt_ratio) < 1e-5:
return
# crop img
height, width = im.shape[0], im.shape[1]
min_x_bnd = max(0, torch.min(gt_boxes_ori[:, 0]) - trim_size)
min_y_bnd = max(0, torch.min(gt_boxes_ori[:, 1]) - trim_size)
max_x_bnd = min(width - trim_size, torch.max(gt_boxes_ori[:, 2]))
max_y_bnd = min(height - trim_size, torch.max(gt_boxes_ori[:, 3]))
sub_index = 0
for sub_index_x in range(aug_factor):
sub_img_path = get_sub_path(sub_index, img_path, out_dir)
sub_annt_path = get_sub_path(sub_index, annt_path, out_dir)
while True:
gt_boxes = gt_boxes_ori.clone()
gt_classes = gt_classes_ori.clone()
x_min_img, y_min_img = random.randint(min_x_bnd,
max_x_bnd), random.randint(
min_y_bnd, max_y_bnd)
gt_boxes[:, ::2] = gt_boxes[:, ::2] - x_min_img
gt_boxes[:, 1::2] = gt_boxes[:, 1::2] - y_min_img
gt_boxes.clamp_(0, trim_size - 1)
overlap = (gt_boxes[:, 2] - gt_boxes[:, 0]) * (
gt_boxes[:, 3] - gt_boxes[:, 1]) / gt_area
ratio = (gt_boxes[:, 2] - gt_boxes[:, 0]) / (gt_boxes[:, 3] -
gt_boxes[:, 1])
boxes_width = gt_boxes[:, 2] - gt_boxes[:, 0]
boxes_height = gt_boxes[:, 3] - gt_boxes[:, 1]
# remove bndbox
not_keep = (overlap < overlap_th) | (ratio > ratio_th) | (
ratio <
(1 / ratio_th)) | (boxes_width < 5) | (boxes_height < 5)
keep = torch.nonzero(not_keep == 0).view(-1)
if keep.numel() == 0:
continue
gt_boxes = gt_boxes[keep]
gt_classes = gt_classes[keep]
annt_info['boxes'] = gt_boxes.numpy().astype(np.uint16)
annt_info['gt_classes_name'] = [
ind_to_class_dict[item]
for item in list(gt_classes.numpy().astype(np.int32))
]
annt_info['ishards'] = np.zeros(len(
annt_info['gt_classes_name'])).astype(np.int32)
img = im[y_min_img:y_min_img + trim_size,
x_min_img:x_min_img + trim_size, :]
break
pil_img = Image.fromarray(img)
ut.save_annt_file(sub_annt_path, annt_info, template_tree)
pil_img.save(sub_img_path)
sub_index += 1