text_plain_output_1.png image_output_1.png
from PIL import Image

from pycocotools.coco import COCO

from torchvision.models.detection.faster_rcnn import FastRCNNPredictor

from torchvision.models.detection.faster_rcnn import FastRCNNPredictor

import os

import os

import torch

import torch

import torchvision

import torchvision

import torchvision.datasets as dset

import torchvision.transforms as transforms

import numpy as np
import pandas as pd
import os

import torchvision
import torchvision.datasets as dset
import torchvision.transforms as transforms
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
path2data = '/kaggle/input/levi9-hack9-2023/train'
path2json = '/kaggle/input/levi9-hack9-2023/train.json'

coco_train = dset.CocoDetection(root=path2data, annFile=path2json, transform=transforms.ToTensor())

img, target = coco_train[0]

import os
import torch
import torch.utils.data
import torchvision
from PIL import Image
from pycocotools.coco import COCO
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
class myOwnDataset(torch.utils.data.Dataset):
    def __init__(self, root, annotation, transforms=None):
        self.root = root
        self.transforms = transforms
        self.coco = COCO(annotation)
        self.ids = list(sorted(self.coco.imgs.keys()))
    def __getitem__(self, index):
        coco = self.coco
        img_id = self.ids[index]
        ann_ids = coco.getAnnIds(imgIds=img_id)
        coco_annotation = coco.loadAnns(ann_ids)
        path = coco.loadImgs(img_id)[0]['file_name']
        img = Image.open(os.path.join(self.root, path))
        num_objs = len(coco_annotation)
        boxes = []
        area = 0
        for i in range(num_objs):
            xmin = coco_annotation[i]['bbox'][0]
            ymin = coco_annotation[i]['bbox'][1]
            xmax = xmin + coco_annotation[i]['bbox'][2]
            ymax = ymin + coco_annotation[i]['bbox'][3]
            area += (xmax - xmin) * (ymax - ymin)
            boxes.append([xmin, ymin, xmax, ymax])
        if num_objs == 0:
            boxes = torch.zeros((0, 4), dtype=torch.float32)
        else:
            boxes = torch.as_tensor(boxes, dtype=torch.float32)
        labels = torch.ones((num_objs,), dtype=torch.int64)
        img_id = torch.tensor([img_id])
        areas = []
        for i in range(num_objs):
            areas.append(coco_annotation[i]['area'])
        area = torch.as_tensor(area, dtype=torch.float32)
        iscrowd = torch.zeros((num_objs,), dtype=torch.int64)
        my_annotation = {}
        my_annotation['boxes'] = boxes
        my_annotation['labels'] = labels
        my_annotation['image_id'] = img_id
        my_annotation['area'] = area
        my_annotation['iscrowd'] = iscrowd
        if self.transforms is not None:
            img = self.transforms(img)
        return (img, my_annotation)
    def __len__(self):
        return len(self.ids)
def get_transform():
    custom_transforms = []
    custom_transforms.append(torchvision.transforms.ToTensor())
    return torchvision.transforms.Compose(custom_transforms)
def collate_fn(batch):
    return tuple(zip(*batch))
def get_model_instance_segmentation(num_classes):
    model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True, progress=True, pretrained_backbone=True)
    in_features = model.roi_heads.box_predictor.cls_score.in_features
    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
    return model

train_batch_size = 1
train_shuffle_dl = True
num_workers_dl = 4
num_classes = 2
num_epochs = 2
lr = 0.005
momentum = 0.9
weight_decay = 0.005

import torch
print('Torch version:', torch.__version__)
my_dataset = myOwnDataset(root=path2data, annotation=path2json, transforms=get_transform())
data_loader = torch.utils.data.DataLoader(my_dataset, batch_size=train_batch_size, shuffle=train_shuffle_dl, num_workers=num_workers_dl, collate_fn=collate_fn)
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
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]
model = get_model_instance_segmentation(num_classes)
model.to(device)
params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params, lr=lr, momentum=momentum, weight_decay=weight_decay)
len_dataloader = len(data_loader)
for epoch in range(num_epochs):
    print(f'Epoch: {epoch}/{num_epochs}')
    model.train()
    i = 0
    for imgs, annotations in data_loader:
        i += 1
        imgs = list((img.to(device) for img in imgs))
        annotations = [{k: v.to(device) for k, v in t.items()} for t in annotations]
        loss_dict = model(imgs, annotations)
        losses = sum((loss for loss in loss_dict.values()))
        optimizer.zero_grad()
        losses.backward()
        optimizer.step()
        print(f'Iteration: {i}/{len_dataloader}, Loss: {losses}')