File size: 2,832 Bytes
22e1f7f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 | import torch
from torchvision.models.detection import FasterRCNN
from torchvision.models.detection.rpn import AnchorGenerator
from torchvision.ops import MultiScaleRoIAlign
from .backbone import ResNet50FPNBackbone
class VisionAN18FasterRCNN(FasterRCNN):
def __init__(
self,
num_classes=3,
pretrained_backbone=True
):
backbone = ResNet50FPNBackbone(
pretrained=pretrained_backbone,
out_channels=256
)
anchor_generator = AnchorGenerator(
sizes=(
(32,),
(64,),
(128,),
(256,),
(512,)
),
aspect_ratios=(
(0.5, 1.0, 2.0),
(0.5, 1.0, 2.0),
(0.5, 1.0, 2.0),
(0.5, 1.0, 2.0),
(0.5, 1.0, 2.0)
)
)
roi_pooler = MultiScaleRoIAlign(
featmap_names=[
"c2",
"c3",
"c4",
"c5",
"pool"
],
output_size=7,
sampling_ratio=2
)
super().__init__(
backbone=backbone,
num_classes=num_classes,
rpn_anchor_generator=anchor_generator,
box_roi_pool=roi_pooler,
min_size=800,
max_size=1333,
rpn_pre_nms_top_n_train=2000,
rpn_pre_nms_top_n_test=1000,
rpn_post_nms_top_n_train=1000,
rpn_post_nms_top_n_test=300,
rpn_nms_thresh=0.7,
rpn_fg_iou_thresh=0.7,
rpn_bg_iou_thresh=0.3,
box_score_thresh=0.05,
box_nms_thresh=0.5,
box_detections_per_img=100
)
self.class_names = {
0: "background",
1: "license_plate",
2: "house_number"
}
def print_architecture(self):
print("=" * 60)
print("VISION-AN18 FASTER R-CNN")
print("=" * 60)
print("Backbone : ResNet50")
print("Backbone : Pretrained ImageNet")
print("Neck : Feature Pyramid Network")
print("RPN : Region Proposal Network")
print("ROI Pooler : MultiScale RoI Align")
print("Detection : Fast R-CNN Head")
print("Classes : 3")
print(" 0 : background")
print(" 1 : license_plate")
print(" 2 : house_number")
print("=" * 60)
def create_model(
num_classes=3,
pretrained_backbone=True
):
model = VisionAN18FasterRCNN(
num_classes=num_classes,
pretrained_backbone=pretrained_backbone
)
return model |