import math
from typing import Optional, Union, Tuple

from torch.nn import Parameter

from torch_geometric.typing import OptTensor, Adj
from torch_geometric.nn.conv import MessagePassing

def glorot(tensor):
    if tensor is not None:
        stdv = math.sqrt(6.0 / (tensor.size(-2) + tensor.size(-1)))
        tensor.data.uniform_(-stdv, stdv)

def zeros(tensor):
    if tensor is not None:
        tensor.data.fill_(0)        

dataset = 'MUTAG'

path = os.path.join(os.getcwd(), 'data', 'Entities')
dataset = Entities(path, dataset)
data = dataset[0]

class RelationalGCNConv(MessagePassing):
    def __init__(self,
                 in_channels: int,
                 out_channels: int,
                 num_relations: int,
                 num_bases: Optional[int]=None,
                 aggr: str='mean',
                 **kwargs):
        super(RelationalGCNConv, self).__init__(aggr=aggr, node_dim=0)

        # aggr, node_dim은 알아서 self의 attribute로 등록해준다. (MessagePassing)
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.num_relations = num_relations
        self.num_bases = num_bases

        # 원본 코드에서 in_channels가 tuple인 이유는
        # src/tgt node가 다른 type인 bipartite graph 구조도 지원하기 위함이다.
        # 예시를 위해 integer로 변경한다.
        self.in_channels = in_channels

        # Define Weights
        if num_bases is not None:
            self.V = Parameter(
                data=Tensor(num_bases, in_channels, out_channels))
            self.a = Parameter(Tensor(num_relations, num_bases))
        else:
            self.V = Parameter(
                Tensor(num_relations, in_channels, out_channels))
            # dummy parameter
            self.register_parameter(name='a', param=None)

        self.root = Parameter(Tensor(in_channels, out_channels))
        self.bias = Parameter(Tensor(out_channels))

        self.reset_parameters()

    def reset_parameters(self):
        glorot(self.V)
        glorot(self.a)
        glorot(self.root)
        zeros(self.bias)

    def forward(self,
                x: OptTensor,
                edge_index: Adj,
                edge_type: OptTensor=None):

        x_l = x
        # node feature가 주어지지 않는다면
        # embedding weight(V) lookup을 위해 아래와 같이 세팅한다.
        if x_l is None:
            x_l = torch.arange(self.in_channels, device=self.V.device)

        x_r = x_l
        size = (x_l.size(0), x_r.size(0))

        # output = (num_nodes, out_channels)
        out = torch.zeros(x_r.size(0), self.out_channels, device=x_r.device)

        V = self.V
        if self.num_bases is not None:
            V = torch.einsum("rb,bio->rio", self.a, V)

        # propagate given relations
        for i in range(self.num_relations):
            # 특정 edge_type에 맞는 edge_index를 선별한다.
            selected_edge_index = masked_edge_index(edge_index, edge_type == i)

            # node_features가 주어지지 않는다면
            if x_l.dtype == torch.long:
                out += self.propagate(selected_edge_index, x=V[i, x_l], size=size)

            # node_features가 주어진다면
            else:
                h = self.propagate(selected_edge_index, x=x_l, size=size)
                out += (h @ V[i])

        out += self.root[x_r] if x_r.dtype == torch.long else x_r @ self.root
        out += self.bias
        return out

from torch_sparse import masked_select_nnz

def masked_edge_index(edge_index, edge_mask):
    """
    :param edge_index: (2, num_edges)
    :param edge_mask: (num_edges) -- source node 기준임
    :return: masked edge_index (edge_mask에 해당하는 Tensor만 가져옴)
    """
    if isinstance(edge_index, Tensor):
        return edge_index[:, edge_mask]
    else:
        # if edge_index == SparseTensor
        return masked_select_nnz(edge_index, edge_mask, layout='coo')

data = data.to(device)

model = RelationalGCNConv(
    in_channels=in_channels, out_channels=out_channels,
    num_relations=num_relations, num_bases=num_bases).to(device)

print(get_num_params(model))

out = model(x=data.x, edge_index=data.edge_index, edge_type=data.edge_type)

_base_ = [
    '../_base_/models/faster_rcnn_r50_fpn.py',
    '../_base_/datasets/coco_detection.py',
    '../_base_/schedules/schedule_1x.py', '../_base_/default_runtime.py'
]

_base_ = './faster_rcnn_r50_fpn_1x_coco.py'
model = dict(
    roi_head=dict(
        bbox_head=dict(
            reg_decoded_bbox=True,
            loss_bbox=dict(type='BoundedIoULoss', loss_weight=10.0))))

# model settings
model = dict(
    type='FasterRCNN',
    pretrained='torchvision://resnet50',
    backbone=dict(
        type='ResNet',
        depth=50,
        num_stages=4,
        out_indices=(0, 1, 2, 3),
        frozen_stages=1,
        norm_cfg=dict(type='BN', requires_grad=True),
        norm_eval=True,
        style='pytorch'),
    neck=dict(
        type='FPN',
        in_channels=[256, 512, 1024, 2048],
        out_channels=256,
        num_outs=5),
    rpn_head=dict(
        type='RPNHead',
        ...
        ),
    roi_head=dict(
        type='StandardRoIHead',
        ...
        ),
    # model training and testing settings
    train_cfg=dict(
        rpn=dict(...),
        rpn_proposal=dict(...),
        rcnn=dict(...),
    test_cfg=dict(
        rpn=dict...),
        rcnn=dict(...)
        # soft-nms is also supported for rcnn testing
        # e.g., nms=dict(type='soft_nms', iou_threshold=0.5, min_score=0.05)
    ))

'images': [
    {
        'file_name': 'COCO_val2014_000000001268.jpg',
        'height': 427,
        'width': 640,
        'id': 1268
    },
    ...
],

'annotations': [
    {
        'segmentation': [[192.81,
            247.09,
            ...
            219.03,
            249.06]],  # if you have mask labels
        'area': 1035.749,
        'iscrowd': 0,
        'image_id': 1268,
        'bbox': [192.81, 224.8, 74.73, 33.43],
        'category_id': 16,
        'id': 42986
    },
    ...
],

'categories': [
    {'id': 0, 'name': 'car'},
    ...
]

# the new config inherits the base configs to highlight the necessary modification
_base_ = './cascade_mask_rcnn_r50_fpn_1x_coco.py'

# 1. dataset settings
dataset_type = 'CocoDataset'
classes = ('a', 'b', 'c', 'd', 'e')
data = dict(
    samples_per_gpu=2,
    workers_per_gpu=2,
    train=dict(
        type=dataset_type,
        # explicitly add your class names to the field `classes`
        classes=classes,
        ann_file='path/to/your/train/annotation_data',
        img_prefix='path/to/your/train/image_data'),
    val=dict(
        type=dataset_type,
        # explicitly add your class names to the field `classes`
        classes=classes,
        ann_file='path/to/your/val/annotation_data',
        img_prefix='path/to/your/val/image_data'),
    test=dict(
        type=dataset_type,
        # explicitly add your class names to the field `classes`
        classes=classes,
        ann_file='path/to/your/test/annotation_data',
        img_prefix='path/to/your/test/image_data'))

# 2. model settings

# explicitly over-write all the `num_classes` field from default 80 to 5.
model = dict(
    roi_head=dict(
        bbox_head=[
            dict(
                type='Shared2FCBBoxHead',
                # explicitly over-write all the `num_classes` field from default 80 to 5.
                num_classes=5),
            dict(
                type='Shared2FCBBoxHead',
                # explicitly over-write all the `num_classes` field from default 80 to 5.
                num_classes=5),
            dict(
                type='Shared2FCBBoxHead',
                # explicitly over-write all the `num_classes` field from default 80 to 5.
                num_classes=5)],
    # explicitly over-write all the `num_classes` field from default 80 to 5.
    mask_head=dict(num_classes=5)))

[
    {
        'filename': 'a.jpg',
        'width': 1280,
        'height': 720,
        'ann': {
            'bboxes': <np.ndarray, float32> (n, 4),
            'labels': <np.ndarray, int64> (n, ),
            'bboxes_ignore': <np.ndarray, float32> (k, 4), (optional field)
            'labels_ignore': <np.ndarray, int64> (k, )     (optional field)
        }
    },
    ...
]

classes = ('person', 'bicycle', 'car')
data = dict(
    train=dict(classes=classes),
    val=dict(classes=classes),
    test=dict(classes=classes))

classes = 'path/to/classes.txt'
data = dict(
    train=dict(classes=classes),
    val=dict(classes=classes),
    test=dict(classes=classes))

img_norm_cfg = dict(
    mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
train_pipeline = [
    dict(type='LoadImageFromFile'),
    dict(type='LoadAnnotations', with_bbox=True),
    dict(type='Resize', img_scale=(1333, 800), keep_ratio=True),
    dict(type='RandomFlip', flip_ratio=0.5),
    dict(type='Normalize', **img_norm_cfg),
    dict(type='Pad', size_divisor=32),
    dict(type='MyTransform'),
    dict(type='DefaultFormatBundle'),
    dict(type='Collect', keys=['img', 'gt_bboxes', 'gt_labels']),
]

import torch.nn as nn
from ..builder import BACKBONES

@BACKBONES.register_module()
class MobileNet(nn.Module):

    def __init__(self, arg1, arg2):
        pass

    def forward(self, x):  # should return a tuple
        pass
    # 구현하는 방법은 이미 구현되어 있는 다른 파일들을 보는 것도 도움이 된다.(그러나 보통 복잡함)

custom_imports = dict(
    imports=['mmdet.models.backbones.mobilenet'],
    allow_failed_imports=False)

model = dict(
    ...
    backbone=dict(
        type='MobileNet',
        arg1=xxx,
        arg2=xxx),
    ...

from ..builder import NECKS

@NECKS.register_module()
class PAFPN(nn.Module):

    def __init__(self,
                in_channels,
                out_channels,
                num_outs,
                start_level=0,
                end_level=-1,
                add_extra_convs=False):
        pass

    def forward(self, inputs):
        # implementation is ignored
        pass

from .pafpn import PAFPN

custom_imports = dict(
    imports=['mmdet.models.necks.pafpn.py'],
    allow_failed_imports=False)

neck=dict(
    type='PAFPN',
    in_channels=[256, 512, 1024, 2048],
    out_channels=256,
    num_outs=5)

custom_imports=dict(
    imports=['mmdet.models.roi_heads.double_roi_head', 
            'mmdet.models.bbox_heads.double_bbox_head',
            'mmdet.models.losses.my_loss'])

loss_bbox=dict(type='MyLoss', loss_weight=1.0))

optimizer = dict(type='Adam', lr=0.0003, weight_decay=0.0001)
# or
optimizer = dict(type='MyOptimizer', a=a_value, b=b_value, c=c_value)

from .registry import OPTIMIZERS
from torch.optim import Optimizer


@OPTIMIZERS.register_module()
class MyOptimizer(Optimizer):

    def __init__(self, a, b, c)

from .my_optimizer import MyOptimizer

custom_imports = dict(imports=['mmdet.core.optimizer.my_optimizer'], allow_failed_imports=False)

_base_ = [
    '../_base_/models/mask_rcnn_r50_fpn.py',
    '../_base_/datasets/cityscapes_instance.py', '../_base_/default_runtime.py'
]

# optimizer
# lr is set for a batch size of 8
optimizer = dict(type='SGD', lr=0.01, momentum=0.9, weight_decay=0.0001)
optimizer_config = dict(grad_clip=None)
# learning policy
lr_config = dict(
    policy='step',
    warmup='linear',
    warmup_iters=500,
    warmup_ratio=0.001,
    step=[7])
# the max_epochs and step in lr_config need specifically tuned for the customized dataset
runner = dict(max_epochs=8)
log_config = dict(interval=100)

load_from = 'https://download.openmmlab.com/mmdetection/v2.0/mask_rcnn/mask_rcnn_r50_caffe_fpn_mstrain-poly_3x_coco/mask_rcnn_r50_caffe_fpn_mstrain-poly_3x_coco_bbox_mAP-0.408__segm_mAP-0.37_20200504_163245-42aa3d00.pth'  # noqa

MessagePassing(aggr="add", flow="source_to_target", node_dim=-2)

class MessagePassing(torch.nn.Module):
    special_args: Set[str] = {
        'edge_index', 'adj_t', 'edge_index_i', 'edge_index_j', 'size',
        'size_i', 'size_j', 'ptr', 'index', 'dim_size'
    }

    def __init__(self, aggr: Optional[str] = "add",
                 flow: str = "source_to_target", node_dim: int = -2):

        super(MessagePassing, self).__init__()

        self.aggr = aggr
        assert self.aggr in ['add', 'mean', 'max', None]

        self.flow = flow
        assert self.flow in ['source_to_target', 'target_to_source']

        self.node_dim = node_dim

        # 함수를 검사하여 인자를 OrderedDict 형태로 취함
        # pop_first=True 이면 첫 인자는 버림
        self.inspector = Inspector(self)
        self.inspector.inspect(self.message)   # message 메서드
        ...

import re
import inspect
from collections import OrderedDict
from typing import Dict, List, Any, Optional, Callable, Set

class Inspector(object):
    def __init__(self, base_class: Any):
        self.base_class: Any = base_class
        self.params: Dict[str, Dict[str, Any]] = {}

    def inspect(self, func: Callable,
                pop_first: bool = False) -> Dict[str, Any]:
        params = inspect.signature(func).parameters
        params = OrderedDict(params)
        if pop_first:
            params.popitem(last=False)
        self.params[func.__name__] = params

    def keys(self, func_names: Optional[List[str]] = None) -> Set[str]:
        keys = []
        for func in func_names or list(self.params.keys()):
            keys += self.params[func].keys()
        return set(keys)

    def __implements__(self, cls, func_name: str) -> bool:
        if cls.__name__ == 'MessagePassing':
            return False
        if func_name in cls.__dict__.keys():
            return True
        return any(self.__implements__(c, func_name) for c in cls.__bases__)

    def implements(self, func_name: str) -> bool:
        return self.__implements__(self.base_class.__class__, func_name)

    def types(self, func_names: Optional[List[str]] = None) -> Dict[str, str]:
        out: Dict[str, str] = {}
        for func_name in func_names or list(self.params.keys()):
            func = getattr(self.base_class, func_name)
            arg_types = parse_types(func)[0][0]
            for key in self.params[func_name].keys():
                if key in out and out[key] != arg_types[key]:
                    raise ValueError(
                        (f'Found inconsistent types for argument {key}. '
                         f'Expected type {out[key]} but found type '
                         f'{arg_types[key]}.'))
                out[key] = arg_types[key]
        return out

    def distribute(self, func_name, kwargs: Dict[str, Any]):
        # func_name = 예) 'message'
        # kwargs = coll_dict
        # inspector.params['message']에 있는 argument들을 불러온 뒤
        # 이들에게 해당하는 데이터를 coll_dict에서 가져옴
        out = {}
        for key, param in self.params[func_name].items():
            data = kwargs.get(key, inspect.Parameter.empty)
            if data is inspect.Parameter.empty:
                if param.default is inspect.Parameter.empty:
                    raise TypeError(f'Required parameter {key} is empty.')
                data = param.default
            out[key] = data
        return out

params = inspect.signature(func).parameters

import inspect
from collections import OrderedDict

def func(a='OH', b=7, *args, **kwargs):
    pass

params = inspect.signature(func).parameters
params = OrderedDict(params)

OrderedDict([('a', <Parameter "a='OH'">), ('b', <Parameter "b=7">), ('args', <Parameter "*args">), ('kwargs', <Parameter "**kwargs">)])

class MessagePassing(torch.nn.Module):
    special_args: Set[str] = {
        'edge_index', 'adj_t', 'edge_index_i', 'edge_index_j', 'size',
        'size_i', 'size_j', 'ptr', 'index', 'dim_size'
    }

    def __init__(self, aggr: Optional[str] = "add",
                 flow: str = "source_to_target", node_dim: int = -2):

        super(MessagePassing, self).__init__()

        self.aggr = aggr
        assert self.aggr in ['add', 'mean', 'max', None]

        self.flow = flow
        assert self.flow in ['source_to_target', 'target_to_source']

        self.node_dim = node_dim

        # 함수를 검사하여 인자를 OrderedDict 형태로 취함
        # pop_first=True 이면 첫 인자는 버림
        self.inspector = Inspector(self)
        self.inspector.inspect(self.message)
        self.inspector.inspect(self.aggregate, pop_first=True)
        self.inspector.inspect(self.message_and_aggregate, pop_first=True)
        self.inspector.inspect(self.update, pop_first=True)

        self.__user_args__ = self.inspector.keys(
            ['message', 'aggregate', 'update']).difference(self.special_args)
        self.__fused_user_args__ = self.inspector.keys(
            ['message_and_aggregate', 'update']).difference(self.special_args)

        # Support for "fused" message passing.
        # message_and_aggregate 메서드를 구현하면 self.fuse = True
        # self.inspector.base_class.__dict__.keys()에서 확인 가능
        self.fuse = self.inspector.implements('message_and_aggregate')

import torch
from torch_geometric.datasets import Planetoid
from torch_geometric.nn import MessagePassing
from torch_geometric.utils import add_self_loops, degree
import torch_geometric.transforms as T

# Load Cora Dataset
dataset = 'Cora'
path = os.path.join(os.getcwd(), 'data', dataset)
dataset = Planetoid(path, dataset, transform=T.NormalizeFeatures())
data = dataset[0]

x, edge_index = data.x, data.edge_index

print(x.shape, edge_index.shape)
# (torch.Size([2708, 1433]), torch.Size([2, 10556]))

class GCNConv(MessagePassing):
    def __init__(self, in_channels, out_channels):
        super(GCNConv, self).__init__(aggr='add')
        self.lin = torch.nn.Linear(in_channels, out_channels)

    def forward(self, x, edge_index):
        # x has shape [N, in_channels]
        # edge_index has shape [2, E]

        # Step 1: Add self-loops to the adjacency matrix.
        edge_index, _ = add_self_loops(edge_index, num_nodes=x.size(0))

        # Step 2: Linearly transform node feature matrix.
        x = self.lin(x)

        # Step 3: Compute normalization.
        row, col = edge_index
        deg = degree(col, x.size(0), dtype=x.dtype)
        deg_inv_sqrt = deg.pow(-0.5)
        deg_inv_sqrt[deg_inv_sqrt == float('inf')] = 0
        norm = deg_inv_sqrt[row] * deg_inv_sqrt[col]

        # Step 4-5: Start propagating messages.
        return self.propagate(edge_index, x=x, norm=norm)

    def message(self, x_j, norm):
        # x_j has shape [E, out_channels]

        # Step 4: Normalize node features.
        return norm.view(-1, 1) * x_j

conv = GCNConv(x.shape[1], 32)

# conv.inspector.params
# = message_passing이 갖고 있는 method의 인자들을 일부(pop_item=True)를 제외하고 취한 것
# Ex) {'message': OrderedDict([('x_j', <Parameger "x_j">)])}
for param, item in conv.inspector.params.items():
    print(param, ': ', item)

# check
print(conv.__user_args__)
print(conv.__fused_user_args__)
print(conv.fuse)

"""
message :  OrderedDict([('x_j', <Parameter "x_j">), ('norm', <Parameter "norm">)])
aggregate :  OrderedDict([('index', <Parameter "index: torch.Tensor">), ('ptr', <Parameter "ptr: Optional[torch.Tensor] = None">), ('dim_size', <Parameter "dim_size: Optional[int] = None">)])
message_and_aggregate :  OrderedDict()
update :  OrderedDict()

# check
{'norm', 'x_j'}
set()
False
"""

def propagate(self, edge_index: Adj, size: Size = None, **kwargs):
    size = self.__check_input__(edge_index, size)

    # Run "fused" message and aggregation (if applicable).
    # (생략)

    # Otherwise, run both functions in separation.
    elif isinstance(edge_index, Tensor) or not self.fuse:
        coll_dict = self.__collect__(
            self.__user_args__, edge_index, size, kwargs)

# Step 4-5: Start propagating messages.
return self.propagate(edge_index, x=x, norm=norm)

{'norm': tensor([0.2500, 0.2236, 0.2500,  ..., 0.5000, 0.2000, 0.2000]),
 'x_j': tensor([[-0.0106, -0.0185, -0.0095,  ...,  0.0051, -0.0180,  0.0261],
                 ...,
                [-0.0148, -0.0149, -0.0153,  ..., -0.0033, -0.0236,  0.0217]],
                 grad_fn=<IndexSelectBackward>),
 'adj_t': None,
 'edge_index': tensor([[   0,    0,    0,  ..., 2705, 2706, 2707],
                       [ 633, 1862, 2582,  ..., 2705, 2706, 2707]]),
 'edge_index_i': tensor([ 633, 1862, 2582,  ..., 2705, 2706, 2707]),
 'edge_index_j': tensor([   0,    0,    0,  ..., 2705, 2706, 2707]),
 'ptr': None,
 'index': tensor([ 633, 1862, 2582,  ..., 2705, 2706, 2707]),
 'size': [2708, None],
 'size_i': 2708,
 'size_j': 2708,
 'dim_size': 2708}

msg_kwargs = self.inspector.distribute('message', coll_dict)
out = self.message(**msg_kwargs)

# For `GNNExplainer`, we require a separate message and aggregate
# (생략)

aggr_kwargs = self.inspector.distribute('aggregate', coll_dict)
out = self.aggregate(out, **aggr_kwargs)

update_kwargs = self.inspector.distribute('update', coll_dict)
output = self.update(out, **update_kwargs)

print({a:b.shape for a, b in msg_kwargs.items()})
# {'x_j': torch.Size([13264, 32]), 'norm': torch.Size([13264])}

from mmdet.apis import init_detector, inference_detector
import mmcv

# Specify the path to model config and checkpoint file
config_file = 'configs/faster_rcnn/faster_rcnn_r50_fpn_1x_coco.py'
checkpoint_file = 'checkpoints/faster_rcnn_r50_fpn_1x_coco_20200130-047c8118.pth'

# build the model from a config file and a checkpoint file
model = init_detector(config_file, checkpoint_file, device='cuda:0')

# test a single image and show the results
img = 'demo/demo.jpg'  # or img = mmcv.imread(img), which will only load it once
result = inference_detector(model, img)
# visualize the results in a new window
model.show_result(img, result)
# or save the visualization results to image files
model.show_result(img, result, out_file='demo/demo_result.jpg')

# test a video and show the results
video = mmcv.VideoReader('demo/demo.mp4')
for frame in video:
    result = inference_detector(model, frame)
    model.show_result(frame, result, wait_time=1)

import os.path as osp
import mmcv

def convert_balloon_to_coco(ann_file, out_file, image_prefix):
    data_infos = mmcv.load(ann_file)
    
    annotations = []
    images = []
    obj_count = 0
    for idx, v in enumerate(mmcv.track_iter_progress(data_infos.values())):
        filename = v['filename']
        img_path = osp.join(image_prefix, filename)
        height, width = mmcv.imread(img_path).shape[:2]
        
        images.append(dict(
            id=idx,
            file_name=filename,
            height=height,
            width=width))
        
        bboxes = []
        labels = []
        masks = []
        for _, obj in v['regions'].items():
            assert not obj['region_attributes']
            obj = obj['shape_attributes']
            px = obj['all_points_x']
            py = obj['all_points_y']
            poly = [(x + 0.5, y + 0.5) for x, y in zip(px, py)]
            poly = [p for x in poly for p in x]
            
            x_min, y_min, x_max, y_max = (
                min(px), min(py), max(px), max(py))
            
            
            data_anno = dict(
                image_id=idx,
                id=obj_count,
                category_id=0,
                bbox=[x_min, y_min, x_max - x_min, y_max - y_min],
                area=(x_max - x_min) * (y_max - y_min),
                segmentation=[poly],
                iscrowd=0)
            annotations.append(data_anno)
            obj_count += 1
    
    coco_format_json = dict(
        images=images,
        annotations=annotations,
        categories=[{'id':0, 'name': 'balloon'}])
    mmcv.dump(coco_format_json, out_file)


convert_balloon_to_coco('train/via_region_data.json',
                        'train/annotation_coco.json',
                        'train')

convert_balloon_to_coco('val/via_region_data.json',
                        'val/annotation_coco.json',
                        'val')

# The new config inherits a base config to highlight the necessary modification
_base_ = '../mask_rcnn/mask_rcnn_r50_caffe_fpn_mstrain-poly_1x_coco.py'

# We also need to change the num_classes in head to match the dataset's annotation
model = dict(
    roi_head=dict(
        bbox_head=dict(num_classes=1),
        mask_head=dict(num_classes=1)))

# Modify dataset related settings
dataset_type = 'COCODataset'
classes = ('balloon',)
data = dict(
    train=dict(
        img_prefix='data/balloon/train/',
        classes=classes,
        ann_file='data/balloon/train/annotation_coco.json'),
    val=dict(
        img_prefix='data/balloon/val/',
        classes=classes,
        ann_file='data/balloon/val/annotation_coco.json'),
    test=dict(
        img_prefix='data/balloon/val/',
        classes=classes,
        ann_file='data/balloon/val/annotation_coco.json'))

# We can use the pre-trained Mask RCNN model to obtain higher performance
load_from = 'checkpoints/mask_rcnn_r50_caffe_fpn_mstrain-poly_3x_coco_bbox_mAP-0.408__segm_mAP-0.37_20200504_163245-42aa3d00.pth'

from ..builder import NECKS

@NECKS.register_module()
class AugFPN(nn.Module):

    def __init__(self,
                in_channels,
                out_channels,
                num_outs,
                start_level=0,
                end_level=-1,
                add_extra_convs=False):
        pass

    def forward(self, inputs):
        # implementation is ignored
        pass

custom_imports = dict(
    imports=['mmdet.models.necks.augfpn.py'],
    allow_failed_imports=False)