fix(dataset): 边界框中心对应的网格计算

Author: zjZSTU

py/lib/models/location_dataset.py

@@ -13,14 +13,13 @@
 from torch.utils.data import DataLoader
 from torch.utils.data import Dataset
 from utils import file
+from utils import util
 import torchvision.transforms as transforms
 
-cate_list = ['cucumber', 'eggplant', 'mushroom']
-
 
 class LocationDataset(Dataset):
 
-    def __init__(self, root_dir, transform=None, S=7, B=2, C=20):
+    def __init__(self, root_dir, cate_list, transform=None, S=7, B=2, C=20):
         """
         保存图像以及标注框性能
         :param root_dir: 根目录
@@ -34,6 +33,7 @@ def __init__(self, root_dir, transform=None, S=7, B=2, C=20):
         self.S = S
         self.B = B
         self.C = C
+        self.cate_list = cate_list
 
         jpeg_path_list = []
         xml_path_list = []
@@ -54,11 +54,13 @@ def __getitem__(self, index):
         assert index < len(self.jpeg_path_list), 'image length: %d' % len(self.jpeg_path_list)
 
         # print(self.jpeg_path_list[index])
-        image = cv2.imread(self.jpeg_path_list[index])
+        img_path = self.jpeg_path_list[index]
+        image = cv2.imread(img_path)
         img_h, img_w = image.shape[:2]
-        ratio_h = 1
-        ratio_w = 1
+        ratio_h = 1.0
+        ratio_w = 1.0
         if self.transform:
+            # [H, W, C] -> [C, H, W]
             image = self.transform(image)
             # 计算图像缩放比例
             dst_img_h, dst_img_w = image.shape[1:3]
@@ -73,7 +75,8 @@ def __getitem__(self, index):
 
         target = torch.zeros((self.S * self.S, self.C + self.B * 5))
         bndboxs, name_list = file.parse_location_xml(self.xml_path_list[index])
-        # 缩放边界框坐标（x, y, w, h）
+        bndboxs = util.bbox_corner_to_center(bndboxs)
+        # 缩放边界框坐标（x_center, y_center, w, h）
         bndboxs[:, 0] = bndboxs[:, 0] * ratio_w
         bndboxs[:, 1] = bndboxs[:, 1] * ratio_h
         bndboxs[:, 2] = bndboxs[:, 2] * ratio_w
@@ -89,6 +92,8 @@ def __getitem__(self, index):
             # 边界框中心位于哪个网格
             grid_x = int(box_x / grid_w)
             grid_y = int(box_y / grid_h)
+            # 行/列从０开始计数
+            print(grid_x + 1, grid_y + 1)
             # 边界框中心相对于网格的比例（0,1）
             x = (box_x % grid_w) / grid_w
             y = (box_y % grid_h) / grid_h
@@ -97,22 +102,24 @@ def __getitem__(self, index):
             h = box_h / img_h
             # 该网格内是否已填充（每个网格1个标注边界框）
             if grid_nums[grid_x, grid_y] > 1:
-                print('网格(%d, %d)已填充：%s' % (grid_x, grid_y, str(target[grid_x, grid_y])))
+                print('网格(%d, %d)已填充：%s' % (grid_x, grid_y, img_path))
             else:
+                grid_nums[grid_x, grid_y] = 1
+
                 # 转换类别和标签
-                cate_idx = cate_list.index(name)
+                cate_idx = self.cate_list.index(name)
+                # 指定网格
+                grid_idx = self.S * grid_y + grid_x
                 # 指定类别概率为1
-                target[grid_x * grid_y, cate_idx] = 1
+                target[grid_idx, cate_idx] = 1
                 for j in range(self.B):
                     # 置信度
-                    target[grid_x * grid_y, self.C + j] = 1
+                    target[grid_idx, self.C + j] = 1
                     # 相应的边界框坐标
-                    target[grid_x * grid_y, self.C + self.B + 4 * j] = x
-                    target[grid_x * grid_y, self.C + self.B + 4 * j + 1] = y
-                    target[grid_x * grid_y, self.C + self.B + 4 * j + 2] = w
-                    target[grid_x * grid_y, self.C + self.B + 4 * j + 3] = h
-
-                grid_nums[grid_x, grid_y] += 1
+                    target[grid_idx, self.C + self.B + 4 * j] = x
+                    target[grid_idx, self.C + self.B + 4 * j + 1] = y
+                    target[grid_idx, self.C + self.B + 4 * j + 2] = w
+                    target[grid_idx, self.C + self.B + 4 * j + 3] = h
 
         return image, target
 
@@ -129,17 +136,18 @@ def __len__(self):
         transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
     ])
 
-    data_set = LocationDataset(root_dir, transform, 7, 2, 3)
-    # print(data_set)
-    # print(len(data_set))
-    #
-    # image, target = data_set.__getitem__(3)
-    # print(image.shape)
-    # print(target.shape)
-    # print(target)
-
-    data_loader = DataLoader(data_set, shuffle=True, batch_size=8, num_workers=8)
-    items = next(iter(data_loader))
-    inputs, labels = items
-    print(inputs.shape)
-    print(labels.shape)
+    cate_list = ['cucumber', 'eggplant', 'mushroom']
+    data_set = LocationDataset(root_dir, cate_list, transform, 7, 2, 3)
+    print(data_set)
+    print(len(data_set))
+
+    image, target = data_set.__getitem__(3)
+    print(image.shape)
+    print(target.shape)
+    print(target)
+
+    # data_loader = DataLoader(data_set, shuffle=True, batch_size=8, num_workers=8)
+    # items = next(iter(data_loader))
+    # inputs, labels = items
+    # print(inputs.shape)
+    # print(labels.shape)

py/lib/utils/file.py

@@ -69,8 +69,7 @@ def parse_location_xml(xml_path):
         else:
             pass
 
-        bndboxs = np.array(bndboxs)
-        return bndboxs, name_list
+        return np.array(bndboxs), name_list
 
 
 def save_model(model, model_save_path):