这个模板拿来就可以直接用,我已经测试过了。
你需要先安装pycocotools库,直接用pip安装就可以了
然后还有一些依赖你就看头文件,缺哪个安装哪个就行
import torch
from torch import nn
import torch.nn.functional as f
import torchvision
import torchvision.transforms as tfs
import torchvision.models as models
from torch.autograd import Variable
from torch.utils.data import dataloader
import numpy as np
from torchvision.datasets import CocoDetection
import cv2
class COCOdataset(CocoDetection):
# 类别名称
CLASSES_NAME = (
'__back_ground__', 'person', 'bicycle', 'car', 'motorcycle',
'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light',
'fire hydrant', 'stop sign', 'parking meter', 'bench',
'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant',
'bear', 'zebra', 'giraffe', 'backpack', 'umbrella',
'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard',
'sports ball', 'kite', 'baseball bat', 'baseball glove',
'skateboard', 'surfboard', 'tennis racket', 'bottle',
'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl',
'banana', 'apple', 'sandwich', 'orange', 'broccoli',
'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair',
'couch', 'potted plant', 'bed', 'dining table', 'toilet',
'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
'microwave', 'oven', 'toaster', 'sink', 'refrigerator',
'book', 'clock', 'vase', 'scissors', 'teddy bear',
'hair drier', 'toothbrush')
prefix = "../data/coco/"
def __init__(self, img_path=None, ann_path=None, resize_size=[800,1333],is_train = True, transform=None)