tkshnkmr

import os
import torch
import torch.utils.data
import torchvision
from PIL import Image
from pycocotools.coco import COCO

class myOwnDataset(torch.utils.data.Dataset):
    def __init__(self, root, annotation, transforms=None):
        self.root = root

tkshnkmr

# In my case, just added ToTensor
def get_transform():
    custom_transforms = []
    custom_transforms.append(torchvision.transforms.ToTensor())
    return torchvision.transforms.Compose(custom_transforms)

tkshnkmr

# path to your own data and coco file
train_data_dir = 'my_data/train'
train_coco = 'my_data/my_train_coco.json'

# create own Dataset
my_dataset = myOwnDataset(root=train_data_dir,
                          annotation=train_coco,
                          transforms=get_transform()
                          )

tkshnkmr

# select device (whether GPU or CPU)
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')

# DataLoader is iterable over Dataset
for imgs, targets in data_loader:
    imgs = list(img.to(device) for img in imgs)
    targets = [{k: v.to(device) for k, v in t.items()} for t in targets]
    print(targets)

tkshnkmr

# select device (whether GPU or CPU)
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')

# DataLoader is iterable over Dataset
for imgs, annotations in data_loader:
    imgs = list(img.to(device) for img in imgs)
    annotations = [{k: v.to(device) for k, v in t.items()} for t in annotations]
    print(annotations)

tkshnkmr

[{'boxes': tensor([[ 67.6667,  28.1667, 427.6667, 451.5000]]), 'labels': tensor([1]), 'image_id': tensor([117]), 'area': tensor([152400.]), 'iscrowd': tensor([0])}]
[{'boxes': tensor([[154.7301,  75.7785, 313.2076, 240.1384]]), 'labels': tensor([1]), 'image_id': tensor([191]), 'area': tensor([26047.3418]), 'iscrowd': tensor([0])}]
[{'boxes': tensor([[ 86.4694,  39.7959, 288.0000, 384.6939]]), 'labels': tensor([1]), 'image_id': tensor([194]), 'area': tensor([69507.5000]), 'iscrowd': tensor([0])}]
[{'boxes': tensor([[121.0870, 343.5652, 762.3913, 905.5217]]), 'labels': tensor([1]), 'image_id': tensor([51]), 'area': tensor([360385.1562]), 'iscrowd': tensor([0])}]

tkshnkmr

from torchvision.models.detection.faster_rcnn import FastRCNNPredictor

def get_model_instance_segmentation(num_classes):
    # load an instance segmentation model pre-trained pre-trained on COCO
    model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=False)
    # get number of input features for the classifier
    in_features = model.roi_heads.box_predictor.cls_score.in_features
    # replace the pre-trained head with a new one
    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)

tkshnkmr

import os
import torch
from torch.utils import data
from PIL import Image
from torchvision import transforms

class simpleDataset(data.Dataset):
    
    # initialise function of class
    def __init__(self, root, filenames, labels):

tkshnkmr

# data directory
root = "my_data"

# assume we have 3 jpg images
filenames = ['img1.jpg', 'img2.jpg', 'img3.jpg']

# the class of image might be ['black cat', 'tabby cat', 'tabby cat']
labels = [0, 1, 1]

# create own Dataset

tkshnkmr

# data loader
batch_size = 1
num_workers = 4

data_loader = torch.utils.data.DataLoader(my_dataset,
                                          batch_size=batch_size,
                                          shuffle=False,
                                          num_workers=num_workers
                                         )