Faster R-CNN ResNet-50 FPN Mask R-CNN ResNet-50 FPN

IoU on Testset

import torch
import torchvision.models.detection as models
import pytorch_lightning as pl

import numpy as np
import albumentations as A
from albumentations.pytorch.transforms import ToTensorV2
from torchvision.transforms import ToTensor, Compose, ToPILImage, Normalize
from detection_nbdev.dataset import XMLDetectionDataModule, XMLDetectionDataset
from detection_nbdev.model import TorchVisionDetector

# Setup models
model_names = ['fasterrcnn_resnet50_fpn', 'retinanet_resnet50_fpn', 'maskrcnn_resnet50_fpn']

# Prepare data
data_dir = '../detection/data/'
image_transform = Compose([
    ToTensor(),
    lambda x: x.unsqueeze(0)])

dm = XMLDetectionDataModule(
    data_dir, 
    image_transform=ToTensor())
dm.setup(mode='use_dir')

# for mn in model_names:
#     model=TorchVisionDetector(mn, pretrained=True)
#     trainer = pl.Trainer(gpus='1', checkpoint_callback=False)
#     trainer.test(model, dm.test_dataloader())

Visualizing Trajectory

import numpy as np
import cv2
import glob
import torchvision
from PIL import Image
from detection_nbdev.utils import visualize_bbox

model=TorchVisionDetector('maskrcnn_resnet50_fpn', pretrained=True)
model.eval();

root = '../detection/data/videoframes/*.jpg'
transform = torchvision.transforms.ToTensor()
files = glob.glob(root)

import os
# Prepare output
os.makedirs('./out', exist_ok=True)

for i, file in enumerate(files):
    im = Image.open(file).convert("RGB")
    x = transform(im).unsqueeze(0)
    y_pred = model(x)
    
    bboxes = [list(map(int, _y)) for _y in y_pred[0]]
    if len(bboxes) > 0:
        im = Image.fromarray(visualize_bbox(
            np.asarray(im),
            bboxes[0],
            'ball'
        ))
    im.save(f'out/{i:06d}.png')

---------------------------------------------------------------------------
KeyboardInterrupt                         Traceback (most recent call last)
<ipython-input-10-f84ad8250feb> in <module>
      2     im = Image.open(file).convert("RGB")
      3     x = transform(im).unsqueeze(0)
----> 4     y_pred = model(x)
      5 
      6     bboxes = [list(map(int, _y)) for _y in y_pred[0]]

~/miniconda3/lib/python3.8/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    725             result = self._slow_forward(*input, **kwargs)
    726         else:
--> 727             result = self.forward(*input, **kwargs)
    728         for hook in itertools.chain(
    729                 _global_forward_hooks.values(),

~/detection-nbdev/detection_nbdev/model.py in forward(self, x)
    125 
    126     def forward(self, x):
--> 127         predictions = self.model(x)
    128         # filter out predictions of sports ball
    129         func = lambda x : x[1] == sports_ball_ID

~/miniconda3/lib/python3.8/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    725             result = self._slow_forward(*input, **kwargs)
    726         else:
--> 727             result = self.forward(*input, **kwargs)
    728         for hook in itertools.chain(
    729                 _global_forward_hooks.values(),

~/miniconda3/lib/python3.8/site-packages/torchvision/models/detection/generalized_rcnn.py in forward(self, images, targets)
     94                                      .format(degen_bb, target_idx))
     95 
---> 96         features = self.backbone(images.tensors)
     97         if isinstance(features, torch.Tensor):
     98             features = OrderedDict([('0', features)])

~/miniconda3/lib/python3.8/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    725             result = self._slow_forward(*input, **kwargs)
    726         else:
--> 727             result = self.forward(*input, **kwargs)
    728         for hook in itertools.chain(
    729                 _global_forward_hooks.values(),

~/miniconda3/lib/python3.8/site-packages/torchvision/models/detection/backbone_utils.py in forward(self, x)
     41 
     42     def forward(self, x):
---> 43         x = self.body(x)
     44         x = self.fpn(x)
     45         return x

~/miniconda3/lib/python3.8/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    725             result = self._slow_forward(*input, **kwargs)
    726         else:
--> 727             result = self.forward(*input, **kwargs)
    728         for hook in itertools.chain(
    729                 _global_forward_hooks.values(),

~/miniconda3/lib/python3.8/site-packages/torchvision/models/_utils.py in forward(self, x)
     61         out = OrderedDict()
     62         for name, module in self.items():
---> 63             x = module(x)
     64             if name in self.return_layers:
     65                 out_name = self.return_layers[name]

~/miniconda3/lib/python3.8/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    725             result = self._slow_forward(*input, **kwargs)
    726         else:
--> 727             result = self.forward(*input, **kwargs)
    728         for hook in itertools.chain(
    729                 _global_forward_hooks.values(),

~/miniconda3/lib/python3.8/site-packages/torch/nn/modules/container.py in forward(self, input)
    115     def forward(self, input):
    116         for module in self:
--> 117             input = module(input)
    118         return input
    119 

~/miniconda3/lib/python3.8/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    725             result = self._slow_forward(*input, **kwargs)
    726         else:
--> 727             result = self.forward(*input, **kwargs)
    728         for hook in itertools.chain(
    729                 _global_forward_hooks.values(),

~/miniconda3/lib/python3.8/site-packages/torchvision/models/resnet.py in forward(self, x)
    102         identity = x
    103 
--> 104         out = self.conv1(x)
    105         out = self.bn1(out)
    106         out = self.relu(out)

~/miniconda3/lib/python3.8/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    725             result = self._slow_forward(*input, **kwargs)
    726         else:
--> 727             result = self.forward(*input, **kwargs)
    728         for hook in itertools.chain(
    729                 _global_forward_hooks.values(),

~/miniconda3/lib/python3.8/site-packages/torch/nn/modules/conv.py in forward(self, input)
    421 
    422     def forward(self, input: Tensor) -> Tensor:
--> 423         return self._conv_forward(input, self.weight)
    424 
    425 class Conv3d(_ConvNd):

~/miniconda3/lib/python3.8/site-packages/torch/nn/modules/conv.py in _conv_forward(self, input, weight)
    417                             weight, self.bias, self.stride,
    418                             _pair(0), self.dilation, self.groups)
--> 419         return F.conv2d(input, weight, self.bias, self.stride,
    420                         self.padding, self.dilation, self.groups)
    421 

KeyboardInterrupt: