pokouqiege/models/line_detect/line_detect.py

import os
from typing import Any, Callable, List, Optional, Tuple, Union
import torch
from torch import nn


from libs.vision_libs import ops
from libs.vision_libs.models import MobileNet_V3_Large_Weights, mobilenet_v3_large, EfficientNet_V2_S_Weights, \
    efficientnet_v2_s, detection, EfficientNet_V2_L_Weights, efficientnet_v2_l, EfficientNet_V2_M_Weights, \
    efficientnet_v2_m
from libs.vision_libs.models.detection.anchor_utils import AnchorGenerator
from libs.vision_libs.models.detection.rpn import RPNHead, RegionProposalNetwork
from libs.vision_libs.models.detection.ssdlite import _mobilenet_extractor
from libs.vision_libs.models.detection.transform import GeneralizedRCNNTransform
from libs.vision_libs.ops import misc as misc_nn_ops, MultiScaleRoIAlign
from libs.vision_libs.transforms._presets import ObjectDetection
from libs.vision_libs.models._api import register_model, Weights, WeightsEnum
from libs.vision_libs.models._meta import _COCO_PERSON_CATEGORIES, _COCO_PERSON_KEYPOINT_NAMES, _COCO_CATEGORIES
from libs.vision_libs.models._utils import _ovewrite_value_param, handle_legacy_interface
from libs.vision_libs.models.resnet import resnet50, ResNet50_Weights, ResNet18_Weights, resnet18
from libs.vision_libs.models.detection._utils import overwrite_eps
from libs.vision_libs.models.detection.backbone_utils import _resnet_fpn_extractor, _validate_trainable_layers, \
    BackboneWithFPN, resnet_fpn_backbone
from libs.vision_libs.models.detection.faster_rcnn import FasterRCNN, TwoMLPHead, FastRCNNPredictor
from .loi_heads import RoIHeads

from .trainer import Trainer
from ..base import backbone_factory
from ..base.backbone_factory import get_convnext_fpn, get_anchor_generator
# from ..base.backbone_factory import get_convnext_fpn, get_anchor_generator
from ..base.base_detection_net import BaseDetectionNet
import torch.nn.functional as F

from ..base.high_reso_resnet import resnet50fpn, resnet18fpn

__all__ = [
    "LineDetect",
    "linedetect_resnet50_fpn",
]

def _default_anchorgen():
    anchor_sizes = ((32,), (64,), (128,), (256,), (512,))
    aspect_ratios = ((0.5, 1.0, 2.0),) * len(anchor_sizes)
    return AnchorGenerator(anchor_sizes, aspect_ratios)


class LineDetect(BaseDetectionNet):


    def __init__(
            self,
            backbone,
            num_classes=None,
            # transform parameters
            min_size=512,
            max_size=2048,
            image_mean=None,
            image_std=None,
            # RPN parameters
            rpn_anchor_generator=None,
            rpn_head=None,
            rpn_pre_nms_top_n_train=2000,
            rpn_pre_nms_top_n_test=1000,
            rpn_post_nms_top_n_train=2000,
            rpn_post_nms_top_n_test=1000,
            rpn_nms_thresh=0.7,
            rpn_fg_iou_thresh=0.7,
            rpn_bg_iou_thresh=0.3,
            rpn_batch_size_per_image=256,
            rpn_positive_fraction=0.5,
            rpn_score_thresh=0.0,
            # Box parameters
            box_roi_pool=None,
            box_head=None,
            box_predictor=None,
            box_score_thresh=0.05,
            box_nms_thresh=0.5,
            box_detections_per_img=100,
            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,
            # keypoint parameters
            line_roi_pool=None,
            line_head=None,
            line_predictor=None,
            num_keypoints=None,
            **kwargs,
    ):

        out_channels = backbone.out_channels

        if rpn_anchor_generator is None:
            rpn_anchor_generator = _default_anchorgen()
        if rpn_head is None:
            rpn_head = RPNHead(out_channels, rpn_anchor_generator.num_anchors_per_location()[0])

        rpn_pre_nms_top_n = dict(training=rpn_pre_nms_top_n_train, testing=rpn_pre_nms_top_n_test)
        rpn_post_nms_top_n = dict(training=rpn_post_nms_top_n_train, testing=rpn_post_nms_top_n_test)

        rpn = RegionProposalNetwork(
            rpn_anchor_generator,
            rpn_head,
            rpn_fg_iou_thresh,
            rpn_bg_iou_thresh,
            rpn_batch_size_per_image,
            rpn_positive_fraction,
            rpn_pre_nms_top_n,
            rpn_post_nms_top_n,
            rpn_nms_thresh,
            score_thresh=rpn_score_thresh,
        )

        if box_roi_pool is None:
            box_roi_pool = MultiScaleRoIAlign(featmap_names=["0", "1", "2", "3"], output_size=7, sampling_ratio=2)

        if box_head is None:
            resolution = box_roi_pool.output_size[0]
            representation_size = 1024
            box_head = TwoMLPHead(out_channels * resolution**2, representation_size)

        if box_predictor is None:
            representation_size = 1024
            box_predictor = ObjectionPredictor(representation_size, num_classes)

        roi_heads = RoIHeads(
            # Box
            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,
        )

        if image_mean is None:
            image_mean = [0.485, 0.456, 0.406]
        if image_std is None:
            image_std = [0.229, 0.224, 0.225]
        transform = GeneralizedRCNNTransform(min_size, max_size, image_mean, image_std, **kwargs)

        super().__init__(backbone, rpn, roi_heads, transform)



        if not isinstance(line_roi_pool, (MultiScaleRoIAlign, type(None))):
            raise TypeError(
                "keypoint_roi_pool should be of type MultiScaleRoIAlign or None instead of {type(keypoint_roi_pool)}"
            )
        if min_size is None:
            min_size = (640, 672, 704, 736, 768, 800)

        if num_keypoints is not None:
            if line_predictor is not None:
                raise ValueError("num_keypoints should be None when keypoint_predictor is specified")
        else:
            num_keypoints = 2


        if line_roi_pool is None:
            line_roi_pool = MultiScaleRoIAlign(featmap_names=["0", "1", "2", "3"], output_size=14, sampling_ratio=2)

        if line_head is None:
            keypoint_layers = tuple(1 for _ in range(8))
            line_head = LineHeads(16, keypoint_layers)

        if line_predictor is None:
            keypoint_dim_reduced = 512  # == keypoint_layers[-1]
            line_predictor = LinePredictor(keypoint_dim_reduced)


        self.roi_heads.line_roi_pool = line_roi_pool
        self.roi_heads.line_head = line_head
        self.roi_heads.line_predictor = line_predictor

    def start_train(self, cfg):
        # cfg = read_yaml(cfg)
        self.trainer = Trainer()
        self.trainer.train_from_cfg(model=self, cfg=cfg)

    def load_weights(self, save_path, device='cuda'):
        if os.path.exists(save_path):
            checkpoint = torch.load(save_path, map_location=device)

            self.load_state_dict(checkpoint['model_state_dict'])
            # if optimizer is not None:
            #     optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
            # epoch = checkpoint['epoch']
            # loss = checkpoint['loss']
            # print(f"Loaded best model from {save_path} at epoch {epoch} with loss {loss:.4f}")
            print(f"Loaded model from {save_path}")
        else:
            print(f"No saved model found at {save_path}")
        return self


class TwoMLPHead(nn.Module):
    """
    Standard heads for FPN-based models

    Args:
        in_channels (int): number of input channels
        representation_size (int): size of the intermediate representation
    """

    def __init__(self, in_channels, representation_size):
        super().__init__()

        self.fc6 = nn.Linear(in_channels, representation_size)
        self.fc7 = nn.Linear(representation_size, representation_size)

    def forward(self, x):
        x = x.flatten(start_dim=1)

        x = F.relu(self.fc6(x))
        x = F.relu(self.fc7(x))

        return x


class ObjectionConvFCHead(nn.Sequential):
    def __init__(
        self,
        input_size: Tuple[int, int, int],
        conv_layers: List[int],
        fc_layers: List[int],
        norm_layer: Optional[Callable[..., nn.Module]] = None,
    ):
        """
        Args:
            input_size (Tuple[int, int, int]): the input size in CHW format.
            conv_layers (list): feature dimensions of each Convolution layer
            fc_layers (list): feature dimensions of each FCN layer
            norm_layer (callable, optional): Module specifying the normalization layer to use. Default: None
        """
        in_channels, in_height, in_width = input_size

        blocks = []
        previous_channels = in_channels
        for current_channels in conv_layers:
            blocks.append(misc_nn_ops.Conv2dNormActivation(previous_channels, current_channels, norm_layer=norm_layer))
            previous_channels = current_channels
        blocks.append(nn.Flatten())
        previous_channels = previous_channels * in_height * in_width
        for current_channels in fc_layers:
            blocks.append(nn.Linear(previous_channels, current_channels))
            blocks.append(nn.ReLU(inplace=True))
            previous_channels = current_channels

        super().__init__(*blocks)
        for layer in self.modules():
            if isinstance(layer, nn.Conv2d):
                nn.init.kaiming_normal_(layer.weight, mode="fan_out", nonlinearity="relu")
                if layer.bias is not None:
                    nn.init.zeros_(layer.bias)


class ObjectionPredictor(nn.Module):
    """
    Standard classification + bounding box regression layers
    for Fast R-CNN.

    Args:
        in_channels (int): number of input channels
        num_classes (int): number of output classes (including background)
    """

    def __init__(self, in_channels, num_classes):
        super().__init__()
        self.cls_score = nn.Linear(in_channels, num_classes)
        self.bbox_pred = nn.Linear(in_channels, num_classes * 4)

    def forward(self, x):
        if x.dim() == 4:
            torch._assert(
                list(x.shape[2:]) == [1, 1],
                f"x has the wrong shape, expecting the last two dimensions to be [1,1] instead of {list(x.shape[2:])}",
            )
        x = x.flatten(start_dim=1)
        scores = self.cls_score(x)
        bbox_deltas = self.bbox_pred(x)

        return scores, bbox_deltas

class LineHeads(nn.Sequential):
    def __init__(self, in_channels, layers):
        d = []
        next_feature = in_channels
        for out_channels in layers:
            d.append(nn.Conv2d(next_feature, out_channels, 3, stride=1, padding=1))
            d.append(nn.ReLU(inplace=True))
            next_feature = out_channels
        super().__init__(*d)
        for m in self.children():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
                nn.init.constant_(m.bias, 0)


class LinePredictor(nn.Module):
    def __init__(self, in_channels, out_channels=1 ):
        super().__init__()
        input_features = in_channels
        deconv_kernel = 4
        self.kps_score_lowres = nn.ConvTranspose2d(
            input_features,
            out_channels,
            deconv_kernel,
            stride=2,
            padding=deconv_kernel // 2 - 1,
        )
        nn.init.kaiming_normal_(self.kps_score_lowres.weight, mode="fan_out", nonlinearity="relu")
        nn.init.constant_(self.kps_score_lowres.bias, 0)
        self.up_scale = 2
        self.out_channels = out_channels

    def forward(self, x):
        print(f'before kps_score_lowres x:{x.shape}')
        x = self.kps_score_lowres(x)
        print(f'kps_score_lowres x:{x.shape}')
        return torch.nn.functional.interpolate(
            x, scale_factor=float(self.up_scale), mode="bilinear", align_corners=False, recompute_scale_factor=False
        )

def linedetect_newresnet18fpn(
        *,

        num_classes: Optional[int] = None,
        num_points:Optional[int] = None,

        **kwargs: Any,
) -> LineDetect:
    # weights = LineNet_ResNet50_FPN_Weights.verify(weights)
    # weights_backbone = ResNet50_Weights.verify(weights_backbone)
    if num_classes is None:
        num_classes = 2
    if num_points is None:
        num_points = 2


    backbone =resnet18fpn()
    featmap_names=['0', '1', '2', '3','pool']
    # print(f'featmap_names:{featmap_names}')
    roi_pooler = MultiScaleRoIAlign(
        featmap_names=featmap_names,
        output_size=7,
        sampling_ratio=2
    )
    num_features=len(featmap_names)
    anchor_sizes = tuple((int(16 * 2 ** i),) for i in range(num_features))  # 自动生成不同大小
    # print(f'anchor_sizes:{anchor_sizes}')
    aspect_ratios = ((0.5, 1.0, 2.0),) * num_features
    # print(f'aspect_ratios:{aspect_ratios}')


    anchor_generator =  AnchorGenerator(sizes=anchor_sizes, aspect_ratios=aspect_ratios)

    model = LineDetect(backbone, num_classes, num_keypoints=num_points,rpn_anchor_generator=anchor_generator,box_roi_pool=roi_pooler, **kwargs)

    return model



def linedetect_resnet18_fpn(
        *,
        num_classes: Optional[int] = None,
        num_points: Optional[int] = None,
        **kwargs: Any,
) -> LineDetect:

    if num_classes is None:
        num_classes = 2
    if num_points is None:
        num_points = 2

    backbone = resnet_fpn_backbone(backbone_name='resnet18',weights=None)
    model = LineDetect(backbone, num_classes, num_keypoints=num_points, **kwargs)

    return model

def linedetect_resnet50_fpn(
        *,
        num_classes: Optional[int] = None,
        num_points: Optional[int] = None,
        **kwargs: Any,
) -> LineDetect:
    if num_classes is None:
        num_classes = 2
    if num_points is None:
        num_points = 2

    backbone = resnet_fpn_backbone(backbone_name='resnet18', weights=None)
    model = LineDetect(backbone, num_classes, num_keypoints=num_points, **kwargs)


    return model