From 65d533f0a1554d8a583b9a2e39eb98fbe00f2795 Mon Sep 17 00:00:00 2001 From: Administrator Date: Wed, 30 Aug 2023 17:37:53 +0800 Subject: [PATCH] tensor --- AI_example.py | 14 ++++++++------ DMPRUtils/DMPR_process.py | 2 +- DMPR_YOLO/jointUtil.py | 38 +++++++++++++++++--------------------- 3 files changed, 26 insertions(+), 28 deletions(-) diff --git a/AI_example.py b/AI_example.py index 824eff3..d9478ea 100644 --- a/AI_example.py +++ b/AI_example.py @@ -45,10 +45,10 @@ def main(): 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' + # outpth = 'images/output' + outpth = 'images/debug_out' folders = os.listdir(impth) for file in folders: imgpath = os.path.join(impth, file) @@ -58,7 +58,7 @@ def main(): # 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)') @@ -69,6 +69,7 @@ def main(): # DMPR process 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)') @@ -95,7 +96,8 @@ def main(): # # 绘制删除满足 在膨胀框内 && 角度差小于90度 的box for *xyxy, conf, cls, flag in reversed(joint_det): if flag == 0: - label = f'{int(cls)} {conf:.2f}' + # label = f'{int(cls)} {conf:.2f}' + label = None plot_one_box(xyxy, img0, label=label, color=rainbows[int(cls)], line_thickness=2) # save diff --git a/DMPRUtils/DMPR_process.py b/DMPRUtils/DMPR_process.py index 671c632..3fa5fb0 100644 --- a/DMPRUtils/DMPR_process.py +++ b/DMPRUtils/DMPR_process.py @@ -18,7 +18,7 @@ MarkingPoint = namedtuple('MarkingPoint', ['x', 'y', 'direction', 'shape']) def plot_points(image, pred_points, line_thickness=3): """Plot marking points on the image.""" - if len(pred_points): + if pred_points.size: tl = line_thickness or round(0.002 * (image.shape[0] + image.shape[1]) / 2) + 1 # line/font thickness tf = max(tl - 1, 1) # font thickness for conf, *point in pred_points: diff --git a/DMPR_YOLO/jointUtil.py b/DMPR_YOLO/jointUtil.py index c854d7f..1a553ec 100644 --- a/DMPR_YOLO/jointUtil.py +++ b/DMPR_YOLO/jointUtil.py @@ -1,39 +1,35 @@ import math +import numpy as np import torch def dmpr_yolo(dmpr_det, yolo_det, img_shape, cls:int): - device_ = yolo_det.device - - # dmpr_det内容为conf, x, y, θ, shape - if dmpr_det.device != device_: - dmpr_det = dmpr_det.to(device_) # 创建yolo_det_clone内容为x1, y1, x2, y2, conf, cls, unlabel (unlabel代表该类是否需要忽略，0：不忽略 其他：忽略) - yolo_det_clone = yolo_det.clone().detach() - tmp_0_tensor = torch.zeros([len(yolo_det), 1], device=device_) - yolo_det_clone = torch.cat([yolo_det_clone, tmp_0_tensor], dim=1) + yolo_det_clone = yolo_det.copy() + tmp_0_tensor = np.zeros([len(yolo_det), 1]) + yolo_det_clone = np.concatenate([yolo_det_clone, tmp_0_tensor], axis=1) # cls为需要计算的类别 yolo_det = yolo_det[yolo_det[:, -1] == cls] # new_yolo_det为膨胀后数据，内容为x1, y1, x2, y2, flag (flag代表膨胀后车位内是否包含角点 且 与角点方向差值小于90度, 其值为第一个满足条件的角点索引) - new_yolo_det = torch.zeros([len(yolo_det), 5], device=device_) + new_yolo_det = np.zeros([len(yolo_det), 5]) # yolo框膨胀，长的边两边各膨胀0.4倍总长，短的边两边各膨胀0.2倍总长 x_length = yolo_det[:, 2] - yolo_det[:, 0] #x2-x1 y_length = yolo_det[:, 3] - yolo_det[:, 1] #y2-y1 # x, y哪个方向差值大哪个方向膨胀的多 - x_dilate_coefficient = ((x_length > y_length).int() + 1)*0.2 - y_dilate_coefficient = ((~(x_length > y_length)).int() + 1)*0.2 + x_dilate_coefficient = ((x_length > y_length) + 1)*0.2 + y_dilate_coefficient = ((~(x_length > y_length)) + 1)*0.2 # 膨胀 - new_yolo_det[:, 0] = torch.round(yolo_det[:, 0] - x_dilate_coefficient * x_length).clamp_(0, img_shape[1]) #x1 膨胀 - new_yolo_det[:, 1] = torch.round(yolo_det[:, 1] - y_dilate_coefficient * y_length).clamp_(0, img_shape[0]) #y1 膨胀 - new_yolo_det[:, 2] = torch.round(yolo_det[:, 2] + x_dilate_coefficient * x_length).clamp_(0, img_shape[1]) #x2 膨胀 - new_yolo_det[:, 3] = torch.round(yolo_det[:, 3] + y_dilate_coefficient * y_length).clamp_(0, img_shape[0]) #y2 膨胀 + new_yolo_det[:, 0] = np.round(yolo_det[:, 0] - x_dilate_coefficient * x_length).clip(0, img_shape[1]) #x1 膨胀 + new_yolo_det[:, 1] = np.round(yolo_det[:, 1] - y_dilate_coefficient * y_length).clip(0, img_shape[0]) #y1 膨胀 + new_yolo_det[:, 2] = np.round(yolo_det[:, 2] + x_dilate_coefficient * x_length).clip(0, img_shape[1]) #x2 膨胀 + new_yolo_det[:, 3] = np.round(yolo_det[:, 3] + y_dilate_coefficient * y_length).clip(0, img_shape[0]) #y2 膨胀 # 判断膨胀后yolo框包含角点关系 && 包含角点的时候计算水平框中心点与角点的角度关系 # for i in range(0, len(new_yolo_det)): @@ -54,19 +50,19 @@ def dmpr_yolo(dmpr_det, yolo_det, img_shape, cls:int): # new_yolo_det[i, 4] = j + 1 # elif (ang_diff < -180) and (360 + ang_diff <= 90): # new_yolo_det[i, 4] = j + 1 - m, n = len(new_yolo_det), len(dmpr_det) + m, n = new_yolo_det.size, dmpr_det.size if not m or not n: return yolo_det_clone, new_yolo_det - new_yolo = new_yolo_det.unsqueeze(dim=1).repeat(1, n, 1) # 扩展为 (m , n, 5) - dmpr_det = dmpr_det.unsqueeze(dim=0).repeat(m, 1, 1) - yolo_dmpr = torch.cat((new_yolo, dmpr_det), dim=2) # (m, n, 10) + new_yolo = new_yolo_det[:, np.newaxis, :].repeat(dmpr_det.shape[0], 1) # 扩展为 (m , n, 5) + dmpr_det = dmpr_det[np.newaxis, ...].repeat(new_yolo_det.shape[0], 0) + yolo_dmpr = np.concatenate((new_yolo, dmpr_det), axis=2) # (m, n, 10) x_p, y_p = yolo_dmpr[..., 6], yolo_dmpr[..., 7] x1, y1, x2, y2 = yolo_dmpr[..., 0], yolo_dmpr[..., 1], yolo_dmpr[..., 2], yolo_dmpr[..., 3] x_c, y_c = (x1+x2)/2, (y1+y2)/2 - direction1 = torch.atan2(y_c - y_p, x_c - x_p) / math.pi * 180 + direction1 = np.arctan2(y_c - y_p, x_c - x_p) / math.pi * 180 direction2 = yolo_dmpr[..., 8] / math.pi * 180 ang_diff = direction1 - direction2 @@ -75,7 +71,7 @@ def dmpr_yolo(dmpr_det, yolo_det, img_shape, cls:int): 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))) - res = torch.sum(mask, dim=1).float() + res = np.sum(mask, axis=1) # 索引两次更新tensor test1 # yolo_det_clone[yolo_det_clone[:, -2] == cls][:, -1] = new_yolo_det[:, 4]