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))