From 9e2bf52fb2f08a4c25999f82b77fbd15272855ff Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=99=BD=E9=B9=AD=E5=85=88=E7=94=9F?= <766529835@qq.com> Date: Fri, 3 Feb 2023 22:12:14 +0800 Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E6=94=B9=E8=AE=AD=E7=BB=83=E8=BF=87?= =?UTF-8?q?=E7=A8=8B=E4=B8=AD=E7=9A=84eval?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- get_map.py | 2 +- train.py | 2 +- utils/callbacks.py | 40 ++++++++++++++++++++++++++++++---------- 3 files changed, 32 insertions(+), 12 deletions(-) diff --git a/get_map.py b/get_map.py index a46eb6d..7b7e502 100644 --- a/get_map.py +++ b/get_map.py @@ -25,7 +25,7 @@ if __name__ == "__main__": # map_mode为3代表仅仅计算VOC_map。 # map_mode为4代表利用COCO工具箱计算当前数据集的0.50:0.95map。需要获得预测结果、获得真实框后并安装pycocotools才行 #-------------------------------------------------------------------------------------------------------------------# - map_mode = 3 + map_mode = 0 #--------------------------------------------------------------------------------------# # 此处的classes_path用于指定需要测量VOC_map的类别 # 一般情况下与训练和预测所用的classes_path一致即可 diff --git a/train.py b/train.py index 6497378..d449b4b 100644 --- a/train.py +++ b/train.py @@ -41,7 +41,7 @@ if __name__ == "__main__": # Cuda 是否使用Cuda # 没有GPU可以设置成False #---------------------------------# - Cuda = False + Cuda = True #---------------------------------------------------------------------# # distributed 用于指定是否使用单机多卡分布式运行 # 终端指令仅支持Ubuntu。CUDA_VISIBLE_DEVICES用于在Ubuntu下指定显卡。 diff --git a/utils/callbacks.py b/utils/callbacks.py index 56191f9..a58aa65 100644 --- a/utils/callbacks.py +++ b/utils/callbacks.py @@ -7,7 +7,7 @@ matplotlib.use('Agg') import scipy.signal from matplotlib import pyplot as plt from torch.utils.tensorboard import SummaryWriter - +from utils.utils_rbox import rbox2poly, poly2hbb import shutil import numpy as np @@ -145,21 +145,31 @@ class EvalCallback(): if results[0] is None: return - top_label = np.array(results[0][:, 6], dtype = 'int32') - top_conf = results[0][:, 4] * results[0][:, 5] - top_boxes = results[0][:, :4] - + top_label = np.array(results[0][:, 7], dtype = 'int32') + top_conf = results[0][:, 5] * results[0][:, 6] + top_rboxes = results[0][:, :5] + top_polys = rbox2poly(top_rboxes) + #---------------------------------------------------------# + # 将归一化的预测结果变为真实的预测框 + #---------------------------------------------------------# + top_polys[..., [0, 2, 4, 6]] *= image_shape[1] + top_polys[..., [1, 3, 5, 7]] *= image_shape[0] + top_hbbs = poly2hbb(top_polys) top_100 = np.argsort(top_conf)[::-1][:self.max_boxes] - top_boxes = top_boxes[top_100] + top_hbbs = top_hbbs[top_100] top_conf = top_conf[top_100] top_label = top_label[top_100] for i, c in list(enumerate(top_label)): predicted_class = self.class_names[int(c)] - box = top_boxes[i] + hbb = top_hbbs[i] score = str(top_conf[i]) - top, left, bottom, right = box + xc, yc, w, h = hbb + left = xc - w/2 + top = yc - h/2 + right = xc + w/2 + bottom = yc + h/2 if predicted_class not in class_names: continue @@ -190,6 +200,12 @@ class EvalCallback(): #------------------------------# gt_boxes = np.array([np.array(list(map(int,box.split(',')))) for box in line[1:]]) #------------------------------# + # 将polygon转换为hbb + #------------------------------# + hbbs = np.zeros((gt_boxes.shape[0], 6)) + hbbs[:, :5] = poly2hbb(gt_boxes[:, :8]) + hbbs[:, 5] = gt_boxes[:, 8] + #------------------------------# # 获得预测txt #------------------------------# self.get_map_txt(image_id, image, self.class_names, self.map_out_path) @@ -198,8 +214,12 @@ class EvalCallback(): # 获得真实框txt #------------------------------# with open(os.path.join(self.map_out_path, "ground-truth/"+image_id+".txt"), "w") as new_f: - for box in gt_boxes: - left, top, right, bottom, obj = box + for hbb in hbbs: + xc, yc, w, h, obj = hbb + left = xc - w/2 + top = yc - h/2 + right = xc + w/2 + bottom = yc + h/2 obj_name = self.class_names[obj] new_f.write("%s %s %s %s %s\n" % (obj_name, left, top, right, bottom)) -- GitLab