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