数据集学习笔记(三):COCO创建dataloader用于训练

安装环境

# 安装Cython
pip install -i https://pypi.tuna.tsinghua.edu.cn/simple Cython
 
# 安装pycocotools
pip install git+https://github.com/philferriere/cocoapi.git#subdirectory=PythonAPI

如果下面这个安装不起，就通过这个下载
然后pip install 即可

COCO训练数据加载代码

import torchvision.datasets as datasets
from torchvision import transforms
import torch,cv2
"""""""""""""""""""""""""""""""""COCO dataloader"""""""""""""""""""""""""""""
train_root = 'E:/dataset/Aquarium/train/'
val_root = 'E:/dataset/Aquarium/valid/'
font = cv2.FONT_HERSHEY_SIMPLEX
train_annFile = 'E:/dataset/Aquarium/annotations/train_annotations.coco.json'
val_annFile = 'E:/dataset/Aquarium/annotations/val_annotations.coco.json'
# 定义 coco collate_fn
def collate_fn_coco(batch):
    return tuple(zip(*batch))

# 创建 coco dataset
data_transform = {
    "train": transforms.Compose([transforms.RandomResizedCrop(224),  # 随机裁剪，在缩放成224*224
                                 # transforms.RandomErasing(p=0.5), # 随机遮挡 概率0.5
                                 transforms.RandomHorizontalFlip(),  # 水平方向随机翻转，概率为0.5
                                 transforms.ToTensor(),
                                 transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])]),
    "val": transforms.Compose([transforms.Resize(256),
                               transforms.CenterCrop(224),
                               transforms.ToTensor(),
                               transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])}
coco_train = datasets.CocoDetection(train_root, train_annFile, transform=data_transform["train"])
coco_val = datasets.CocoDetection(val_root, val_annFile, transform=data_transform["val"])

# 创建 dataloader
train_loader = torch.utils.data.DataLoader(coco_train, batch_size=16 ,shuffle=True,num_workers=0,pin_memory=True,collate_fn=collate_fn_coco,drop_last=True)
val_loader = torch.utils.data.DataLoader(coco_val, batch_size=16 ,shuffle=False,num_workers=0,pin_memory=True,collate_fn=collate_fn_coco,drop_last=False)

# 可视化
for imgs, labels in train_loader:
    for i in range(len(imgs)):
        bboxes = []
        ids = []
        img = imgs[i]
        labels_ = labels[i]
        for label in labels_:
            bboxes.append([label['bbox'][0],
                           label['bbox'][1],
                           label['bbox'][0] + label['bbox'][2],
                           label['bbox'][1] + label['bbox'][3]
                           ])
            ids.append(label['category_id'])

        img = img.permute(1, 2, 0).numpy()
        img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
        for box, id_ in zip(bboxes, ids):
            x1 = int(box[0])
            y1 = int(box[1])
            x2 = int(box[2])
            y2 = int(box[3])
            cv2.rectangle(img, (x1, y1), (x2, y2), (0, 0, 255), thickness=2)
            cv2.putText(img, text=str(id_), org=(x1 + 5, y1 + 5), fontFace=font, fontScale=1,
                        thickness=2, lineType=cv2.LINE_AA, color=(0, 255, 0))
        cv2.imshow('test', img)
        cv2.waitKey()