├── .idea
├── Repulsion_Loss.iml
├── misc.xml
├── modules.xml
├── vcs.xml
└── workspace.xml
├── README.md
├── anchors.py
├── csv_eval.py
├── dataloader.py
├── img
├── 1.png
└── 2.png
├── lib
├── __init__.py
├── __pycache__
│ └── __init__.cpython-36.pyc
├── build.sh
└── nms
│ ├── __init__.py
│ ├── __pycache__
│ ├── __init__.cpython-36.pyc
│ └── pth_nms.cpython-36.pyc
│ ├── _ext
│ ├── __init__.py
│ ├── __pycache__
│ │ └── __init__.cpython-36.pyc
│ └── nms
│ │ ├── __init__.py
│ │ ├── __pycache__
│ │ └── __init__.cpython-36.pyc
│ │ └── _nms.so
│ ├── build.py
│ ├── pth_nms.py
│ └── src
│ ├── cuda
│ ├── nms_kernel.cu
│ ├── nms_kernel.cu.o
│ └── nms_kernel.h
│ ├── nms.c
│ ├── nms.h
│ ├── nms_cuda.c
│ └── nms_cuda.h
├── losses.py
├── model.py
├── train.py
└── utils.py
/.idea/Repulsion_Loss.iml:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
5 |
6 |
7 |
8 |
9 |
10 |
11 |
--------------------------------------------------------------------------------
/.idea/misc.xml:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
--------------------------------------------------------------------------------
/.idea/modules.xml:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
5 |
6 |
7 |
8 |
--------------------------------------------------------------------------------
/.idea/vcs.xml:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
5 |
6 |
--------------------------------------------------------------------------------
/.idea/workspace.xml:
--------------------------------------------------------------------------------
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 |
51 |
52 |
53 |
54 |
55 | true
56 | DEFINITION_ORDER
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 |
2 | #include
3 |
4 | int cpu_nms(THLongTensor * keep_out, THLongTensor * num_out, THFloatTensor * boxes, THLongTensor * order, THFloatTensor * areas, float nms_overlap_thresh) {
5 | // boxes has to be sorted
6 | THArgCheck(THLongTensor_isContiguous(keep_out), 0, "keep_out must be contiguous");
7 | THArgCheck(THLongTensor_isContiguous(boxes), 2, "boxes must be contiguous");
8 | THArgCheck(THLongTensor_isContiguous(order), 3, "order must be contiguous");
9 | THArgCheck(THLongTensor_isContiguous(areas), 4, "areas must be contiguous");
10 | // Number of ROIs
11 | long boxes_num = THFloatTensor_size(boxes, 0);
12 | long boxes_dim = THFloatTensor_size(boxes, 1);
13 |
14 | long * keep_out_flat = THLongTensor_data(keep_out);
15 | float * boxes_flat = THFloatTensor_data(boxes);
16 | long * order_flat = THLongTensor_data(order);
17 | float * areas_flat = THFloatTensor_data(areas);
18 |
19 | THByteTensor* suppressed = THByteTensor_newWithSize1d(boxes_num);
20 | THByteTensor_fill(suppressed, 0);
21 | unsigned char * suppressed_flat = THByteTensor_data(suppressed);
22 |
23 | // nominal indices
24 | int i, j;
25 | // sorted indices
26 | int _i, _j;
27 | // temp variables for box i's (the box currently under consideration)
28 | float ix1, iy1, ix2, iy2, iarea;
29 | // variables for computing overlap with box j (lower scoring box)
30 | float xx1, yy1, xx2, yy2;
31 | float w, h;
32 | float inter, ovr;
33 |
34 | long num_to_keep = 0;
35 | for (_i=0; _i < boxes_num; ++_i) {
36 | i = order_flat[_i];
37 | if (suppressed_flat[i] == 1) {
38 | continue;
39 | }
40 | keep_out_flat[num_to_keep++] = i;
41 | ix1 = boxes_flat[i * boxes_dim];
42 | iy1 = boxes_flat[i * boxes_dim + 1];
43 | ix2 = boxes_flat[i * boxes_dim + 2];
44 | iy2 = boxes_flat[i * boxes_dim + 3];
45 | iarea = areas_flat[i];
46 | for (_j = _i + 1; _j < boxes_num; ++_j) {
47 | j = order_flat[_j];
48 | if (suppressed_flat[j] == 1) {
49 | continue;
50 | }
51 | xx1 = fmaxf(ix1, boxes_flat[j * boxes_dim]);
52 | yy1 = fmaxf(iy1, boxes_flat[j * boxes_dim + 1]);
53 | xx2 = fminf(ix2, boxes_flat[j * boxes_dim + 2]);
54 | yy2 = fminf(iy2, boxes_flat[j * boxes_dim + 3]);
55 | w = fmaxf(0.0, xx2 - xx1 + 1);
56 | h = fmaxf(0.0, yy2 - yy1 + 1);
57 | inter = w * h;
58 | ovr = inter / (iarea + areas_flat[j] - inter);
59 | if (ovr >= nms_overlap_thresh) {
60 | suppressed_flat[j] = 1;
61 | }
62 | }
63 | }
64 |
65 | long *num_out_flat = THLongTensor_data(num_out);
66 | *num_out_flat = num_to_keep;
67 | THByteTensor_free(suppressed);
68 | return 1;
69 | }
--------------------------------------------------------------------------------
/lib/nms/src/nms.h:
--------------------------------------------------------------------------------
1 | int cpu_nms(THLongTensor * keep_out, THLongTensor * num_out, THFloatTensor * boxes, THLongTensor * order, THFloatTensor * areas, float nms_overlap_thresh);
--------------------------------------------------------------------------------
/lib/nms/src/nms_cuda.c:
--------------------------------------------------------------------------------
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 | #include
8 | #include |
9 | #include
10 | #include
11 |
12 | #include "cuda/nms_kernel.h"
13 |
14 |
15 | extern THCState *state;
16 |
17 | int gpu_nms(THLongTensor * keep, THLongTensor* num_out, THCudaTensor * boxes, float nms_overlap_thresh) {
18 | // boxes has to be sorted
19 | THArgCheck(THLongTensor_isContiguous(keep), 0, "boxes must be contiguous");
20 | THArgCheck(THCudaTensor_isContiguous(state, boxes), 2, "boxes must be contiguous");
21 | // Number of ROIs
22 | int boxes_num = THCudaTensor_size(state, boxes, 0);
23 | int boxes_dim = THCudaTensor_size(state, boxes, 1);
24 |
25 | float* boxes_flat = THCudaTensor_data(state, boxes);
26 |
27 | const int col_blocks = DIVUP(boxes_num, threadsPerBlock);
28 | THCudaLongTensor * mask = THCudaLongTensor_newWithSize2d(state, boxes_num, col_blocks);
29 | unsigned long long* mask_flat = THCudaLongTensor_data(state, mask);
30 |
31 | _nms(boxes_num, boxes_flat, mask_flat, nms_overlap_thresh);
32 |
33 | THLongTensor * mask_cpu = THLongTensor_newWithSize2d(boxes_num, col_blocks);
34 | THLongTensor_copyCuda(state, mask_cpu, mask);
35 | THCudaLongTensor_free(state, mask);
36 |
37 | unsigned long long * mask_cpu_flat = THLongTensor_data(mask_cpu);
38 |
39 | THLongTensor * remv_cpu = THLongTensor_newWithSize1d(col_blocks);
40 | unsigned long long* remv_cpu_flat = THLongTensor_data(remv_cpu);
41 | THLongTensor_fill(remv_cpu, 0);
42 |
43 | long * keep_flat = THLongTensor_data(keep);
44 | long num_to_keep = 0;
45 |
46 | int i, j;
47 | for (i = 0; i < boxes_num; i++) {
48 | int nblock = i / threadsPerBlock;
49 | int inblock = i % threadsPerBlock;
50 |
51 | if (!(remv_cpu_flat[nblock] & (1ULL << inblock))) {
52 | keep_flat[num_to_keep++] = i;
53 | unsigned long long *p = &mask_cpu_flat[0] + i * col_blocks;
54 | for (j = nblock; j < col_blocks; j++) {
55 | remv_cpu_flat[j] |= p[j];
56 | }
57 | }
58 | }
59 |
60 | long * num_out_flat = THLongTensor_data(num_out);
61 | * num_out_flat = num_to_keep;
62 |
63 | THLongTensor_free(mask_cpu);
64 | THLongTensor_free(remv_cpu);
65 |
66 | return 1;
67 | }
68 |
--------------------------------------------------------------------------------
/lib/nms/src/nms_cuda.h:
--------------------------------------------------------------------------------
1 | int gpu_nms(THLongTensor * keep_out, THLongTensor* num_out, THCudaTensor * boxes, float nms_overlap_thresh);
--------------------------------------------------------------------------------
/losses.py:
--------------------------------------------------------------------------------
1 | import numpy as np
2 | import torch
3 | import torch.nn as nn
4 | import random
5 |
6 |
7 | def calc_iou(a, b):
8 | area = (b[:, 2] - b[:, 0]) * (b[:, 3] - b[:, 1])
9 |
10 | iw = torch.min(torch.unsqueeze(a[:, 2], dim=1), b[:, 2]) - torch.max(torch.unsqueeze(a[:, 0], 1), b[:, 0])
11 | ih = torch.min(torch.unsqueeze(a[:, 3], dim=1), b[:, 3]) - torch.max(torch.unsqueeze(a[:, 1], 1), b[:, 1])
12 |
13 | iw = torch.clamp(iw, min=0)
14 | ih = torch.clamp(ih, min=0)
15 |
16 | ua = torch.unsqueeze((a[:, 2] - a[:, 0]) * (a[:, 3] - a[:, 1]), dim=1) + area - iw * ih
17 |
18 | ua = torch.clamp(ua, min=1e-8)
19 |
20 | intersection = iw * ih
21 |
22 | IoU = intersection / ua
23 |
24 | return IoU
25 |
26 |
27 | def IoG(box_a, box_b):
28 |
29 | inter_xmin = torch.max(box_a[:, 0], box_b[:, 0])
30 | inter_ymin = torch.max(box_a[:, 1], box_b[:, 1])
31 | inter_xmax = torch.min(box_a[:, 2], box_b[:, 2])
32 | inter_ymax = torch.min(box_a[:, 3], box_b[:, 3])
33 | Iw = torch.clamp(inter_xmax - inter_xmin, min=0)
34 | Ih = torch.clamp(inter_ymax - inter_ymin, min=0)
35 | I = Iw * Ih
36 | G = (box_a[:, 2] - box_a[:, 0]) * (box_a[:, 3] - box_a[:, 1])
37 | return I / G
38 |
39 |
40 | def smooth_ln(x, smooth):
41 | return torch.where(
42 | torch.le(x, smooth),
43 | -torch.log(1 - x),
44 | ((x - smooth) / (1 - smooth)) - np.log(1 - smooth)
45 | )
46 |
47 |
48 | class FocalLoss(nn.Module):
49 | # def __init__(self):
50 |
51 | def forward(self, classifications, regressions, anchors, annotations, ignores):
52 | alpha = 0.25
53 | gamma = 2.0
54 | batch_size = classifications.shape[0]
55 | classification_losses = []
56 | regression_losses = []
57 | RepGT_losses = []
58 | RepBox_losses = []
59 |
60 | anchor = anchors[0, :, :]
61 |
62 | anchor_widths = anchor[:, 2] - anchor[:, 0]
63 | anchor_heights = anchor[:, 3] - anchor[:, 1]
64 | anchor_ctr_x = anchor[:, 0] + 0.5 * anchor_widths
65 | anchor_ctr_y = anchor[:, 1] + 0.5 * anchor_heights
66 |
67 | for j in range(batch_size):
68 |
69 | classification = classifications[j, :, :]
70 | regression = regressions[j, :, :]
71 |
72 | bbox_annotation = annotations[j, :, :]
73 | bbox_annotation = bbox_annotation[bbox_annotation[:, 4] != -1]
74 |
75 | if bbox_annotation.shape[0] == 0:
76 | regression_losses.append(torch.tensor(0).float().cuda())
77 | classification_losses.append(torch.tensor(0).float().cuda())
78 | RepGT_losses.append(torch.tensor(0).float().cuda())
79 |
80 | continue
81 |
82 | classification = torch.clamp(classification, 1e-4, 1.0 - 1e-4)
83 |
84 | ignore = ignores[j, :, :]
85 | ignore = ignore[ignore[:, 4] != -1]
86 | if ignore.shape[0] > 0:
87 | iou_igno = calc_iou(anchor, ignore)
88 | iou_igno_max, iou_igno_argmax = torch.max(iou_igno, dim=1)
89 | index_igno = torch.lt(iou_igno_max, 0.5)
90 | anchor_keep = anchor[index_igno, :]
91 | classification = classification[index_igno, :]
92 | regression = regression[index_igno, :]
93 | anchor_widths_keep = anchor_widths[index_igno]
94 | anchor_heights_keep = anchor_heights[index_igno]
95 | anchor_ctr_x_keep = anchor_ctr_x[index_igno]
96 | anchor_ctr_y_keep = anchor_ctr_y[index_igno]
97 | else:
98 | anchor_keep = anchor
99 | anchor_widths_keep = anchor_widths
100 | anchor_heights_keep = anchor_heights
101 | anchor_ctr_x_keep = anchor_ctr_x
102 | anchor_ctr_y_keep = anchor_ctr_y
103 |
104 | IoU = calc_iou(anchor_keep, bbox_annotation[:, :4]) # num_anchors x num_annotations
105 |
106 | IoU_max, IoU_argmax = torch.max(IoU, dim=1) # num_anchors x 1
107 |
108 | # compute the loss for classification
109 | targets = torch.ones(classification.shape) * -1
110 | targets = targets.cuda()
111 |
112 | targets[torch.lt(IoU_max, 0.4), :] = 0
113 |
114 | positive_indices = torch.ge(IoU_max, 0.5)
115 |
116 | num_positive_anchors = positive_indices.sum()
117 |
118 | assigned_annotations = bbox_annotation[IoU_argmax, :] #which gt the anchor matches
119 |
120 | targets[positive_indices, :] = 0
121 | targets[positive_indices, assigned_annotations[positive_indices, 4].long()] = 1
122 |
123 | alpha_factor = torch.ones(targets.shape).cuda() * alpha
124 |
125 | alpha_factor = torch.where(torch.eq(targets, 1.), alpha_factor, 1. - alpha_factor)
126 | focal_weight = torch.where(torch.eq(targets, 1.), 1. - classification, classification)
127 | focal_weight = alpha_factor * torch.pow(focal_weight, gamma)
128 |
129 | bce = -(targets * torch.log(classification) + (1.0 - targets) * torch.log(1.0 - classification))
130 |
131 | # cls_loss = focal_weight * torch.pow(bce, gamma)
132 | cls_loss = focal_weight * bce
133 |
134 | cls_loss = torch.where(torch.ne(targets, -1.0), cls_loss, torch.zeros(cls_loss.shape).cuda())
135 |
136 | classification_losses.append(cls_loss.sum() / torch.clamp(num_positive_anchors.float(), min=1.0))
137 |
138 | # compute the loss for regression
139 |
140 | if positive_indices.sum() > 0:
141 | assigned_annotations = assigned_annotations[positive_indices, :] #num_pos_anchors * 5
142 |
143 | anchor_widths_pi = anchor_widths_keep[positive_indices]
144 | anchor_heights_pi = anchor_heights_keep[positive_indices]
145 | anchor_ctr_x_pi = anchor_ctr_x_keep[positive_indices]
146 | anchor_ctr_y_pi = anchor_ctr_y_keep[positive_indices]
147 |
148 | gt_widths = assigned_annotations[:, 2] - assigned_annotations[:, 0]
149 | gt_heights = assigned_annotations[:, 3] - assigned_annotations[:, 1]
150 | gt_ctr_x = assigned_annotations[:, 0] + 0.5 * gt_widths
151 | gt_ctr_y = assigned_annotations[:, 1] + 0.5 * gt_heights
152 |
153 | # clip widths to 1
154 | gt_widths = torch.clamp(gt_widths, min=1)
155 | gt_heights = torch.clamp(gt_heights, min=1)
156 |
157 | targets_dx = (gt_ctr_x - anchor_ctr_x_pi) / anchor_widths_pi
158 | targets_dy = (gt_ctr_y - anchor_ctr_y_pi) / anchor_heights_pi
159 | targets_dw = torch.log(gt_widths / anchor_widths_pi)
160 | targets_dh = torch.log(gt_heights / anchor_heights_pi)
161 |
162 | targets = torch.stack((targets_dx, targets_dy, targets_dw, targets_dh))
163 | targets = targets.t()
164 |
165 | targets = targets / torch.Tensor([[0.1, 0.1, 0.2, 0.2]]).cuda()
166 |
167 | negative_indices = 1 - positive_indices
168 |
169 | regression_diff = torch.abs(targets - regression[positive_indices, :])
170 |
171 | regression_loss = torch.where(
172 | torch.le(regression_diff, 1.0 / 9.0),
173 | 0.5 * 9.0 * torch.pow(regression_diff, 2),
174 | regression_diff - 0.5 / 9.0
175 | )
176 | regression_loss = regression_loss.mean()
177 | regression_losses.append(regression_loss)
178 |
179 |
180 | # predict regression to boxes that are positive
181 | if bbox_annotation.shape[0] == 1:
182 | RepGT_losses.append(torch.tensor(0).float().cuda())
183 | # RepBox_losses.append(torch.tensor(0).float().cuda())
184 | else:
185 | regression_pos = regression[positive_indices, :]
186 | regression_pos_dx = regression_pos[:, 0]
187 | regression_pos_dy = regression_pos[:, 1]
188 | regression_pos_dw = regression_pos[:, 2]
189 | regression_pos_dh = regression_pos[:, 3]
190 | predict_w = torch.exp(regression_pos_dw) * anchor_widths_pi
191 | predict_h = torch.exp(regression_pos_dh) * anchor_heights_pi
192 | predict_x = regression_pos_dx * anchor_widths_pi + anchor_ctr_x_pi
193 | predict_y = regression_pos_dy * anchor_heights_pi + anchor_ctr_y_pi
194 | predict_xmin = predict_x - 0.5 * predict_w
195 | predict_ymin = predict_y - 0.5 * predict_h
196 | predict_xmax = predict_x + 0.5 * predict_w
197 | predict_ymax = predict_y + 0.5 * predict_h
198 | predict_boxes = torch.stack((predict_xmin, predict_ymin, predict_xmax, predict_ymax)).t()
199 |
200 | # add RepGT_losses
201 | IoU_pos = IoU[positive_indices, :]
202 | IoU_max_keep, IoU_argmax_keep = torch.max(IoU_pos, dim=1, keepdim=True) # num_anchors x 1
203 | for idx in range(IoU_argmax_keep.shape[0]):
204 | IoU_pos[idx, IoU_argmax_keep[idx]] = -1
205 | IoU_sec, IoU_argsec = torch.max(IoU_pos, dim=1)
206 |
207 | assigned_annotations_sec = bbox_annotation[IoU_argsec, :] # which gt the anchor iou second num_anchors * 5
208 |
209 | IoG_to_minimize = IoG(assigned_annotations_sec, predict_boxes)
210 | RepGT_loss = smooth_ln(IoG_to_minimize, 0.5)
211 | RepGT_loss = RepGT_loss.mean()
212 | RepGT_losses.append(RepGT_loss)
213 |
214 | # add PepBox losses
215 | IoU_argmax_pos = IoU_argmax[positive_indices].float()
216 | IoU_argmax_pos = IoU_argmax_pos.unsqueeze(0).t()
217 | predict_boxes = torch.cat([predict_boxes, IoU_argmax_pos], dim=1)
218 | predict_boxes_np = predict_boxes.detach().cpu().numpy()
219 | num_gt = bbox_annotation.shape[0]
220 | predict_boxes_sampled = []
221 | for id in range(num_gt):
222 | index = np.where(predict_boxes_np[:, 4]==id)[0]
223 | if index.shape[0]:
224 | idx = random.choice(range(index.shape[0]))
225 | predict_boxes_sampled.append(predict_boxes[index[idx], :4])
226 | predict_boxes_sampled = torch.stack(predict_boxes_sampled)
227 | iou_repbox = calc_iou(predict_boxes_sampled, predict_boxes_sampled)
228 | mask = torch.lt(iou_repbox, 1.).float()
229 | iou_repbox = iou_repbox * mask
230 | RepBox_loss = smooth_ln(iou_repbox, 0.5)
231 | RepBox_loss = RepBox_loss.sum() / torch.clamp(torch.sum(torch.gt(iou_repbox, 0)).float(), min=1.0)
232 | RepBox_losses.append(RepBox_loss)
233 |
234 | else:
235 | regression_losses.append(torch.tensor(0).float().cuda())
236 | RepGT_losses.append(torch.tensor(0).float().cuda())
237 | RepBox_losses.append(torch.tensor(0).float().cuda())
238 |
239 | return torch.stack(classification_losses).mean(dim=0, keepdim=True), \
240 | torch.stack(regression_losses).mean(dim=0, keepdim=True), \
241 | torch.stack(RepGT_losses).mean(dim=0, keepdim=True), \
242 | torch.stack(RepBox_losses).mean(dim=0, keepdim=True)
243 |
244 |
245 |
--------------------------------------------------------------------------------
/model.py:
--------------------------------------------------------------------------------
1 | import torch.nn as nn
2 | import torch
3 | import math
4 | import time
5 | import os
6 | import numpy as np
7 | import cv2
8 | import matplotlib.pyplot as plt
9 | import torch.utils.model_zoo as model_zoo
10 | from torch.nn import init
11 | from utils import BasicBlock, Bottleneck, BBoxTransform, ClipBoxes
12 | from anchors import Anchors
13 | import losses
14 | from lib.nms.pth_nms import pth_nms
15 | from dataloader import UnNormalizer
16 | unnormalize = UnNormalizer()
17 |
18 |
19 | def nms(dets, thresh):
20 | "Dispatch to either CPU or GPU NMS implementations.\
21 | Accept dets as tensor"""
22 | return pth_nms(dets, thresh)
23 |
24 |
25 |
26 | class PyramidFeatures(nn.Module):
27 | def __init__(self, C3_size, C4_size, C5_size, feature_size=256):
28 | super(PyramidFeatures, self).__init__()
29 |
30 | # upsample C5 to get P5 from the FPN paper
31 | self.P5_1 = nn.Conv2d(C5_size, feature_size, kernel_size=1, stride=1, padding=0)
32 | self.P5_upsampled = nn.Upsample(scale_factor=2, mode='nearest')
33 | self.P5_2 = nn.Conv2d(feature_size, feature_size, kernel_size=3, stride=1, padding=1)
34 |
35 | # add P5 elementwise to C4
36 | self.P4_1 = nn.Conv2d(C4_size, feature_size, kernel_size=1, stride=1, padding=0)
37 | self.P4_upsampled = nn.Upsample(scale_factor=2, mode='nearest')
38 | self.P4_2 = nn.Conv2d(feature_size, feature_size, kernel_size=3, stride=1, padding=1)
39 |
40 | # add P4 elementwise to C3
41 | self.P3_1 = nn.Conv2d(C3_size, feature_size, kernel_size=1, stride=1, padding=0)
42 | self.P3_2 = nn.Conv2d(feature_size, feature_size, kernel_size=3, stride=1, padding=1)
43 |
44 | # "P6 is obtained via a 3x3 stride-2 conv on C5"
45 | self.P6 = nn.Conv2d(C5_size, feature_size, kernel_size=3, stride=2, padding=1)
46 |
47 | # "P7 is computed by applying ReLU followed by a 3x3 stride-2 conv on P6"
48 | self.P7_1 = nn.ReLU()
49 | self.P7_2 = nn.Conv2d(feature_size, feature_size, kernel_size=3, stride=2, padding=1)
50 |
51 | def forward(self, inputs):
52 | C3, C4, C5 = inputs
53 |
54 | P5_x = self.P5_1(C5)
55 | P5_upsampled_x = self.P5_upsampled(P5_x)
56 | P5_x = self.P5_2(P5_x)
57 |
58 | P4_x = self.P4_1(C4)
59 | P4_x = P5_upsampled_x + P4_x
60 | P4_upsampled_x = self.P4_upsampled(P4_x)
61 | P4_x = self.P4_2(P4_x)
62 |
63 | P3_x = self.P3_1(C3)
64 | P3_x = P3_x + P4_upsampled_x
65 | P3_x = self.P3_2(P3_x)
66 |
67 | P6_x = self.P6(C5)
68 |
69 | P7_x = self.P7_1(P6_x)
70 | P7_x = self.P7_2(P7_x)
71 |
72 | return [P3_x, P4_x, P5_x, P6_x, P7_x]
73 |
74 |
75 | class RegressionModel(nn.Module):
76 | def __init__(self, num_features_in, num_anchors=9, feature_size=256):
77 | super(RegressionModel, self).__init__()
78 |
79 | self.conv1 = nn.Conv2d(num_features_in, feature_size, kernel_size=3, padding=1)
80 | self.act1 = nn.ReLU()
81 |
82 | self.conv2 = nn.Conv2d(feature_size, feature_size, kernel_size=3, padding=1)
83 | self.act2 = nn.ReLU()
84 |
85 | self.conv3 = nn.Conv2d(feature_size, feature_size, kernel_size=3, padding=1)
86 | self.act3 = nn.ReLU()
87 |
88 | self.conv4 = nn.Conv2d(feature_size, feature_size, kernel_size=3, padding=1)
89 | self.act4 = nn.ReLU()
90 |
91 | self.output = nn.Conv2d(feature_size, num_anchors * 4, kernel_size=3, padding=1)
92 |
93 | def forward(self, x):
94 | out = self.conv1(x)
95 | out = self.act1(out)
96 |
97 | out = self.conv2(out)
98 | out = self.act2(out)
99 |
100 | out = self.conv3(out)
101 | out = self.act3(out)
102 |
103 | out = self.conv4(out)
104 | out = self.act4(out)
105 |
106 | out = self.output(out)
107 |
108 | # out is B x C x W x H, with C = 4*num_anchors
109 | out = out.permute(0, 2, 3, 1)
110 |
111 | return out.contiguous().view(out.shape[0], -1, 4)
112 |
113 |
114 | class ClassificationModel(nn.Module):
115 | def __init__(self, num_features_in, num_anchors=9, num_classes=80, prior=0.01, feature_size=256):
116 | super(ClassificationModel, self).__init__()
117 |
118 | self.num_classes = num_classes
119 | self.num_anchors = num_anchors
120 |
121 | self.conv1 = nn.Conv2d(num_features_in, feature_size, kernel_size=3, padding=1)
122 | self.act1 = nn.ReLU()
123 |
124 | self.conv2 = nn.Conv2d(feature_size, feature_size, kernel_size=3, padding=1)
125 | self.act2 = nn.ReLU()
126 |
127 | self.conv3 = nn.Conv2d(feature_size, feature_size, kernel_size=3, padding=1)
128 | self.act3 = nn.ReLU()
129 |
130 | self.conv4 = nn.Conv2d(feature_size, feature_size, kernel_size=3, padding=1)
131 | self.act4 = nn.ReLU()
132 |
133 | self.output = nn.Conv2d(feature_size, num_anchors * num_classes, kernel_size=3, padding=1)
134 | self.output_act = nn.Sigmoid()
135 |
136 | def forward(self, x):
137 | out = self.conv1(x)
138 | out = self.act1(out)
139 |
140 | out = self.conv2(out)
141 | out = self.act2(out)
142 |
143 | out = self.conv3(out)
144 | out = self.act3(out)
145 |
146 | out = self.conv4(out)
147 | out = self.act4(out)
148 |
149 | out = self.output(out)
150 | out = self.output_act(out)
151 |
152 | # out is B x C x W x H, with C = n_classes + n_anchors
153 | out1 = out.permute(0, 2, 3, 1)
154 |
155 | batch_size, width, height, channels = out1.shape
156 |
157 | out2 = out1.view(batch_size, width, height, self.num_anchors, self.num_classes)
158 |
159 | return out2.contiguous().view(x.shape[0], -1, self.num_classes)
160 |
161 |
162 | class ResNet(nn.Module):
163 |
164 | def __init__(self, num_classes, block, layers):
165 | self.inplanes = 64
166 | super(ResNet, self).__init__()
167 | self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)
168 | self.bn1 = nn.BatchNorm2d(64)
169 | self.relu = nn.ReLU(inplace=True)
170 | self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
171 | self.layer1 = self._make_layer(block, 64, layers[0])
172 | self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
173 | self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
174 | self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
175 |
176 | if block == BasicBlock:
177 | fpn_sizes = [self.layer2[layers[1] - 1].conv2.out_channels, self.layer3[layers[2] - 1].conv2.out_channels,
178 | self.layer4[layers[3] - 1].conv2.out_channels]
179 | elif block == Bottleneck:
180 | fpn_sizes = [self.layer2[layers[1] - 1].conv3.out_channels, self.layer3[layers[2] - 1].conv3.out_channels,
181 | self.layer4[layers[3] - 1].conv3.out_channels]
182 |
183 | self.fpn = PyramidFeatures(fpn_sizes[0], fpn_sizes[1], fpn_sizes[2])
184 |
185 | self.regressionModel = RegressionModel(256)
186 | self.classificationModel = ClassificationModel(256, num_classes=num_classes)
187 |
188 | self.anchors = Anchors()
189 |
190 | self.regressBoxes = BBoxTransform()
191 |
192 | self.clipBoxes = ClipBoxes()
193 |
194 | self.focalLoss = losses.FocalLoss()
195 |
196 | for m in self.modules():
197 | if isinstance(m, nn.Conv2d):
198 | n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
199 | m.weight.data.normal_(0, math.sqrt(2. / n))
200 | # init.xavier_normal(m.weight)
201 | elif isinstance(m, nn.BatchNorm2d):
202 | m.weight.data.fill_(1)
203 | m.bias.data.zero_()
204 |
205 | prior = 0.01
206 |
207 | self.classificationModel.output.weight.data.fill_(0)
208 | self.classificationModel.output.bias.data.fill_(-math.log((1.0 - prior) / prior))
209 |
210 | self.regressionModel.output.weight.data.fill_(0)
211 | self.regressionModel.output.bias.data.fill_(0)
212 |
213 | self.freeze_bn()
214 |
215 | def _make_layer(self, block, planes, blocks, stride=1):
216 | downsample = None
217 | if stride != 1 or self.inplanes != planes * block.expansion:
218 | downsample = nn.Sequential(
219 | nn.Conv2d(self.inplanes, planes * block.expansion,
220 | kernel_size=1, stride=stride, bias=False),
221 | nn.BatchNorm2d(planes * block.expansion),
222 | )
223 |
224 | layers = []
225 | layers.append(block(self.inplanes, planes, stride, downsample))
226 | self.inplanes = planes * block.expansion
227 | for i in range(1, blocks):
228 | layers.append(block(self.inplanes, planes))
229 |
230 | return nn.Sequential(*layers)
231 |
232 | def freeze_bn(self):
233 | '''Freeze BatchNorm layers.'''
234 | for layer in self.modules():
235 | if isinstance(layer, nn.BatchNorm2d):
236 | layer.eval()
237 |
238 | def forward(self, inputs):
239 |
240 | if self.training:
241 | img_batch, annotations, ignores = inputs
242 | else:
243 | img_batch = inputs
244 |
245 | x = self.conv1(img_batch)
246 | x = self.bn1(x)
247 | x = self.relu(x)
248 | x = self.maxpool(x)
249 |
250 | x1 = self.layer1(x)
251 | x2 = self.layer2(x1)
252 | x3 = self.layer3(x2)
253 | x4 = self.layer4(x3)
254 |
255 | features = self.fpn([x2, x3, x4])
256 |
257 | regression = torch.cat([self.regressionModel(feature) for feature in features], dim=1)
258 |
259 | classification = torch.cat([self.classificationModel(feature) for feature in features], dim=1)
260 |
261 | anchors = self.anchors(img_batch)
262 |
263 | if self.training:
264 | return self.focalLoss(classification, regression, anchors, annotations, ignores)
265 | else:
266 |
267 | transformed_anchors = self.regressBoxes(anchors, regression)
268 | transformed_anchors = self.clipBoxes(transformed_anchors, img_batch)
269 |
270 | scores = torch.max(classification, dim=2, keepdim=True)[0]
271 | scores_over_thresh = (scores > 0.05)[0, :, 0]
272 |
273 | if scores_over_thresh.sum() == 0:
274 | # no boxes to NMS, just return
275 | # return [torch.zeros(0), torch.zeros(0), torch.zeros(0, 4)]
276 | return [None, None, None]
277 |
278 | classification = classification[:, scores_over_thresh, :]
279 | transformed_anchors = transformed_anchors[:, scores_over_thresh, :]
280 | scores = scores[:, scores_over_thresh, :]
281 |
282 | anchors_nms_idx = nms(torch.cat([transformed_anchors, scores], dim=2)[0, :, :], 0.5)
283 | nms_scores, nms_class = classification[0, anchors_nms_idx, :].max(dim=1)
284 | return [nms_scores, nms_class, transformed_anchors[0, anchors_nms_idx, :]]
285 |
286 |
287 | def resnet18(num_classes, **kwargs):
288 | """Constructs a ResNet-18 model.
289 | Args:
290 | pretrained (bool): If True, returns a model pre-trained on ImageNet
291 | """
292 | model = ResNet(num_classes, BasicBlock, [2, 2, 2, 2], **kwargs)
293 | return model
294 |
295 |
296 | def resnet34(num_classes, **kwargs):
297 | """Constructs a ResNet-34 model.
298 | Args:
299 | pretrained (bool): If True, returns a model pre-trained on ImageNet
300 | """
301 | model = ResNet(num_classes, BasicBlock, [3, 4, 6, 3], **kwargs)
302 | return model
303 |
304 |
305 | def resnet50(num_classes, **kwargs):
306 | """Constructs a ResNet-50 model.
307 | Args:
308 | pretrained (bool): If True, returns a model pre-trained on ImageNet
309 | """
310 | model = ResNet(num_classes, Bottleneck, [3, 4, 6, 3], **kwargs)
311 | return model
312 |
313 |
314 | def resnet101(num_classes, **kwargs):
315 | """Constructs a ResNet-101 model.
316 | Args:
317 | pretrained (bool): If True, returns a model pre-trained on ImageNet
318 | """
319 | model = ResNet(num_classes, Bottleneck, [3, 4, 23, 3], **kwargs)
320 | return model
321 |
322 |
323 | def resnet152(num_classes, **kwargs):
324 | """Constructs a ResNet-152 model.
325 | Args:
326 | pretrained (bool): If True, returns a model pre-trained on ImageNet
327 | """
328 | model = ResNet(num_classes, Bottleneck, [3, 8, 36, 3], **kwargs)
329 | return model
--------------------------------------------------------------------------------
/train.py:
--------------------------------------------------------------------------------
1 | import time
2 | import os
3 | import copy
4 | import argparse
5 | import pdb
6 | import collections
7 | import sys
8 |
9 | import numpy as np
10 |
11 | import torch
12 | import torch.nn as nn
13 | import torch.optim as optim
14 | from torch.optim import lr_scheduler
15 | from torch.autograd import Variable
16 | from torchvision import datasets, models, transforms
17 | import torchvision
18 | from tensorboardX import SummaryWriter
19 |
20 | import model
21 | from anchors import Anchors
22 | import losses
23 | from dataloader import CSVDataset, collater, Resizer, AspectRatioBasedSampler, Augmenter, UnNormalizer, Normalizer
24 | from torch.utils.data import Dataset, DataLoader
25 |
26 | import csv_eval
27 | import cv2
28 | assert torch.__version__.split('.')[1] == '4'
29 |
30 | print('CUDA available: {}'.format(torch.cuda.is_available()))
31 |
32 | ckpt = False
33 | def main(args=None):
34 |
35 | parser = argparse.ArgumentParser(description='Simple training script for training a RetinaNet network.')
36 |
37 | parser.add_argument('--csv_train', help='Path to file containing training annotations (see readme)')
38 | parser.add_argument('--csv_classes', help='Path to file containing class list (see readme)')
39 | parser.add_argument('--csv_val', help='Path to file containing validation annotations (optional, see readme)')
40 |
41 | parser.add_argument('--depth', help='Resnet depth, must be one of 18, 34, 50, 101, 152', type=int, default=50)
42 | parser.add_argument('--epochs', help='Number of epochs', type=int, default=50)
43 |
44 | parser.add_argument('--model_name', help='name of the model to save')
45 | parser.add_argument('--pretrained', help='pretrained model name')
46 |
47 | parser = parser.parse_args(args)
48 |
49 | # Create the data loaders
50 | dataset_train = CSVDataset(train_file=parser.csv_train, class_list=parser.csv_classes, transform=transforms.Compose([Resizer(), Augmenter(), Normalizer()]))
51 |
52 | if parser.csv_val is None:
53 | dataset_val = None
54 | print('No validation annotations provided.')
55 | else:
56 | dataset_val = CSVDataset(train_file=parser.csv_val, class_list=parser.csv_classes, transform=transforms.Compose([Resizer(), Normalizer()]))
57 |
58 | sampler = AspectRatioBasedSampler(dataset_train, batch_size=2, drop_last=False)
59 | dataloader_train = DataLoader(dataset_train, num_workers=16, collate_fn=collater, batch_sampler=sampler)
60 | #dataloader_train = DataLoader(dataset_train, num_workers=16, collate_fn=collater, batch_size=8, shuffle=True)
61 |
62 | if dataset_val is not None:
63 | sampler_val = AspectRatioBasedSampler(dataset_val, batch_size=2, drop_last=False)
64 | dataloader_val = DataLoader(dataset_val, num_workers=16, collate_fn=collater, batch_sampler=sampler_val)
65 | #dataloader_val = DataLoader(dataset_train, num_workers=16, collate_fn=collater, batch_size=8, shuffle=True)
66 |
67 | # Create the model_pose_level_attention
68 | if parser.depth == 18:
69 | retinanet = model.resnet18(num_classes=dataset_train.num_classes())
70 | elif parser.depth == 34:
71 | retinanet = model.resnet34(num_classes=dataset_train.num_classes())
72 | elif parser.depth == 50:
73 | retinanet = model.resnet50(num_classes=dataset_train.num_classes())
74 | elif parser.depth == 101:
75 | retinanet = model.resnet101(num_classes=dataset_train.num_classes())
76 | elif parser.depth == 152:
77 | retinanet = model.resnet152(num_classes=dataset_train.num_classes())
78 | else:
79 | raise ValueError('Unsupported model depth, must be one of 18, 34, 50, 101, 152')
80 |
81 | if ckpt:
82 | retinanet = torch.load('')
83 | print('load ckpt')
84 | else:
85 | retinanet_dict = retinanet.state_dict()
86 | pretrained_dict = torch.load('./weight/' + parser.pretrained)
87 | pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in retinanet_dict}
88 | retinanet_dict.update(pretrained_dict)
89 | retinanet.load_state_dict(retinanet_dict)
90 | print('load pretrained backbone')
91 |
92 | print(retinanet)
93 | retinanet = torch.nn.DataParallel(retinanet, device_ids=[0])
94 | retinanet.cuda()
95 |
96 | retinanet.training = True
97 |
98 | optimizer = optim.Adam(retinanet.parameters(), lr=1e-5)
99 | #optimizer = optim.SGD(retinanet.parameters(), lr=1e-3, momentum=0.9, weight_decay=1e-4)
100 |
101 | scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, patience=3, verbose=True)
102 | #scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=20, gamma=0.1)
103 |
104 | loss_hist = collections.deque(maxlen=500)
105 |
106 | retinanet.train()
107 | retinanet.module.freeze_bn()
108 |
109 | print('Num training images: {}'.format(len(dataset_train)))
110 | f_map = open('./mAP_txt/' + parser.model_name + '.txt', 'a')
111 | writer = SummaryWriter(log_dir='./summary')
112 | iters = 0
113 | for epoch_num in range(0, parser.epochs):
114 |
115 | retinanet.train()
116 | retinanet.module.freeze_bn()
117 |
118 | epoch_loss = []
119 | #scheduler.step()
120 |
121 | for iter_num, data in enumerate(dataloader_train):
122 |
123 | iters += 1
124 |
125 | optimizer.zero_grad()
126 |
127 | classification_loss, regression_loss, repgt_loss, repbox_loss = retinanet([data['img'].cuda().float(), data['annot'], data['ignore']])
128 |
129 | classification_loss = classification_loss.mean()
130 | regression_loss = regression_loss.mean()
131 | repgt_loss = repgt_loss.mean()
132 | repbox_loss = repbox_loss.mean()
133 |
134 | loss = classification_loss + regression_loss + 0.5 * repgt_loss + 0.5 * repbox_loss
135 |
136 | if bool(loss == 0):
137 | continue
138 |
139 | loss.backward()
140 |
141 | torch.nn.utils.clip_grad_norm_(retinanet.parameters(), 0.1)
142 |
143 | optimizer.step()
144 |
145 | loss_hist.append(float(loss))
146 |
147 | epoch_loss.append(float(loss))
148 |
149 | print('Epoch: {} | Iteration: {} | Classification loss: {:1.5f} | Regression loss: {:1.5f} | RepGT loss {:1.5f} | RepBox loss {:1.5f} | Running loss: {:1.5f}'.format(epoch_num, iter_num, float(classification_loss), float(regression_loss), float(repgt_loss), float(repbox_loss), np.mean(loss_hist)))
150 |
151 | writer.add_scalar('classification_loss', classification_loss, iters)
152 | writer.add_scalar('regression_loss', regression_loss, iters)
153 | writer.add_scalar('loss', loss, iters)
154 |
155 | del classification_loss
156 | del regression_loss
157 |
158 |
159 | if parser.csv_val is not None:
160 |
161 | print('Evaluating dataset')
162 |
163 | mAP = csv_eval.evaluate(dataset_val, retinanet)
164 | f_map.write('mAP:{}, epoch:{}'.format(mAP[0][0], epoch_num))
165 | f_map.write('\n')
166 |
167 | scheduler.step(np.mean(epoch_loss))
168 |
169 | torch.save(retinanet.module, './ckpt/' + parser.model_name + '_{}.pt'.format(epoch_num))
170 |
171 | retinanet.eval()
172 |
173 | writer.export_scalars_to_json("./summary/' + parser.pretrained + 'all_scalars.json")
174 | f_map.close()
175 | writer.close()
176 |
177 | if __name__ == '__main__':
178 | main()
179 |
--------------------------------------------------------------------------------
/utils.py:
--------------------------------------------------------------------------------
1 | import torch
2 | import torch.nn as nn
3 | import numpy as np
4 |
5 |
6 | def conv3x3(in_planes, out_planes, stride=1):
7 | """3x3 convolution with padding"""
8 | return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
9 | padding=1, bias=False)
10 |
11 | class BasicBlock(nn.Module):
12 | expansion = 1
13 |
14 | def __init__(self, inplanes, planes, stride=1, downsample=None):
15 | super(BasicBlock, self).__init__()
16 | self.conv1 = conv3x3(inplanes, planes, stride)
17 | self.bn1 = nn.BatchNorm2d(planes)
18 | self.relu = nn.ReLU(inplace=True)
19 | self.conv2 = conv3x3(planes, planes)
20 | self.bn2 = nn.BatchNorm2d(planes)
21 | self.downsample = downsample
22 | self.stride = stride
23 |
24 | def forward(self, x):
25 | residual = x
26 |
27 | out = self.conv1(x)
28 | out = self.bn1(out)
29 | out = self.relu(out)
30 |
31 | out = self.conv2(out)
32 | out = self.bn2(out)
33 |
34 | if self.downsample is not None:
35 | residual = self.downsample(x)
36 |
37 | out += residual
38 | out = self.relu(out)
39 |
40 | return out
41 |
42 |
43 | class Bottleneck(nn.Module):
44 | expansion = 4
45 |
46 | def __init__(self, inplanes, planes, stride=1, downsample=None):
47 | super(Bottleneck, self).__init__()
48 | self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
49 | self.bn1 = nn.BatchNorm2d(planes)
50 | self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
51 | padding=1, bias=False)
52 | self.bn2 = nn.BatchNorm2d(planes)
53 | self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
54 | self.bn3 = nn.BatchNorm2d(planes * 4)
55 | self.relu = nn.ReLU(inplace=True)
56 | self.downsample = downsample
57 | self.stride = stride
58 |
59 | def forward(self, x):
60 | residual = x
61 |
62 | out = self.conv1(x)
63 | out = self.bn1(out)
64 | out = self.relu(out)
65 |
66 | out = self.conv2(out)
67 | out = self.bn2(out)
68 | out = self.relu(out)
69 |
70 | out = self.conv3(out)
71 | out = self.bn3(out)
72 |
73 | if self.downsample is not None:
74 | residual = self.downsample(x)
75 |
76 | out += residual
77 | out = self.relu(out)
78 |
79 | return out
80 |
81 |
82 | class SELayer(nn.Module):
83 | def __init__(self, channel, reduction=16):
84 | super(SELayer, self).__init__()
85 | self.avg_pool = nn.AdaptiveAvgPool2d(1)
86 | self.fc = nn.Sequential(
87 | nn.Linear(channel, channel // reduction),
88 | nn.ReLU(inplace=True),
89 | nn.Linear(channel // reduction, channel),
90 | nn.Sigmoid()
91 | )
92 |
93 | def forward(self, x):
94 | b, c, _, _ = x.size()
95 | y = self.avg_pool(x).view(b, c)
96 | y = self.fc(y).view(b, c, 1, 1)
97 | return x * y
98 |
99 | class BottleneckSE(nn.Module):
100 | expansion = 4
101 |
102 | def __init__(self, inplanes, planes, stride=1, downsample=None, reduction=16):
103 | super(BottleneckSE, self).__init__()
104 | self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
105 | self.bn1 = nn.BatchNorm2d(planes)
106 | self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
107 | padding=1, bias=False)
108 | self.bn2 = nn.BatchNorm2d(planes)
109 | self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
110 | self.bn3 = nn.BatchNorm2d(planes * 4)
111 | self.relu = nn.ReLU(inplace=True)
112 | self.se = SELayer(planes * 4, reduction)
113 | self.downsample = downsample
114 | self.stride = stride
115 |
116 | def forward(self, x):
117 | residual = x
118 |
119 | out = self.conv1(x)
120 | out = self.bn1(out)
121 | out = self.relu(out)
122 |
123 | out = self.conv2(out)
124 | out = self.bn2(out)
125 | out = self.relu(out)
126 |
127 | out = self.conv3(out)
128 | out = self.bn3(out)
129 | out = self.se(out)
130 |
131 | if self.downsample is not None:
132 | residual = self.downsample(x)
133 |
134 | out += residual
135 | out = self.relu(out)
136 |
137 | return out
138 |
139 |
140 | class CBAM_Module(nn.Module):
141 |
142 | def __init__(self, channels, reduction):
143 | super(CBAM_Module, self).__init__()
144 | self.avg_pool = nn.AdaptiveAvgPool2d(1)
145 | self.max_pool = nn.AdaptiveMaxPool2d(1)
146 | self.fc1 = nn.Conv2d(channels, channels // reduction, kernel_size=1,
147 | padding=0)
148 | self.relu = nn.ReLU(inplace=True)
149 | self.fc2 = nn.Conv2d(channels // reduction, channels, kernel_size=1,
150 | padding=0)
151 | self.sigmoid_channel = nn.Sigmoid()
152 | self.conv_after_concat = nn.Conv2d(2, 1, kernel_size = 7, stride=1, padding = 3)
153 | self.sigmoid_spatial = nn.Sigmoid()
154 |
155 | def forward(self, x):
156 | module_input = x
157 | avg = self.avg_pool(x)
158 | mx = self.max_pool(x)
159 | avg = self.fc1(avg)
160 | mx = self.fc1(mx)
161 | avg = self.relu(avg)
162 | mx = self.relu(mx)
163 | avg = self.fc2(avg)
164 | mx = self.fc2(mx)
165 | x = avg + mx
166 | x = self.sigmoid_channel(x)
167 | x = module_input * x
168 | module_input = x
169 | avg = torch.mean(x, 1, True)
170 | mx, _ = torch.max(x, 1, True)
171 | x = torch.cat((avg, mx), 1)
172 | x = self.conv_after_concat(x)
173 | x = self.sigmoid_spatial(x)
174 | x = module_input * x
175 | return x
176 |
177 | class BottleneckCBAM(nn.Module):
178 | expansion = 4
179 |
180 | def __init__(self, inplanes, planes, stride=1, downsample=None, reduction=16):
181 | super(BottleneckCBAM, self).__init__()
182 | self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
183 | self.bn1 = nn.BatchNorm2d(planes)
184 | self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
185 | padding=1, bias=False)
186 | self.bn2 = nn.BatchNorm2d(planes)
187 | self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
188 | self.bn3 = nn.BatchNorm2d(planes * 4)
189 | self.relu = nn.ReLU(inplace=True)
190 | self.se = CBAM_Module(planes * 4, reduction)
191 | self.downsample = downsample
192 | self.stride = stride
193 |
194 | def forward(self, x):
195 | residual = x
196 |
197 | out = self.conv1(x)
198 | out = self.bn1(out)
199 | out = self.relu(out)
200 |
201 | out = self.conv2(out)
202 | out = self.bn2(out)
203 | out = self.relu(out)
204 |
205 | out = self.conv3(out)
206 | out = self.bn3(out)
207 | out = self.se(out)
208 |
209 | if self.downsample is not None:
210 | residual = self.downsample(x)
211 |
212 | out += residual
213 | out = self.relu(out)
214 |
215 | return out
216 |
217 |
218 | class BBoxTransform(nn.Module):
219 |
220 | def __init__(self, mean=None, std=None):
221 | super(BBoxTransform, self).__init__()
222 | if mean is None:
223 | self.mean = torch.from_numpy(np.array([0, 0, 0, 0]).astype(np.float32)).cuda()
224 | else:
225 | self.mean = mean
226 | if std is None:
227 | self.std = torch.from_numpy(np.array([0.1, 0.1, 0.2, 0.2]).astype(np.float32)).cuda()
228 | else:
229 | self.std = std
230 |
231 | def forward(self, boxes, deltas):
232 |
233 | widths = boxes[:, :, 2] - boxes[:, :, 0]
234 | heights = boxes[:, :, 3] - boxes[:, :, 1]
235 | ctr_x = boxes[:, :, 0] + 0.5 * widths
236 | ctr_y = boxes[:, :, 1] + 0.5 * heights
237 |
238 | dx = deltas[:, :, 0] * self.std[0] + self.mean[0]
239 | dy = deltas[:, :, 1] * self.std[1] + self.mean[1]
240 | dw = deltas[:, :, 2] * self.std[2] + self.mean[2]
241 | dh = deltas[:, :, 3] * self.std[3] + self.mean[3]
242 |
243 | pred_ctr_x = ctr_x + dx * widths
244 | pred_ctr_y = ctr_y + dy * heights
245 | pred_w = torch.exp(dw) * widths
246 | pred_h = torch.exp(dh) * heights
247 |
248 | pred_boxes_x1 = pred_ctr_x - 0.5 * pred_w
249 | pred_boxes_y1 = pred_ctr_y - 0.5 * pred_h
250 | pred_boxes_x2 = pred_ctr_x + 0.5 * pred_w
251 | pred_boxes_y2 = pred_ctr_y + 0.5 * pred_h
252 |
253 | pred_boxes = torch.stack([pred_boxes_x1, pred_boxes_y1, pred_boxes_x2, pred_boxes_y2], dim=2)
254 |
255 | return pred_boxes
256 |
257 |
258 | class ClipBoxes(nn.Module):
259 |
260 | def __init__(self, width=None, height=None):
261 | super(ClipBoxes, self).__init__()
262 |
263 | def forward(self, boxes, img):
264 |
265 | batch_size, num_channels, height, width = img.shape
266 |
267 | boxes[:, :, 0] = torch.clamp(boxes[:, :, 0], min=0)
268 | boxes[:, :, 1] = torch.clamp(boxes[:, :, 1], min=0)
269 |
270 | boxes[:, :, 2] = torch.clamp(boxes[:, :, 2], max=width)
271 | boxes[:, :, 3] = torch.clamp(boxes[:, :, 3], max=height)
272 |
273 | return boxes
274 |
--------------------------------------------------------------------------------
| |