diff --git a/.idea/AI.iml b/.idea/AI.iml index 5bfdd80..513351e 100644 --- a/.idea/AI.iml +++ b/.idea/AI.iml @@ -4,7 +4,7 @@ - + \ No newline at end of file diff --git a/.idea/deployment.xml b/.idea/deployment.xml index f3eb157..dd2f1c8 100644 --- a/.idea/deployment.xml +++ b/.idea/deployment.xml @@ -2,6 +2,13 @@ + + + + + + + diff --git a/.idea/misc.xml b/.idea/misc.xml index 729b1a3..da088d4 100644 --- a/.idea/misc.xml +++ b/.idea/misc.xml @@ -1,4 +1,4 @@ - + \ No newline at end of file diff --git a/AI_example.py b/AI_example.py index d9478ea..d565713 100644 --- a/AI_example.py +++ b/AI_example.py @@ -6,6 +6,7 @@ import torch from DMPRUtils.DMPR_process import DMPR_process, plot_points from DMPRUtils.model.detector import DirectionalPointDetector +from DMPRUtils.yolo_net import Model from DMPR_YOLO.jointUtil import dmpr_yolo from conf import config from models.experimental import attempt_load @@ -16,12 +17,12 @@ from utils.torch_utils import select_device def main(): ##预先设置的参数 - device_ = '0' ##选定模型，可选 cpu,'0','1' + device_ = '1' ##选定模型，可选 cpu,'0','1' ##以下参数目前不可改 Detweights = 'weights/urbanManagement/yolo/best.pt' seg_nclass = 2 - DMPRweights = "weights/urbanManagement/DMPR/dp_detector_499.pth" + DMPRweights = "weights/urbanManagement/DMPR/dp_detector_299.pth" conf_thres, iou_thres, classes = 0.25, 0.45, 3 labelnames = "weights/yolov5/class5/labelnames.json" rainbows = [[0, 0, 255], [0, 255, 0], [255, 0, 0], [255, 0, 255], [255, 255, 0], [255, 129, 0], [255, 0, 127], @@ -41,26 +42,29 @@ def main(): # DMPR model args = config.get_parser_for_inference().parse_args() - DMPRmodel = DirectionalPointDetector(3, args.depth_factor, config.NUM_FEATURE_MAP_CHANNEL).to(device) + # DMPRmodel = DirectionalPointDetector(3, args.depth_factor, config.NUM_FEATURE_MAP_CHANNEL).to(device) + # DMPRmodel.load_state_dict(torch.load(DMPRweights)) + DMPRmodel = Model(args.cfg, ch=3).to(device) DMPRmodel.load_state_dict(torch.load(DMPRweights)) # 图像测试 - # impth = 'images/input' - impth = 'images/debug' - # outpth = 'images/output' - outpth = 'images/debug_out' + impth = 'images/input' + # impth = 'images/debug' + # impth = '/home/thsw/WJ/zjc/AI_old/images/input_0' + outpth = 'images/output' + # outpth = 'images/debug_out' folders = os.listdir(impth) for file in folders: imgpath = os.path.join(impth, file) img0 = cv2.imread(imgpath) assert img0 is not None, 'Image Not Found ' + imgpath - # t_start = time.time() + t_start = time.time() # yolo process det0 = yolo_process(img0, model, device, args, half) det0 = det0.cpu().detach().numpy() - # t_yolo = time.time() - # print(f't_yolo. ({t_yolo - t_start:.3f}s)') + t_yolo = time.time() + print(f't_yolo. ({t_yolo - t_start:.3f}s)') # plot所有box # for *xyxy, conf, cls in reversed(det0): @@ -71,8 +75,8 @@ def main(): det1 = DMPR_process(img0, DMPRmodel, device, args) det1 = det1.cpu().detach().numpy() - # t_dmpr = time.time() - # print(f't_dmpr. ({t_dmpr - t_yolo:.3f}s)') + t_dmpr = time.time() + print(f't_dmpr. ({t_dmpr - t_yolo:.3f}s)') # 绘制角点 plot_points(img0, det1) @@ -82,10 +86,10 @@ def main(): # yolo joint DMPR cls = 0 #需要过滤的box类别 - joint_det, dilate_box = dmpr_yolo(det1, det0, img0.shape, cls) + joint_det, dilate_box = dmpr_yolo(det1, det0, img0.shape, cls, args.scale_ratio) - # t_joint = time.time() - # print(f't_joint. ({t_joint - t_dmpr:.3f}s)') + t_joint = time.time() + print(f't_joint. ({t_joint - t_dmpr:.3f}s)') # t_end = time.time() # print(f'Done. ({t_end - t_start:.3f}s)') diff --git a/DMPRUtils/DMPR_process.py b/DMPRUtils/DMPR_process.py index 3fa5fb0..0409deb 100644 --- a/DMPRUtils/DMPR_process.py +++ b/DMPRUtils/DMPR_process.py @@ -1,5 +1,6 @@ import math import os +import time from collections import namedtuple import cv2 @@ -132,7 +133,12 @@ def get_predicted_points2(prediction, thresh): def detect_marking_points(detector, image, thresh, device): """Given image read from opencv, return detected marking points.""" + t1 = time.time() + torch.cuda.synchronize(device) prediction = detector(preprocess_image(image).to(device)) + torch.cuda.synchronize(device) + t2 = time.time() + print(f'detector: {t2 - t1:.3f}s') return get_predicted_points2(prediction[0], thresh) def scale_coords2(img1_shape, coords, img0_shape, ratio_pad=None): diff --git a/DMPR_YOLO/jointUtil.py b/DMPR_YOLO/jointUtil.py index 1a553ec..05a69ef 100644 --- a/DMPR_YOLO/jointUtil.py +++ b/DMPR_YOLO/jointUtil.py @@ -4,7 +4,7 @@ import numpy as np import torch -def dmpr_yolo(dmpr_det, yolo_det, img_shape, cls:int): +def dmpr_yolo(dmpr_det, yolo_det, img_shape, cls:int, scale_ratio): # 创建yolo_det_clone内容为x1, y1, x2, y2, conf, cls, unlabel (unlabel代表该类是否需要忽略，0：不忽略 其他：忽略) yolo_det_clone = yolo_det.copy() @@ -22,8 +22,8 @@ def dmpr_yolo(dmpr_det, yolo_det, img_shape, cls:int): y_length = yolo_det[:, 3] - yolo_det[:, 1] #y2-y1 # x, y哪个方向差值大哪个方向膨胀的多 - x_dilate_coefficient = ((x_length > y_length) + 1)*0.2 - y_dilate_coefficient = ((~(x_length > y_length)) + 1)*0.2 + x_dilate_coefficient = ((x_length > y_length) + 1)*scale_ratio + y_dilate_coefficient = ((~(x_length > y_length)) + 1)*scale_ratio # 膨胀 new_yolo_det[:, 0] = np.round(yolo_det[:, 0] - x_dilate_coefficient * x_length).clip(0, img_shape[1]) #x1 膨胀 @@ -64,12 +64,24 @@ def dmpr_yolo(dmpr_det, yolo_det, img_shape, cls:int): direction1 = np.arctan2(y_c - y_p, x_c - x_p) / math.pi * 180 direction2 = yolo_dmpr[..., 8] / math.pi * 180 + # direction3 = (direction2 + 90) if (direction2 + 90) <= 180 else (direction2 - 270) + direction3 = direction2 + 90 # L形角点另外一个方向 + direction3[direction3 > 180] -= 360 ang_diff = direction1 - direction2 + ang_diff2 = direction1 - direction3 # 判断膨胀后yolo框包含角点关系 & & 包含角点的时候计算水平框中心点与角点的角度关系 # direction ∈ （-180， 180） 若角差大于180，需算补角 + # T形角点比较一个方向，L形角点比较两个方向 mask = (x_p >= x1) & (x_p <= x2) & (y_p >= y1) & (y_p <= y2) & \ - (((ang_diff >= -90) & (ang_diff <= 90)) | ((ang_diff > 180) & ((360 - ang_diff) <= 90)) | (((ang_diff) < -180) & ((360 + ang_diff) <= 90))) + (((yolo_dmpr[..., 9] <= 0.5) & # T形角点情况 + (((ang_diff >= -90) & (ang_diff <= 90)) | ((ang_diff > 180) & ((360 - ang_diff) <= 90)) | + (((ang_diff) < -180) & ((360 + ang_diff) <= 90)))) | + ((yolo_dmpr[..., 9] > 0.5) & # L形角点情况 + (((ang_diff >= -90) & (ang_diff <= 90)) | ((ang_diff > 180) & ((360 - ang_diff) <= 90)) | + (((ang_diff) < -180) & ((360 + ang_diff) <= 90))) & + (((ang_diff2 >= -90) & (ang_diff2 <= 90)) | ((ang_diff2 > 180) & ((360 - ang_diff2) <= 90)) | + (((ang_diff2) < -180) & ((360 + ang_diff2) <= 90))))) res = np.sum(mask, axis=1) diff --git a/conf/config.py b/conf/config.py index 1a5c801..7689197 100644 --- a/conf/config.py +++ b/conf/config.py @@ -5,11 +5,14 @@ import argparse # 5: sin(direction) NUM_FEATURE_MAP_CHANNEL = 6 - def add_common_arguments(parser): """Add common arguments for training and inference.""" parser.add_argument('--detector_weights', default=r'E:\pycharmProject\DMPR-PS\weights\dp_detector_499.pth', help="The weights of pretrained detector.") + parser.add_argument('--cfg', type=str, default='models/yolov5s.yaml', + help='model.yaml path') + parser.add_argument('--hyp', type=str, default='conf/hyp.scratch.yaml', + help='hyperparameters path') parser.add_argument('--depth_factor', type=int, default=32, help="Depth factor.") parser.add_argument('--disable_cuda', action='store_true', @@ -41,6 +44,8 @@ def get_parser_for_inference(): help='object confidence threshold') parser.add_argument('--iou-thres', type=float, default=0.45, help='IOU threshold for NMS') + parser.add_argument('--scale-ratio', type=float, default=0.5, + help='detected box scale ratio') parser.add_argument('--ovlap-thres', type=float, default=0.6, help='overlap threshold for OBS') parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS') diff --git a/weights/urbanManagement/DMPR/dp_detector_250.pth b/weights/urbanManagement/DMPR/dp_detector_250.pth deleted file mode 100644 index e092608..0000000 Binary files a/weights/urbanManagement/DMPR/dp_detector_250.pth and /dev/null differ diff --git a/weights/urbanManagement/DMPR/dp_detector_499.pth b/weights/urbanManagement/DMPR/dp_detector_499.pth deleted file mode 100644 index e9fc784..0000000 Binary files a/weights/urbanManagement/DMPR/dp_detector_499.pth and /dev/null differ diff --git a/weights/urbanManagement/yolo/best.pt b/weights/urbanManagement/yolo/best.pt index 4542fb3..53efa68 100644 Binary files a/weights/urbanManagement/yolo/best.pt and b/weights/urbanManagement/yolo/best.pt differ diff --git a/weights/urbanManagement/yolo/last.pt b/weights/urbanManagement/yolo/last.pt index b1cad8c..b7d3b04 100644 Binary files a/weights/urbanManagement/yolo/last.pt and b/weights/urbanManagement/yolo/last.pt differ