def __getitem__(self, index):
img = ImageHelper.read_image(self.img_list[index],
tool=self.configer.get('data', 'image_tool'),
mode=self.configer.get('data', 'input_mode'))
img_size = ImageHelper.get_size(img)
bboxes, labels = self.__read_json_file(self.json_list[index])
if self.aug_transform is not None:
img, bboxes, labels = self.aug_transform(img, bboxes=bboxes, labels=labels)
img_scale = ImageHelper.get_size(img)[0] / img_size[0]
labels = torch.from_numpy(labels).long()
bboxes = torch.from_numpy(bboxes).float()
meta = dict(
ori_img_size=img_size,
border_size=ImageHelper.get_size(img),
img_scale=img_scale,
)
if self.img_transform is not None:
img = self.img_transform(img)
return dict(
img=DataContainer(img, stack=True),
bboxes=DataContainer(bboxes, stack=False),
labels=DataContainer(labels, stack=False),
meta=DataContainer(meta, stack=False, cpu_only=True)
)
def __getitem__(self, index):
img = ImageHelper.read_image(
self.img_list[index],
tool=self.configer.get('data', 'image_tool'),
mode=self.configer.get('data', 'input_mode'))
ori_img_size = ImageHelper.get_size(img)
if self.aug_transform is not None:
img = self.aug_transform(img)
border_hw = ImageHelper.get_size(img)[::-1]
if self.img_transform is not None:
img = self.img_transform(img)
meta = dict(ori_img_size=ori_img_size,
border_hw=border_hw,
img_path=self.img_list[index])
return dict(img=DataContainer(img,
stack=True,
return_dc=True,
samples_per_gpu=True),
meta=DataContainer(meta,
stack=False,
cpu_only=True,
return_dc=True,
samples_per_gpu=True))
def make_input(self, image=None, input_size=None,
min_side_length=None, max_side_length=None, scale=None):
in_width, in_height = None, None
if input_size is None and min_side_length is None and max_side_length is None:
in_width, in_height = ImageHelper.get_size(image)
elif input_size is not None and min_side_length is None and max_side_length is None:
in_width, in_height = input_size
elif input_size is None and min_side_length is not None and max_side_length is None:
width, height = ImageHelper.get_size(image)
scale_ratio = min_side_length / min(width, height)
w_scale_ratio, h_scale_ratio = scale_ratio, scale_ratio
in_width, in_height = int(round(width * w_scale_ratio)), int(round(height * h_scale_ratio))
elif input_size is None and min_side_length is None and max_side_length is not None:
width, height = ImageHelper.get_size(image)
scale_ratio = max_side_length / max(width, height)
w_scale_ratio, h_scale_ratio = scale_ratio, scale_ratio
in_width, in_height = int(round(width * w_scale_ratio)), int(round(height * h_scale_ratio))
else:
Log.error('Incorrect target size setting.')
exit(1)
if not isinstance(scale, (list, tuple)):
image = ImageHelper.resize(image, (int(in_width * scale), int(in_height * scale)), interpolation='linear')
img_tensor = ToTensor()(image)
img_tensor = Normalize(div_value=self.configer.get('normalize', 'div_value'),
mean=self.configer.get('normalize', 'mean'),
std=self.configer.get('normalize', 'std'))(img_tensor)
img_tensor = img_tensor.unsqueeze(0).to(torch.device('cpu' if self.configer.get('gpu') is None else 'cuda'))
return img_tensor
else:
img_tensor_list = []
for s in scale:
image = ImageHelper.resize(image, (int(in_width * s), int(in_height * s)), interpolation='linear')
img_tensor = ToTensor()(image)
img_tensor = Normalize(div_value=self.configer.get('normalize', 'div_value'),
mean=self.configer.get('normalize', 'mean'),
std=self.configer.get('normalize', 'std'))(img_tensor)
img_tensor = img_tensor.unsqueeze(0).to(
torch.device('cpu' if self.configer.get('gpu') is None else 'cuda'))
img_tensor_list.append(img_tensor)
return img_tensor_list
def __test_img(self, image_path, json_path, raw_path, vis_path):
Log.info('Image Path: {}'.format(image_path))
img = ImageHelper.read_image(
image_path,
tool=self.configer.get('data', 'image_tool'),
mode=self.configer.get('data', 'input_mode'))
ori_img_bgr = ImageHelper.get_cv2_bgr(img,
mode=self.configer.get(
'data', 'input_mode'))
img, scale = BoundResize()(img)
inputs = self.blob_helper.make_input(img, scale=1.0)
with torch.no_grad():
# Forward pass.
test_group = self.det_net(inputs, scale)
test_indices_and_rois, test_roi_locs, test_roi_scores, test_rois_num = test_group
batch_detections = self.decode(test_roi_locs, test_roi_scores,
test_indices_and_rois,
test_rois_num, self.configer,
ImageHelper.get_size(img))
json_dict = self.__get_info_tree(batch_detections[0],
ori_img_bgr,
scale=scale)
image_canvas = self.det_parser.draw_bboxes(
ori_img_bgr.copy(),
json_dict,
conf_threshold=self.configer.get('vis', 'conf_threshold'))
cv2.imwrite(vis_path, image_canvas)
cv2.imwrite(raw_path, ori_img_bgr)
Log.info('Json Path: {}'.format(json_path))
JsonHelper.save_file(json_dict, json_path)
return json_dict
def __getitem__(self, index):
img = ImageHelper.read_image(
self.img_list[index],
tool=self.configer.get('data', 'image_tool'),
mode=self.configer.get('data', 'input_mode'))
img_size = ImageHelper.get_size(img)
bboxes, labels = self.__read_json_file(self.json_list[index])
if self.aug_transform is not None:
img, bboxes, labels = self.aug_transform(img,
bboxes=bboxes,
labels=labels)
labels = torch.from_numpy(labels).long()
bboxes = torch.from_numpy(bboxes).float()
scale1 = 600 / min(img_size)
scale2 = 1000 / max(img_size)
scale = min(scale1, scale2)
if self.img_transform is not None:
img = self.img_transform(img)
return img, scale, bboxes, labels
def __test_img(self, image_path, save_path):
Log.info('Image Path: {}'.format(image_path))
ori_image = ImageHelper.read_image(
image_path,
tool=self.configer.get('data', 'image_tool'),
mode=self.configer.get('data', 'input_mode'))
ori_width, ori_height = ImageHelper.get_size(ori_image)
ori_img_bgr = ImageHelper.get_cv2_bgr(ori_image,
mode=self.configer.get(
'data', 'input_mode'))
heatmap_avg = np.zeros(
(ori_height, ori_width, self.configer.get('network',
'heatmap_out')))
for i, scale in enumerate(self.configer.get('test', 'scale_search')):
image = self.blob_helper.make_input(ori_image,
input_size=self.configer.get(
'test', 'input_size'),
scale=scale)
with torch.no_grad():
heatmap_out_list = self.pose_net(image)
heatmap_out = heatmap_out_list[-1]
# extract outputs, resize, and remove padding
heatmap = heatmap_out.squeeze(0).cpu().numpy().transpose(
1, 2, 0)
heatmap = cv2.resize(heatmap, (ori_width, ori_height),
interpolation=cv2.INTER_CUBIC)
heatmap_avg = heatmap_avg + heatmap / len(
self.configer.get('test', 'scale_search'))
all_peaks = self.__extract_heatmap_info(heatmap_avg)
image_canvas = self.__draw_key_point(all_peaks, ori_img_bgr)
ImageHelper.save(image_canvas, save_path)
def ms_test(self, ori_image):
ori_width, ori_height = ImageHelper.get_size(ori_image)
total_logits = np.zeros(
(ori_height, ori_width, self.configer.get('data', 'num_classes')),
np.float32)
for scale in self.configer.get('test', 'scale_search'):
image, border_hw = self._get_blob(ori_image, scale=scale)
results = self._predict(image)
results = cv2.resize(results[:border_hw[0], :border_hw[1]],
(ori_width, ori_height),
interpolation=cv2.INTER_CUBIC)
total_logits += results
if self.configer.get('data', 'image_tool') == 'cv2':
mirror_image = cv2.flip(ori_image, 1)
else:
mirror_image = ori_image.transpose(Image.FLIP_LEFT_RIGHT)
image, border_hw = self._get_blob(mirror_image, scale=1.0)
results = self._predict(image)
results = results[:border_hw[0], :border_hw[1]]
results = cv2.resize(results[:, ::-1], (ori_width, ori_height),
interpolation=cv2.INTER_CUBIC)
total_logits += results
return total_logits
def __list_dirs(self, root_dir, dataset):
img_list = list()
label_list = list()
size_list = list()
image_dir = os.path.join(root_dir, dataset, 'image')
label_dir = os.path.join(root_dir, dataset, 'label')
img_extension = os.listdir(image_dir)[0].split('.')[-1]
for file_name in os.listdir(label_dir):
image_name = '.'.join(file_name.split('.')[:-1])
img_path = os.path.join(image_dir,
'{}.{}'.format(image_name, img_extension))
label_path = os.path.join(label_dir, file_name)
if not os.path.exists(label_path) or not os.path.exists(img_path):
Log.error('Label Path: {} not exists.'.format(label_path))
continue
img_list.append(img_path)
label_list.append(label_path)
img = ImageHelper.read_image(
img_path,
tool=self.configer.get('data', 'image_tool'),
mode=self.configer.get('data', 'input_mode'))
size_list.append(ImageHelper.get_size(img))
if dataset == 'train' and self.configer.get('data', 'include_val'):
image_dir = os.path.join(root_dir, 'val/image')
label_dir = os.path.join(root_dir, 'val/label')
for file_name in os.listdir(label_dir):
image_name = '.'.join(file_name.split('.')[:-1])
img_path = os.path.join(
image_dir, '{}.{}'.format(image_name, img_extension))
label_path = os.path.join(label_dir, file_name)
if not os.path.exists(label_path) or not os.path.exists(
img_path):
Log.error('Label Path: {} not exists.'.format(label_path))
continue
img_list.append(img_path)
label_list.append(label_path)
img = ImageHelper.read_image(
img_path,
tool=self.configer.get('data', 'image_tool'),
mode=self.configer.get('data', 'input_mode'))
size_list.append(ImageHelper.get_size(img))
return img_list, label_list, size_list
def __test_img(self, image_path, label_path, vis_path, raw_path):
Log.info('Image Path: {}'.format(image_path))
ori_image = ImageHelper.read_image(image_path,
tool=self.configer.get('data', 'image_tool'),
mode=self.configer.get('data', 'input_mode'))
ori_width, ori_height = ImageHelper.get_size(ori_image)
total_logits = np.zeros((ori_height, ori_width, self.configer.get('data', 'num_classes')), np.float32)
for scale in self.configer.get('test', 'scale_search'):
image = self.blob_helper.make_input(image=ori_image,
input_size=self.configer.get('test', 'input_size'),
scale=scale)
if self.configer.get('test', 'crop_test'):
crop_size = self.configer.get('test', 'crop_size')
if image.size()[3] > crop_size[0] and image.size()[2] > crop_size[1]:
results = self._crop_predict(image, crop_size)
else:
results = self._predict(image)
else:
results = self._predict(image)
results = cv2.resize(results, (ori_width, ori_height), interpolation=cv2.INTER_LINEAR)
total_logits += results
if self.configer.get('test', 'mirror'):
if self.configer.get('data', 'image_tool') == 'cv2':
image = cv2.flip(ori_image, 1)
else:
image = ori_image.transpose(Image.FLIP_LEFT_RIGHT)
image = self.blob_helper.make_input(image, input_size=self.configer.get('test', 'input_size'), scale=1.0)
if self.configer.get('test', 'crop_test'):
crop_size = self.configer.get('test', 'crop_size')
if image.size()[3] > crop_size[0] and image.size()[2] > crop_size[1]:
results = self._crop_predict(image, crop_size)
else:
results = self._predict(image)
else:
results = self._predict(image)
results = cv2.resize(results[:, ::-1], (ori_width, ori_height), interpolation=cv2.INTER_LINEAR)
total_logits += results
label_map = np.argmax(total_logits, axis=-1)
label_img = np.array(label_map, dtype=np.uint8)
image_bgr = cv2.cvtColor(np.array(ori_image), cv2.COLOR_RGB2BGR)
image_canvas = self.seg_parser.colorize(label_img, image_canvas=image_bgr)
ImageHelper.save(image_canvas, save_path=vis_path)
ImageHelper.save(ori_image, save_path=raw_path)
if not self.configer.is_empty('data', 'label_list'):
label_img = self.__relabel(label_img)
label_img = Image.fromarray(label_img, 'P')
Log.info('Label Path: {}'.format(label_path))
ImageHelper.save(label_img, label_path)
def ss_test(self, ori_image):
ori_width, ori_height = ImageHelper.get_size(ori_image)
total_logits = np.zeros((ori_height, ori_width, self.configer.get('data', 'num_classes')), np.float32)
image, border_hw = self._get_blob(ori_image, scale=1.0)
results = self._predict(image)
results = cv2.resize(results[:border_hw[0], :border_hw[1]],
(ori_width, ori_height), interpolation=cv2.INTER_CUBIC)
total_logits += results
return total_logits
def __call__(self, img):
img_size = ImageHelper.get_size(img)
scale1 = self.resize_bound[0] / min(img_size)
scale2 = self.resize_bound[1] / max(img_size)
scale = min(scale1, scale2)
target_size = [int(round(i * scale)) for i in img_size]
img = ImageHelper.resize(img,
target_size=target_size,
interpolation='cubic')
return img, scale
def sscrop_test(self, ori_image):
ori_width, ori_height = ImageHelper.get_size(ori_image)
total_logits = np.zeros((ori_height, ori_width, self.configer.get('data', 'num_classes')), np.float32)
image, _ = self._get_blob(ori_image, scale=1.0)
crop_size = self.configer.get('test', 'crop_size')
if image.size()[3] > crop_size[0] and image.size()[2] > crop_size[1]:
results = self._crop_predict(image, crop_size)
else:
results = self._predict(image)
results = cv2.resize(results, (ori_width, ori_height), interpolation=cv2.INTER_CUBIC)
total_logits += results
return total_logits
def __getitem__(self, index):
img = ImageHelper.read_image(
self.img_list[index],
tool=self.configer.get('data', 'image_tool'),
mode=self.configer.get('data', 'input_mode'))
img_size = ImageHelper.get_size(img)
labelmap = ImageHelper.read_image(self.label_list[index],
tool=self.configer.get(
'data', 'image_tool'),
mode='P')
if self.configer.exists('data', 'label_list'):
labelmap = self._encode_label(labelmap)
if self.configer.exists('data', 'reduce_zero_label'):
labelmap = self._reduce_zero_label(labelmap)
ori_target = ImageHelper.tonp(labelmap)
ori_target[ori_target == 255] = -1
if self.aug_transform is not None:
img, labelmap = self.aug_transform(img, labelmap=labelmap)
border_size = ImageHelper.get_size(img)
if self.img_transform is not None:
img = self.img_transform(img)
if self.label_transform is not None:
labelmap = self.label_transform(labelmap)
meta = dict(ori_img_size=img_size,
border_size=border_size,
ori_target=ori_target)
return dict(
img=DataContainer(img, stack=True),
labelmap=DataContainer(labelmap, stack=True),
meta=DataContainer(meta, stack=False, cpu_only=True),
)
def make_input(self, image=None, input_size=None,
min_side_length=None, max_side_length=None, scale=None):
if input_size is not None and min_side_length is None and max_side_length is None:
if input_size[0] == -1 and input_size[1] == -1:
in_width, in_height = ImageHelper.get_size(image)
elif input_size[0] != -1 and input_size[1] != -1:
in_width, in_height = input_size
elif input_size[0] == -1 and input_size[1] != -1:
width, height = ImageHelper.get_size(image)
scale_ratio = input_size[1] / height
w_scale_ratio, h_scale_ratio = scale_ratio, scale_ratio
in_width, in_height = int(round(width * w_scale_ratio)), int(round(height * h_scale_ratio))
else:
assert input_size[0] != -1 and input_size[1] == -1
width, height = ImageHelper.get_size(image)
scale_ratio = input_size[0] / width
w_scale_ratio, h_scale_ratio = scale_ratio, scale_ratio
in_width, in_height = int(round(width * w_scale_ratio)), int(round(height * h_scale_ratio))
elif input_size is None and min_side_length is not None and max_side_length is None:
width, height = ImageHelper.get_size(image)
scale_ratio = min_side_length / min(width, height)
w_scale_ratio, h_scale_ratio = scale_ratio, scale_ratio
in_width, in_height = int(round(width * w_scale_ratio)), int(round(height * h_scale_ratio))
elif input_size is None and min_side_length is None and max_side_length is not None:
width, height = ImageHelper.get_size(image)
scale_ratio = max_side_length / max(width, height)
w_scale_ratio, h_scale_ratio = scale_ratio, scale_ratio
in_width, in_height = int(round(width * w_scale_ratio)), int(round(height * h_scale_ratio))
elif input_size is None and min_side_length is not None and max_side_length is not None:
width, height = ImageHelper.get_size(image)
scale_ratio = min_side_length / min(width, height)
bound_scale_ratio = max_side_length / max(width, height)
scale_ratio = min(scale_ratio, bound_scale_ratio)
w_scale_ratio, h_scale_ratio = scale_ratio, scale_ratio
in_width, in_height = int(round(width * w_scale_ratio)), int(round(height * h_scale_ratio))
else:
in_width, in_height = ImageHelper.get_size(image)
image = ImageHelper.resize(image, (int(in_width * scale), int(in_height * scale)), interpolation='cubic')
img_tensor = ToTensor()(image)
img_tensor = Normalize(div_value=self.configer.get('normalize', 'div_value'),
mean=self.configer.get('normalize', 'mean'),
std=self.configer.get('normalize', 'std'))(img_tensor)
img_tensor = img_tensor.unsqueeze(0).to(torch.device('cpu' if self.configer.get('gpu') is None else 'cuda'))
return img_tensor
def __test_img(self, image_path, json_path, raw_path, vis_path):
Log.info('Image Path: {}'.format(image_path))
ori_image = ImageHelper.read_image(image_path,
tool=self.configer.get('data', 'image_tool'),
mode=self.configer.get('data', 'input_mode'))
ori_width, ori_height = ImageHelper.get_size(ori_image)
ori_img_bgr = ImageHelper.get_cv2_bgr(ori_image, mode=self.configer.get('data', 'input_mode'))
heatmap_avg = np.zeros((ori_height, ori_width, self.configer.get('network', 'heatmap_out')))
paf_avg = np.zeros((ori_height, ori_width, self.configer.get('network', 'paf_out')))
multiplier = [scale * self.configer.get('test', 'input_size')[1] / ori_height
for scale in self.configer.get('test', 'scale_search')]
stride = self.configer.get('network', 'stride')
for i, scale in enumerate(multiplier):
image, border_hw = self._get_blob(ori_image, scale=scale)
with torch.no_grad():
paf_out_list, heatmap_out_list = self.pose_net(image)
paf_out = paf_out_list[-1]
heatmap_out = heatmap_out_list[-1]
# extract outputs, resize, and remove padding
heatmap = heatmap_out.squeeze(0).cpu().numpy().transpose(1, 2, 0)
heatmap = cv2.resize(heatmap, None, fx=stride, fy=stride, interpolation=cv2.INTER_CUBIC)
heatmap = cv2.resize(heatmap[:border_hw[0], :border_hw[1]],
(ori_width, ori_height), interpolation=cv2.INTER_CUBIC)
paf = paf_out.squeeze(0).cpu().numpy().transpose(1, 2, 0)
paf = cv2.resize(paf, None, fx=stride, fy=stride, interpolation=cv2.INTER_CUBIC)
paf = cv2.resize(paf[:border_hw[0], :border_hw[1]],
(ori_width, ori_height), interpolation=cv2.INTER_CUBIC)
heatmap_avg = heatmap_avg + heatmap / len(multiplier)
paf_avg = paf_avg + paf / len(multiplier)
all_peaks = self.__extract_heatmap_info(heatmap_avg)
special_k, connection_all = self.__extract_paf_info(ori_img_bgr, paf_avg, all_peaks)
subset, candidate = self.__get_subsets(connection_all, special_k, all_peaks)
json_dict = self.__get_info_tree(ori_img_bgr, subset, candidate)
image_canvas = self.pose_parser.draw_points(ori_img_bgr.copy(), json_dict)
image_canvas = self.pose_parser.link_points(image_canvas, json_dict)
ImageHelper.save(image_canvas, vis_path)
ImageHelper.save(ori_img_bgr, raw_path)
Log.info('Json Save Path: {}'.format(json_path))
JsonHelper.save_file(json_dict, json_path)
def __getitem__(self, index):
img = ImageHelper.read_image(
self.img_list[index],
tool=self.configer.get('data', 'image_tool'),
mode=self.configer.get('data', 'input_mode'))
if os.path.exists(self.mask_list[index]):
maskmap = ImageHelper.read_image(self.mask_list[index],
tool=self.configer.get(
'data', 'image_tool'),
mode='P')
else:
maskmap = np.ones((img.size[1], img.size[0]), dtype=np.uint8)
if self.configer.get('data', 'image_tool') == 'pil':
maskmap = ImageHelper.np2img(maskmap)
kpts, bboxes = self.__read_json_file(self.json_list[index])
if self.aug_transform is not None and len(bboxes) > 0:
img, maskmap, kpts, bboxes = self.aug_transform(img,
maskmap=maskmap,
kpts=kpts,
bboxes=bboxes)
elif self.aug_transform is not None:
img, maskmap, kpts = self.aug_transform(img,
maskmap=maskmap,
kpts=kpts)
width, height = ImageHelper.get_size(maskmap)
maskmap = ImageHelper.resize(
maskmap, (width // self.configer.get('network', 'stride'),
height // self.configer.get('network', 'stride')),
interpolation='nearest')
maskmap = torch.from_numpy(np.array(maskmap, dtype=np.float32))
maskmap = maskmap.unsqueeze(0)
kpts = torch.from_numpy(kpts).float()
heatmap = self.heatmap_generator(kpts, [width, height], maskmap)
vecmap = self.paf_generator(kpts, [width, height], maskmap)
if self.img_transform is not None:
img = self.img_transform(img)
return dict(img=DataContainer(img, stack=True),
heatmap=DataContainer(heatmap, stack=True),
maskmap=DataContainer(maskmap, stack=True),
vecmap=DataContainer(vecmap, stack=True))
def __getitem__(self, index):
img = ImageHelper.read_image(
self.img_list[index],
tool=self.configer.get('data', 'image_tool'),
mode=self.configer.get('data', 'input_mode'))
kpts, bboxes = self.__read_json_file(self.json_list[index])
if self.aug_transform is not None:
img, kpts, bboxes = self.aug_transform(img,
kpts=kpts,
bboxes=bboxes)
kpts = torch.from_numpy(kpts).float()
heatmap = self.heatmap_generator(kpts, ImageHelper.get_size(img))
if self.img_transform is not None:
img = self.img_transform(img)
return img, heatmap
def __test_img(self, image_path, json_path, raw_path, vis_path):
Log.info('Image Path: {}'.format(image_path))
image = ImageHelper.read_image(
image_path,
tool=self.configer.get('data', 'image_tool'),
mode=self.configer.get('data', 'input_mode'))
ori_img_bgr = ImageHelper.get_cv2_bgr(image,
mode=self.configer.get(
'data', 'input_mode'))
width, height = ImageHelper.get_size(image)
scale1 = self.configer.get('test', 'resize_bound')[0] / min(
width, height)
scale2 = self.configer.get('test', 'resize_bound')[1] / max(
width, height)
scale = min(scale1, scale2)
inputs = self.blob_helper.make_input(image, scale=scale)
b, c, h, w = inputs.size()
border_wh = [w, h]
if self.configer.exists('test', 'fit_stride'):
stride = self.configer.get('test', 'fit_stride')
pad_w = 0 if (w % stride == 0) else stride - (w % stride) # right
pad_h = 0 if (h % stride == 0) else stride - (h % stride) # down
expand_image = torch.zeros(
(b, c, h + pad_h, w + pad_w)).to(inputs.device)
expand_image[:, :, 0:h, 0:w] = inputs
inputs = expand_image
data_dict = dict(
img=inputs,
meta=DataContainer([[
dict(ori_img_size=ImageHelper.get_size(ori_img_bgr),
aug_img_size=border_wh,
img_scale=scale,
input_size=[inputs.size(3),
inputs.size(2)])
]],
cpu_only=True))
with torch.no_grad():
# Forward pass.
test_group = self.det_net(data_dict)
test_indices_and_rois, test_roi_locs, test_roi_scores, test_rois_num = test_group
batch_detections = self.decode(test_roi_locs, test_roi_scores,
test_indices_and_rois, test_rois_num,
self.configer,
DCHelper.tolist(data_dict['meta']))
json_dict = self.__get_info_tree(batch_detections[0],
ori_img_bgr,
scale=scale)
image_canvas = self.det_parser.draw_bboxes(
ori_img_bgr.copy(),
json_dict,
conf_threshold=self.configer.get('res', 'vis_conf_thre'))
cv2.imwrite(vis_path, image_canvas)
cv2.imwrite(raw_path, ori_img_bgr)
Log.info('Json Path: {}'.format(json_path))
JsonHelper.save_file(json_dict, json_path)
return json_dict
def __getitem__(self, index):
image = ImageHelper.read_image(
self.img_list[index],
tool=self.configer.get('data', 'image_tool'),
mode=self.configer.get('data', 'input_mode'))
img_size = ImageHelper.get_size(image)
if self.configer.exists('test', 'input_size'):
input_size = self.configer.get('test', 'input_size')
if input_size[0] == -1 and input_size[1] == -1:
in_width, in_height = ImageHelper.get_size(image)
elif input_size[0] != -1 and input_size[1] != -1:
in_width, in_height = input_size
elif input_size[0] == -1 and input_size[1] != -1:
width, height = ImageHelper.get_size(image)
scale_ratio = input_size[1] / height
w_scale_ratio, h_scale_ratio = scale_ratio, scale_ratio
in_width, in_height = int(round(width * w_scale_ratio)), int(
round(height * h_scale_ratio))
else:
assert input_size[0] != -1 and input_size[1] == -1
width, height = ImageHelper.get_size(image)
scale_ratio = input_size[0] / width
w_scale_ratio, h_scale_ratio = scale_ratio, scale_ratio
in_width, in_height = int(round(width * w_scale_ratio)), int(
round(height * h_scale_ratio))
elif self.configer.exists(
'test', 'min_side_length') and not self.configer.exists(
'test', 'max_side_length'):
width, height = ImageHelper.get_size(image)
scale_ratio = self.configer.get('test', 'min_side_length') / min(
width, height)
w_scale_ratio, h_scale_ratio = scale_ratio, scale_ratio
in_width, in_height = int(round(width * w_scale_ratio)), int(
round(height * h_scale_ratio))
elif not self.configer.exists(
'test', 'min_side_length') and self.configer.exists(
'test', 'max_side_length'):
width, height = ImageHelper.get_size(image)
scale_ratio = self.configer.get('test', 'max_side_length') / max(
width, height)
w_scale_ratio, h_scale_ratio = scale_ratio, scale_ratio
in_width, in_height = int(round(width * w_scale_ratio)), int(
round(height * h_scale_ratio))
elif self.configer.exists('test',
'min_side_length') and self.configer.exists(
'test', 'max_side_length'):
width, height = ImageHelper.get_size(image)
scale_ratio = self.configer.get('test', 'min_side_length') / min(
width, height)
bound_scale_ratio = self.configer.get(
'test', 'max_side_length') / max(width, height)
scale_ratio = min(scale_ratio, bound_scale_ratio)
w_scale_ratio, h_scale_ratio = scale_ratio, scale_ratio
in_width, in_height = int(round(width * w_scale_ratio)), int(
round(height * h_scale_ratio))
else:
in_width, in_height = ImageHelper.get_size(image)
img = ImageHelper.resize(image, (int(in_width), int(in_height)),
interpolation='linear')
if self.img_transform is not None:
img = self.img_transform(img)
meta = dict(ori_img_size=img_size,
border_hw=[in_height, in_width],
img_path=self.img_list[index])
return dict(img=DataContainer(img,
stack=True,
return_dc=True,
samples_per_gpu=True),
meta=DataContainer(meta,
stack=False,
cpu_only=True,
return_dc=True,
samples_per_gpu=True))
