def predict(self, img):
# # Preprocessing
(rows, cols, channels) = img.shape
image = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
torch.cuda.synchronize()
start = time.time()
x = image.astype(np.float32)
x = (x / 255 - self.means) / self.stds
x = x.astype(np.float32)
x = x[:, :, ::-1].copy()
x = torch.from_numpy(x).permute(2, 0, 1)
x = Variable(x.unsqueeze(0))
if self.cuda_use:
x = x.cuda()
contours, output = self.detector.detect(x)
torch.cuda.synchronize()
end = time.time()
print "Text Detection Time : {}".format(end - start)
image, contours = rescale_result(image, contours, rows, cols)
img_viz = visualize_detection(image, contours)
return img_viz, contours
def inference(detector, test_loader, output_dir):
total_time = 0.
if cfg.exp_name != "MLT2017":
osmkdir(output_dir)
else:
if not os.path.exists(output_dir):
mkdirs(output_dir)
for i, (image, meta) in enumerate(test_loader):
image = to_device(image)
torch.cuda.synchronize()
idx = 0 # test mode can only run with batch_size == 1
# visualization
img_show = image[idx].permute(1, 2, 0).cpu().numpy()
img_show = ((img_show * cfg.stds + cfg.means) * 255).astype(np.uint8)
img_show = cv2.cvtColor(img_show, cv2.COLOR_BGR2RGB)
# get detection result
contours, output = detector.detect(image, img_show)
tr_pred, tcl_pred = output['tr'], output['tcl']
torch.cuda.synchronize()
print('detect {} / {} images: {}.'.format(i + 1, len(test_loader),
meta['image_id'][idx]))
pred_vis = visualize_detection(img_show, contours, tr_pred[1],
tcl_pred[1])
path = os.path.join(cfg.vis_dir, '{}_test'.format(cfg.exp_name),
meta['image_id'][idx])
cv2.imwrite(path, pred_vis)
H, W = meta['Height'][idx].item(), meta['Width'][idx].item()
img_show, contours = rescale_result(img_show, contours, H, W)
# write to file
if cfg.exp_name == "Icdar2015":
fname = "res_" + meta['image_id'][idx].replace('jpg', 'txt')
contours = data_transfer_ICDAR(contours)
write_to_file(contours, os.path.join(output_dir, fname))
elif cfg.exp_name == "TD500":
fname = "res_" + meta['image_id'][idx].replace('JPG', 'txt')
im_show = data_transfer_TD500(contours,
os.path.join(output_dir, fname),
img_show)
id_img = meta['image_id'][idx].replace("img_",
"").replace("JPG", "jpg")
path = os.path.join(cfg.vis_dir, '{}_test'.format(cfg.exp_name),
id_img)
cv2.imwrite(path, im_show)
else:
fname = meta['image_id'][idx].replace('jpg', 'txt')
write_to_file(contours, os.path.join(output_dir, fname))
def inference(detector, test_loader, output_dir):
total_time = 0.
for i, (image, train_mask, tr_mask, tcl_mask, radius_map, sin_map, cos_map,
meta) in enumerate(test_loader):
image, train_mask, tr_mask, tcl_mask, radius_map, sin_map, cos_map = to_device(
image, train_mask, tr_mask, tcl_mask, radius_map, sin_map, cos_map)
torch.cuda.synchronize()
start = time.time()
idx = 0 # test mode can only run with batch_size == 1
# get detection result
contours, output = detector.detect(image)
torch.cuda.synchronize()
end = time.time()
total_time += end - start
fps = (i + 1) / total_time
print('detect {} / {} images: {}. ({:.2f} fps)'.format(
i + 1, len(test_loader), meta['image_id'][idx], fps))
# visualization
tr_pred, tcl_pred = output['tr'], output['tcl']
img_show = image[idx].permute(1, 2, 0).cpu().numpy()
img_show = ((img_show * cfg.stds + cfg.means) * 255).astype(np.uint8)
pred_vis = visualize_detection(img_show, contours, tr_pred[1],
tcl_pred[1])
gt_contour = []
for annot, n_annot in zip(meta['annotation'][idx],
meta['n_annotation'][idx]):
if n_annot.item() > 0:
gt_contour.append(annot[:n_annot].int().cpu().numpy())
gt_vis = visualize_detection(img_show, gt_contour,
tr_mask[idx].cpu().numpy(),
tcl_mask[idx].cpu().numpy())
im_vis = np.concatenate([pred_vis, gt_vis], axis=0)
path = os.path.join(cfg.vis_dir, '{}_test'.format(cfg.exp_name),
meta['image_id'][idx])
cv2.imwrite(path, im_vis)
H, W = meta['Height'][idx].item(), meta['Width'][idx].item()
img_show, contours = rescale_result(img_show, contours, H, W)
# write to file
mkdirs(output_dir)
write_to_file(
contours,
os.path.join(output_dir,
meta['image_id'][idx].replace('jpg', 'txt')))
def inference(detector, test_loader, output_dir):
total_time = 0.0
for i, (image, reg_mask, meta) in enumerate(test_loader):
image, reg_mask = to_device(image, reg_mask)
torch.cuda.synchronize()
start = time.time()
index = 0
contours, aster_text, output = detector.detect(image)
torch.cuda.synchronize()
end = time.time()
total_time += end - start
fps = (i + 1) / total_time
print('detect {} | {} images: {}. ({:.2f} fps)'.format(i + 1, len(test_loader), meta['image_id'][index], fps))
# visualization
pred_mask = output['reg']
img_show = image[index].permute(1, 2, 0).cpu().numpy()
img_show = ((img_show * cfg.stds + cfg.means) * 255).astype(np.uint8)
if (cfg.spotter):
pred_vis = visualize_detection_end_to_end(img_show, contours, aster_text, pred_mask)
else:
pred_vis = visualize_detection(img_show, contours, pred_mask)
gt_contour = []
for annot, n_annot in zip(meta['annotation'][index], meta['n_annotation'][index]):
if n_annot.item() > 0:
gt_contour.append(annot[:n_annot].int().cpu().numpy())
gt_vis = visualize_detection(img_show, gt_contour, reg_mask[index].cpu().numpy())
im_vis = np.concatenate([pred_vis, gt_vis], axis=0)
path = os.path.join(cfg.visualization_directory, '{0}_{1}_test'.format(cfg.dataset_name, cfg.backbone), meta['image_id'][index])
cv2.imwrite(path.replace('.gif', '.png'), im_vis)
H, W = meta['Height'][index].item(), meta['Width'][index].item()
img_show, contours = rescale_result(img_show, contours, H, W)
mkdirs(output_dir)
write_to_file(contours, aster_text,
os.path.join(output_dir, meta['image_id'][index].replace('ts_', '')
.replace('.jpg', '.txt').replace('.JPG', '.txt').replace('.png', '.txt').replace('.gif', '.txt')))
def inference(detector, test_loader, output_dir):
total_time = 0.
for i, (image, meta) in enumerate(test_loader):
# print (image)
image = to_device(image)
torch.cuda.synchronize()
start = time.time()
idx = 0 # test mode can only run with batch_size == 1
# get detection result
contours, output = detector.detect(image)
torch.cuda.synchronize()
end = time.time()
total_time += end - start
fps = (i + 1) / total_time
print('detect {} / {} images: {}. ({:.2f} fps)'.format(
i, len(test_loader), meta['image_id'][idx], fps))
# visualization
tr_pred, tcl_pred = output['tr'], output['tcl']
img_show = image[idx].permute(1, 2, 0).cpu().numpy()
img_show = ((img_show * cfg.stds + cfg.means) * 255).astype(np.uint8)
# print (meta)
H, W = meta['Height'][idx].item(), meta['Width'][idx].item()
img_show, contours = rescale_result(img_show, contours, H, W)
# print (contours)
pred_vis = visualize_detection(img_show, contours)
path = os.path.join(cfg.vis_dir, '{}_deploy'.format(cfg.exp_name),
meta['image_id'][idx])
cv2.imwrite(path, pred_vis)
# write to file
mkdirs(output_dir)
write_to_file(
contours,
os.path.join(output_dir,
meta['image_id'][idx].replace('jpg', 'txt')))
def inference(detector, test_loader, output_dir):
total_time = 0.
post_all_time =0.
net_all_time = 0.
backbone_all_time = 0.
IM_all_time = 0.
detach_all_time =0.
if cfg.exp_name != "MLT2017":
osmkdir(output_dir)
else:
if not os.path.exists(output_dir):
mkdirs(output_dir)
for i, (image, train_mask, tr_mask, meta) in enumerate(test_loader):
image, train_mask, tr_mask = to_device(image, train_mask, tr_mask)
torch.cuda.synchronize()
idx = 0 # test mode can only run with batch_size == 1
# visualization
img_show = image[idx].permute(1, 2, 0).cpu().numpy()
img_show = ((img_show * cfg.stds + cfg.means) * 255).astype(np.uint8)
# compute time
start = time.time()
# get detection result
contours, output, net_time, post_time = detector.detect(image, img_show)
end = time.time()
#total_time += end - start
total_time += (net_time + post_time)
post_all_time += post_time
net_all_time += net_time
backbone_all_time+= output["backbone_time"]
IM_all_time += output["IM_time"]
detach_all_time += output["detach_time"]
fps = (i + 1) / total_time
print('detect {} / {} images: {}. ({:.2f} fps); backbone-time:{:.2f}, IM-time:{:.2f}, post-time:{:0.2f}, Transfer-time:{:.2f}'.format(i + 1, len(test_loader), meta['image_id'][idx], fps, backbone_all_time*1000/(i+1), IM_all_time*1000/(i+1), post_all_time*1000/(i+1), detach_all_time*1000/(i+1)))
if cfg.exp_name == "Icdar2015" or cfg.exp_name == "MLT2017" or cfg.exp_name == "TD500":
pred_vis = visualize_detection(img_show, output['bbox'], output['tr'])
else:
pred_vis = visualize_detection(img_show, contours, output['tr'])
gt_contour = []
for annot, n_annot in zip(meta['annotation'][idx], meta['n_annotation'][idx]):
if n_annot.item() > 0:
gt_contour.append(annot[:n_annot].int().cpu().numpy())
gt_vis = visualize_gt(img_show, gt_contour, tr_mask[idx].cpu().numpy())
im_vis = np.concatenate([pred_vis, gt_vis], axis=0)
path = os.path.join(cfg.vis_dir, '{}_test'.format(cfg.exp_name), meta['image_id'][idx].split(".")[0]+".jpg")
cv2.imwrite(path, im_vis)
H, W = meta['Height'][idx].item(), meta['Width'][idx].item()
img_show, contours = rescale_result(img_show, contours, H, W)
# write to file
if cfg.exp_name == "Icdar2015":
fname = "res_" + meta['image_id'][idx].replace('jpg', 'txt')
contours = data_transfer_ICDAR(contours)
write_to_file(contours, os.path.join(output_dir, fname))
elif cfg.exp_name == "MLT2017":
out_dir = os.path.join(output_dir, str(cfg.checkepoch))
if not os.path.exists(out_dir):
mkdirs(out_dir)
fname = meta['image_id'][idx].split("/")[-1].replace('ts', 'res')
fname = fname.split(".")[0] + ".txt"
data_transfer_MLT2017(contours, os.path.join(out_dir, fname))
elif cfg.exp_name == "TD500":
fname = "res_" + meta['image_id'][idx].split(".")[0]+".txt"
data_transfer_TD500(contours, os.path.join(output_dir, fname))
else:
fname = meta['image_id'][idx].replace('jpg', 'txt')
write_to_file(contours, os.path.join(output_dir, fname))
def inference(detector, test_loader, output_dir):
total_time = 0.
if cfg.exp_name != "MLT2017":
osmkdir(output_dir)
else:
if not os.path.exists(output_dir):
mkdirs(output_dir)
for i, (image, train_mask, tr_mask, tcl_mask, radius_map, sin_map, cos_map, meta) in enumerate(test_loader):
image, train_mask, tr_mask, tcl_mask, radius_map, sin_map, cos_map = to_device(
image, train_mask, tr_mask, tcl_mask, radius_map, sin_map, cos_map)
torch.cuda.synchronize()
start = time.time()
idx = 0 # test mode can only run with batch_size == 1
# visualization
img_show = image[idx].permute(1, 2, 0).cpu().numpy()
img_show = ((img_show * cfg.stds + cfg.means) * 255).astype(np.uint8)
# get detection result
contours, output = detector.detect(image, img_show)
tr_pred, tcl_pred = output['tr'], output['tcl']
torch.cuda.synchronize()
end = time.time()
total_time += end - start
fps = (i + 1) / total_time
print('detect {} / {} images: {}. ({:.2f} fps)'.format(i + 1, len(test_loader), meta['image_id'][idx], fps))
pred_vis = visualize_detection(img_show, contours, tr_pred[1], tcl_pred[1])
gt_contour = []
for annot, n_annot in zip(meta['annotation'][idx], meta['n_annotation'][idx]):
if n_annot.item() > 0:
gt_contour.append(annot[:n_annot].int().cpu().numpy())
gt_vis = visualize_gt(img_show, gt_contour,
tr_mask[idx].cpu().numpy(), tcl_mask[idx, :, :, 0].cpu().numpy())
im_vis = np.concatenate([pred_vis, gt_vis], axis=0)
# path = os.path.join(cfg.vis_dir, '{}_test'.format(cfg.exp_name), meta['image_id'][idx])
# cv2.imwrite(path, im_vis)
H, W = meta['Height'][idx].item(), meta['Width'][idx].item()
img_show, contours = rescale_result(img_show, contours, H, W)
# write to file
if cfg.exp_name == "Icdar2015":
fname = "res_" + meta['image_id'][idx].replace('jpg', 'txt')
contours = data_transfer_ICDAR(contours)
write_to_file(contours, os.path.join(output_dir, fname))
elif cfg.exp_name == "MLT2017":
path = os.path.join(cfg.vis_dir, '{}_test'.format(cfg.exp_name),
meta['image_id'][idx].split("/")[-1])
cv2.imwrite(path, im_vis)
out_dir = os.path.join(output_dir, str(cfg.checkepoch))
if not os.path.exists(out_dir):
mkdirs(out_dir)
fname = meta['image_id'][idx].split("/")[-1].replace('ts', 'res')
fname = fname.split(".")[0] + ".txt"
data_transfer_MLT2017(contours, os.path.join(out_dir, fname))
elif cfg.exp_name == "TD500":
fname = "res_img_" + meta['image_id'][idx].replace('jpg', 'txt')
data_transfer_TD500(contours, os.path.join(output_dir, fname))
else:
fname = meta['image_id'][idx].replace('jpg', 'txt')
write_to_file(contours, os.path.join(output_dir, fname))
def Predict(self,
image_path,
output_img_path="output.jpg",
output_txt_path="output.txt",
tr_thresh=0.4,
tcl_thresh=0.4):
cfg = self.system_dict["local"]["cfg"]
model = self.system_dict["local"]["model"]
start = time.time()
image = pil_load_img(image_path)
transform = BaseTransform(size=cfg.input_size,
mean=cfg.means,
std=cfg.stds)
H, W, _ = image.shape
image, polygons = transform(image)
# to pytorch channel sequence
image = image.transpose(2, 0, 1)
meta = {
'image_id': 0,
'image_path': image_path,
'Height': H,
'Width': W
}
image = torch.from_numpy(np.expand_dims(image, axis=0))
image = to_device(image)
if (self.system_dict["local"]["cfg"].cuda):
torch.cuda.synchronize()
end = time.time()
print("Image loading time: {}".format(end - start))
start = time.time()
detector = TextDetector(model,
tr_thresh=tr_thresh,
tcl_thresh=tcl_thresh)
# get detection result
contours, output = detector.detect(image)
torch.cuda.synchronize()
end = time.time()
print("Inference time - {}".format(end - start))
start = time.time()
tr_pred, tcl_pred = output['tr'], output['tcl']
img_show = image[0].permute(1, 2, 0).cpu().numpy()
img_show = ((img_show * cfg.stds + cfg.means) * 255).astype(np.uint8)
img_show, contours = rescale_result(img_show, contours, H, W)
pred_vis = visualize_detection(img_show, contours)
cv2.imwrite(output_img_path, pred_vis)
# write to file
self.write_to_file(contours, output_txt_path)
end = time.time()
print("Writing output time - {}".format(end - start))
