fix(detect): 中心坐标和网格对应关系

Author: zjZSTU

py/detector.py

@@ -43,22 +43,32 @@ def load_data(img_path, xml_path):
     ratio_h = scale_h / h
     ratio_w = scale_w / w
 
-    scale_bndboxs = torch.from_numpy(bndboxs).float()
-    scale_bndboxs[:, 0] = scale_bndboxs[:, 0] * ratio_w
-    scale_bndboxs[:, 1] = scale_bndboxs[:, 1] * ratio_h
-    scale_bndboxs[:, 2] = scale_bndboxs[:, 2] * ratio_w
-    scale_bndboxs[:, 3] = scale_bndboxs[:, 3] * ratio_h
-    scale_bndboxs = scale_bndboxs.int().numpy()
-
+    # [C, H, W] -> [N, C, H, W]
     img = img.unsqueeze(0)
+
     data_dict = {}
     data_dict['src'] = src
     data_dict['src_size'] = (h, w)
     data_dict['bndboxs'] = bndboxs
+    data_dict['name_list'] = name_list
+
     data_dict['img'] = img
     data_dict['scale_size'] = (scale_h, scale_w)
     data_dict['ratio'] = (ratio_h, ratio_w)
-    return img, scale_bndboxs, name_list, data_dict
+
+    return img, data_dict
+
+
+def load_model():
+    model_path = './models/checkpoint_yolo_v1_24.pth'
+    model = YOLO_v1(S=7, B=2, C=3)
+    model.load_state_dict(torch.load(model_path))
+    model.eval()
+    for param in model.parameters():
+        param.requires_grad = False
+    model = model.to(device)
+
+    return model
 
 
 def deform_bboxs(pred_bboxs, data_dict):
@@ -76,8 +86,8 @@ def deform_bboxs(pred_bboxs, data_dict):
         col = int(i % S)
 
         x_center, y_center, box_w, box_h = pred_bboxs[i]
-        bboxs[i, 0] = (row + x_center) * grid_w
-        bboxs[i, 1] = (col + y_center) * grid_h
+        bboxs[i, 0] = (col + x_center) * grid_w
+        bboxs[i, 1] = (row + y_center) * grid_h
         bboxs[i, 2] = box_w * scale_w
         bboxs[i, 3] = box_h * scale_h
     # (x_center, y_center, w, h) -> (xmin, ymin, xmax, ymax)
@@ -103,17 +113,9 @@ def deform_bboxs(pred_bboxs, data_dict):
     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
     # device = "cpu"
 
-    img, bndboxs, name_list, data_dict = load_data('../imgs/cucumber_9.jpg', '../imgs/cucumber_9.xml')
-
-    model_path = './models/checkpoint_yolo_v1_24.pth'
-    model = YOLO_v1(S=7, B=2, C=3)
-    model.load_state_dict(torch.load(model_path))
-    model.eval()
-    for param in model.parameters():
-        param.requires_grad = False
-    model = model.to(device)
-
-    # 缩放图像
+    img, data_dict = load_data('../imgs/cucumber_9.jpg', '../imgs/cucumber_9.xml')
+    model = load_model()
+    # 计算
     outputs = model.forward(img.to(device)).cpu().squeeze(0)
     print(outputs.shape)
 
@@ -140,5 +142,5 @@ def deform_bboxs(pred_bboxs, data_dict):
     # 预测边界框的缩放，回到原始图像
     pred_bboxs = deform_bboxs(pred_cate_bboxs, data_dict)
     # 在原图绘制标注边界框和预测边界框
-    dst = draw.plot_bboxs(data_dict['src'], data_dict['bndboxs'], name_list, pred_bboxs, pred_cates, pred_cate_probs)
-    draw.show(dst)
+    dst = draw.plot_bboxs(data_dict['src'], data_dict['bndboxs'], data_dict['name_list'], pred_bboxs, pred_cates, pred_cate_probs)
+    draw.show(dst)