def _init_test_roi_heads_faster_rcnn(self):
out_channels = 256
num_classes = 91
box_fg_iou_thresh = 0.5
box_bg_iou_thresh = 0.5
box_batch_size_per_image = 512
box_positive_fraction = 0.25
bbox_reg_weights = None
box_score_thresh = 0.05
box_nms_thresh = 0.5
box_detections_per_img = 100
box_roi_pool = ops.MultiScaleRoIAlign(
featmap_names=['0', '1', '2', '3'],
output_size=7,
sampling_ratio=2)
resolution = box_roi_pool.output_size[0]
representation_size = 1024
box_head = TwoMLPHead(out_channels * resolution**2,
representation_size)
representation_size = 1024
box_predictor = FastRCNNPredictor(representation_size, num_classes)
roi_heads = RoIHeads(box_roi_pool, box_head, box_predictor,
box_fg_iou_thresh, box_bg_iou_thresh,
box_batch_size_per_image, box_positive_fraction,
bbox_reg_weights, box_score_thresh,
box_nms_thresh, box_detections_per_img)
return roi_heads
def create_model(num_classes):
backbone_network = cfg.backbone
if backbone_network == 'mobilenet':
backbone = MobileNetV2(weights_path=cfg.pretrained_weights).features
backbone.out_channels = 1280
anchor_generator = AnchorsGenerator(sizes=cfg.anchor_size,
aspect_ratios=cfg.anchor_ratio)
roi_pooler = ops.MultiScaleRoIAlign(
featmap_names=['0'], # 在哪些特征层上进行roi pooling
output_size=[7, 7], # roi_pooling输出特征矩阵尺寸
sampling_ratio=2) # 采样率
model = FasterRCNN(backbone=backbone,
num_classes=num_classes,
rpn_anchor_generator=anchor_generator,
box_roi_pool=roi_pooler)
return model
def __init__(self):
super(TransformModule, self).__init__()
self.model = ops.MultiScaleRoIAlign(['feat1', 'feat2'], 3, 2)
self.image_sizes = [(512, 512)]
def __init__(self):
super().__init__()
self.model = ops.MultiScaleRoIAlign(["feat1", "feat2"], 3, 2)
self.image_sizes = [(512, 512)]
def create_model(num_classes):
global backbone, model
backbone_network = cfg.backbone
anchor_sizes = tuple((f,) for f in cfg.anchor_size)
aspect_ratios = tuple((f,) for f in cfg.anchor_ratio) * len(anchor_sizes)
anchor_generator = AnchorsGenerator(sizes=anchor_sizes,
aspect_ratios=aspect_ratios)
if backbone_network == 'mobilenet':
backbone = MobileNetV2(weights_path=cfg.backbone_pretrained_weights).features
backbone.out_channels = 1280
roi_pooler = ops.MultiScaleRoIAlign(featmap_names=['0'], # roi pooling in which resolution feature
output_size=cfg.roi_out_size, # roi_pooling output feature size
sampling_ratio=cfg.roi_sample_rate) # sampling_ratio
model = FasterRCNN(backbone=backbone, num_classes=num_classes,
# transform parameters
min_size=cfg.min_size, max_size=cfg.max_size,
image_mean=cfg.image_mean, image_std=cfg.image_std,
# rpn parameters
rpn_anchor_generator=anchor_generator, box_roi_pool=roi_pooler,
rpn_pre_nms_top_n_train=cfg.rpn_pre_nms_top_n_train,
rpn_pre_nms_top_n_test=cfg.rpn_pre_nms_top_n_test,
rpn_post_nms_top_n_train=cfg.rpn_post_nms_top_n_train,
rpn_post_nms_top_n_test=cfg.rpn_post_nms_top_n_test,
rpn_nms_thresh=cfg.rpn_nms_thresh,
rpn_fg_iou_thresh=cfg.rpn_fg_iou_thresh,
rpn_bg_iou_thresh=cfg.rpn_bg_iou_thresh,
rpn_batch_size_per_image=cfg.rpn_batch_size_per_image,
rpn_positive_fraction=cfg.rpn_positive_fraction,
# Box parameters
box_head=None, box_predictor=None,
# remove low threshold target
box_score_thresh=cfg.box_score_thresh,
box_nms_thresh=cfg.box_nms_thresh,
box_detections_per_img=cfg.box_detections_per_img,
box_fg_iou_thresh=cfg.box_fg_iou_thresh,
box_bg_iou_thresh=cfg.box_bg_iou_thresh,
box_batch_size_per_image=cfg.box_batch_size_per_image,
box_positive_fraction=cfg.box_positive_fraction,
bbox_reg_weights=cfg.bbox_reg_weights
)
elif backbone_network == 'resnet50_fpn':
backbone = resnet50_fpn_backbone()
roi_pooler = ops.MultiScaleRoIAlign(
featmap_names=['0', '1', '2', '3'],
output_size=cfg.roi_out_size,
sampling_ratio=cfg.roi_sample_rate)
model = FasterRCNN(backbone=backbone, num_classes=num_classes,
# transform parameters
min_size=cfg.min_size, max_size=cfg.max_size,
image_mean=cfg.image_mean, image_std=cfg.image_std,
# rpn parameters
rpn_anchor_generator=anchor_generator, box_roi_pool=roi_pooler,
rpn_pre_nms_top_n_train=cfg.rpn_pre_nms_top_n_train,
rpn_pre_nms_top_n_test=cfg.rpn_pre_nms_top_n_test,
rpn_post_nms_top_n_train=cfg.rpn_post_nms_top_n_train,
rpn_post_nms_top_n_test=cfg.rpn_post_nms_top_n_test,
rpn_nms_thresh=cfg.rpn_nms_thresh,
rpn_fg_iou_thresh=cfg.rpn_fg_iou_thresh,
rpn_bg_iou_thresh=cfg.rpn_bg_iou_thresh,
rpn_batch_size_per_image=cfg.rpn_batch_size_per_image,
rpn_positive_fraction=cfg.rpn_positive_fraction,
# Box parameters
box_head=None, box_predictor=None,
# remove low threshold target
box_score_thresh=cfg.box_score_thresh,
box_nms_thresh=cfg.box_nms_thresh,
box_detections_per_img=cfg.box_detections_per_img,
box_fg_iou_thresh=cfg.box_fg_iou_thresh,
box_bg_iou_thresh=cfg.box_bg_iou_thresh,
box_batch_size_per_image=cfg.box_batch_size_per_image,
box_positive_fraction=cfg.box_positive_fraction,
bbox_reg_weights=cfg.bbox_reg_weights
)
# weights_dict = torch.load(cfg.pretrained_weights)
# missing_keys, unexpected_keys = model.load_state_dict(weights_dict, strict=False)
# if len(missing_keys) != 0 or len(unexpected_keys) != 0:
# print("missing_keys: ", missing_keys)
# print("unexpected_keys: ", unexpected_keys)
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
return model
