tensor

AI_example.py

@@ -48,11 +48,11 @@ def main():
     DMPRmodel.load_state_dict(torch.load(DMPRweights))
 
     # 图像测试
-    impth = 'images/input'
-    # impth = 'images/debug'
+    # impth = 'images/input'
+    impth = 'images/debug'
     # impth = '/home/thsw/WJ/zjc/AI_old/images/input_0'
-    outpth = 'images/output'
-    # outpth = 'images/debug_out'
+    # outpth = 'images/output'
+    outpth = 'images/debug_out'
     folders = os.listdir(impth)
     for file in folders:
         imgpath = os.path.join(impth, file)

DMPR_YOLO/jointUtil.py

@@ -15,7 +15,7 @@ def dmpr_yolo(dmpr_det, yolo_det, img_shape, cls:int, scale_ratio):
     yolo_det = yolo_det[yolo_det[:, -1] == cls]
 
     # new_yolo_det为膨胀后数据，内容为x1, y1, x2, y2, flag (flag代表膨胀后车位内是否包含角点 且 与角点方向差值小于90度, 其值为第一个满足条件的角点索引)
-    new_yolo_det = np.zeros([len(yolo_det), 5])
+    new_yolo_det = np.zeros([len(yolo_det), 7])
 
     # yolo框膨胀，长的边两边各膨胀0.4倍总长，短的边两边各膨胀0.2倍总长
     x_length = yolo_det[:, 2] - yolo_det[:, 0]      #x2-x1
@@ -25,6 +25,10 @@ def dmpr_yolo(dmpr_det, yolo_det, img_shape, cls:int, scale_ratio):
     x_dilate_coefficient = ((x_length > y_length) + 1)*scale_ratio
     y_dilate_coefficient = ((~(x_length > y_length)) + 1)*scale_ratio
 
+    # 原始框中心点x_c, y_c
+    new_yolo_det[:, 5] = (yolo_det[:, 0] + yolo_det[:, 2]) / 2
+    new_yolo_det[:, 6] = (yolo_det[:, 1] + yolo_det[:, 3]) / 2
+
     # 膨胀
     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 膨胀
@@ -58,12 +62,12 @@ def dmpr_yolo(dmpr_det, yolo_det, img_shape, cls:int, scale_ratio):
     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]
+    x_p, y_p = yolo_dmpr[..., 8], yolo_dmpr[..., 9]
     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
+    x_c, y_c = yolo_dmpr[..., 5], yolo_dmpr[..., 6]
 
     direction1 = np.arctan2(y_c - y_p, x_c - x_p) / math.pi * 180
-    direction2 = yolo_dmpr[..., 8] / math.pi * 180
+    direction2 = yolo_dmpr[..., 10] / math.pi * 180
     # direction3 = (direction2 + 90) if (direction2 + 90) <= 180 else (direction2 - 270)
     direction3 = direction2 + 90                    # L形角点另外一个方向
     direction3[direction3 > 180] -= 360
@@ -74,10 +78,10 @@ def dmpr_yolo(dmpr_det, yolo_det, img_shape, cls:int, scale_ratio):
     # direction ∈ （-180， 180） 若角差大于180，需算补角
     # T形角点比较一个方向，L形角点比较两个方向
     mask = (x_p >= x1) & (x_p <= x2) & (y_p >= y1) & (y_p <= y2) & \
-           (((yolo_dmpr[..., 9] <= 0.5) &                          # T形角点情况
+           (((yolo_dmpr[..., 11] <= 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形角点情况
+            ((yolo_dmpr[..., 11] > 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)) |