DEVELOP_GUIDE.md 3.52 KB
Newer Older
xinghao's avatar
xinghao committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
# Develop Guide

## Custom Dataset
Add your custom dataset is simple and flexible.
For example, create `ssd/data/datasets/my_dataset.py`:
```python
import torch.utils.data

from ssd.structures.container import Container

class MyDataset(torch.utils.data.Dataset):
    def __init__(self, ..., transform=None, target_transform=None):
        # as you would do normally
        ...
        self.transform = transform
        self.target_transform = target_transform

    def __getitem__(self, index):
        # load the image as a PIL Image
        image = ...

        # load the bounding boxes in x1, y1, x2, y2 order.
        boxes = np.array((N, 4), dtype=np.float32)
        # and labels
        labels = np.array((N, ), dtype=np.int64)

        if self.transform:
            image, boxes, labels = self.transform(image, boxes, labels)
        if self.target_transform:
            boxes, labels = self.target_transform(boxes, labels)
        targets = Container(
            boxes=boxes,
            labels=labels,
        )
        # return the image, the targets and the index in your dataset
        return image, targets, index
```

in `ssd/data/datasets/__init__.py`
```python
from .my_dataset import MyDataset

_DATASETS = {
    'VOCDataset': VOCDataset,
    'COCODataset': COCODataset,
    'MyDataset': MyDataset,
}
```

in `ssd/config/path_catlog.py`:
```python
DATASETS = {
    ...
    'my_custom_dataset': {
        "arg1": "your/arg",
        "arg2": "your/arg",
    },
    ...
}

@staticmethod
def get(name):
    ...
    if name == 'my_custom_dataset':
        attrs = DatasetCatalog.DATASETS[name]
        return dict(factory="MyDataset", args=attrs)
    ...
```

in your `config.ymal`:
```yaml
DATASETS:
  TRAIN: ("my_custom_dataset", )
  TEST: ("my_custom_test_dataset", )
```

### Test
While the aforementioned example should work for training, it's also easy to add your custom test code:
in `ssd/data/datasets/evaluation/__init__.py`
```python
if isinstance(dataset, MyDataset):
    return my_own_evaluation(**args)
```

## Custom Backbone

It very simple to add your own backbone for SSD.
For example, create `ssd/modeling/backbone/my_backbone.py`:
```python
import torch.nn as nn

from ssd.modeling import registry
from ssd.utils.model_zoo import load_state_dict_from_url


class MyBackbone(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        ...

    def forward(self, x):
        features = []
        
        # forward your network
        
        # add arbitrary feature you want to do prediction upon it.
        
        features.append(feature1)
        features.append(feature2)
        features.append(feature3)
        features.append(feature4)

        # return them as a tuple
        return tuple(features)

@registry.BACKBONES.register('my_backbone')
def my_backbone(cfg, pretrained=True):
    model = MyBackbone(cfg)
    model_url = 'you_model_url'
    if pretrained:
        model.init_from_pretrain(load_state_dict_from_url(model_url))
    return model
```
in `ssd/modeling/backbone/__init__.py`:
```python
from .my_backbone import MyBackbone
```

in your `config.ymal`:
```yaml
MODEL:
  BACKBONE:
    NAME: 'my_backbone'
    OUT_CHANNELS: (-, -, -, -) # should match feature1 - feature4's out_channels in MyBackbone
  PRIORS:
    FEATURE_MAPS: [-, -, -, -] # feature1 - feature4's size
    STRIDES: [-, -, -, -] # feature1 - feature4's output stride
    MIN_SIZES: [21, 45, 99, 153] # your custom anchor settings
    MAX_SIZES: [45, 99, 153, 207]
    ASPECT_RATIOS: [[2, 3], [2, 3], [2, 3], [2, 3]]
    BOXES_PER_LOCATION: [6, 6, 6, 6]
```