def _eval(self):
## todo
self.model.eval()
# torch.cuda.empty_cache() # speed up evaluating after training finished
img_path = os.path.join(self.test_path, 'img')
gt_path = os.path.join(self.test_path, 'gt')
result_save_path = os.path.join(self.save_dir, 'result')
if os.path.exists(result_save_path):
shutil.rmtree(result_save_path, ignore_errors=True)
if not os.path.exists(result_save_path):
os.makedirs(result_save_path)
short_size = 736
# 预测所有测试图片
img_paths = [os.path.join(img_path, x) for x in os.listdir(img_path)]
for img_path in tqdm(img_paths, desc='test models'):
img_name = os.path.basename(img_path).split('.')[0]
save_name = os.path.join(result_save_path,
'res_' + img_name + '.txt')
assert os.path.exists(img_path), 'file is not exists'
img = cv2.imread(img_path)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
h, w = img.shape[:2]
scale = short_size / min(h, w)
img = cv2.resize(img, None, fx=scale, fy=scale)
# 将图片由(w,h)变为(1, img_channel, h, w)
tensor = transforms.ToTensor()(img)
tensor = tensor.unsqueeze_(0)
tensor = tensor.to(self.device)
with torch.no_grad():
torch.cuda.synchronize(self.device)
preds = self.model(tensor)[0]
torch.cuda.synchronize(self.device)
preds, boxes_list = decode(preds)
scale = (preds.shape[1] / w, preds.shape[0] / h)
if len(boxes_list):
boxes_list = boxes_list / scale
np.savetxt(save_name,
boxes_list.reshape(-1, 8),
delimiter=',',
fmt='%d')
# 开始计算 recall precision f1
result_dict = cal_recall_precison_f1(gt_path=gt_path,
result_path=result_save_path)
return result_dict['recall'], result_dict['precision'], result_dict[
'hmean']
def cal_f1(self):
self.logger.info('Calculate Recall Precision F1...')
result_save_path = os.path.join(self.save_dir, 'result')
if os.path.exists(result_save_path):
shutil.rmtree(result_save_path, ignore_errors=True)
if not os.path.exists(result_save_path):
os.makedirs(result_save_path)
self.logger.info("Load predict result from {}".format(self.npy_path))
pred_result = np.load(self.npy_path).tolist()
for_cal_map = pred_result["for_cal_map"]
for i, img_name in enumerate(sorted(os.listdir(self.test_img_path))):
predict_boxes = for_cal_map[img_name]['pred_coord']
seq = []
if predict_boxes is not None:
seq.extend([','.join([str(int(b)) for b in box[:-1]]) + '\n' for box in predict_boxes])
with open(os.path.join(result_save_path, str(os.path.basename(img_name).split('.')[0]) + '.txt'), 'w') as f:
f.writelines(seq)
res_dict = cal_recall_precison_f1(gt_path=self.test_gt_path, result_path=result_save_path)
precision, recall, hmean = res_dict['precision'], res_dict['recall'], res_dict['hmean']
self.logger.info('iou=0.5 ##### precision: {:.6f}, recall: {:.6f}, f1: {:.6f}'.format(precision, recall, hmean))
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
if __name__ == '__main__':
os.environ['CUDA_VISIBLE_DEVICES'] = str('0')
model_path = r'output/PAN_shufflenetv2_FPEM_FFM.pth'
img_path = r'/mnt/e/zj/dataset/icdar2015/test/img'
gt_path = r'/mnt/e/zj/dataset/icdar2015/test/gt'
save_path = './output/result'#model_path.replace('checkpoint/best_model.pth', 'result/')
gpu_id = 0
save_path = main(model_path, img_path, save_path, gpu_id=gpu_id)
result = cal_recall_precison_f1(gt_path=gt_path, result_path=save_path)
print(result)
def _eval_map(self):
self.logger.info('Enter evaluating...')
self.model.eval()
result_save_path = os.path.join(self.save_dir, 'result')
if os.path.exists(result_save_path):
shutil.rmtree(result_save_path, ignore_errors=True)
if not os.path.exists(result_save_path):
os.makedirs(result_save_path)
predict_embeddings, joint_boxes, all_gt_boxes = [], [], []
qbs_words, qbs_queries, qbs_targets, db_targets, gt_targets = [], [], [], [], []
overlaps, used_test_word = [], []
# Compute a mapping from class string to class id...
self.label_encoder.fit([word for word in self.test_words])
# Create queries...
test_unique_words, counts = np.unique(self.test_words,
return_counts=True)
for idx, test_word in enumerate(self.test_words):
gt_targets.extend(self.label_encoder.transform([test_word]))
if test_word not in used_test_word and test_word in test_unique_words:
qbs_words.append(test_word)
qbs_queries.append(self.words_embeddings[test_word])
qbs_targets.extend(self.label_encoder.transform([test_word]))
used_test_word.append(test_word)
for i, (img_file, gt_file) in enumerate(
zip(self.test_img_files, self.test_gt_files)):
self.logger.info('Evaluating {} image: {}'.format(i, img_file))
# Get test gt boxes & gt words...
gt_boxes, gt_words = [], []
with open(gt_file, mode='r', encoding='utf-8') as f:
lines = f.readlines()
for line in lines:
line = line.strip().rstrip('\n').lstrip(
'\ufeff').strip().split(',', maxsplit=8)
gt_boxes.append([int(ver) for ver in line[:8]])
gt_words.append(str(line[-1]).strip().lower())
# Get img...
im = Image.open(img_file)
im = im.convert("RGB")
im, ratio_w, ratio_h = resize_img(
im, long_size=self.config['trainer']['long_size'])
with torch.no_grad():
if str(self.device).__contains__('cuda'):
torch.cuda.synchronize(self.device)
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean=(0.5, 0.5, 0.5),
std=(0.5, 0.5, 0.5))
])
im = transform(im).unsqueeze(0)
im = im.to(self.device)
(predict_score, predict_geo), predict_embed = self.model(im)
if str(self.device).__contains__('cuda'):
torch.cuda.synchronize(self.device)
# Predicting boxes...
predict_boxes, _ = get_boxes(
score=predict_score.squeeze(0).cpu().numpy(),
geo=predict_geo.squeeze(0).cpu().numpy(),
cls_score_thresh=self.cls_score_thresh,
bbox_nms_overlap=self.bbox_nms_overlap)
predict_embed = predict_embed.squeeze(0).cpu().numpy()
if predict_boxes is None:
continue
self.logger.info(
'Idx: {0} ===> Predict result [predict_boxes: {1}; gt_boxes: {2}]'
.format(i, predict_boxes.shape, len(gt_boxes)))
for predict_box in predict_boxes:
min_x = min(predict_box[0], predict_box[2], predict_box[4],
predict_box[6])
max_x = max(predict_box[0], predict_box[2], predict_box[4],
predict_box[6])
min_y = min(predict_box[1], predict_box[3], predict_box[5],
predict_box[7])
max_y = max(predict_box[1], predict_box[3], predict_box[5],
predict_box[7])
w, h = max_x - min_x, max_y - min_y
differ = h * 0.2 if h < w else w * 0.2
min_x, max_x = int((min_x + differ) / 4), int(
(max_x - differ) / 4)
min_y, max_y = int((min_y + differ) / 4), int(
(max_y - differ) / 4)
if min_x > max_x or min_y > max_y:
continue
predict_embeddings.append(
np.mean(predict_embed[:, min_y:max_y, min_x:max_x],
axis=(1, 2)))
predict_boxes = adjust_ratio(predict_boxes, ratio_w, ratio_h)
seq = []
if predict_boxes is not None:
seq.extend([
','.join([str(int(b)) for b in box[:-1]]) + '\n'
for box in predict_boxes
])
with open(
os.path.join(
result_save_path,
str(os.path.basename(img_file).split('.')[0]) +
'.txt'), 'w') as f:
f.writelines(seq)
joint_boxes.extend(predict_boxes[:, :8])
all_gt_boxes.extend(gt_boxes)
gt_boxes = np.array(gt_boxes)
# Calculate overlap...
overlap = cal_overlap(predict_boxes, gt_boxes)
overlaps.append(overlap)
inds = overlap.argmax(axis=1)
db_targets.extend(
self.label_encoder.transform([gt_words[idx] for idx in inds]))
# End evaluate...
db = np.vstack(predict_embeddings) if len(
predict_embeddings) != 0 else np.array(predict_embeddings)
all_overlaps = np.zeros((len(joint_boxes), len(all_gt_boxes)),
dtype=np.float32)
x, y = 0, 0
for o in overlaps:
all_overlaps[y:y + o.shape[0], x:x + o.shape[1]] = o
y += o.shape[0]
x += o.shape[1]
db_targets, qbs_targets, qbs_words = np.array(db_targets), np.array(
qbs_targets), np.array(qbs_words)
qbs_queries, joint_boxes = np.array(qbs_queries), np.array(joint_boxes)
assert (qbs_queries.shape[0] == qbs_targets.shape[0])
assert (db.shape[0] == db_targets.shape[0])
self.logger.info('Calculate mAP...')
mAP_qbs, mR_qbs = cal_map(qbs_queries,
qbs_targets,
db,
db_targets,
gt_targets,
joint_boxes,
all_overlaps,
self.query_nms_overlap,
self.overlap_thresh,
qbs_words,
num_workers=0)
mAP_qbs, mR_qbs = np.mean(mAP_qbs * 100), np.mean(mR_qbs * 100)
# Calculate recall precision f1
res_dict = cal_recall_precison_f1(gt_path=self.test_gt_path,
result_path=result_save_path)
return res_dict['recall'], res_dict['precision'], res_dict[
'hmean'], mAP_qbs, mR_qbs