dir self.S = S # S: 每个cell预测的边界框数量 self.B = B # B: 每个边界框预测的坐标数 self.C = C # C: 物体类别数 self.transform = transform def __len__(self): return len(self.annotations) def __getitem__(self, index): label_path = os.path.join(self.label_dir, self.annotations.iloc[index, 1]) boxes = [] with open(label_path) as f: for label in f.readlines(): class_label, x, y, width, height = [float(x) if float(x) != int(float(x)) else int(x) for x in label.replace("\n", "").split()] boxes.append([class_label, x, y, width, height]) img_path = os.path.join(self.img_dir, self.annotations.iloc[index, 0]) image = Image.open(img_path) boxes = torch.tensor(boxes) if self.transform: image, boxes = self.transform(image, boxes) label_matrix = torch.zeros((self.S, self.S, self.C + 5 * self.B)) for box in boxes: class_label, x, y, width, height = box class_label = int(class_label) # i,j represents the cell row and cell column i, j = int(self.S * y), int(self.S * x) x_cell, y_cell = self.S * x - j, self.S * y - i # width and height of cell width_cell, height_cell = (width * self.S, height * self.S) if label_matrix[i, j, 20] == 0: # Set that there exists an object # Coordinates of bounding box box_coordinates = torch.tensor([x_cell, y_cell, width_cell, height_cell]) label_matrix[i, j, 20:25] = box_coordinates # Set the class label_matrix[i, j, class_label] = 1 # Set no-object confidence to 0 label_matrix[i, j, 25:] = torch.tensor([0] * 5) return image, label_matrix
这是一个 PyTorch 中的数据集类，用于加载 Pascal VOC 数据集的图片和标签。在 `__init__` 方法中，我们读取了 CSV 文件，设置了图片和标签的路径，以及每个 cell 预测的边界框数、每个边界框预测的坐标数和物体类别数。在 `__getitem__` 方法中，我们首先读取了标签文件，解析出每个边界框的类别、坐标和大小，然后加载对应的图片，对图片和标签进行了预处理，并最终返回图片和标签矩阵。其中，标签矩阵的大小为 `(S, S, C + 5B)`，其中 `S` 表示每个 cell 的大小，`C` 表示物体类别数，`B` 表示每个边界框预测的坐标数，这里为 4 个坐标值和 1 个置信度值。在标签矩阵中，每个 cell 中包含了表示是否存在物体的置信度值和对应边界框的坐标和类别信息。
扫一扫
2208
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
