├── LICENSE
├── README.md
├── configs
├── config.py
├── hsd_res101_coco_512.yaml
├── hsd_resx101_coco_512.yaml
├── hsd_vgg_coco_320.yaml
└── hsd_vgg_coco_512.yaml
├── data
├── __init__.py
├── coco.py
├── data_augment.py
├── scripts
│ ├── VOC2007.sh
│ └── VOC2012.sh
├── voc0712.py
└── voc_eval.py
├── demo.py
├── eval.py
├── hsd.jpg
├── layers
├── __init__.py
├── functions
│ ├── __init__.py
│ ├── detection.py
│ ├── prior_box.py
│ └── prior_layer.py
└── modules
│ ├── __init__.py
│ ├── focal_loss_sigmoid.py
│ ├── focal_loss_softmax.py
│ ├── hsd_multibox_loss.py
│ ├── multibox_loss.py
│ ├── weight_smooth_l1_loss.py
│ └── weight_softmax_loss.py
├── make.sh
├── models
├── attention.py
├── deform
│ ├── .gitignore
│ ├── README.md
│ ├── build.py
│ ├── functions
│ │ ├── __init__.py
│ │ └── deform_conv.py
│ ├── make.sh
│ ├── modules
│ │ ├── __init__.py
│ │ └── deform_conv.py
│ └── src
│ │ ├── deform_conv.c
│ │ ├── deform_conv.h
│ │ ├── deform_conv_cuda.c
│ │ ├── deform_conv_cuda.h
│ │ ├── deform_conv_cuda_kernel.cu
│ │ └── deform_conv_cuda_kernel.h
├── hsd_res.py
├── hsd_resx.py
├── hsd_vgg.py
├── model_builder.py
└── model_helper.py
├── train.py
└── utils
├── __init__.py
├── augmentations.py
├── averageMeter.py
├── box_utils.py
├── build.py
├── collections.py
├── get_class_map.py
├── nms
├── __init__.py
├── cpu_nms.c
├── cpu_nms.pyx
├── gpu_nms.cpp
├── gpu_nms.hpp
├── gpu_nms.pyx
├── nms_kernel.cu
└── py_cpu_nms.py
├── nms_wrapper.py
└── timer.py
/LICENSE:
--------------------------------------------------------------------------------
1 | MIT License
2 |
3 | Copyright (c) 2019 JialeCao001
4 |
5 | Permission is hereby granted, free of charge, to any person obtaining a copy
6 | of this software and associated documentation files (the "Software"), to deal
7 | in the Software without restriction, including without limitation the rights
8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9 | copies of the Software, and to permit persons to whom the Software is
10 | furnished to do so, subject to the following conditions:
11 |
12 | The above copyright notice and this permission notice shall be included in all
13 | copies or substantial portions of the Software.
14 |
15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21 | SOFTWARE.
22 |
--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
1 | # HSD
2 |
3 | This code is a official implementation of "[*Hierarchical Shot Detector (ICCV2019)*](https://openaccess.thecvf.com/content_ICCV_2019/papers/Cao_Hierarchical_Shot_Detector_ICCV_2019_paper.pdf)" on COCO object detection with Pytorch.
4 |
5 | ## Introduction
6 | We propose a novel pipeline for accurate object detection (called ROC). Instead of simultaneous classification and regression, ROC firstly conducts box regression, secondly predicts the feature sampling locations for box classification, and finally classifies regressed boxes with the features of offset locations. To achieve the better detection accuracy, a hierarchical shot detector is proposed by stacking two ROC modules. Meanwhile, the contextual information is also incorporated to enrich the features of the second ROC module.
7 |
8 |
9 |
10 |
HSD detection pipeline.
11 |
2 |
3 | int deform_conv_forward(THFloatTensor *input, THFloatTensor *offset,
4 | THFloatTensor *output)
5 | {
6 | // if (!THFloatTensor_isSameSizeAs(input1, input2))
7 | // return 0;
8 | // THFloatTensor_resizeAs(output, input);
9 | // THFloatTensor_cadd(output, input1, 1.0, input2);
10 | return 1;
11 | }
12 |
13 | int deform_conv_backward(THFloatTensor *grad_output, THFloatTensor *grad_input,
14 | THFloatTensor *grad_offset)
15 | {
16 | // THFloatTensor_resizeAs(grad_input, grad_output);
17 | // THFloatTensor_fill(grad_input, 1);
18 | return 1;
19 | }
20 |
--------------------------------------------------------------------------------
/models/deform/src/deform_conv.h:
--------------------------------------------------------------------------------
1 | int deform_conv_forward(THFloatTensor *input, THFloatTensor *offset,
2 | THFloatTensor *output);
3 | int deform_conv_backward(THFloatTensor *grad_output, THFloatTensor *grad_input,
4 | THFloatTensor *grad_offset);
5 |
--------------------------------------------------------------------------------
/models/deform/src/deform_conv_cuda.h:
--------------------------------------------------------------------------------
1 | int deform_conv_forward_cuda(THCudaTensor *input,
2 | THCudaTensor *weight, /*THCudaTensor * bias, */
3 | THCudaTensor *offset, THCudaTensor *output,
4 | THCudaTensor *columns, THCudaTensor *ones, int kW,
5 | int kH, int dW, int dH, int padW, int padH,
6 | int dilationH, int dilationW,
7 | int deformable_group);
8 |
9 | int deform_conv_backward_input_cuda(
10 | THCudaTensor *input, THCudaTensor *offset, THCudaTensor *gradOutput,
11 | THCudaTensor *gradInput, THCudaTensor *gradOffset, THCudaTensor *weight,
12 | THCudaTensor *columns, int kW, int kH, int dW, int dH, int padW, int padH,
13 | int dilationH, int dilationW, int deformable_group);
14 |
15 | int deform_conv_backward_parameters_cuda(
16 | THCudaTensor *input, THCudaTensor *offset, THCudaTensor *gradOutput,
17 | THCudaTensor *gradWeight, /*THCudaTensor *gradBias, */
18 | THCudaTensor *columns, THCudaTensor *ones, int kW, int kH, int dW, int dH,
19 | int padW, int padH, int dilationH, int dilationW, int deformable_group,
20 | float scale);
21 |
--------------------------------------------------------------------------------
/models/deform/src/deform_conv_cuda_kernel.h:
--------------------------------------------------------------------------------
1 | template
2 | void deformable_im2col(cudaStream_t stream, const DType *data_im,
3 | const DType *data_offset, const int channels,
4 | const int height, const int width, const int ksize_h,
5 | const int ksize_w, const int pad_h, const int pad_w,
6 | const int stride_h, const int stride_w,
7 | const int dilation_h, const int dilation_w,
8 | const int deformable_group, DType *data_col);
9 |
10 | template
11 | void deformable_col2im(cudaStream_t stream, const DType *data_col,
12 | const DType *data_offset, const int channels,
13 | const int height, const int width, const int ksize_h,
14 | const int ksize_w, const int pad_h, const int pad_w,
15 | const int stride_h, const int stride_w,
16 | const int dilation_h, const int dilation_w,
17 | const int deformable_group, DType *grad_im);
18 |
19 | template
20 | void deformable_col2im_coord(cudaStream_t stream, const DType *data_col,
21 | const DType *data_im, const DType *data_offset,
22 | const int channels, const int height,
23 | const int width, const int ksize_h,
24 | const int ksize_w, const int pad_h,
25 | const int pad_w, const int stride_h,
26 | const int stride_w, const int dilation_h,
27 | const int dilation_w, const int deformable_group,
28 | DType *grad_offset);
29 |
--------------------------------------------------------------------------------
/models/hsd_res.py:
--------------------------------------------------------------------------------
1 | # -*- coding: utf-8 -*-
2 | # Written by yq_yao
3 |
4 | import torch
5 | import torch.nn as nn
6 | import torch.nn.functional as F
7 | from torch.autograd import Variable
8 | from models.attention import PAM_Module
9 | from models.model_helper import weights_init
10 |
11 |
12 | def conv3x3(in_planes, out_planes, stride=1):
13 | "3x3 convolution with padding"
14 | return nn.Conv2d(
15 | in_planes,
16 | out_planes,
17 | kernel_size=3,
18 | stride=stride,
19 | padding=1,
20 | bias=False)
21 |
22 |
23 | class BasicBlock(nn.Module):
24 | expansion = 1
25 |
26 | def __init__(self, inplanes, planes, stride=1, downsample=None):
27 | super(BasicBlock, self).__init__()
28 | self.conv1 = conv3x3(inplanes, planes, stride)
29 | self.bn1 = nn.BatchNorm2d(planes)
30 | self.relu = nn.ReLU(inplace=True)
31 | self.conv2 = conv3x3(planes, planes)
32 | self.bn2 = nn.BatchNorm2d(planes)
33 | self.downsample = downsample
34 | self.stride = stride
35 |
36 | def forward(self, x):
37 | residual = x
38 |
39 | out = self.conv1(x)
40 | out = self.bn1(out)
41 | out = self.relu(out)
42 |
43 | out = self.conv2(out)
44 | out = self.bn2(out)
45 |
46 | if self.downsample is not None:
47 | residual = self.downsample(x)
48 | out += residual
49 | out = self.relu(out)
50 |
51 | return out
52 |
53 |
54 | class Bottleneck(nn.Module):
55 | expansion = 4
56 |
57 | def __init__(self, inplanes, planes, stride=1, downsample=None):
58 | super(Bottleneck, self).__init__()
59 | self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
60 | self.bn1 = nn.BatchNorm2d(planes)
61 | self.conv2 = nn.Conv2d(
62 | planes,
63 | planes,
64 | kernel_size=3,
65 | stride=stride,
66 | padding=1,
67 | bias=False)
68 | self.bn2 = nn.BatchNorm2d(planes)
69 | self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
70 | self.bn3 = nn.BatchNorm2d(planes * 4)
71 | self.relu = nn.ReLU(inplace=True)
72 | self.downsample = downsample
73 | self.stride = stride
74 |
75 | def forward(self, x):
76 | residual = x
77 |
78 | out = self.conv1(x)
79 | out = self.bn1(out)
80 | out = self.relu(out)
81 |
82 | out = self.conv2(out)
83 | out = self.bn2(out)
84 | out = self.relu(out)
85 |
86 | out = self.conv3(out)
87 | out = self.bn3(out)
88 |
89 | if self.downsample is not None:
90 | residual = self.downsample(x)
91 |
92 | out += residual
93 | out = self.relu(out)
94 |
95 | return out
96 | class BasicConv(nn.Module):
97 |
98 | def __init__(self, in_planes, out_planes, kernel_size, stride=1, padding=0, dilation=1, groups=1, relu=True, bn=True, bias=False):
99 | super(BasicConv, self).__init__()
100 | self.out_channels = out_planes
101 | self.conv = nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias)
102 | self.bn = nn.BatchNorm2d(out_planes,eps=1e-5, momentum=0.01, affine=True) if bn else None
103 | self.relu = nn.ReLU(inplace=True) if relu else None
104 |
105 | def forward(self, x):
106 | x = self.conv(x)
107 | if self.bn is not None:
108 | x = self.bn(x)
109 | if self.relu is not None:
110 | x = self.relu(x)
111 | return x
112 | class FEModule(nn.Module):
113 | def __init__(self, in_channels, out_channels, stride=1, norm_layer=nn.BatchNorm2d):
114 | super(FEModule, self).__init__()
115 | self.out_channels = out_channels
116 | inter_channels = in_channels // 4
117 | self.brancha = nn.Sequential(nn.Conv2d(in_channels, inter_channels, 3, padding=1, bias=False),
118 | norm_layer(inter_channels),
119 | nn.ReLU())
120 | self.sa = PAM_Module(inter_channels)
121 | self.brancha1 = nn.Sequential(nn.Conv2d(inter_channels, inter_channels, 3, padding=1, bias=False),
122 | norm_layer(inter_channels),
123 | nn.ReLU())
124 |
125 | self.sl = nn.Sequential(
126 | BasicConv(in_channels, inter_channels, kernel_size=1, stride=1),
127 | BasicConv(inter_channels, inter_channels, kernel_size=3, stride=1, padding=1)
128 | )
129 | # self.sl = BasicConv(in_channels, inter_channels+inter_channels, kernel_size=3, padding=1, stride=1)
130 | self.sn = nn.Sequential(
131 | BasicConv(in_channels, inter_channels, kernel_size=1, stride=1),
132 | BasicConv(inter_channels, inter_channels, kernel_size=3, stride=1, padding=3, dilation=3)
133 | )
134 | self.fuse = nn.Sequential(nn.Dropout2d(0.1, False),
135 | nn.Conv2d(inter_channels + inter_channels + inter_channels, out_channels,
136 | kernel_size=3, stride=stride, padding=1, bias=False),
137 | norm_layer(out_channels),
138 | nn.ReLU())
139 |
140 |
141 |
142 | def forward(self, x):
143 | sa_feat = self.sa(self.brancha(x))
144 | sa_conv = self.brancha1(sa_feat)
145 | sl_output = self.sl(x)
146 | sn_output = self.sn(x)
147 | feat_cat = torch.cat([sa_conv, sl_output, sn_output], dim=1)
148 | output = self.fuse(feat_cat)
149 | return output
150 | #
151 | def trans_head():
152 | arm_trans = []
153 | arm_trans += [BasicConv(512, 256, kernel_size=3, stride=1, padding=1)]
154 | arm_trans += [BasicConv(1024, 256, kernel_size=3, stride=1, padding=1)]
155 | arm_trans += [BasicConv(2048, 256, kernel_size=3, stride=1, padding=1)]
156 | arm_trans += [BasicConv(2048, 256, kernel_size=3, stride=1, padding=1)]
157 |
158 | orm_trans = []
159 | orm_trans += [BasicConv(256, 256, kernel_size=3, stride=1, padding=1)]
160 | orm_trans += [BasicConv(256, 256, kernel_size=3, stride=1, padding=1)]
161 | orm_trans += [BasicConv(256, 256, kernel_size=3, stride=1, padding=1)]
162 | orm_trans += [BasicConv(256, 256, kernel_size=3, stride=1, padding=1)]
163 |
164 | return arm_trans, orm_trans
165 |
166 | class HSDResnet(nn.Module):
167 | def __init__(self, block, num_blocks, size):
168 | super(HSDResnet, self).__init__()
169 | self.inplanes = 64
170 |
171 | self.conv1 = nn.Conv2d(
172 | 3, 64, kernel_size=7, stride=2, padding=3, bias=False)
173 | self.bn1 = nn.BatchNorm2d(64)
174 | # Bottom-up layers
175 | self.layer1 = self._make_layer(block, 64, num_blocks[0], stride=1)
176 | self.layer2 = self._make_layer(block, 128, num_blocks[1], stride=2)
177 | self.layer3 = self._make_layer(block, 256, num_blocks[2], stride=2)
178 | self.layer4 = self._make_layer(block, 512, num_blocks[3], stride=2)
179 | self.extras0 = self._make_layer(block, 512, 2, stride=2)
180 |
181 | self.fe1 = FEModule(512,256)
182 | self.fe2 = FEModule(512,256)
183 | self.fe3 = FEModule(512,256)
184 | self.arm_trans = nn.ModuleList(trans_head()[0])
185 | self.orm_trans = nn.ModuleList(trans_head()[1])
186 |
187 | self._init_modules()
188 |
189 | def _make_layer(self, block, planes, blocks, stride=1):
190 | downsample = None
191 | if stride != 1 or self.inplanes != planes * block.expansion:
192 | downsample = nn.Sequential(
193 | nn.Conv2d(
194 | self.inplanes,
195 | planes * block.expansion,
196 | kernel_size=1,
197 | stride=stride,
198 | bias=False),
199 | nn.BatchNorm2d(planes * block.expansion),
200 | )
201 |
202 | layers = []
203 | layers.append(block(self.inplanes, planes, stride, downsample))
204 | self.inplanes = planes * block.expansion
205 | for i in range(1, blocks):
206 | layers.append(block(self.inplanes, planes))
207 |
208 | return nn.Sequential(*layers)
209 |
210 | def _init_modules(self):
211 | self.extras0.apply(weights_init)
212 | self.arm_trans.apply(weights_init)
213 | self.orm_trans.apply(weights_init)
214 | self.fe1.apply(weights_init)
215 | self.fe2.apply(weights_init)
216 | self.fe3.apply(weights_init)
217 |
218 | def forward(self, x):
219 | c1 = F.relu(self.bn1(self.conv1(x)))
220 | c1 = F.max_pool2d(c1, kernel_size=3, stride=2, padding=1)
221 | c2 = self.layer1(c1)
222 | c3 = self.layer2(c2)
223 | c4 = self.layer3(c3)
224 | c5 = self.layer4(c4)
225 | c6 = self.extras0(c5)
226 |
227 | c3_0 = self.arm_trans[0](c3)
228 | c4_0 = self.arm_trans[1](c4)
229 | c5_0 = self.arm_trans[2](c5)
230 | c6_0 = self.arm_trans[3](c6)
231 |
232 | arm_sources = [c3_0, c4_0, c5_0, c6_0]
233 |
234 | odm_sources = []
235 | up = F.upsample(arm_sources[1], size=arm_sources[0].size()[2:], mode='bilinear')
236 | odm_sources.append(self.fe1(torch.cat([up, arm_sources[0]], dim = 1)))
237 | up = F.upsample(arm_sources[2], size=arm_sources[1].size()[2:], mode='bilinear')
238 | odm_sources.append(self.fe2(torch.cat([up, arm_sources[1]], dim=1)))
239 | up = F.upsample(arm_sources[3], size=arm_sources[2].size()[2:], mode='bilinear')
240 | odm_sources.append(self.fe3(torch.cat([up, arm_sources[2]], dim=1)))
241 | odm_sources.append(self.orm_trans[3](arm_sources[3]))
242 |
243 | return arm_sources, odm_sources
244 |
245 |
246 | def HSDResnet50(size):
247 | return HSDResnet(Bottleneck, [3, 4, 6, 3], size)
248 |
249 |
250 | def HSDResnet101(size):
251 | return HSDResnet(Bottleneck, [3, 4, 23, 3], size)
252 |
253 |
254 | def HSDResnet152(size):
255 | return HSDResnet(Bottleneck, [3, 8, 36, 3], size)
256 |
--------------------------------------------------------------------------------
/models/hsd_vgg.py:
--------------------------------------------------------------------------------
1 | # -*- coding: utf-8 -*-
2 | # Written by yq_yao
3 |
4 | import torch
5 | import torch.nn as nn
6 | import torch.nn.functional as F
7 | from torch.autograd import Variable
8 | import torch.nn.init as init
9 | from models.model_helper import weights_init
10 | from models.attention import PAM_Module
11 |
12 |
13 | class L2Norm(nn.Module):
14 | def __init__(self, n_channels, scale):
15 | super(L2Norm, self).__init__()
16 | self.n_channels = n_channels
17 | self.gamma = scale or None
18 | self.eps = 1e-10
19 | self.weight = nn.Parameter(torch.Tensor(self.n_channels))
20 | self.reset_parameters()
21 |
22 | def reset_parameters(self):
23 | init.constant_(self.weight, self.gamma)
24 |
25 | def forward(self, x):
26 | norm = x.pow(2).sum(dim=1, keepdim=True).sqrt() + self.eps
27 | x = x / norm
28 | out = self.weight.unsqueeze(0).unsqueeze(2).unsqueeze(3).expand_as(
29 | x) * x
30 | return out
31 |
32 | class BasicBlock(nn.Module):
33 | def __init__(self, in_planes, out_planes, stride=1):
34 | super(BasicBlock, self).__init__()
35 | self.out_channels = out_planes
36 | inter_planes = in_planes // 4
37 | self.single_branch = nn.Sequential(
38 | BasicConv(in_planes, inter_planes, kernel_size=(3, 3), stride=stride, padding=(1, 1)),
39 | BasicConv(inter_planes, inter_planes, kernel_size=3, stride=1, padding=2, dilation=2),
40 | BasicConv(inter_planes, out_planes, kernel_size=(3, 3), stride=1, padding=(1, 1))
41 | )
42 |
43 | def forward(self, x):
44 | out = self.single_branch(x)
45 | return out
46 | # This function is derived from torchvision VGG make_layers()
47 | # https://github.com/pytorch/vision/blob/master/torchvision/models/vgg.py
48 |
49 | class BasicConv(nn.Module):
50 |
51 | def __init__(self, in_planes, out_planes, kernel_size, stride=1, padding=0, dilation=1, groups=1, relu=True, bn=True, bias=False):
52 | super(BasicConv, self).__init__()
53 | self.out_channels = out_planes
54 | self.conv = nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias)
55 | self.bn = nn.BatchNorm2d(out_planes,eps=1e-5, momentum=0.01, affine=True) if bn else None
56 | self.relu = nn.ReLU(inplace=True) if relu else None
57 |
58 | def forward(self, x):
59 | x = self.conv(x)
60 | if self.bn is not None:
61 | x = self.bn(x)
62 | if self.relu is not None:
63 | x = self.relu(x)
64 | return x
65 |
66 |
67 |
68 | def vgg(cfg, i, batch_norm=False):
69 | layers = []
70 | in_channels = i
71 | for v in cfg:
72 | if v == 'M':
73 | layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
74 | elif v == 'C':
75 | layers += [nn.MaxPool2d(kernel_size=2, stride=2, ceil_mode=True)]
76 | else:
77 | conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1)
78 | if batch_norm:
79 | layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)]
80 | else:
81 | layers += [conv2d, nn.ReLU(inplace=True)]
82 | in_channels = v
83 | pool5 = nn.MaxPool2d(kernel_size=3, stride=1, padding=1)
84 | conv6 = nn.Conv2d(512, 1024, kernel_size=3, padding=3, dilation=3)
85 | conv7 = nn.Conv2d(1024, 1024, kernel_size=1)
86 | layers += [
87 | pool5, conv6,
88 | nn.ReLU(inplace=True), conv7,
89 | nn.ReLU(inplace=True)
90 | ]
91 | return layers
92 |
93 |
94 | base = {
95 | '300': [
96 | 64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'C', 512, 512, 512, 'M',
97 | 512, 512, 512
98 | ],
99 | '512': [
100 | 64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'C', 512, 512, 512, 'M',
101 | 512, 512, 512
102 | ],
103 | }
104 |
105 |
106 | def add_extras(size):
107 | layers = []
108 | layers += [BasicBlock(1024, 256, stride=2)]
109 | layers += [BasicBlock(256, 256, stride=2)]
110 | return layers
111 |
112 | #
113 |
114 | class SELayer(nn.Module):
115 | def __init__(self, channel, reduction=8):
116 | super(SELayer, self).__init__()
117 | self.avg_pool = nn.AdaptiveAvgPool2d(1)
118 | self.fc = nn.Sequential(
119 | nn.Linear(channel, channel // reduction),
120 | nn.ReLU(inplace=True),
121 | nn.Linear(channel // reduction, channel),
122 | nn.Sigmoid()
123 | )
124 |
125 | def forward(self, x):
126 | b, c, _, _ = x.size()
127 | y = self.avg_pool(x).view(b, c)
128 | y = self.fc(y).view(b, c, 1, 1)
129 | return x * (1+y)
130 |
131 | class FEModule(nn.Module):
132 | def __init__(self, in_channels, out_channels, stride=1, norm_layer=nn.BatchNorm2d):
133 | super(FEModule, self).__init__()
134 | self.out_channels = out_channels
135 | inter_channels = in_channels // 4
136 | self.brancha = nn.Sequential(nn.Conv2d(in_channels, inter_channels, 3, padding=1, bias=False),
137 | norm_layer(inter_channels),
138 | nn.ReLU())
139 |
140 | self.sa = PAM_Module(inter_channels)
141 | self.brancha1 = nn.Sequential(nn.Conv2d(inter_channels, inter_channels, 3, padding=1, bias=False),
142 | norm_layer(inter_channels),
143 | nn.ReLU())
144 |
145 | # aspp
146 | self.sl = nn.Sequential(
147 | BasicConv(in_channels, inter_channels, kernel_size=1, stride=1),
148 | BasicConv(inter_channels, inter_channels, kernel_size=3, stride=1, padding=1)
149 | )
150 | self.sn = nn.Sequential(
151 | BasicConv(in_channels, inter_channels, kernel_size=1, stride=1),
152 | BasicConv(inter_channels, inter_channels, kernel_size=3, stride=1, padding=3, dilation=3)
153 | )
154 |
155 | self.fuse = nn.Sequential(nn.Dropout2d(0.1, False),
156 | nn.Conv2d(inter_channels + inter_channels + inter_channels, out_channels,
157 | kernel_size=3, stride=stride, padding=1, bias=False),
158 | norm_layer(out_channels),
159 | nn.ReLU())
160 |
161 |
162 |
163 | def forward(self, x):
164 | sa_feat = self.sa(self.brancha(x))
165 | sa_conv = self.brancha1(sa_feat)
166 |
167 | sl_output = self.sl(x)
168 | sn_output = self.sn(x)
169 |
170 | feat_cat = torch.cat([sa_conv, sl_output, sn_output], dim=1)
171 | sasc_output = self.fuse(feat_cat)
172 |
173 | return sasc_output
174 |
175 | def trans_head():
176 | arm_trans = []
177 | arm_trans += [BasicConv(512, 256, kernel_size=3, stride=1, padding=1)]
178 | arm_trans += [BasicConv(1024, 256, kernel_size=3, stride=1, padding=1)]
179 | arm_trans += [BasicConv(256, 256, kernel_size=3, stride=1, padding=1)]
180 | arm_trans += [BasicConv(256, 256, kernel_size=3, stride=1, padding=1)]
181 |
182 | orm_trans = []
183 | orm_trans += [BasicConv(256, 512, kernel_size=3, stride=1, padding=1)]
184 | orm_trans += [BasicConv(256, 512, kernel_size=3, stride=1, padding=1)]
185 | orm_trans += [BasicConv(256, 512, kernel_size=3, stride=1, padding=1)]
186 | orm_trans += [BasicConv(256, 256, kernel_size=3, stride=1, padding=1)]
187 |
188 | return arm_trans, orm_trans
189 |
190 | class VGG16Extractor(nn.Module):
191 | def __init__(self, size, channel_size='48'):
192 | super(VGG16Extractor, self).__init__()
193 | self.vgg = nn.ModuleList(vgg(base[str(size)], 3))
194 | self.extras = nn.ModuleList(add_extras(str(size)))
195 |
196 | self.fe1 = FEModule(512,256)
197 | self.fe2 = FEModule(512,256)
198 | self.fe3 = FEModule(512,256)
199 | self.arm_trans = nn.ModuleList(trans_head()[0])
200 | self.orm_trans = nn.ModuleList(trans_head()[1])
201 |
202 | self._init_modules()
203 |
204 | def _init_modules(self):
205 | self.extras.apply(weights_init)
206 | self.orm_trans.apply(weights_init)
207 | self.arm_trans.apply(weights_init)
208 | self.fe1.apply(weights_init)
209 | self.fe2.apply(weights_init)
210 | self.fe3.apply(weights_init)
211 |
212 |
213 | def forward(self, x):
214 | """Applies network layers and ops on input image(s) x.
215 | Args:
216 | x: input image or batch of images. Shape: [batch,3*batch,300,300].
217 | Return:
218 | Depending on phase:
219 | test:
220 | Variable(tensor) of output class label predictions,
221 | confidence score, and corresponding location predictions for
222 | each object detected. Shape: [batch,topk,7]
223 | train:
224 | list of concat outputs from:
225 | 1: confidence layers, Shape: [batch*num_priors,num_classes]
226 | 2: localization layers, Shape: [batch,num_priors*4]
227 | 3: priorbox layers, Shape: [2,num_priors*4]
228 | """
229 | arm_sources = list()
230 |
231 | for i in range(23):
232 | x = self.vgg[i](x)
233 | #38x38
234 | c2 = x
235 | c2 = self.arm_trans[0](c2)
236 | arm_sources.append(c2)
237 |
238 | for k in range(23, len(self.vgg)):
239 | x = self.vgg[k](x)
240 | #19x19
241 | c3 = x
242 | c3 = self.arm_trans[1](c3)
243 | arm_sources.append(c3)
244 |
245 | # 10x10
246 | x = self.extras[0](x)
247 |
248 | # c4 = x
249 | c4 = self.arm_trans[2](x)
250 | arm_sources.append(c4)
251 |
252 | # 5x5
253 | x = self.extras[1](x)
254 | # c5 = x
255 | c5 = self.arm_trans[3](x)
256 | arm_sources.append(c5)
257 |
258 | odm_sources = []
259 | up = F.upsample(arm_sources[1], size=arm_sources[0].size()[2:], mode='bilinear')
260 | odm_sources.append(self.fe1(torch.cat([up, arm_sources[0]], dim = 1)))
261 | up = F.upsample(arm_sources[2], size=arm_sources[1].size()[2:], mode='bilinear')
262 | odm_sources.append(self.fe2(torch.cat([up, arm_sources[1]], dim=1)))
263 | up = F.upsample(arm_sources[3], size=arm_sources[2].size()[2:], mode='bilinear')
264 | odm_sources.append(self.fe3(torch.cat([up, arm_sources[2]], dim=1)))
265 | odm_sources.append(self.orm_trans[3](arm_sources[3]))
266 |
267 |
268 | return arm_sources, odm_sources
269 |
270 |
271 | def hsd_vgg(size):
272 | return VGG16Extractor(size)
--------------------------------------------------------------------------------
/models/model_builder.py:
--------------------------------------------------------------------------------
1 | import torch
2 | import torch.nn as nn
3 | import torch.nn.functional as F
4 | from torch.autograd import Variable
5 | from layers import *
6 | import os
7 | from models.model_helper import weights_init, weights_init1
8 | import importlib
9 | from layers.functions.prior_layer import PriorLayer
10 | from models.deform.modules import ConvOffset2d
11 |
12 | def get_func(func_name):
13 | """Helper to return a function object by name. func_name must identify a
14 | function in this module or the path to a function relative to the base
15 | 'modeling' module.
16 | """
17 | if func_name == '':
18 | return None
19 | try:
20 | parts = func_name.split('.')
21 | # Refers to a function in this module
22 | if len(parts) == 1:
23 | return globals()[parts[0]]
24 | # Otherwise, assume we're referencing a module under modeling
25 | module_name = 'models.' + '.'.join(parts[:-1])
26 | module = importlib.import_module(module_name)
27 | return getattr(module, parts[-1])
28 | except Exception:
29 | print('Failed to find function: %s', func_name)
30 | raise
31 |
32 |
33 | class SSD(nn.Module):
34 | """Single Shot Multibox Architecture
35 | The network is composed of a base VGG network followed by the
36 | added multibox conv layers. Each multibox layer branches into
37 | 1) conv2d for class conf scores
38 | 2) conv2d for localization predictions
39 | 3) associated priorbox layer to produce default bounding
40 | boxes specific to the layer's feature map size.
41 | See: https://arxiv.org/pdf/1512.02325.pdf for more details.
42 |
43 | Args:
44 | phase: (string) Can be "test" or "train"
45 | base: VGG16 layers for input, size of either 300 or 500
46 | extras: extra layers that feed to multibox loc and conf layers
47 | head: "multibox head" consists of loc and conf conv layers
48 | """
49 |
50 | def _init_modules(self):
51 | self.arm_loc.apply(weights_init)
52 | self.arm_conf.apply(weights_init)
53 | self.deform_conv_c.apply(weights_init)
54 | self.deform_conv_l1.apply(weights_init)
55 | self.deform_conv_c1.apply(weights_init)
56 | self.offset.apply(weights_init1)
57 | self.offset1.apply(weights_init1)
58 | self.offlat.apply(weights_init)
59 | self.offlat1.apply(weights_init)
60 |
61 | if self.cfg.MODEL.LOAD_PRETRAINED_WEIGHTS:
62 | weights = torch.load(self.cfg.MODEL.PRETRAIN_WEIGHTS)
63 | print("load pretrain model {}".format(
64 | self.cfg.MODEL.PRETRAIN_WEIGHTS))
65 | if self.cfg.MODEL.TYPE.split('_')[-1] == 'vgg':
66 | self.extractor.vgg.load_state_dict(weights)
67 | elif self.cfg.MODEL.TYPE.split('_')[-1] == 'res':
68 | self.extractor.load_state_dict(weights, strict=False)
69 | else:
70 | self.extractor.load_state_dict(weights['state_dict'], strict=False)
71 | print(weights['state_dict'])
72 | def __init__(self, cfg):
73 | super(SSD, self).__init__()
74 | self.cfg = cfg
75 | self.size = cfg.MODEL.SIZE
76 | if self.size == '300':
77 | size_cfg = cfg.SMALL
78 | else:
79 | size_cfg = cfg.BIG
80 | self.num_classes = cfg.MODEL.NUM_CLASSES
81 | self.prior_layer = PriorLayer(cfg)
82 | self.priorbox = PriorBox(cfg)
83 | self.priors = self.priorbox.forward()
84 | self.extractor = get_func(cfg.MODEL.CONV_BODY)(self.size)
85 | if cfg.MODEL.CASCADE:
86 | self.odm_channels = size_cfg.ODM_CHANNELS
87 | self.arm_num_classes = 2
88 |
89 | self.arm_loc = nn.ModuleList()
90 | self.arm_conf = nn.ModuleList()
91 | self.arm_channels = size_cfg.ARM_CHANNELS
92 | self.num_anchors = size_cfg.NUM_ANCHORS
93 | self.input_fixed = size_cfg.INPUT_FIXED
94 | self.arm_loc = nn.ModuleList()
95 | self.arm_conf = nn.ModuleList()
96 | self.deform_conv_c = nn.ModuleList()
97 | self.deform_conv_l1 = nn.ModuleList()
98 | self.deform_conv_c1 = nn.ModuleList()
99 | self.offset = nn.ModuleList()
100 | self.offset1 = nn.ModuleList()
101 | self.offlat1 = nn.ModuleList()
102 | self.offlat = nn.ModuleList()
103 |
104 | for i in range(len(self.arm_channels)):
105 | if cfg.MODEL.CASCADE:
106 | self.arm_loc += [nn.Conv2d(self.arm_channels[i], self.num_anchors[i] * 4, kernel_size=3, padding=1)]
107 | self.arm_conf += [nn.Conv2d(self.arm_channels[i], self.num_anchors[i] * self.arm_num_classes, kernel_size=3, padding=1)]
108 |
109 | self.deform_conv_c += [ConvOffset2d(256*3, self.num_anchors[i] * self.arm_num_classes, 3, stride=1, padding=1, num_deformable_groups=3)]
110 | self.deform_conv_c1 += [ConvOffset2d(256*3, self.num_anchors[i] * self.num_classes, 3, stride=1, padding=1, num_deformable_groups=3)]
111 | self.deform_conv_l1 += [ConvOffset2d(256*3, self.num_anchors[i] * 4, 3, stride=1, padding=1, num_deformable_groups=3)]
112 | self.offset += [nn.Conv2d(36, 2 * 3 * 3 * 3, kernel_size=1, stride=1, padding=0, bias=False, groups = 3)]
113 | self.offset1 += [nn.Conv2d(36, 2 * 3 * 3 * 3, kernel_size=1, stride=1, padding=0, bias=False, groups=3)]
114 | self.offlat += [nn.Sequential(nn.Conv2d(12, 36, kernel_size=1, stride=1, padding=0, groups = 3), nn.ReLU())]
115 | self.offlat1 += [nn.Sequential(nn.Conv2d(12, 36, kernel_size=1, stride=1, padding=0, groups = 3), nn.ReLU())]
116 |
117 | else:
118 | self.arm_loc += [nn.Conv2d(self.arm_channels[i], self.num_anchors[i] * 4, kernel_size=3, padding=1)]
119 | self.arm_conf += [nn.Conv2d(self.arm_channels[i], self.num_anchors[i] * self.num_classes, kernel_size=3, padding=1)]
120 | self.deform_conv_c += [ConvOffset2d(256 * 3, self.num_anchors[i] * self.num_classes, 3, stride=1, padding=1, num_deformable_groups=3)]
121 | self.offset += [nn.Conv2d(36, 2 * 3 * 3 * 3, kernel_size=1, stride=1, padding=0, bias=False, groups=3)]
122 | self.offlat += [nn.Sequential(nn.Conv2d(12, 36, kernel_size=1, stride=1, padding=0, groups=3), nn.ReLU())]
123 |
124 | if cfg.TRAIN.TRAIN_ON:
125 | self._init_modules()
126 |
127 | def forward(self, x):
128 |
129 | arm_loc = list()
130 | arm_conf = list()
131 | if self.cfg.MODEL.CASCADE:
132 | odm_loc = list()
133 | odm_conf = list()
134 | arm_xs, odm_xs = self.extractor(x)
135 |
136 | for (x, x0, l0, c0, deform_conv_c, deform_conv_l1, deform_conv_c1, offset, offset1, offlat, offlat1) in \
137 | zip(odm_xs, arm_xs, self.arm_loc, self.arm_conf,
138 | self.deform_conv_c, self.deform_conv_l1, self.deform_conv_c1,
139 | self.offset, self.offset1, self.offlat, self.offlat1):
140 |
141 | # s1
142 | x_l = l0(x0)
143 | arm_loc.append(x_l.permute(0, 2, 3, 1).contiguous())
144 | x_offset = offlat(x_l.detach())
145 | x_offset = offset(x_offset)
146 | x03 = torch.cat([x0, x0, x0], dim=1)
147 | x_c = deform_conv_c(x03, x_offset)
148 | arm_conf.append(x_c.permute(0, 2, 3, 1).contiguous())
149 | # s2
150 | x3 = torch.cat([x, x, x], dim=1)
151 | xl = deform_conv_l1(x3, x_offset)
152 | x_offset1 = offlat1(xl.detach())
153 | x_offset1 = x_offset+offset1(x_offset1)
154 | xc = deform_conv_c1(x3, x_offset1)
155 |
156 | odm_loc.append(xl.permute(0, 2, 3, 1).contiguous())
157 | odm_conf.append(xc.permute(0, 2, 3, 1).contiguous())
158 |
159 | arm_loc = torch.cat([o.view(o.size(0), -1) for o in arm_loc], 1)
160 | arm_conf = torch.cat([o.view(o.size(0), -1) for o in arm_conf], 1)
161 | odm_loc = torch.cat([o.view(o.size(0), -1) for o in odm_loc], 1)
162 | odm_conf = torch.cat([o.view(o.size(0), -1) for o in odm_conf], 1)
163 | else:
164 | arm_xs, arm_xs1 = self.extractor(x)
165 |
166 | for (x0, x0_c, l0, c0, deform_conv_c, offset, offlat) in \
167 | zip(arm_xs, arm_xs1, self.arm_loc, self.arm_conf, self.deform_conv_c, self.offset, self.offlat):
168 |
169 | x_l = l0(x0)
170 | arm_loc.append(x_l.permute(0, 2, 3, 1).contiguous())
171 |
172 | x_offset = offlat(x_l.detach())
173 | x_offset = offset(x_offset)
174 | x03 = torch.cat([x0_c, x0_c, x0_c], dim=1)
175 | x_c = deform_conv_c(x03, x_offset)
176 | arm_conf.append(x_c.permute(0, 2, 3, 1).contiguous())
177 |
178 |
179 | arm_loc = torch.cat([o.view(o.size(0), -1) for o in arm_loc], 1)
180 | arm_conf = torch.cat([o.view(o.size(0), -1) for o in arm_conf], 1)
181 |
182 | img_wh = (x.size(3), x.size(2))
183 | feature_maps_wh = [(t.size(3), t.size(2)) for t in arm_xs]
184 |
185 | if self.cfg.MODEL.CASCADE:
186 | output = (arm_loc.view(arm_loc.size(0), -1, 4),
187 | arm_conf.view(
188 | arm_conf.size(0), -1, self.arm_num_classes),
189 | odm_loc.view(odm_loc.size(0), -1, 4),
190 | odm_conf.view(odm_conf.size(0), -1, self.num_classes),
191 | self.priors if self.input_fixed else self.prior_layer(
192 | img_wh, feature_maps_wh))
193 | else:
194 | output = (arm_loc.view(arm_loc.size(0), -1, 4),
195 | arm_conf.view(arm_conf.size(0), -1, self.num_classes),
196 | self.priors if self.input_fixed else self.prior_layer(
197 | img_wh, feature_maps_wh))
198 | return output
199 |
--------------------------------------------------------------------------------
/models/model_helper.py:
--------------------------------------------------------------------------------
1 | # -*- coding: utf-8 -*-
2 | # Written by yq_yao
3 |
4 | import torch
5 | import torch.nn as nn
6 | import torch.nn.functional as F
7 | from torch.autograd import Variable
8 | import torch.nn.init as init
9 |
10 |
11 | def xavier(param):
12 | init.xavier_uniform_(param)
13 |
14 |
15 | # def weights_init(m):
16 | # if isinstance(m, nn.Conv2d):
17 | # xavier(m.weight.data)
18 | # m.bias.data.zero_()
19 |
20 |
21 | def weights_init(m):
22 | for key in m.state_dict():
23 | if key.split('.')[-1] == 'weight':
24 | if 'conv' in key:
25 | init.kaiming_normal(m.state_dict()[key], mode='fan_out')
26 | if 'offset' in key:
27 | #init.kaiming_normal(m.state_dict()[key], mode='fan_out')
28 | m.state_dict()[key][...] = 0
29 |
30 | if 'bn' in key:
31 | m.state_dict()[key][...] = 1
32 | elif key.split('.')[-1] == 'bias':
33 | m.state_dict()[key][...] = 0
34 | def weights_init1(m):
35 | for key in m.state_dict():
36 | print(key)
37 | if key.split('.')[-1] == 'weight':
38 | init.constant_(m.state_dict()[key], 0.0)
39 | elif key.split('.')[-1] == 'bias':
40 | m.state_dict()[key][...] = 0
41 |
42 |
43 |
44 | def trans_layers(block, fpn_num):
45 | layers = list()
46 | for i in range(fpn_num):
47 | layers += [
48 | nn.Sequential(
49 | nn.Conv2d(block[i], 256, kernel_size=3, stride=1, padding=1),
50 | nn.ReLU(inplace=True),
51 | nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1))
52 | ]
53 |
54 | return layers
55 |
56 |
57 | def trans_layers_2(raw_channels, inner_channels):
58 | layers = list()
59 | fpn_num = len(raw_channels)
60 | for i in range(fpn_num):
61 | layers += [
62 | nn.Sequential(
63 | nn.Conv2d(
64 | raw_channels[i],
65 | inner_channels[i],
66 | kernel_size=3,
67 | stride=1,
68 | padding=1), nn.ReLU(inplace=True),
69 | nn.Conv2d(
70 | inner_channels[i],
71 | inner_channels[i],
72 | kernel_size=3,
73 | stride=1,
74 | padding=1))
75 | ]
76 |
77 | return layers
78 |
79 |
80 | def latent_layers(fpn_num):
81 | layers = []
82 | for i in range(fpn_num):
83 | layers += [nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1)]
84 | return layers
85 |
86 |
87 | def up_layers(fpn_num):
88 | layers = []
89 | for i in range(fpn_num - 1):
90 | layers += [nn.Upsample(scale_factor=2, mode='bilinear')]
91 | return layers
92 |
93 |
94 | class FpnAdapter(nn.Module):
95 | def __init__(self, block, fpn_num):
96 | super(FpnAdapter, self).__init__()
97 | self.trans_layers = nn.ModuleList(trans_layers(block, fpn_num))
98 | self.up_layers = nn.ModuleList(up_layers(fpn_num))
99 | self.latent_layers = nn.ModuleList(latent_layers(fpn_num))
100 | self._init_modules()
101 |
102 | def _init_modules(self):
103 | self.trans_layers.apply(weights_init)
104 | self.latent_layers.apply(weights_init)
105 |
106 | def forward(self, x):
107 | trans_layers_list = list()
108 | fpn_out = list()
109 | for (p, t) in zip(x, self.trans_layers):
110 | trans_layers_list.append(t(p))
111 | last = F.relu(
112 | self.latent_layers[-1](trans_layers_list[-1]), inplace=True)
113 | # last layer
114 | fpn_out.append(last)
115 | _up = self.up_layers[-1](last)
116 | for i in range(len(trans_layers_list) - 2, -1, -1):
117 | q = F.relu(trans_layers_list[i] + _up, inplace=True)
118 | q = F.relu(self.latent_layers[i](q), inplace=True)
119 | fpn_out.append(q)
120 | if i > 0:
121 | _up = self.up_layers[i - 1](q)
122 | fpn_out = fpn_out[::-1]
123 | return fpn_out
124 |
125 |
126 | class ConvPool(nn.Module):
127 | def __init__(self, inplane, plane):
128 | super(ConvPool, self).__init__()
129 | self.conv = nn.Conv2d(inplane, plane, kernel_size=1, stride=1)
130 | self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
131 | self._init_modules()
132 |
133 | def _init_modules(self):
134 | self.conv.apply(weights_init)
135 |
136 | def forward(self, x):
137 | out = self.conv(x)
138 | out = self.pool(out)
139 | return x, out
140 |
141 |
142 | class ConvUpsample(nn.Module):
143 | def __init__(self, inplace, plane):
144 | super(ConvUpsample, self).__init__()
145 | self.conv = nn.Conv2d(inplace, plane, kernel_size=1, stride=1)
146 | self.up_sample = nn.Upsample(scale_factor=2, mode='bilinear')
147 | self.smooth_conv = nn.Conv2d(plane, plane, kernel_size=1, stride=1)
148 | self._init_modules()
149 |
150 | def _init_modules(self):
151 | self.conv.apply(weights_init)
152 | self.smooth_conv.apply(weights_init)
153 |
154 | def forward(self, x):
155 | out = self.conv(x)
156 | out = self.up_sample(out)
157 | out = self.smooth_conv(out)
158 | return x, out
159 |
160 |
161 | class ConvPoolUpsample(nn.Module):
162 | def __init__(self, inplace, plane):
163 | super(ConvPoolUpsample, self).__init__()
164 | self.up_conv = nn.Conv2d(inplace, plane, kernel_size=1, stride=1)
165 | self.pool_conv = nn.Conv2d(inplace, plane, kernel_size=1, stride=1)
166 | self.up_sample = nn.Upsample(scale_factor=2, mode='bilinear')
167 | self.smooth_conv = nn.Conv2d(plane, plane, kernel_size=1, stride=1)
168 | self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
169 |
170 | self._init_modules()
171 |
172 | def _init_modules(self):
173 | self.up_conv.apply(weights_init)
174 | self.smooth_conv.apply(weights_init)
175 | self.pool_conv.apply(weights_init)
176 |
177 | def forward(self, x):
178 | up_out = self.up_conv(x)
179 | pool_out = self.pool_conv(x)
180 | up_out = self.up_sample(up_out)
181 | up_out = self.smooth_conv(up_out)
182 | pool_out = self.pool(pool_out)
183 | return x, pool_out, up_out
184 |
185 |
186 | def weave_layers(block, weave_num):
187 | layers = list()
188 | add_channel = 32
189 | for i in range(weave_num):
190 | if i == 0:
191 | layers += [ConvPool(block[i], add_channel)]
192 | elif i == weave_num - 1:
193 | layers += [ConvUpsample(block[i], add_channel)]
194 | else:
195 | layers += [ConvPoolUpsample(block[i], add_channel)]
196 | return layers
197 |
198 |
199 | class WeaveBlock(nn.Module):
200 | def __init__(self, raw_channel, weave_add_channel, dense_num):
201 | super(WeaveBlock, self).__init__()
202 | layers = list()
203 | for j in range(dense_num):
204 | layers += [
205 | nn.Conv2d(
206 | raw_channel, weave_add_channel[j], kernel_size=1, stride=1)
207 | ]
208 | self.weave_layers = nn.ModuleList(layers)
209 | self._init_modules()
210 |
211 | def _init_modules(self):
212 | self.weave_layers.apply(weights_init)
213 |
214 | def forward(self, x):
215 | out = list()
216 | out.append(x)
217 | for i in range(len(self.weave_layers)):
218 | out.append(self.weave_layers[i](x))
219 | return out
220 |
221 |
222 | def weave_layers_2(raw_channels, weave_add_channels):
223 | layers = list()
224 | num = 2
225 | weave_num = len(raw_channels)
226 | for i in range(weave_num):
227 | if i == 0 or i == weave_num - 1:
228 | layers += [
229 | WeaveBlock(raw_channels[i], weave_add_channels[i], num - 1)
230 | ]
231 | else:
232 | layers += [WeaveBlock(raw_channels[i], weave_add_channels[i], num)]
233 | return layers
234 |
235 |
236 | def weave_concat_layers_2(raw_channels, weave_add_channels, weave_channels):
237 | layers = list()
238 | weave_num = len(raw_channels)
239 | for i in range(weave_num):
240 | if i == 0:
241 | add_channel = weave_add_channels[i + 1][0]
242 | elif i == weave_num - 1:
243 | add_channel = weave_add_channels[i - 1][1]
244 | else:
245 | add_channel = weave_add_channels[i - 1][1] + weave_add_channels[
246 | i + 1][0]
247 | layers += [
248 | nn.Conv2d(
249 | raw_channels[i] + add_channel,
250 | weave_channels[i],
251 | kernel_size=1,
252 | stride=1)
253 | ]
254 | return layers
255 |
256 |
257 | def weave_concat_layers(block, weave_num, channel):
258 | layers = list()
259 | for i in range(weave_num):
260 | if i == 0 or i == weave_num - 1:
261 | add_channel = channel
262 | else:
263 | add_channel = channel * 2
264 | layers += [
265 | nn.Conv2d(block[i] + add_channel, 256, kernel_size=1, stride=1)
266 | ]
267 | return layers
268 |
269 |
270 | def adaptive_upsample(x, size):
271 | return F.upsample(x, size, mode='bilinear')
272 |
273 |
274 | def adaptive_pool(x, size):
275 | return F.adaptive_max_pool2d(x, size)
276 |
277 |
278 | class WeaveAdapter2(nn.Module):
279 | def __init__(self, raw_channels, weave_add_channels, weave_channels):
280 | super(WeaveAdapter2, self).__init__()
281 | self.trans_layers = nn.ModuleList(
282 | trans_layers_2(raw_channels, weave_channels))
283 | self.weave_layers = nn.ModuleList(
284 | weave_layers_2(weave_channels, weave_add_channels))
285 | self.weave_concat_layers = nn.ModuleList(
286 | weave_concat_layers_2(weave_channels, weave_add_channels,
287 | weave_channels))
288 | self.weave_num = len(raw_channels)
289 | self._init_modules()
290 |
291 | def _init_modules(self):
292 | self.trans_layers.apply(weights_init)
293 | self.weave_concat_layers.apply(weights_init)
294 |
295 | def forward(self, x):
296 | trans_layers_list = list()
297 | weave_out = list()
298 | for (p, t) in zip(x, self.trans_layers):
299 | trans_layers_list.append(t(p))
300 | weave_list = list()
301 | for (t, w) in zip(trans_layers_list, self.weave_layers):
302 | weave_list.append(w(t))
303 |
304 | for i in range(self.weave_num):
305 | b, c, h, w = weave_list[i][0].size()
306 | if i == 0:
307 | up = adaptive_upsample(weave_list[i + 1][1], (h, w))
308 | weave = torch.cat((up, weave_list[i][0]), 1)
309 | elif i == self.weave_num - 1:
310 | pool = adaptive_pool(weave_list[i - 1][-1], (h, w))
311 | weave = torch.cat((pool, weave_list[i][0]), 1)
312 | else:
313 | up = adaptive_upsample(weave_list[i + 1][1], (h, w))
314 | pool = adaptive_pool(weave_list[i - 1][-1], (h, w))
315 | weave = torch.cat((up, pool, weave_list[i][0]), 1)
316 | weave = F.relu(self.weave_concat_layers[i](weave), inplace=True)
317 | weave_out.append(weave)
318 | return weave_out
319 |
320 |
321 | class WeaveAdapter(nn.Module):
322 | def __init__(self, block, weave_num):
323 | super(WeaveAdapter, self).__init__()
324 | self.trans_layers = nn.ModuleList(trans_layers(block, weave_num))
325 | self.weave_layers = nn.ModuleList(
326 | weave_layers([256, 256, 256, 256], weave_num))
327 | self.weave_concat_layers = nn.ModuleList(
328 | weave_concat_layers([256, 256, 256, 256], weave_num, 48))
329 | self.weave_num = weave_num
330 | self._init_modules()
331 |
332 | def _init_modules(self):
333 | self.trans_layers.apply(weights_init)
334 | self.weave_concat_layers.apply(weights_init)
335 |
336 | def forward(self, x):
337 | trans_layers_list = list()
338 | weave_out = list()
339 | for (p, t) in zip(x, self.trans_layers):
340 | trans_layers_list.append(t(p))
341 | weave_list = list()
342 | for (t, w) in zip(trans_layers_list, self.weave_layers):
343 | weave_list.append(w(t))
344 |
345 | for i in range(self.weave_num):
346 | if i == 0:
347 | weave = torch.cat((weave_list[i][0], weave_list[i + 1][-1]), 1)
348 | elif i == self.weave_num - 1:
349 | weave = torch.cat((weave_list[i][0], weave_list[i - 1][1]), 1)
350 | else:
351 | weave = torch.cat((weave_list[i][0], weave_list[i - 1][1],
352 | weave_list[i + 1][-1]), 1)
353 | weave = F.relu(self.weave_concat_layers[i](weave), inplace=True)
354 | weave_out.append(weave)
355 | return weave_out
356 |
--------------------------------------------------------------------------------
/train.py:
--------------------------------------------------------------------------------
1 | import os
2 | import torch
3 | torch.set_num_threads(2)
4 | import torch.nn as nn
5 | import torch.optim as optim
6 | import torch.backends.cudnn as cudnn
7 | import torch.nn.init as init
8 | import argparse
9 | from torch.autograd import Variable
10 | import torch.utils.data as data
11 | from data import COCODetection, VOCDetection, BaseTransform, preproc
12 | from layers.modules import MultiBoxLoss, HSDMultiBoxLoss
13 | from layers.functions import Detect
14 | from utils.nms_wrapper import nms, soft_nms
15 | from configs.config import cfg, cfg_from_file
16 | import numpy as np
17 | import time
18 | import os
19 | import sys
20 | import pickle
21 | import datetime
22 | from models.model_builder import SSD
23 | import yaml
24 |
25 |
26 | def arg_parse():
27 | parser = argparse.ArgumentParser(description='HSD Training')
28 | parser.add_argument(
29 | '--cfg',
30 | dest='cfg_file',
31 | required=True,
32 | help='Config file for training (and optionally testing)')
33 | parser.add_argument(
34 | '--num_workers',
35 | default=8,
36 | type=int,
37 | help='Number of workers used in dataloading')
38 | parser.add_argument('--ngpu', default=2, type=int, help='gpus')
39 | parser.add_argument(
40 | '--resume_net', default=None, help='resume net for retraining')
41 | parser.add_argument(
42 | '--resume_epoch',
43 | default=0,
44 | type=int,
45 | help='resume iter for retraining')
46 |
47 | parser.add_argument(
48 | '--save_folder',
49 | default='./weights/hsd',
50 | help='Location to save checkpoint models')
51 | args = parser.parse_args()
52 | return args
53 |
54 | def detection_collate(batch):
55 | """Custom collate fn for dealing with batches of images that have a different
56 | number of associated object annotations (bounding boxes).
57 |
58 | Arguments:
59 | batch: (tuple) A tuple of tensor images and lists of annotations
60 |
61 | Return:
62 | A tuple containing:
63 | 1) (tensor) batch of images stacked on their 0 dim
64 | 2) (list of tensors) annotations for a given image are stacked on 0 dim
65 | """
66 | targets = []
67 | imgs = []
68 | img_info = []
69 | for sample in batch:
70 | imgs.append(sample[0])
71 | targets.append(torch.FloatTensor(sample[1]))
72 | img_info.append(torch.FloatTensor(sample[2]))
73 | return torch.stack(imgs, 0), targets, img_info
74 |
75 |
76 | def adjust_learning_rate(optimizer, epoch, step_epoch, gamma, epoch_size,
77 | iteration):
78 | """Sets the learning rate
79 | # Adapted from PyTorch Imagenet example:
80 | # https://github.com/pytorch/examples/blob/master/imagenet/main.py
81 | """
82 | ## warmup
83 | if epoch <= cfg.TRAIN.WARMUP_EPOCH:
84 | if cfg.TRAIN.WARMUP:
85 | iteration += (epoch_size * (epoch - 1))
86 | lr = 1e-6 + (cfg.SOLVER.BASE_LR - 1e-6) * iteration / (
87 | epoch_size * cfg.TRAIN.WARMUP_EPOCH)
88 | else:
89 | lr = cfg.SOLVER.BASE_LR
90 | else:
91 | div = 0
92 | if epoch > step_epoch[-1]:
93 | div = len(step_epoch) - 1
94 | else:
95 | for idx, v in enumerate(step_epoch):
96 | if epoch > step_epoch[idx] and epoch <= step_epoch[idx + 1]:
97 | div = idx
98 | break
99 | lr = cfg.SOLVER.BASE_LR * (gamma**div)
100 |
101 | for param_group in optimizer.param_groups:
102 | param_group['lr'] = lr
103 | return lr
104 |
105 | def train(train_loader, net, criterion, optimizer, epoch, epoch_step, gamma,
106 | end_epoch, cfg):
107 | net.train()
108 | begin = time.time()
109 | epoch_size = len(train_loader)
110 | for iteration, (imgs, targets, _) in enumerate(train_loader):
111 | t0 = time.time()
112 | lr = adjust_learning_rate(optimizer, epoch, epoch_step, gamma,
113 | epoch_size, iteration)
114 | imgs = imgs.cuda()
115 | imgs.requires_grad_()
116 | with torch.no_grad():
117 | targets = [anno.cuda() for anno in targets]
118 | output = net(imgs)
119 | optimizer.zero_grad()
120 | if not cfg.MODEL.CASCADE:
121 | ssd_criterion = criterion[0]
122 | loss_l, loss_c = ssd_criterion(output, targets)
123 | loss = loss_l + loss_c
124 | else:
125 | arm_criterion = criterion[0]
126 | odm_criterion = criterion[1]
127 | arm_loss_l, arm_loss_c = arm_criterion(output, targets)
128 | odm_loss_l, odm_loss_c = odm_criterion(
129 | output, targets, use_arm=True, filter_object=True)
130 | loss = arm_loss_l + arm_loss_c + odm_loss_l + odm_loss_c
131 | loss.backward()
132 | optimizer.step()
133 | t1 = time.time()
134 | iteration_time = t1 - t0
135 | all_time = ((end_epoch - epoch) * epoch_size +
136 | (epoch_size - iteration)) * iteration_time
137 | eta = str(datetime.timedelta(seconds=int(all_time)))
138 | if iteration % 10 == 0:
139 | if not cfg.MODEL.CASCADE:
140 | print('Epoch:' + repr(epoch) + ' || epochiter: ' +
141 | repr(iteration % epoch_size) + '/' + repr(epoch_size) +
142 | ' || L: %.4f C: %.4f||' %
143 | (loss_l.item(), loss_c.item()) +
144 | 'iteration time: %.4f sec. ||' % (t1 - t0) +
145 | 'LR: %.5f' % (lr) + ' || eta time: {}'.format(eta))
146 | else:
147 | print('Epoch:' + repr(epoch) + ' || epochiter: ' +
148 | repr(iteration % epoch_size) + '/' + repr(epoch_size) +
149 | '|| 1st_L: %.4f 1st_C: %.4f||' %
150 | (arm_loss_l.item(), arm_loss_c.item()) +
151 | ' 2rd_L: %.4f 2rd_C: %.4f||' %
152 | (odm_loss_l.item(), odm_loss_c.item()) +
153 | ' loss: %.4f||' % (loss.item()) +
154 | 'iteration time: %.4f sec. ||' % (t1 - t0) +
155 | 'LR: %.5f' % (lr) + ' || eta time: {}'.format(eta))
156 |
157 |
158 | def save_checkpoint(net, epoch, size, optimizer):
159 | save_name = os.path.join(
160 | args.save_folder,
161 | cfg.MODEL.TYPE + "_epoch_{}_{}".format(str(epoch), str(size)) + '.pth')
162 | torch.save({
163 | 'epoch': epoch,
164 | 'size': size,
165 | 'batch_size': cfg.TRAIN.BATCH_SIZE,
166 | 'model': net.state_dict(),
167 | 'optimizer': optimizer.state_dict()
168 | }, save_name)
169 |
170 |
171 | def eval_net(val_dataset,
172 | val_loader,
173 | net,
174 | detector,
175 | cfg,
176 | transform,
177 | max_per_image=300,
178 | thresh=0.01,
179 | batch_size=1):
180 | net.eval()
181 | num_images = len(val_dataset)
182 | num_classes = cfg.MODEL.NUM_CLASSES
183 | eval_save_folder = "./eval/"
184 | if not os.path.exists(eval_save_folder):
185 | os.mkdir(eval_save_folder)
186 | all_boxes = [[[] for _ in range(num_images)] for _ in range(num_classes)]
187 | det_file = os.path.join(eval_save_folder, 'detections.pkl')
188 | st = time.time()
189 | for idx, (imgs, _, img_info) in enumerate(val_loader):
190 | with torch.no_grad():
191 | t1 = time.time()
192 | x = imgs
193 | x = x.cuda()
194 | output = net(x)
195 | t4 = time.time()
196 | boxes, scores = detector.forward(output)
197 | t2 = time.time()
198 | for k in range(boxes.size(0)):
199 | i = idx * batch_size + k
200 | boxes_ = boxes[k]
201 | scores_ = scores[k]
202 | boxes_ = boxes_.cpu().numpy()
203 | scores_ = scores_.cpu().numpy()
204 | img_wh = img_info[k]
205 | scale = np.array([img_wh[0], img_wh[1], img_wh[0], img_wh[1]])
206 | boxes_ *= scale
207 | for j in range(1, num_classes):
208 | inds = np.where(scores_[:, j] > thresh)[0]
209 | if len(inds) == 0:
210 | all_boxes[j][i] = np.empty([0, 5], dtype=np.float32)
211 | continue
212 | c_bboxes = boxes_[inds]
213 | c_scores = scores_[inds, j]
214 | c_dets = np.hstack((c_bboxes,
215 | c_scores[:, np.newaxis])).astype(
216 | np.float32, copy=False)
217 | keep = nms(c_dets, cfg.TEST.NMS_OVERLAP, force_cpu=True)
218 | keep = keep[:50]
219 | c_dets = c_dets[keep, :]
220 | all_boxes[j][i] = c_dets
221 | t3 = time.time()
222 | detect_time = t2 - t1
223 | nms_time = t3 - t2
224 | forward_time = t4 - t1
225 | if idx % 10 == 0:
226 | print('im_detect: {:d}/{:d} {:.3f}s {:.3f}s {:.3f}s'.format(
227 | i + 1, num_images, forward_time, detect_time, nms_time))
228 | print("detect time: ", time.time() - st)
229 | with open(det_file, 'wb') as f:
230 | pickle.dump(all_boxes, f, pickle.HIGHEST_PROTOCOL)
231 | print('Evaluating detections')
232 | val_dataset.evaluate_detections(all_boxes, eval_save_folder)
233 |
234 |
235 |
236 | def main():
237 | global args
238 | args = arg_parse()
239 | cfg_from_file(args.cfg_file)
240 | save_folder = args.save_folder
241 | batch_size = cfg.TRAIN.BATCH_SIZE
242 | bgr_means = cfg.TRAIN.BGR_MEAN
243 | p = 0.6
244 | gamma = cfg.SOLVER.GAMMA
245 | momentum = cfg.SOLVER.MOMENTUM
246 | weight_decay = cfg.SOLVER.WEIGHT_DECAY
247 | size = cfg.MODEL.SIZE
248 | thresh = cfg.TEST.CONFIDENCE_THRESH
249 | if cfg.DATASETS.DATA_TYPE == 'VOC':
250 | trainvalDataset = VOCDetection
251 | top_k = 200
252 | else:
253 | trainvalDataset = COCODetection
254 | top_k = 300
255 | dataset_name = cfg.DATASETS.DATA_TYPE
256 | dataroot = cfg.DATASETS.DATAROOT
257 | trainSet = cfg.DATASETS.TRAIN_TYPE
258 | valSet = cfg.DATASETS.VAL_TYPE
259 | num_classes = cfg.MODEL.NUM_CLASSES
260 | start_epoch = args.resume_epoch
261 | epoch_step = cfg.SOLVER.EPOCH_STEPS
262 | end_epoch = cfg.SOLVER.END_EPOCH
263 | if not os.path.exists(save_folder):
264 | os.mkdir(save_folder)
265 | torch.set_default_tensor_type('torch.cuda.FloatTensor')
266 | net = SSD(cfg)
267 | print(net)
268 | if cfg.MODEL.SIZE == '300':
269 | size_cfg = cfg.SMALL
270 | else:
271 | size_cfg = cfg.BIG
272 | optimizer = optim.SGD(
273 | net.parameters(),
274 | lr=cfg.SOLVER.BASE_LR,
275 | momentum=momentum,
276 | weight_decay=weight_decay)
277 | if args.resume_net != None:
278 | checkpoint = torch.load(args.resume_net)
279 | state_dict = checkpoint['model']
280 | from collections import OrderedDict
281 | new_state_dict = OrderedDict()
282 | for k, v in state_dict.items():
283 | head = k[:7]
284 | if head == 'module.':
285 | name = k[7:] # remove `module.`
286 | else:
287 | name = k
288 | new_state_dict[name] = v
289 | net.load_state_dict(new_state_dict)
290 | optimizer.load_state_dict(checkpoint['optimizer'])
291 | print('Loading resume network...')
292 | if args.ngpu > 1:
293 | net = torch.nn.DataParallel(net)
294 | net.cuda()
295 | cudnn.benchmark = True
296 |
297 | criterion = list()
298 | if cfg.MODEL.CASCADE:
299 | detector = Detect(cfg)
300 | arm_criterion = HSDMultiBoxLoss(cfg, 2)
301 | odm_criterion = HSDMultiBoxLoss(cfg, cfg.MODEL.NUM_CLASSES)
302 | criterion.append(arm_criterion)
303 | criterion.append(odm_criterion)
304 | else:
305 | detector = Detect(cfg)
306 | ssd_criterion = MultiBoxLoss(cfg)
307 | criterion.append(ssd_criterion)
308 |
309 | TrainTransform = preproc(size_cfg.IMG_WH, bgr_means, p)
310 | ValTransform = BaseTransform(size_cfg.IMG_WH, bgr_means, (2, 0, 1))
311 |
312 | val_dataset = trainvalDataset(dataroot, valSet, ValTransform, dataset_name)
313 | val_loader = data.DataLoader(
314 | val_dataset,
315 | batch_size,
316 | shuffle=False,
317 | num_workers=args.num_workers,
318 | collate_fn=detection_collate)
319 |
320 | for epoch in range(start_epoch + 1, end_epoch + 1):
321 | train_dataset = trainvalDataset(dataroot, trainSet, TrainTransform,
322 | dataset_name)
323 | epoch_size = len(train_dataset)
324 | train_loader = data.DataLoader(
325 | train_dataset,
326 | batch_size,
327 | shuffle=True,
328 | num_workers=args.num_workers,
329 | collate_fn=detection_collate)
330 | train(train_loader, net, criterion, optimizer, epoch, epoch_step,
331 | gamma, end_epoch, cfg)
332 | if (epoch % 10 == 0) or (epoch % 10 == 0 and epoch >= 110):
333 | save_checkpoint(net, epoch, size, optimizer)
334 | if (epoch >= 50 and epoch % 10 == 0):
335 | eval_net(
336 | val_dataset,
337 | val_loader,
338 | net,
339 | detector,
340 | cfg,
341 | ValTransform,
342 | top_k,
343 | thresh=thresh,
344 | batch_size=batch_size)
345 | save_checkpoint(net, end_epoch, size, optimizer)
346 |
347 |
348 | if __name__ == '__main__':
349 | main()
350 |
--------------------------------------------------------------------------------
/utils/__init__.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/JialeCao001/HSD/8abcf78db5f313266a3bb3f85b9424927fe59a2d/utils/__init__.py
--------------------------------------------------------------------------------
/utils/averageMeter.py:
--------------------------------------------------------------------------------
1 | class AverageMeter(object):
2 | """Computes and stores the average and current value"""
3 |
4 | def __init__(self):
5 | self.reset()
6 |
7 | def reset(self):
8 | self.val = 0
9 | self.avg = 0
10 | self.sum = 0
11 | self.count = 0
12 |
13 | def update(self, val, n=1):
14 | self.val = val
15 | self.sum += val * n
16 | self.count += n
17 | self.avg = self.sum / self.count
--------------------------------------------------------------------------------
/utils/build.py:
--------------------------------------------------------------------------------
1 | # --------------------------------------------------------
2 | # Fast R-CNN
3 | # Copyright (c) 2015 Microsoft
4 | # Licensed under The MIT License [see LICENSE for details]
5 | # Written by Ross Girshick
6 | # --------------------------------------------------------
7 |
8 | import os
9 | from os.path import join as pjoin
10 | import numpy as np
11 | from distutils.core import setup
12 | from distutils.extension import Extension
13 | from Cython.Distutils import build_ext
14 |
15 |
16 | def find_in_path(name, path):
17 | "Find a file in a search path"
18 | # adapted fom http://code.activestate.com/recipes/52224-find-a-file-given-a-search-path/
19 | for dir in path.split(os.pathsep):
20 | binpath = pjoin(dir, name)
21 | if os.path.exists(binpath):
22 | return os.path.abspath(binpath)
23 | return None
24 |
25 |
26 | def locate_cuda():
27 | """Locate the CUDA environment on the system
28 |
29 | Returns a dict with keys 'home', 'nvcc', 'include', and 'lib64'
30 | and values giving the absolute path to each directory.
31 |
32 | Starts by looking for the CUDAHOME env variable. If not found, everything
33 | is based on finding 'nvcc' in the PATH.
34 | """
35 |
36 | # first check if the CUDAHOME env variable is in use
37 | if 'CUDAHOME' in os.environ:
38 | home = os.environ['CUDAHOME']
39 | nvcc = pjoin(home, 'bin', 'nvcc')
40 | else:
41 | # otherwise, search the PATH for NVCC
42 | default_path = pjoin(os.sep, 'usr', 'local', 'cuda', 'bin')
43 | nvcc = find_in_path('nvcc',
44 | os.environ['PATH'] + os.pathsep + default_path)
45 | if nvcc is None:
46 | raise EnvironmentError(
47 | 'The nvcc binary could not be '
48 | 'located in your $PATH. Either add it to your path, or set $CUDAHOME'
49 | )
50 | home = os.path.dirname(os.path.dirname(nvcc))
51 |
52 | cudaconfig = {
53 | 'home': home,
54 | 'nvcc': nvcc,
55 | 'include': pjoin(home, 'include'),
56 | 'lib64': pjoin(home, 'lib64')
57 | }
58 | for k, v in cudaconfig.items():
59 | if not os.path.exists(v):
60 | raise EnvironmentError(
61 | 'The CUDA %s path could not be located in %s' % (k, v))
62 |
63 | return cudaconfig
64 |
65 |
66 | CUDA = locate_cuda()
67 |
68 | # Obtain the numpy include directory. This logic works across numpy versions.
69 | try:
70 | numpy_include = np.get_include()
71 | except AttributeError:
72 | numpy_include = np.get_numpy_include()
73 |
74 |
75 | def customize_compiler_for_nvcc(self):
76 | """inject deep into distutils to customize how the dispatch
77 | to gcc/nvcc works.
78 |
79 | If you subclass UnixCCompiler, it's not trivial to get your subclass
80 | injected in, and still have the right customizations (i.e.
81 | distutils.sysconfig.customize_compiler) run on it. So instead of going
82 | the OO route, I have this. Note, it's kindof like a wierd functional
83 | subclassing going on."""
84 |
85 | # tell the compiler it can processes .cu
86 | self.src_extensions.append('.cu')
87 |
88 | # save references to the default compiler_so and _comple methods
89 | default_compiler_so = self.compiler_so
90 | super = self._compile
91 |
92 | # now redefine the _compile method. This gets executed for each
93 | # object but distutils doesn't have the ability to change compilers
94 | # based on source extension: we add it.
95 | def _compile(obj, src, ext, cc_args, extra_postargs, pp_opts):
96 | print(extra_postargs)
97 | if os.path.splitext(src)[1] == '.cu':
98 | # use the cuda for .cu files
99 | self.set_executable('compiler_so', CUDA['nvcc'])
100 | # use only a subset of the extra_postargs, which are 1-1 translated
101 | # from the extra_compile_args in the Extension class
102 | postargs = extra_postargs['nvcc']
103 | else:
104 | postargs = extra_postargs['gcc']
105 |
106 | super(obj, src, ext, cc_args, postargs, pp_opts)
107 | # reset the default compiler_so, which we might have changed for cuda
108 | self.compiler_so = default_compiler_so
109 |
110 | # inject our redefined _compile method into the class
111 | self._compile = _compile
112 |
113 |
114 | # run the customize_compiler
115 | class custom_build_ext(build_ext):
116 | def build_extensions(self):
117 | customize_compiler_for_nvcc(self.compiler)
118 | build_ext.build_extensions(self)
119 |
120 |
121 | ext_modules = [
122 | Extension(
123 | "nms.cpu_nms", ["nms/cpu_nms.pyx"],
124 | extra_compile_args={'gcc': ["-Wno-cpp", "-Wno-unused-function"]},
125 | include_dirs=[numpy_include]),
126 | Extension(
127 | 'nms.gpu_nms',
128 | ['nms/nms_kernel.cu', 'nms/gpu_nms.pyx'],
129 | library_dirs=[CUDA['lib64']],
130 | libraries=['cudart'],
131 | language='c++',
132 | runtime_library_dirs=[CUDA['lib64']],
133 | # this syntax is specific to this build system
134 | # we're only going to use certain compiler args with nvcc and not with gcc
135 | # the implementation of this trick is in customize_compiler() below
136 | extra_compile_args={
137 | 'gcc': ["-Wno-unused-function"],
138 | 'nvcc': [
139 | '-arch=sm_61', '--ptxas-options=-v', '-c',
140 | '--compiler-options', "'-fPIC'"
141 | ]
142 | },
143 | include_dirs=[numpy_include, CUDA['include']])
144 | ]
145 |
146 | setup(
147 | name='mot_utils',
148 | ext_modules=ext_modules,
149 | # inject our custom trigger
150 | cmdclass={'build_ext': custom_build_ext},
151 | )
152 |
--------------------------------------------------------------------------------
/utils/collections.py:
--------------------------------------------------------------------------------
1 | # Copyright (c) 2017-present, Facebook, Inc.
2 | #
3 | # Licensed under the Apache License, Version 2.0 (the "License");
4 | # you may not use this file except in compliance with the License.
5 | # You may obtain a copy of the License at
6 | #
7 | # http://www.apache.org/licenses/LICENSE-2.0
8 | #
9 | # Unless required by applicable law or agreed to in writing, software
10 | # distributed under the License is distributed on an "AS IS" BASIS,
11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | # See the License for the specific language governing permissions and
13 | # limitations under the License.
14 | ##############################################################################
15 | """A simple attribute dictionary used for representing configuration options."""
16 |
17 | from __future__ import absolute_import
18 | from __future__ import division
19 | from __future__ import print_function
20 | from __future__ import unicode_literals
21 |
22 |
23 | class AttrDict(dict):
24 |
25 | IMMUTABLE = '__immutable__'
26 |
27 | def __init__(self, *args, **kwargs):
28 | super(AttrDict, self).__init__(*args, **kwargs)
29 | self.__dict__[AttrDict.IMMUTABLE] = False
30 |
31 | def __getattr__(self, name):
32 | if name in self.__dict__:
33 | return self.__dict__[name]
34 | elif name in self:
35 | return self[name]
36 | else:
37 | raise AttributeError(name)
38 |
39 | def __setattr__(self, name, value):
40 | if not self.__dict__[AttrDict.IMMUTABLE]:
41 | if name in self.__dict__:
42 | self.__dict__[name] = value
43 | else:
44 | self[name] = value
45 | else:
46 | raise AttributeError(
47 | 'Attempted to set "{}" to "{}", but AttrDict is immutable'.
48 | format(name, value))
49 |
50 | def immutable(self, is_immutable):
51 | """Set immutability to is_immutable and recursively apply the setting
52 | to all nested AttrDicts.
53 | """
54 | self.__dict__[AttrDict.IMMUTABLE] = is_immutable
55 | # Recursively set immutable state
56 | for v in self.__dict__.values():
57 | if isinstance(v, AttrDict):
58 | v.immutable(is_immutable)
59 | for v in self.values():
60 | if isinstance(v, AttrDict):
61 | v.immutable(is_immutable)
62 |
63 | def is_immutable(self):
64 | return self.__dict__[AttrDict.IMMUTABLE]
65 |
--------------------------------------------------------------------------------
/utils/get_class_map.py:
--------------------------------------------------------------------------------
1 | import numpy as np
2 | import os
3 | import argparse
4 | import os.path as osp
5 |
6 |
7 | def check_size(submit_file):
8 | max_size = 60 * 1024 * 1024
9 | if osp.getsize(submit_file) > max_size:
10 | raise (
11 | IOError,
12 | "File size exceeds the specified maximum size, which is 60M for the server."
13 | )
14 |
15 |
16 | def parse_submission(submit_file):
17 | with open(submit_file, 'r') as f:
18 | lines = f.readlines()
19 | submit_dict = dict()
20 | final_dict = dict()
21 | splitlines = [x.strip().split(' ') for x in lines]
22 | for idx, val in enumerate(splitlines):
23 | cls = str(int(float(val[1])))
24 | if cls not in submit_dict:
25 | submit_dict[cls] = list()
26 | final_dict[cls] = dict()
27 | submit_dict[cls].append(
28 | [val[0], val[2], val[3], val[4], val[5], val[6]])
29 | for k, v in submit_dict.items():
30 | image_ids = [x[0] for x in v]
31 | confidence = np.array([float(x[1]) for x in v])
32 | BB = np.array([[float(z) for z in x[2:]] for x in v])
33 | sorted_ind = np.argsort(-confidence)
34 | sorted_scores = np.sort(-confidence)
35 | BB = BB[sorted_ind, :]
36 | image_ids = [image_ids[x] for x in sorted_ind]
37 | final_dict[k]["image_ids"] = image_ids
38 | final_dict[k]["BB"] = np.array(BB)
39 | return final_dict
40 |
41 |
42 | def parse_gt_annotation(gt_file):
43 | with open(gt_file, 'r') as f:
44 | lines = f.readlines()
45 | info = [x.strip().split() for x in lines]
46 | gt = {}
47 | for item in info:
48 | img_id = item[0]
49 | obj_struct = {}
50 | obj_struct['class'] = item[1]
51 | obj_struct['bbox'] = [
52 | int(item[2]),
53 | int(item[3]),
54 | int(item[4]),
55 | int(item[5])
56 | ]
57 | if img_id not in gt:
58 | gt[img_id] = list()
59 | gt[img_id].append(obj_struct)
60 | return gt
61 |
62 |
63 | def get_class_recs(recs, classname):
64 | npos = 0
65 | class_recs = {}
66 | for key in recs.keys():
67 | R = [obj for obj in recs[key] if obj['class'] == classname]
68 | bbox = np.array([x['bbox'] for x in R])
69 | det = [False] * len(R)
70 | npos += len(R)
71 | class_recs[key] = {'bbox': bbox, 'det': det}
72 | return class_recs, npos
73 |
74 |
75 | def compute_ap(rec, prec):
76 | mrec = np.concatenate(([0.], rec, [1.]))
77 | mpre = np.concatenate(([0.], prec, [0.]))
78 | for i in range(mpre.size - 1, 0, -1):
79 | mpre[i - 1] = np.maximum(mpre[i - 1], mpre[i])
80 | i = np.where(mrec[1:] != mrec[:-1])[0]
81 | ap = np.sum((mrec[i + 1] - mrec[i]) * mpre[i + 1])
82 | return ap
83 |
84 |
85 | def eval(submit_file, gt_file, ovthresh, classname):
86 | recs = parse_gt_annotation(gt_file)
87 | submit_result = parse_submission(submit_file)
88 | # get one class result
89 | class_recs, npos = get_class_recs(recs, classname)
90 | image_ids = submit_result[classname]["image_ids"]
91 | BB = submit_result[classname]["BB"]
92 | nd = len(image_ids)
93 | tp = np.zeros(nd)
94 | fp = np.zeros(nd)
95 | for d in range(nd):
96 | if image_ids[d] not in recs.keys():
97 | raise KeyError(
98 | "Can not find image {} in the groundtruth file, did you submit the result file for the right dataset?"
99 | .format(image_ids[d]))
100 | for d in range(nd):
101 | R = class_recs[image_ids[d]]
102 | bb = BB[d, :].astype(float)
103 | ovmax = -np.inf
104 | BBGT = R['bbox'].astype(float)
105 | if BBGT.size > 0:
106 | ixmin = np.maximum(BBGT[:, 0], bb[0])
107 | iymin = np.maximum(BBGT[:, 1], bb[1])
108 | ixmax = np.minimum(BBGT[:, 2], bb[2])
109 | iymax = np.minimum(BBGT[:, 3], bb[3])
110 | iw = np.maximum(ixmax - ixmin + 1., 0.)
111 | ih = np.maximum(iymax - iymin + 1., 0.)
112 | inters = iw * ih
113 | uni = ((bb[2] - bb[0] + 1.) * (bb[3] - bb[1] + 1.) +
114 | (BBGT[:, 2] - BBGT[:, 0] + 1.) *
115 | (BBGT[:, 3] - BBGT[:, 1] + 1.) - inters)
116 | overlaps = inters / uni
117 | ovmax = np.max(overlaps)
118 | jmax = np.argmax(overlaps)
119 | if ovmax > ovthresh:
120 | if not R['det'][jmax]:
121 | tp[d] = 1.
122 | R['det'][jmax] = 1
123 | else:
124 | fp[d] = 1.
125 | else:
126 | fp[d] = 1.
127 | fp = np.cumsum(fp)
128 | tp = np.cumsum(tp)
129 | rec = tp / float(npos)
130 | prec = tp / np.maximum(tp + fp, np.finfo(np.float64).eps)
131 | ap = compute_ap(rec, prec)
132 | return ap
133 |
134 |
135 | def result_eval(submit_file, gt, class_list):
136 | ove_aap = []
137 | for ove in np.arange(0.5, 1.0, 0.05):
138 | cls_aap = []
139 | for cls in class_list:
140 | ap = eval(submit_file, gt, ove, cls)
141 | cls_aap.append(ap)
142 | cls_mAP = np.average(cls_aap)
143 | print("thresh", round(ove, 3), "map", round(cls_mAP * 100, 3))
144 | ove_aap.append(cls_mAP)
145 | mAP = np.average(ove_aap) * 100
146 | return round(mAP, 3)
147 |
148 |
149 | if __name__ == '__main__':
150 | '''
151 | submit_file: image_id, class, score, xmin, ymin, xmax, ymax
152 | gt_file: image_id, class, xmin, ymin, xmax, ymax
153 | '''
154 | class_list = []
155 | for i in range(1, 61):
156 | class_list.append(str(i))
157 | submit_file = "./results/fpn_dcn_result.csv"
158 | gt_file = "./results/val_label.txt"
159 | check_size(submit_file)
160 | mAP = result_eval(submit_file, gt_file, class_list)
161 | out = {'Average AP': str(round(mAP, 3))}
162 | print(out)
--------------------------------------------------------------------------------
/utils/nms/__init__.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/JialeCao001/HSD/8abcf78db5f313266a3bb3f85b9424927fe59a2d/utils/nms/__init__.py
--------------------------------------------------------------------------------
/utils/nms/cpu_nms.pyx:
--------------------------------------------------------------------------------
1 | # --------------------------------------------------------
2 | # Fast R-CNN
3 | # Copyright (c) 2015 Microsoft
4 | # Licensed under The MIT License [see LICENSE for details]
5 | # Written by Ross Girshick
6 | # --------------------------------------------------------
7 |
8 | import numpy as np
9 | cimport numpy as np
10 |
11 | cdef inline np.float32_t max(np.float32_t a, np.float32_t b):
12 | return a if a >= b else b
13 |
14 | cdef inline np.float32_t min(np.float32_t a, np.float32_t b):
15 | return a if a <= b else b
16 |
17 | cdef inline np.float32_t abs(np.float32_t a, np.float32_t b):
18 | return a - b if a >= b else b - a
19 |
20 | def get_iou_weights(np.ndarray[np.float32_t, ndim=1] ious, np.float threshold, float init_weight):
21 |
22 | cdef:
23 | int num = ious.shape[0]
24 | # np.ndarray[np.float32_t, ndim=1] out = np.zeros(num, dtype=np.float)
25 | int idx
26 | float iou
27 | float weight
28 |
29 | for idx, iou in enumerate(ious):
30 | weight = init_weight
31 | if iou > 0.0:
32 | if iou > threshold + 0.1:
33 | weight += 1.0
34 | elif iou < threshold - 0.1:
35 | weight += 1.0
36 | else:
37 | weight += 0.0
38 | ious[idx] = weight
39 | return ious
40 |
41 | def get_mask(np.ndarray[np.float32_t, ndim=1] ious, np.float threshold):
42 | cdef:
43 | int num = ious.shape[0]
44 | int idx = 0
45 | float distance
46 | float iou
47 | np.ndarray[np.int64_t, ndim=1] out = np.zeros((num), dtype=np.int64)
48 | for idx, iou in enumerate(ious):
49 | # if iou >= threshold:
50 | # distance = iou - threshold
51 | # if distance < 0.1:
52 | # out[idx] = 0
53 | # elif distance < 0.2:
54 | # out[idx] = 1
55 | # else:
56 | # out[idx] = 2
57 | # else:
58 | # distance = threshold - iou
59 | # if distance < 0.1:
60 | # out[idx] = 2
61 | # elif distance < 0.2:
62 | # out[idx] = 1
63 | # else:
64 | # out[idx] = 0
65 | distance = abs(iou, threshold)
66 | if distance < 0.1:
67 | # out[:,2] = 1
68 | out[idx] = 2
69 | elif distance < 0.2:
70 | # out[:,1] = 1
71 | out[idx] = 1
72 | else:
73 | # out[:,0] = 0
74 | out[idx] = 0
75 | return out
76 |
77 | def cpu_nms(np.ndarray[np.float32_t, ndim=2] dets, np.float thresh):
78 | cdef np.ndarray[np.float32_t, ndim=1] x1 = dets[:, 0]
79 | cdef np.ndarray[np.float32_t, ndim=1] y1 = dets[:, 1]
80 | cdef np.ndarray[np.float32_t, ndim=1] x2 = dets[:, 2]
81 | cdef np.ndarray[np.float32_t, ndim=1] y2 = dets[:, 3]
82 | cdef np.ndarray[np.float32_t, ndim=1] scores = dets[:, 4]
83 |
84 | cdef np.ndarray[np.float32_t, ndim=1] areas = (x2 - x1 + 1) * (y2 - y1 + 1)
85 | cdef np.ndarray[np.int_t, ndim=1] order = scores.argsort()[::-1]
86 |
87 | cdef int ndets = dets.shape[0]
88 | cdef np.ndarray[np.int_t, ndim=1] suppressed = \
89 | np.zeros((ndets), dtype=np.int)
90 |
91 | # nominal indices
92 | cdef int _i, _j
93 | # sorted indices
94 | cdef int i, j
95 | # temp variables for box i's (the box currently under consideration)
96 | cdef np.float32_t ix1, iy1, ix2, iy2, iarea
97 | # variables for computing overlap with box j (lower scoring box)
98 | cdef np.float32_t xx1, yy1, xx2, yy2
99 | cdef np.float32_t w, h
100 | cdef np.float32_t inter, ovr
101 |
102 | keep = []
103 | for _i in range(ndets):
104 | i = order[_i]
105 | if suppressed[i] == 1:
106 | continue
107 | keep.append(i)
108 | ix1 = x1[i]
109 | iy1 = y1[i]
110 | ix2 = x2[i]
111 | iy2 = y2[i]
112 | iarea = areas[i]
113 | for _j in range(_i + 1, ndets):
114 | j = order[_j]
115 | if suppressed[j] == 1:
116 | continue
117 | xx1 = max(ix1, x1[j])
118 | yy1 = max(iy1, y1[j])
119 | xx2 = min(ix2, x2[j])
120 | yy2 = min(iy2, y2[j])
121 | w = max(0.0, xx2 - xx1 + 1)
122 | h = max(0.0, yy2 - yy1 + 1)
123 | inter = w * h
124 | ovr = inter / (iarea + areas[j] - inter)
125 | if ovr >= thresh:
126 | suppressed[j] = 1
127 |
128 | return keep
129 |
130 | def cpu_soft_nms(np.ndarray[float, ndim=2] boxes, float sigma=0.5, float Nt=0.3, float threshold=0.001, unsigned int method=0):
131 | cdef unsigned int N = boxes.shape[0]
132 | cdef float iw, ih, box_area
133 | cdef float ua
134 | cdef int pos = 0
135 | cdef float maxscore = 0
136 | cdef int maxpos = 0
137 | cdef float x1,x2,y1,y2,tx1,tx2,ty1,ty2,ts,area,weight,ov
138 |
139 | for i in range(N):
140 | maxscore = boxes[i, 4]
141 | maxpos = i
142 |
143 | tx1 = boxes[i,0]
144 | ty1 = boxes[i,1]
145 | tx2 = boxes[i,2]
146 | ty2 = boxes[i,3]
147 | ts = boxes[i,4]
148 |
149 | pos = i + 1
150 | # get max box
151 | while pos < N:
152 | if maxscore < boxes[pos, 4]:
153 | maxscore = boxes[pos, 4]
154 | maxpos = pos
155 | pos = pos + 1
156 |
157 | # add max box as a detection
158 | boxes[i,0] = boxes[maxpos,0]
159 | boxes[i,1] = boxes[maxpos,1]
160 | boxes[i,2] = boxes[maxpos,2]
161 | boxes[i,3] = boxes[maxpos,3]
162 | boxes[i,4] = boxes[maxpos,4]
163 |
164 | # swap ith box with position of max box
165 | boxes[maxpos,0] = tx1
166 | boxes[maxpos,1] = ty1
167 | boxes[maxpos,2] = tx2
168 | boxes[maxpos,3] = ty2
169 | boxes[maxpos,4] = ts
170 |
171 | tx1 = boxes[i,0]
172 | ty1 = boxes[i,1]
173 | tx2 = boxes[i,2]
174 | ty2 = boxes[i,3]
175 | ts = boxes[i,4]
176 |
177 | pos = i + 1
178 | # NMS iterations, note that N changes if detection boxes fall below threshold
179 | while pos < N:
180 | x1 = boxes[pos, 0]
181 | y1 = boxes[pos, 1]
182 | x2 = boxes[pos, 2]
183 | y2 = boxes[pos, 3]
184 | s = boxes[pos, 4]
185 |
186 | area = (x2 - x1 + 1) * (y2 - y1 + 1)
187 | iw = (min(tx2, x2) - max(tx1, x1) + 1)
188 | if iw > 0:
189 | ih = (min(ty2, y2) - max(ty1, y1) + 1)
190 | if ih > 0:
191 | ua = float((tx2 - tx1 + 1) * (ty2 - ty1 + 1) + area - iw * ih)
192 | ov = iw * ih / ua #iou between max box and detection box
193 |
194 | if method == 1: # linear
195 | if ov > Nt:
196 | weight = 1 - ov
197 | else:
198 | weight = 1
199 | elif method == 2: # gaussian
200 | weight = np.exp(-(ov * ov)/sigma)
201 | else: # original NMS
202 | if ov > Nt:
203 | weight = 0
204 | else:
205 | weight = 1
206 |
207 | boxes[pos, 4] = weight*boxes[pos, 4]
208 |
209 | # if box score falls below threshold, discard the box by swapping with last box
210 | # update N
211 | if boxes[pos, 4] < threshold:
212 | boxes[pos,0] = boxes[N-1, 0]
213 | boxes[pos,1] = boxes[N-1, 1]
214 | boxes[pos,2] = boxes[N-1, 2]
215 | boxes[pos,3] = boxes[N-1, 3]
216 | boxes[pos,4] = boxes[N-1, 4]
217 | N = N - 1
218 | pos = pos - 1
219 |
220 | pos = pos + 1
221 |
222 | keep = [i for i in range(N)]
223 | return keep
224 |
--------------------------------------------------------------------------------
/utils/nms/gpu_nms.hpp:
--------------------------------------------------------------------------------
1 | void _nms(int* keep_out, int* num_out, const float* boxes_host, int boxes_num,
2 | int boxes_dim, float nms_overlap_thresh, int device_id);
3 |
--------------------------------------------------------------------------------
/utils/nms/gpu_nms.pyx:
--------------------------------------------------------------------------------
1 | # --------------------------------------------------------
2 | # Faster R-CNN
3 | # Copyright (c) 2015 Microsoft
4 | # Licensed under The MIT License [see LICENSE for details]
5 | # Written by Ross Girshick
6 | # --------------------------------------------------------
7 |
8 | import numpy as np
9 | cimport numpy as np
10 |
11 | assert sizeof(int) == sizeof(np.int32_t)
12 |
13 | cdef extern from "gpu_nms.hpp":
14 | void _nms(np.int32_t*, int*, np.float32_t*, int, int, float, int)
15 |
16 | def gpu_nms(np.ndarray[np.float32_t, ndim=2] dets, np.float thresh,
17 | np.int32_t device_id=0):
18 | cdef int boxes_num = dets.shape[0]
19 | cdef int boxes_dim = dets.shape[1]
20 | cdef int num_out
21 | cdef np.ndarray[np.int32_t, ndim=1] \
22 | keep = np.zeros(boxes_num, dtype=np.int32)
23 | cdef np.ndarray[np.float32_t, ndim=1] \
24 | scores = dets[:, 4]
25 | cdef np.ndarray[np.int_t, ndim=1] \
26 | order = scores.argsort()[::-1]
27 | cdef np.ndarray[np.float32_t, ndim=2] \
28 | sorted_dets = dets[order, :]
29 | _nms(&keep[0], &num_out, &sorted_dets[0, 0], boxes_num, boxes_dim, thresh, device_id)
30 | keep = keep[:num_out]
31 | return list(order[keep])
32 |
--------------------------------------------------------------------------------
/utils/nms/nms_kernel.cu:
--------------------------------------------------------------------------------
1 | // ------------------------------------------------------------------
2 | // Faster R-CNN
3 | // Copyright (c) 2015 Microsoft
4 | // Licensed under The MIT License [see fast-rcnn/LICENSE for details]
5 | // Written by Shaoqing Ren
6 | // ------------------------------------------------------------------
7 |
8 | #include "gpu_nms.hpp"
9 | #include
10 | #include
11 |
12 | #define CUDA_CHECK(condition) \
13 | /* Code block avoids redefinition of cudaError_t error */ \
14 | do { \
15 | cudaError_t error = condition; \
16 | if (error != cudaSuccess) { \
17 | std::cout << cudaGetErrorString(error) << std::endl; \
18 | } \
19 | } while (0)
20 |
21 | #define DIVUP(m,n) ((m) / (n) + ((m) % (n) > 0))
22 | int const threadsPerBlock = sizeof(unsigned long long) * 8;
23 |
24 | __device__ inline float devIoU(float const * const a, float const * const b) {
25 | float left = max(a[0], b[0]), right = min(a[2], b[2]);
26 | float top = max(a[1], b[1]), bottom = min(a[3], b[3]);
27 | float width = max(right - left + 1, 0.f), height = max(bottom - top + 1, 0.f);
28 | float interS = width * height;
29 | float Sa = (a[2] - a[0] + 1) * (a[3] - a[1] + 1);
30 | float Sb = (b[2] - b[0] + 1) * (b[3] - b[1] + 1);
31 | return interS / (Sa + Sb - interS);
32 | }
33 |
34 | __global__ void nms_kernel(const int n_boxes, const float nms_overlap_thresh,
35 | const float *dev_boxes, unsigned long long *dev_mask) {
36 | const int row_start = blockIdx.y;
37 | const int col_start = blockIdx.x;
38 |
39 | // if (row_start > col_start) return;
40 |
41 | const int row_size =
42 | min(n_boxes - row_start * threadsPerBlock, threadsPerBlock);
43 | const int col_size =
44 | min(n_boxes - col_start * threadsPerBlock, threadsPerBlock);
45 |
46 | __shared__ float block_boxes[threadsPerBlock * 5];
47 | if (threadIdx.x < col_size) {
48 | block_boxes[threadIdx.x * 5 + 0] =
49 | dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 5 + 0];
50 | block_boxes[threadIdx.x * 5 + 1] =
51 | dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 5 + 1];
52 | block_boxes[threadIdx.x * 5 + 2] =
53 | dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 5 + 2];
54 | block_boxes[threadIdx.x * 5 + 3] =
55 | dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 5 + 3];
56 | block_boxes[threadIdx.x * 5 + 4] =
57 | dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 5 + 4];
58 | }
59 | __syncthreads();
60 |
61 | if (threadIdx.x < row_size) {
62 | const int cur_box_idx = threadsPerBlock * row_start + threadIdx.x;
63 | const float *cur_box = dev_boxes + cur_box_idx * 5;
64 | int i = 0;
65 | unsigned long long t = 0;
66 | int start = 0;
67 | if (row_start == col_start) {
68 | start = threadIdx.x + 1;
69 | }
70 | for (i = start; i < col_size; i++) {
71 | if (devIoU(cur_box, block_boxes + i * 5) > nms_overlap_thresh) {
72 | t |= 1ULL << i;
73 | }
74 | }
75 | const int col_blocks = DIVUP(n_boxes, threadsPerBlock);
76 | dev_mask[cur_box_idx * col_blocks + col_start] = t;
77 | }
78 | }
79 |
80 | void _set_device(int device_id) {
81 | int current_device;
82 | CUDA_CHECK(cudaGetDevice(¤t_device));
83 | if (current_device == device_id) {
84 | return;
85 | }
86 | // The call to cudaSetDevice must come before any calls to Get, which
87 | // may perform initialization using the GPU.
88 | CUDA_CHECK(cudaSetDevice(device_id));
89 | }
90 |
91 | void _nms(int* keep_out, int* num_out, const float* boxes_host, int boxes_num,
92 | int boxes_dim, float nms_overlap_thresh, int device_id) {
93 | _set_device(device_id);
94 |
95 | float* boxes_dev = NULL;
96 | unsigned long long* mask_dev = NULL;
97 |
98 | const int col_blocks = DIVUP(boxes_num, threadsPerBlock);
99 |
100 | CUDA_CHECK(cudaMalloc(&boxes_dev,
101 | boxes_num * boxes_dim * sizeof(float)));
102 | CUDA_CHECK(cudaMemcpy(boxes_dev,
103 | boxes_host,
104 | boxes_num * boxes_dim * sizeof(float),
105 | cudaMemcpyHostToDevice));
106 |
107 | CUDA_CHECK(cudaMalloc(&mask_dev,
108 | boxes_num * col_blocks * sizeof(unsigned long long)));
109 |
110 | dim3 blocks(DIVUP(boxes_num, threadsPerBlock),
111 | DIVUP(boxes_num, threadsPerBlock));
112 | dim3 threads(threadsPerBlock);
113 | nms_kernel<<>>(boxes_num,
114 | nms_overlap_thresh,
115 | boxes_dev,
116 | mask_dev);
117 |
118 | std::vector mask_host(boxes_num * col_blocks);
119 | CUDA_CHECK(cudaMemcpy(&mask_host[0],
120 | mask_dev,
121 | sizeof(unsigned long long) * boxes_num * col_blocks,
122 | cudaMemcpyDeviceToHost));
123 |
124 | std::vector remv(col_blocks);
125 | memset(&remv[0], 0, sizeof(unsigned long long) * col_blocks);
126 |
127 | int num_to_keep = 0;
128 | for (int i = 0; i < boxes_num; i++) {
129 | int nblock = i / threadsPerBlock;
130 | int inblock = i % threadsPerBlock;
131 |
132 | if (!(remv[nblock] & (1ULL << inblock))) {
133 | keep_out[num_to_keep++] = i;
134 | unsigned long long *p = &mask_host[0] + i * col_blocks;
135 | for (int j = nblock; j < col_blocks; j++) {
136 | remv[j] |= p[j];
137 | }
138 | }
139 | }
140 | *num_out = num_to_keep;
141 |
142 | CUDA_CHECK(cudaFree(boxes_dev));
143 | CUDA_CHECK(cudaFree(mask_dev));
144 | }
145 |
--------------------------------------------------------------------------------
/utils/nms/py_cpu_nms.py:
--------------------------------------------------------------------------------
1 | # --------------------------------------------------------
2 | # Fast R-CNN
3 | # Copyright (c) 2015 Microsoft
4 | # Licensed under The MIT License [see LICENSE for details]
5 | # Written by Ross Girshick
6 | # --------------------------------------------------------
7 |
8 | import numpy as np
9 |
10 | def py_cpu_nms(dets, thresh):
11 | """Pure Python NMS baseline."""
12 | x1 = dets[:, 0]
13 | y1 = dets[:, 1]
14 | x2 = dets[:, 2]
15 | y2 = dets[:, 3]
16 | scores = dets[:, 4]
17 |
18 | areas = (x2 - x1 + 1) * (y2 - y1 + 1)
19 | order = scores.argsort()[::-1]
20 |
21 | keep = []
22 | while order.size > 0:
23 | i = order[0]
24 | keep.append(i)
25 | xx1 = np.maximum(x1[i], x1[order[1:]])
26 | yy1 = np.maximum(y1[i], y1[order[1:]])
27 | xx2 = np.minimum(x2[i], x2[order[1:]])
28 | yy2 = np.minimum(y2[i], y2[order[1:]])
29 |
30 | w = np.maximum(0.0, xx2 - xx1 + 1)
31 | h = np.maximum(0.0, yy2 - yy1 + 1)
32 | inter = w * h
33 | ovr = inter / (areas[i] + areas[order[1:]] - inter)
34 |
35 | inds = np.where(ovr <= thresh)[0]
36 | order = order[inds + 1]
37 |
38 | return keep
39 |
--------------------------------------------------------------------------------
/utils/nms_wrapper.py:
--------------------------------------------------------------------------------
1 | # --------------------------------------------------------
2 | # Fast R-CNN
3 | # Copyright (c) 2015 Microsoft
4 | # Licensed under The MIT License [see LICENSE for details]
5 | # Written by Ross Girshick
6 | # --------------------------------------------------------
7 |
8 | from .nms.cpu_nms import cpu_nms, cpu_soft_nms
9 | from .nms.gpu_nms import gpu_nms
10 |
11 | # def nms(dets, thresh, force_cpu=False):
12 | # """Dispatch to either CPU or GPU NMS implementations."""
13 |
14 | # if dets.shape[0] == 0:
15 | # return []
16 | # if cfg.USE_GPU_NMS and not force_cpu:
17 | # return gpu_nms(dets, thresh, device_id=cfg.GPU_ID)
18 | # else:
19 | # return cpu_nms(dets, thresh)
20 |
21 |
22 | def nms(dets, thresh, force_cpu=False):
23 | """Dispatch to either CPU or GPU NMS implementations."""
24 |
25 | if dets.shape[0] == 0:
26 | return []
27 | if force_cpu:
28 | #return cpu_soft_nms(dets, thresh, method = 0)
29 | return cpu_nms(dets, thresh)
30 | return gpu_nms(dets, thresh)
31 |
32 |
33 | def soft_nms(dets, Nt=0.3, sigma=0.5, thresh=0.001, method=1):
34 | """Dispatch to either CPU or GPU NMS implementations."""
35 |
36 | if dets.shape[0] == 0:
37 | return []
38 | return cpu_soft_nms(dets, sigma, Nt, thresh, method)
--------------------------------------------------------------------------------
/utils/timer.py:
--------------------------------------------------------------------------------
1 | # --------------------------------------------------------
2 | # Fast R-CNN
3 | # Copyright (c) 2015 Microsoft
4 | # Licensed under The MIT License [see LICENSE for details]
5 | # Written by Ross Girshick
6 | # --------------------------------------------------------
7 |
8 | import time
9 |
10 |
11 | class Timer(object):
12 | """A simple timer."""
13 |
14 | def __init__(self):
15 | self.total_time = 0.
16 | self.calls = 0
17 | self.start_time = 0.
18 | self.diff = 0.
19 | self.average_time = 0.
20 |
21 | def tic(self):
22 | # using time.time instead of time.clock because time time.clock
23 | # does not normalize for multithreading
24 | self.start_time = time.time()
25 |
26 | def toc(self, average=True):
27 | self.diff = time.time() - self.start_time
28 | self.total_time += self.diff
29 | self.calls += 1
30 | self.average_time = self.total_time / self.calls
31 | if average:
32 | return self.average_time
33 | else:
34 | return self.diff
35 |
36 | def clear(self):
37 | self.total_time = 0.
38 | self.calls = 0
39 | self.start_time = 0.
40 | self.diff = 0.
41 | self.average_time = 0.
42 |
--------------------------------------------------------------------------------
|