tensor

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

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:

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]