def main(model_path, img_folder, save_path, gpu_id):
if os.path.exists(save_path):
shutil.rmtree(save_path, ignore_errors=True)
if not os.path.exists(save_path):
os.makedirs(save_path)
save_img_folder = os.path.join(save_path, 'img')
if not os.path.exists(save_img_folder):
os.makedirs(save_img_folder)
save_txt_folder = os.path.join(save_path, 'result')
if not os.path.exists(save_txt_folder):
os.makedirs(save_txt_folder)
img_paths = [os.path.join(img_folder, x) for x in os.listdir(img_folder)]
model = Pytorch_model(model_path, gpu_id=gpu_id)
total_frame = 0.0
total_time = 0.0
for img_path in tqdm(img_paths):
img_name = os.path.basename(img_path).split('.')[0]
save_name = os.path.join(save_txt_folder, 'res_' + img_name + '.txt')
_, boxes_list, t = model.predict(img_path)
total_frame += 1
total_time += t
img = draw_bbox(img_path, boxes_list, color=(0, 0, 255))
cv2.imwrite(os.path.join(save_img_folder, '{}.jpg'.format(img_name)), img)
np.savetxt(save_name, boxes_list.reshape(-1, 8), delimiter=',', fmt='%d')
print('fps:{}'.format(total_frame / total_time))
return save_txt_folder
def main(model_path, img_folder, save_path, gpu_id):
if os.path.exists(save_path):
shutil.rmtree(save_path, ignore_errors=True)
if not os.path.exists(save_path):
os.makedirs(save_path)
save_img_folder = os.path.join(save_path, 'img')
if not os.path.exists(save_img_folder):
os.makedirs(save_img_folder)
save_txt_folder = os.path.join(save_path, 'result')
if not os.path.exists(save_txt_folder):
os.makedirs(save_txt_folder)
img_paths = [os.path.join(img_folder, x) for x in os.listdir(img_folder)]
model = Pytorch_model(model_path, gpu_id=gpu_id)
total_frame = 0.0
total_time = 0.0
for img_path in tqdm(img_paths):
img_name = os.path.basename(img_path).split('.')[0]
save_name = os.path.join(save_txt_folder, 'res_' + img_name + '.txt')
_, boxes_list, t = model.predict(img_path, short_size=640)
total_frame += 1
total_time += t
#img = draw_bbox(img_path, boxes_list, color=(0, 0, 255))
#cv2.imwrite(os.path.join(save_img_folder, '{}.jpg'.format(img_name)), img)
file_writer = open(save_name, 'w', encoding='utf-8')
for box in boxes_list:
line = ','.join(list(map(str, box))) + '\n'
file_writer.write(line)
file_writer.close()
#np.savetxt(save_name, boxes_list.reshape(-1, 8), delimiter=',', fmt='%d')
print('fps:{}'.format(total_frame / total_time))
return save_txt_folder
def main(model_path, path, save_path, gpu_id):
if os.path.exists(save_path):
shutil.rmtree(save_path, ignore_errors=True)
if not os.path.exists(save_path):
os.makedirs(save_path)
img_paths = [os.path.join(path, x) for x in os.listdir(path)]
net = PSENet(backbone='resnet152', pretrained=False, result_num=config.n)
model = Pytorch_model(model_path, net=net, scale=1, gpu_id=gpu_id)
pbar = tqdm(total=len(img_paths))
for img_path in img_paths:
img_name = os.path.basename(img_path).split('.')[0]
save_name = os.path.join(save_path, 'res_' + img_name + '.txt')
_, preds = model.predict(img_path)
np.savetxt(save_name, preds.reshape(-1, 8), delimiter=',', fmt='%d')
pbar.update(1)
pbar.close()
def main(model_path, backbone, scale, path, save_path, gpu_id):
device = torch.device("cuda:" + str(gpu_id))
logger = setup_logger(os.path.join(config.output_dir, 'test_log'))
logger.info(config.print())
if os.path.exists(save_path):
shutil.rmtree(save_path, ignore_errors=True)
if not os.path.exists(save_path):
os.makedirs(save_path)
save_img_folder = os.path.join(save_path, 'img')
if not os.path.exists(save_img_folder):
os.makedirs(save_img_folder)
save_txt_folder = os.path.join(save_path, 'result')
if not os.path.exists(save_txt_folder):
os.makedirs(save_txt_folder)
img_paths = [os.path.join(path, x) for x in os.listdir(path)]
net = PSENet(backbone=backbone,
pretrained=config.pretrained,
result_num=config.n)
model = Pytorch_model(model_path, net=net, scale=scale, gpu_id=gpu_id)
num_gpus = torch.cuda.device_count()
if num_gpus > 1:
model = nn.DataParallel(model)
model = model.to(device)
model = nn.DataParallel(model)
recall, precision, f1 = merge_eval(model=model,
save_path=os.path.join(
config.output_dir, 'output'),
test_path=config.testroot,
device=device,
base_path=config.base_path,
use_sub=config.use_sub)
logger.info('test: recall: {:.6f}, precision: {:.6f}, f1: {:.6f}'.format(
recall, precision, f1))
# total_frame = 0.0
# total_time = 0.0
# for img_path in tqdm(img_paths):
# img_name = os.path.basename(img_path).split('.')[0]
# save_name = os.path.join(save_txt_folder, 'res_' + img_name + '.txt')
# _, boxes_list, t = model.predict(img_path)
# total_frame += 1
# total_time += t
# # img = draw_bbox(img_path, boxes_list, color=(0, 0, 255))
# # cv2.imwrite(os.path.join(save_img_folder, '{}.jpg'.format(img_name)), img)
# np.savetxt(save_name, boxes_list.reshape(-1, 8), delimiter=',', fmt='%d')
# print('fps:{}'.format(total_frame / total_time))
return save_txt_folder