in cvnets/models/detection/ssd.py [0:0]
def __init__(self, opts, encoder: BaseEncoder):
super(SingleShotDetector, self).__init__(opts=opts, encoder=encoder)
# delete layers that are not required in detection network
self.encoder.classifier = None
self.encoder.conv_1x1_exp = None
output_strides = getattr(opts, "model.detection.ssd.output_strides", [16, 32, 64, 128, 256, -1 ])
n_os = len(output_strides)
anchors_aspect_ratio = getattr(opts, "model.detection.ssd.anchors_aspect_ratio", [[2, 3]] * len(output_strides))
proj_channels = getattr(opts, "model.detection.ssd.proj_channels", [512, 256, 256, 128, 128, 64])
anchors_aspect_ratio = anchors_aspect_ratio + [[2]] * (n_os - len(anchors_aspect_ratio))
proj_channels = proj_channels + [128] * (n_os - len(proj_channels))
if len(output_strides) != len(anchors_aspect_ratio) != len(proj_channels):
logger.error(
"SSD model requires anchors to be defined for feature maps from each output stride. Also"
"len(anchors_aspect_ratio) == len(output_strides) == len(proj_channels). "
"Got len(output_strides)={}, len(anchors_aspect_ratio)={}, len(proj_channels)={}."
" Please specify correct arguments using following arguments: "
"\n--model.detection.ssd.anchors-aspect-ratio "
"\n--model.detection.ssd.output-strides"
"\n--model.detection.ssd.proj-channels".format(
len(output_strides),
len(anchors_aspect_ratio),
len(proj_channels)
)
)
extra_layers = {}
enc_channels_list = []
in_channels = self.enc_l5_channels
extra_proj_list = [256] * (len(output_strides) - len(proj_channels))
proj_channels = proj_channels + extra_proj_list
for idx, os in enumerate(output_strides):
out_channels = proj_channels[idx]
if os == 8:
enc_channels_list.append(self.enc_l3_channels)
elif os == 16:
enc_channels_list.append(self.enc_l4_channels)
elif os == 32:
enc_channels_list.append(self.enc_l5_channels)
elif os > 32 and os != -1:
extra_layers["os_{}".format(os)] = SeparableConv(
opts=opts, in_channels=in_channels, out_channels=out_channels, kernel_size=3, use_act=True,
use_norm=True, stride=2
)
enc_channels_list.append(out_channels)
in_channels = out_channels
elif os == -1:
extra_layers["os_{}".format(os)] = nn.Sequential(
AdaptiveAvgPool2d(output_size=1),
ConvLayer(opts=opts, in_channels=in_channels, out_channels=out_channels, kernel_size=1,
use_act=True, use_norm=False)
)
enc_channels_list.append(out_channels)
in_channels = out_channels
else:
raise NotImplementedError
self.extra_layers = None if not extra_layers else nn.ModuleDict(extra_layers)
if self.extra_layers is not None:
self.reset_layers(module=self.extra_layers)
# Anchor box related parameters
self.conf_threshold = getattr(opts, "model.detection.ssd.conf_threshold", 0.01)
self.nms_threshold = getattr(opts, "model.detection.ssd.nms_iou_threshold", 0.3)
self.top_k = getattr(opts, "model.detection.ssd.num_objects_per_class", 200)
self.anchor_box_generator = SSDAnchorGenerator(
output_strides=output_strides,
aspect_ratios=anchors_aspect_ratio,
min_ratio=getattr(opts, "model.detection.ssd.min_box_size", 0.1),
max_ratio=getattr(opts, "model.detection.ssd.max_box_size", 1.05)
)
anchors_aspect_ratio = self.anchor_box_generator.num_anchors_per_os()
self.ssd_heads = nn.ModuleList()
# Create SSD detection and classification heads
for os, in_dim, proj_dim, n_anchors in zip(output_strides, enc_channels_list, proj_channels, anchors_aspect_ratio):
self.ssd_heads += [
SSDHead(opts=opts,
in_channels=in_dim,
n_classes=self.n_classes,
n_coordinates=self.coordinates,
n_anchors=n_anchors,
proj_channels=proj_dim,
kernel_size=3 if os != -1 else 1)
]
self.anchors_aspect_ratio = anchors_aspect_ratio
self.output_strides = output_strides