DETR纯代码分享（三）coco_panoptic.py(datasets)
一、Python模块和库
import json
from pathlib import Path
import numpy as np
import torch
from PIL import Image
from panopticapi.utils import rgb2id
from util.box_ops import masks_to_boxes
from .coco import make_coco_transforms
这段代码包含了一系列导入语句，用于导入所需的Python库、模块和函数。
-
import json: 导入Python的JSON模块，用于处理JSON格式的数据。 -
from pathlib import Path: 从Python的pathlib模块中导入Path类，用于处理文件路径和目录。 -
import numpy as np: 导入NumPy库，并将其重命名为np，用于进行数值计算和数组操作。 -
import torch: 导入PyTorch深度学习框架，用于神经网络模型和张量操作。 -
from PIL import Image: 从Python Imaging Library（PIL）中导入Image模块，用于图像处理。 -
from panopticapi.utils import rgb2id: 从panopticapi.utils模块中导入rgb2id函数，该函数用于将RGB图像转换为分割图像的标识符（ID）。 -
from util.box_ops import masks_to_boxes: 从自定义的util.box_ops模块中导入masks_to_boxes函数，该函数用于将掩码（masks）转换为边界框（boxes）。 -
from .coco import make_coco_transforms: 导入当前目录下的coco模块中的make_coco_transforms函数。注意，这里的.表示当前目录，因此代码尝试从当前目录导入coco模块中的make_coco_transforms函数。
这些导入语句用于为后续代码提供所需的功能和模块。根据代码的上下文，这些库、模块和函数可能在处理图像和COCO数据集时使用。
二、CocoPanoptic类
class CocoPanoptic:
def __init__(self, img_folder, ann_folder, ann_file, transforms=None, return_masks=True):
with open(ann_file, 'r') as f:
self.coco = json.load(f)
# sort 'images' field so that they are aligned with 'annotations'
# i.e., in alphabetical order
self.coco['images'] = sorted(self.coco['images'], key=lambda x: x['id'])
# sanity check
if "annotations" in self.coco:
for img, ann in zip(self.coco['images'], self.coco['annotations']):
assert img['file_name'][:-4] == ann['file_name'][:-4]
self.img_folder = img_folder
self.ann_folder = ann_folder
self.ann_file = ann_file
self.transforms = transforms
self.return_masks = return_masks
def __getitem__(self, idx):
ann_info = self.coco['annotations'][idx] if "annotations" in self.coco else self.coco['images'][idx]
img_path = Path(self.img_folder) / ann_info['file_name'].replace('.png', '.jpg')
ann_path = Path(self.ann_folder) / ann_info['file_name']
img = Image.open(img_path).convert('RGB')
w, h = img.size
if "segments_info" in ann_info:
masks = np.asarray(Image.open(ann_path), dtype=np.uint32)
masks = rgb2id(masks)
ids = np.array([ann['id'] for ann in ann_info['segments_info']])
masks = masks == ids[:, None, None]
masks = torch.as_tensor(masks, dtype=torch.uint8)
labels = torch.tensor([ann['category_id'] for ann in ann_info['segments_info']], dtype=torch.int64)
target = {}
target['image_id'] = torch.tensor([ann_info['image_id'] if "image_id" in ann_info else ann_info["id"]])
if self.return_masks:
target['masks'] = masks
target['labels'] = labels
target["boxes"] = masks_to_boxes(masks)
target['size'] = torch.as_tensor([int(h), int(w)])
target['orig_size'] = torch.as_tensor([int(h), int(w)])
if "segments_info" in ann_info:
for name in ['iscrowd', 'area']:
target[name] = torch.tensor([ann[name] for ann in ann_info['segments_info']])
if self.transforms is not None:
img, target = self.transforms(img, target)
return img, target
def __len__(self):
return len(self.coco['images'])
def get_height_and_width(self, idx):
img_info = self.coco['images'][idx]
height = img_info['height']
width = img_info['width']
return height, width提供了一个名为CocoPanoptic的Python类，它用于处理COCO数据集中的图像和注释信息。该类的主要目的是帮助有效地加载和处理COCO数据集，以便用于深度学习模型的训练或评估。它将图像、掩码、标签和边界框等信息打包成可供模型使用的格式。
1、初始化对象
class CocoPanoptic:
def __init__(self, img_folder, ann_folder, ann_file, transforms=None, return_masks=True):
with open(ann_file, 'r') as f:
self.coco = json.load(f)
# sort 'images' field so that they are aligned with 'annotations'
# i.e., in alphabetical order
self.coco['images'] = sorted(self.coco['images'], key=lambda x: x['id'])
# sanity check
if "annotations" in self.coco:
for img, ann in zip(self.coco['images'], self.coco['annotations']):
assert img['file_name'][:-4] == ann['file_name'][:-4]
self.img_folder = img_folder
self.ann_folder = ann_folder
self.ann_file = ann_file
self.transforms = transforms
self.return_masks = return_masks这是一个名为CocoPanoptic的Python类的构造函数，用于初始化对象。以下是构造函数的主要功能和参数解释：
-
__init__(self, img_folder, ann_folder, ann_file, transforms=None, return_masks=True)：这是类的构造函数，用于创建一个CocoPanoptic对象。它接受以下参数：-
img_folder：一个字符串，表示存储图像文件的文件夹路径。 -
ann_folder：一个字符串，表示存储注释文件的文件夹路径。 -
ann_file：一个字符串，表示COCO注释文件的路径，通常是一个JSON文件。 -
transforms：一个可选参数，表示数据转换（transform）操作的函数。如果不提供，则默认为None。 -
return_masks：一个可选参数，是一个布尔值，控制是否返回掩码信息。默认值为True。
-
构造函数的主要任务包括：
-
打开并加载指定的COCO注释文件（
ann_file）并将其存储在类属性self.coco中。 -
对COCO数据集中的'images'字段进行排序，以确保图像信息与注释信息对齐，排序是根据图像字典中的'id'字段进行的，以便按字母顺序排列。
sorted()函数用于对列表进行排序。在这里，key=lambda x: x['id']是排序的关键函数，它告诉sorted函数按照图像的ID字段进行排序。这将按照图像ID的升序对图像进行排序，以便它们以递增顺序排列。 -
执行数据一致性检查，比较每个图像文件名的前缀（去除文件扩展名）是否与对应的注释文件名前缀相匹配。这是一种基本的一致性检查，以确保图像和注释之间的正确对应关系。
[:-4]是一个断言语句，用于检查每个图像的文件名（去掉文件扩展名）是否与相应注释的文件名（去掉文件扩展名）匹配。如果不匹配，这将引发AssertionError，指示存在问题。 -
存储传递给构造函数的各种参数作为对象的属性，包括图像文件夹路径（
img_folder）、注释文件夹路径（ann_folder）、注释文件路径（ann_file）、数据转换函数（transforms）以及是否返回掩码信息（return_masks）。
该类的构造函数用于准备COCO数据集的基本信息，以便在后续的方法中使用，例如__getitem__方法用于获取图像和目标信息，__len__方法用于获取图像总数等。
2、__getitem__方法，用于获取COCO数据集中指定索引（idx）的图像和目标信息
def __getitem__(self, idx):
ann_info = self.coco['annotations'][idx] if "annotations" in self.coco else self.coco['images'][idx]
img_path = Path(self.img_folder) / ann_info['file_name'].replace('.png', '.jpg')
ann_path = Path(self.ann_folder) / ann_info['file_name']
img = Image.open(img_path).convert('RGB')
w, h = img.size
if "segments_info" in ann_info:
masks = np.asarray(Image.open(ann_path), dtype=np.uint32)
masks = rgb2id(masks)
ids = np.array([ann['id'] for ann in ann_info['segments_info']])
masks = masks == ids[:, None, None]
masks = torch.as_tensor(masks, dtype=torch.uint8)
labels = torch.tensor([ann['category_id'] for ann in ann_info['segments_info']], dtype=torch.int64)
target = {}
target['image_id'] = torch.tensor([ann_info['image_id'] if "image_id" in ann_info else ann_info["id"]])
if self.return_masks:
target['masks'] = masks
target['labels'] = labels
target["boxes"] = masks_to_boxes(masks)
target['size'] = torch.as_tensor([int(h), int(w)])
target['orig_size'] = torch.as_tensor([int(h), int(w)])
if "segments_info" in ann_info:
for name in ['iscrowd', 'area']:
target[name] = torch.tensor([ann[name] for ann in ann_info['segments_info']])
if self.transforms is not None:
img, target = self.transforms(img, target)
return img, target（1）加载的图像和相应的注释文件的位置
def __getitem__(self, idx):
ann_info = self.coco['annotations'][idx] if "annotations" in self.coco else self.coco['images'][idx]
img_path = Path(self.img_folder) / ann_info['file_name'].replace('.png', '.jpg')
ann_path = Path(self.ann_folder) / ann_info['file_name']这段代码是CocoPanoptic类的__getitem__方法的一部分。在这部分代码中，它执行以下操作：
-
ann_info = self.coco['annotations'][idx] if "annotations" in self.coco else self.coco['images'][idx]：根据索引（idx）获取与图像或注释相关的信息。这行代码首先检查self.coco字典中是否存在'annotations'字段，如果存在，则根据索引获取相应的注释信息；如果不存在'annotations'字段，则根据索引获取相应的图像信息。这是因为COCO数据集中包含两种类型的数据：图像和注释。 -
img_path = Path(self.img_folder) / ann_info['file_name'].replace('.png', '.jpg')：构建图像文件的路径。它使用Path对象将self.img_folder（图像文件夹路径）与ann_info['file_name']（文件名）拼接在一起，并将.png扩展名替换为.jpg，以确保正确找到图像文件。这是因为通常情况下，COCO数据集中的图像以.jpg格式存储。 -
ann_path = Path(self.ann_folder) / ann_info['file_name']：构建注释文件的路径。它使用Path对象将self.ann_folder（注释文件夹路径）与ann_info['file_name']（文件名）拼接在一起，以获取注释文件的完整路径。
这些操作帮助确定要加载的图像和相应的注释文件的位置。在处理COCO数据集时，通常需要根据图像信息和注释信息来加载相应的图像和注释，以便进行后续的处理和分析。
（2）用于处理图像和分割信息（掩码）
img = Image.open(img_path).convert('RGB')
w, h = img.size
if "segments_info" in ann_info:
masks = np.asarray(Image.open(ann_path), dtype=np.uint32)
masks = rgb2id(masks)
ids = np.array([ann['id'] for ann in ann_info['segments_info']])
masks = masks == ids[:, None, None]
masks = torch.as_tensor(masks, dtype=torch.uint8)
labels = torch.tensor([ann['category_id'] for ann in ann_info['segments_info']], dtype=torch.int64)这段代码是CocoPanoptic类的__getitem__方法的一部分，用于处理图像和分割信息（掩码）。
-
img = Image.open(img_path).convert('RGB')：这行代码使用PIL库打开图像文件，然后将其转换为RGB模式。img_path是之前构建的图像文件路径，convert('RGB')确保图像以RGB格式加载，以便后续的处理和分析。 -
w, h = img.size：这行代码获取加载的图像的宽度（w）和高度（h）信息，这些信息将用于后续的处理。 -
if "segments_info" in ann_info:：这行代码检查ann_info字典中是否包含'segments_info'字段，以确定是否存在分割信息（掩码）。 -
如果存在分割信息：
-
masks = np.asarray(Image.open(ann_path), dtype=np.uint32)：这行代码使用PIL库打开注释文件，然后将其转换为NumPy数组。ann_path是之前构建的注释文件路径，并且dtype=np.uint32指定数据类型为uint32，以适应掩码数据。 -
masks = rgb2id(masks)：调用rgb2id函数，将RGB格式的掩码转换为分割图像的标识符（ID）。这个过程将每个像素的颜色映射到一个唯一的ID，用于表示不同的分割对象。 -
ids = np.array([ann['id'] for ann in ann_info['segments_info']])：这行代码从分割信息中提取ID，将其存储为NumPy数组。 -
masks = masks == ids[:, None, None]：这行代码将掩码与ID进行比较，创建一个布尔掩码，其中每个像素的值表示是否属于相应的分割对象。它使用广播（broadcasting）来实现对每个掩码的逐元素比较。 -
masks = torch.as_tensor(masks, dtype=torch.uint8)：最后，将布尔掩码转换为PyTorch张量，并将数据类型指定为uint8。这个张量将用于存储分割掩码信息。 -
labels = torch.tensor([ann['category_id'] for ann in ann_info['segments_info']], dtype=torch.int64)：这行代码从分割信息中提取类别ID，将其存储为PyTorch张量，并将数据类型指定为int64。这个张量将用于存储每个分割对象的类别标签。
-
这部分代码的目标是加载和处理图像以及与之相关的分割信息（掩码和类别标签），以便在后续的处理中使用。分割信息通常用于图像分割任务，其中需要将图像中的不同对象分割成单独的区域，并为每个区域分配一个类别标签。
（3）target字典
target = {}
target['image_id'] = torch.tensor([ann_info['image_id'] if "image_id" in ann_info else ann_info["id"]])
if self.return_masks:
target['masks'] = masks
target['labels'] = labels
target["boxes"] = masks_to_boxes(masks)
target['size'] = torch.as_tensor([int(h), int(w)])
target['orig_size'] = torch.as_tensor([int(h), int(w)])这段代码用于创建名为target的字典，其中包含了从图像和分割信息中提取的目标信息。以下是代码的详细解释：
-
target = {}：创建一个空的字典，用于存储目标信息。 -
target['image_id']：将图像ID存储在target字典中。如果ann_info中存在'image_id'字段，它将被用作图像ID；否则，将使用'ann_info'中的'id'字段。这是因为COCO数据集中的图像和注释可能具有不同的ID字段。 -
if self.return_masks:：这是一个条件语句，检查是否应该返回分割掩码（masks）。如果self.return_masks为True，将执行以下操作；否则，不会将分割掩码添加到target字典中。 -
target['masks']：如果条件满足（即self.return_masks为True），将分割掩码存储在target字典中。这是可选的，取决于构造函数中的return_masks参数。 -
target['labels']：将分割对象的类别标签存储在target字典中。这些标签是从分割信息中提取的，存储在labels张量中。 -
target["boxes"]：调用masks_to_boxes函数，将分割掩码转换为边界框（boxes），并将边界框信息存储在target字典中。这些边界框表示每个分割对象的位置和大小。masks_to_boxes函数的实现见DETR纯代码分享（六）box_ops.py -
target['size']：存储加载的图像的尺寸（宽度和高度），将它们转换为PyTorch张量。这个张量表示图像的大小。 -
target['orig_size']：存储原始图像的尺寸，将它们转换为PyTorch张量。通常情况下，这与加载的图像大小相同。
这部分代码的主要任务是构建包含图像ID、分割掩码（可选）、类别标签、边界框和图像尺寸信息的target字典。这些信息将在训练或评估深度学习模型时用于目标检测或图像分割任务。
（4）__getitem__方法的余下部分
if "segments_info" in ann_info:
for name in ['iscrowd', 'area']:
target[name] = torch.tensor([ann[name] for ann in ann_info['segments_info']])
if self.transforms is not None:
img, target = self.transforms(img, target)
return img, target这段代码是CocoPanoptic类的__getitem__方法的余下部分，它包含了以下关键步骤：
-
如果
"segments_info"存在于ann_info中，代码进入一个循环，遍历名为['iscrowd', 'area']的属性名称列表。这是因为COCO注释中通常包含有关分割对象的其他属性信息，如是否是拥挤对象（iscrowd）和对象的面积（area）。- 对于每个属性名称，例如
'iscrowd'或'area'，它使用列表解析从ann_info['segments_info']中提取相应属性的值，然后将这些值存储为PyTorch张量，并将其添加到target字典中。这样，target字典中包含了额外的属性信息。
- 对于每个属性名称，例如
-
接下来，代码检查是否定义了数据转换函数（
self.transforms is not None）。数据转换函数通常用于执行数据增强或预处理操作，以便在训练或评估深度学习模型时改进数据质量或适应模型需求。- 如果定义了数据转换函数，它将应用该函数，并将图像和目标作为参数传递给函数。数据转换函数通常接受图像和目标，对它们进行一些操作，然后返回修改后的图像和目标。
-
最后，方法返回包含图像和目标信息的元组
(img, target)。这是在处理完图像和目标后，将它们打包成一个元组，以便在数据加载器中返回。
总之，这段代码完成了以下任务：
- 提取了分割对象的其他属性信息（例如拥挤度和面积）并将其存储在
target字典中。 - 可选地应用了数据转换函数来修改图像和目标。
- 返回了图像和目标的元组，以供后续使用。
3、构建COCO数据集的CocoPanoptic实例
def build(image_set, args):
img_folder_root = Path(args.coco_path)
ann_folder_root = Path(args.coco_panoptic_path)
assert img_folder_root.exists(), f'provided COCO path {img_folder_root} does not exist'
assert ann_folder_root.exists(), f'provided COCO path {ann_folder_root} does not exist'
mode = 'panoptic'
PATHS = {
"train": ("train2017", Path("annotations") / f'{mode}_train2017.json'),
"val": ("val2017", Path("annotations") / f'{mode}_val2017.json'),
}
img_folder, ann_file = PATHS[image_set]
img_folder_path = img_folder_root / img_folder
ann_folder = ann_folder_root / f'{mode}_{img_folder}'
ann_file = ann_folder_root / ann_file
dataset = CocoPanoptic(img_folder_path, ann_folder, ann_file,
transforms=make_coco_transforms(image_set), return_masks=args.masks)
return dataset这段代码定义了一个名为build的函数，该函数用于构建COCO数据集的CocoPanoptic实例。这个函数的主要任务是根据传递的参数和命令行参数构建COCO数据集的适当子集，并返回一个数据集实例，以便在训练或评估深度学习模型时使用。
（1）构建COCO数据集的子集所需的路径和信息
def build(image_set, args):
img_folder_root = Path(args.coco_path)
ann_folder_root = Path(args.coco_panoptic_path)
assert img_folder_root.exists(), f'provided COCO path {img_folder_root} does not exist'
assert ann_folder_root.exists(), f'provided COCO path {ann_folder_root} does not exist'
mode = 'panoptic'
PATHS = {
"train": ("train2017", Path("annotations") / f'{mode}_train2017.json'),
"val": ("val2017", Path("annotations") / f'{mode}_val2017.json'),
}这段代码定义了一个名为build的函数，该函数用于构建COCO数据集的子集。以下是这个函数的主要功能和解释：
-
img_folder_root = Path(args.coco_path)和ann_folder_root = Path(args.coco_panoptic_path)：这两行代码使用Path对象从命令行参数args中获取图像文件夹根路径和注释文件夹根路径。这些根路径是通过命令行参数--coco_path和--coco_panoptic_path传递给函数的。 -
assert img_folder_root.exists(), f'provided COCO path {img_folder_root} does not exist'和assert ann_folder_root.exists(), f'provided COCO path {ann_folder_root} does not exist'：这两行代码用于检查图像文件夹根路径和注释文件夹根路径是否存在。如果路径不存在，将引发断言错误。 -
mode = 'panoptic'：设置数据集的模式为'panoptic'，这通常表示COCO数据集的带有分割信息的模式。 -
PATHS字典：定义了两种模式（"train"和"val"）的图像文件夹和注释文件的路径。这些路径是根据数据集的模式和图像集名称构建的。例如，"train"对应训练集，"val"对应验证集。
这段代码的主要目标是为了准备构建COCO数据集的子集所需的路径和信息，以便在后续的步骤中使用这些路径来加载图像和注释文件，以构建相应的数据集。
（2）构建COCO数据集的子集
img_folder, ann_file = PATHS[image_set]
img_folder_path = img_folder_root / img_folder
ann_folder = ann_folder_root / f'{mode}_{img_folder}'
ann_file = ann_folder_root / ann_file
dataset = CocoPanoptic(img_folder_path, ann_folder, ann_file,
transforms=make_coco_transforms(image_set), return_masks=args.masks)
return dataset这段代码用于根据之前定义的路径和信息构建COCO数据集的子集。以下是代码的详细解释：
-
img_folder, ann_file = PATHS[image_set]：根据传递给函数的image_set参数，从PATHS字典中选择相应的图像文件夹和注释文件路径，并将它们分配给img_folder和ann_file变量。这将根据所需的数据集集合（"train"或"val"）选择相应的文件夹和文件路径。 -
img_folder_path = img_folder_root / img_folder：构建完整的图像文件夹路径，通过将img_folder_root（图像文件夹根路径）与img_folder（子文件夹路径）拼接在一起。这将是用于加载图像文件的完整路径。 -
ann_folder = ann_folder_root / f'{mode}_{img_folder}'：构建完整的注释文件夹路径，通过将ann_folder_root（注释文件夹根路径）与mode（模式）和img_folder（子文件夹路径）拼接在一起。这将是用于加载注释文件的完整路径。 -
ann_file = ann_folder_root / ann_file：构建完整的注释文件路径，通过将ann_folder_root（注释文件夹根路径）与ann_file（注释文件名）拼接在一起。这将是用于加载注释文件的完整路径。 -
创建
CocoPanoptic数据集实例，传递了图像文件夹路径、注释文件夹路径、注释文件路径，以及一些其他参数：transforms=make_coco_transforms(image_set)：使用名为make_coco_transforms的函数创建数据转换，这些转换将在数据加载时应用到图像和目标上。return_masks=args.masks：根据命令行参数args.masks确定是否返回分割掩码。这将传递给CocoPanoptic类的构造函数。
-
返回构建的
CocoPanoptic数据集实例，该实例包含了所需的图像和注释信息，以便在后续的训练或评估中使用。
总之，这段代码的主要任务是根据传递的参数和预定义的路径信息，构建COCO数据集的子集，并返回一个数据集实例，以供后续使用。这个实例包含了图像、注释和其他相关信息，以便进行深度学习任务。