diff --git a/get_map.py b/get_map.py index a46eb6dfcd28c58f5a7ecd5ed4b8b319b73c713a..7b7e502e39e504b6eeb7232ca17641a79329aede 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 64973789e9908e078e22c2be573543f1ae94015a..d449b4bf42b9afb7009bc83dda7b9a231364df08 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 56191f93d2e41f6412340fbf7d91b62e66c30601..a58aa65e06e58f862c3279e04d16882dd24cae95 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))