引言
同样,数据集也是需要build模型的。类CocoDataset是CustomDataset类的子类,而CustomDataset是Dataset
的子类。(之前的创建模型,都是torch.nn.module的子类,数据集的创建就不是咯,注意一下)
关于CustomDataset的定义和其类方法的代码解读,可以看下面这篇文章:
类CocoDataset依然作为形参,添加到@DATASETS.register_module中,作用就是将其保存到注册表的module中。共有四个方法:
load_annotations()
:加载标注文件中的annotation字典,返回图片信息,比如:info{"filename":"284193,faa9000f2678b5e.jpg"}
。get_ann_info(self,idx)
:获得annotation的信息,其实是调用了_parse_ann_info()
;它的形参是指定的图片id,返回值是个字典:bboxes
,bboxes_ignore
,labels
,masks
,mask_polys
,poly_lens
._filter_imgs(self, min_size)
:过滤图片,去除没有annotation标注文件的图片,以及图片尺寸小于min-size的图片。_parse_ann_info(self, ann_info, with_mask=True)
:解析一张图片的annotation的信息,主要是bbox和mask信息,返回值为:bboxes
,bboxes_ignore
,labels
,masks
,mask_polys
,poly_lens
.(如果没用mask分支,就没用后面的三个返回值了)
在父类中custom.py有其初始化,还有另外的几个重要的函数,比如prepare_train_img()、prepare_test_img()等。因为其代码行数太长,所以在上面的链接里写一篇讲解。
以下是coco.py的代码,如有错误的地方,还请指出,后面博主也会继续修改,增加对各个代码段的理解。
coco.py代码注释
import numpy as np
from pycocotools.coco import COCO
from .custom import CustomDataset
from .registry import DATASETS
@DATASETS.register_module
class CocoDataset(CustomDataset):
""" coco api
self.coco = COCO(ann_file)
create class members
# 一维数组,值为对应原coco数据集的annotation/images/categories信息
coco.anns = anns
coco.imgs = imgs
coco.cats = cats
# 两个默认value为list的字典 比如:imgToAnns{"1":[ann1,ann2,ann3,....]},'1'为image_id = 1,ann1时其为 1 的annotation。
coco.imgToAnns = imgToAnns
coco.catToImgs = catToImgs
以上五个members,在实例化COCO时,被创建
"""
CLASSES = ('person', 'mask')
def load_annotations(self, ann_file):
self.coco = COCO(ann_file)
#函数 getImgIds()、getCatIds()、getAnnIds(),返回值为integer array of img/cat/ann ids,形参为过滤条件
self.cat_ids = self.coco.getCatIds() # [1,2] , integer array of cat ids
self.cat2label = {
# dict
cat_id: i + 1 # cat_id : cat_id + 1 &#x