from collections import OrderedDict import torchvision from torchvision.models import maxvit_t from torchvision.models.detection.backbone_utils import BackboneWithFPN from libs.vision_libs import models from libs.vision_libs.models import mobilenet_v3_large, EfficientNet_V2_S_Weights, efficientnet_v2_s, \ EfficientNet_V2_M_Weights, efficientnet_v2_m, EfficientNet_V2_L_Weights, efficientnet_v2_l, ConvNeXt_Base_Weights from libs.vision_libs.models._utils import _ovewrite_value_param, handle_legacy_interface from libs.vision_libs.models.detection import FasterRCNN from libs.vision_libs.models.detection.anchor_utils import AnchorGenerator from libs.vision_libs.models.detection.ssdlite import _mobilenet_extractor from libs.vision_libs.models.resnet import resnet50, ResNet50_Weights, resnet18 from libs.vision_libs.models.detection.backbone_utils import _resnet_fpn_extractor, _validate_trainable_layers from libs.vision_libs.ops import misc as misc_nn_ops, MultiScaleRoIAlign from torch import nn import torch from models.base.high_reso_swin import swin_v2_t def get_resnet50_fpn(): is_trained = False trainable_backbone_layers = _validate_trainable_layers(is_trained, None, 5, 3) norm_layer = misc_nn_ops.FrozenBatchNorm2d if is_trained else nn.BatchNorm2d backbone = resnet50(weights=None, progress=True, norm_layer=norm_layer) backbone = _resnet_fpn_extractor(backbone, trainable_backbone_layers) return backbone def get_resnet18_fpn(): is_trained = False trainable_backbone_layers = _validate_trainable_layers(is_trained, None, 5, 3) norm_layer = misc_nn_ops.FrozenBatchNorm2d if is_trained else nn.BatchNorm2d backbone = resnet18(weights=None, progress=True, norm_layer=norm_layer) backbone = _resnet_fpn_extractor(backbone, trainable_backbone_layers) return backbone def get_mobilenet_v3_large_fpn(): is_trained = False trainable_backbone_layers = _validate_trainable_layers(is_trained, None, 6, 3) norm_layer = misc_nn_ops.FrozenBatchNorm2d if is_trained else nn.BatchNorm2d backbone = mobilenet_v3_large(weights=None, progress=True, norm_layer=norm_layer) backbone = _mobilenet_extractor(backbone, True, trainable_backbone_layers) return backbone def get_convnext_fpn(): convnext = models.convnext_base(weights=ConvNeXt_Base_Weights.DEFAULT) # convnext = models.convnext_small(pretrained=True) # convnext = models.convnext_large(pretrained=True) in_channels_list = [128, 256, 512, 1024] backbone_with_fpn = BackboneWithFPN( convnext.features, return_layers={'1': '0', '3': '1', '5': '2', '7': '3'}, # 确保这些键对应到实际的层 in_channels_list=in_channels_list, out_channels=256 ) return backbone_with_fpn def get_maxvit_fpn(input_size=(224*7,224*7)): maxvit = MaxVitBackbone(input_size=input_size) # print(maxvit.named_children()) # for i,layer in enumerate(maxvit.named_children()): # print(f'layer:{i}:{layer}') test_input = torch.randn(1, 3, 224 * 7, 224 * 7) in_channels_list = [64, 64, 128, 256, 512] featmap_names = ['0', '1', '2', '3', '4', 'pool'] # print(f'featmap_names:{featmap_names}') roi_pooler = MultiScaleRoIAlign( featmap_names=featmap_names, output_size=7, sampling_ratio=2 ) backbone_with_fpn = BackboneWithFPN( maxvit, return_layers={'stem': '0', 'block0': '1', 'block1': '2', 'block2': '3', 'block3': '4'}, # 确保这些键对应到实际的层 in_channels_list=in_channels_list, out_channels=256 ) rpn_anchor_generator = get_anchor_generator(backbone_with_fpn, test_input=test_input), return backbone_with_fpn,rpn_anchor_generator,roi_pooler def get_efficientnetv2_fpn(name='efficientnet_v2_m', pretrained=True): # 加载EfficientNetV2模型 if name == 'efficientnet_v2_s': weights = EfficientNet_V2_S_Weights.IMAGENET1K_V1 if pretrained else None backbone = efficientnet_v2_s(weights=weights).features if name == 'efficientnet_v2_m': weights = EfficientNet_V2_M_Weights.IMAGENET1K_V1 if pretrained else None backbone = efficientnet_v2_m(weights=weights).features if name == 'efficientnet_v2_l': weights = EfficientNet_V2_L_Weights.IMAGENET1K_V1 if pretrained else None backbone = efficientnet_v2_l(weights=weights).features # 定义返回的层索引和名称 return_layers = {"1":"0", "2": "1", "3": "2", "4": "3", "6": "4"} input=torch.randn(1, 3, 512, 512) # out=backbone(input) # print(f'out:{out}') # 获取每个层输出通道数 in_channels_list = [] for layer_idx in [1,2, 3, 4, 6]: module = backbone[layer_idx] # print(f'efficientnet:{backbone}') if hasattr(module, 'out_channels'): in_channels_list.append(module.out_channels) elif hasattr(module[-1], 'out_channels'): # 如果module本身没有out_channels，检查最后一个子模块 in_channels_list.append(module[-1].out_channels) else: raise ValueError(f"Cannot determine out_channels for layer {layer_idx}") # 使用BackboneWithFPN包装backbone print(f'in_channels_list: {in_channels_list}') backbone_with_fpn = BackboneWithFPN( backbone=backbone, return_layers=return_layers, in_channels_list=in_channels_list, out_channels=256 ) out=backbone_with_fpn(input) print(f'out0:{out['0'].shape}') print(f'out1:{out['1'].shape}') print(f'out2:{out['2'].shape}') print(f'out3:{out['3'].shape}') print(f'out4:{out['4'].shape}') return backbone_with_fpn # 加载 ConvNeXt 模型 # convnext = models.convnext_base(pretrained=True) # convnext = models.convnext_tiny(pretrained=True) # convnext = models.convnext_small(pretrained=True) # print(convnext) # # 打印模型的所有命名层 # for name, _ in convnext.features[5].named_children(): # print(name) # 修改 ConvNeXt 以适应 Faster R-CNN # 修改 ConvNeXt 以适应 Faster R-CNN def get_anchor_generator(backbone, test_input): features = backbone(test_input) # 获取 backbone 输出的所有特征图 featmap_names = list(features.keys()) print(f'featmap_names:{featmap_names}') num_features = len(features) # 特征图数量 print(f'num_features:{num_features}') # num_features=num_features-1 # # 定义每层的 anchor 尺寸和比例 # base_sizes = [32, 64, 128] # 支持最多 4 层 # sizes = tuple((size,) for size in base_sizes[:num_features]) 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}') return AnchorGenerator(sizes=anchor_sizes , aspect_ratios=aspect_ratios) class MaxVitBackbone(torch.nn.Module): def __init__(self,input_size=(224*7,224*7)): super(MaxVitBackbone, self).__init__() # 提取MaxVit的部分层作为特征提取器 maxvit_model =maxvit_t(pretrained=False,input_size=input_size) self.stem = maxvit_model.stem # Stem层 self.block0= maxvit_model.blocks[0] self.block1 = maxvit_model.blocks[1] self.block2 = maxvit_model.blocks[2] self.block3 = maxvit_model.blocks[3] def forward(self, x): print("Input size:", x.shape) x = self.stem(x) print("After stem size:", x.shape) x = self.block0(x) print("After block0 size:", x.shape) x = self.block1(x) print("After block1 size:", x.shape) x = self.block2(x) print("After block2 size:", x.shape) x = self.block3(x) print("After block3 size:", x.shape) return x from torchvision.models.feature_extraction import create_feature_extractor def get_swin_transformer_fpn(type='t'): class Trans(nn.Module): def __init__(self): super().__init__() def forward(self,x): x=x.permute(0, 3, 2, 1).contiguous() return x class SwinTransformer(nn.Module): def __init__(self,type='t'): super().__init__() swin = swin_v2_t(weights=None) if type=='t': # 加载 Swin Transformer v2 Tiny # swin = torchvision.models.swin_v2_t(weights=None) swin =swin_v2_t(weights=None) if type=='s': swin=torchvision.models.swin_v2_s(weights=None) if type=='b': swin=torchvision.models.swin_v2_b(weights=None) # for i,layer in enumerate(swin.named_children()): # print(f'layer{i}:,{layer}') # 保存需要提取的层 self.layer0 = swin.features[0] # 第0层 patch embedding self.layer1 =nn.Sequential(swin.features[1],Trans()) # 第1层 stage1 self.layer2 =nn.Sequential(Trans(),swin.features[2]) # 第2层 downsample self.layer3 =nn.Sequential(swin.features[3], Trans()) # 第3层 stage2 self.layer4 =nn.Sequential( Trans(),swin.features[4]) # 第4层 downsample self.layer5 =nn.Sequential(swin.features[5], Trans()) # 第5层 stage3 self.layer6 =nn.Sequential(Trans(),swin.features[6]) # 第6层 downsample self.layer7 =nn.Sequential(swin.features[7], Trans()) # 第7层 stage4 def forward(self, x): x = self.layer0(x) # [B, C, H, W] -> [B, H_, W_, C] print(f'x0:{x.shape}') x = self.layer1(x) print(f'x1:{x.shape}') x = self.layer2(x) x = self.layer3(x) print(f'x2:{x.shape}') x = self.layer4(x) x = self.layer5(x) print(f'x3:{x.shape}') x = self.layer6(x) x = self.layer7(x) print(f'x4:{x.shape}') return x backbone = SwinTransformer(type=type) input=torch.randn(1,3,512,512) out=backbone(input) # print(f'out:{out.keys()}') # for i,layer in enumerate(swin.features.named_children()): # print(f'layer:{i}:{layer}') # out=swin(input) # print(f'out shape:{out.shape}') # channels_list = [96, 192, 384, 768] if type=='t': channels_list = [96, 192, 384, 768] if type=='s': channels_list = [96, 192, 384, 768] if type=='b': channels_list = [128, 256, 512, 1024] backbone_with_fpn = BackboneWithFPN( # swin.features, backbone, return_layers={'layer1': '0', 'layer3': '1', 'layer5': '2', 'layer7': '3'}, in_channels_list=channels_list, out_channels=128 ) 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 ) # out=backbone_with_fpn(input) anchor_generator = get_anchor_generator(backbone_with_fpn, test_input=input) # print(f'out:{out}') return backbone_with_fpn,roi_pooler,anchor_generator if __name__ == '__main__': # backbone_with_fpn, roi_pooler, anchor_generator=get_swin_transformer_fpn(type='t') # model=FasterRCNN(backbone=backbone_with_fpn,num_classes=3,box_roi_pool=roi_pooler,rpn_anchor_generator=anchor_generator) # input=torch.randn(3,3,512,512,device='cuda') # model.eval() # model.to('cuda') # out=model(input) # out=backbone_with_fpn(input) # print(f'out:{out.shape}') backbone=get_efficientnetv2_fpn(name='efficientnet_v2_l') print(backbone) # # maxvit = models.maxvit_t(pretrained=True) # maxvit=MaxVitBackbone() # # print(maxvit.named_children()) # # for i,layer in enumerate(maxvit.named_children()): # print(f'layer:{i}:{layer}') # # in_channels_list = [64,64,128, 256, 512] # backbone_with_fpn = BackboneWithFPN( # maxvit, # return_layers={'stem': '0','block0':'1','block1':'2','block2':'3','block3':'4'}, # 确保这些键对应到实际的层 # in_channels_list=in_channels_list, # out_channels=256 # ) # model = FasterRCNN( # backbone=backbone_with_fpn, # num_classes=91, # COCO 数据集有 91 类 # # rpn_anchor_generator=anchor_generator, # # box_roi_pool=roi_pooler # ) # # test_input = torch.randn(1, 3, 896, 896) # # with torch.no_grad(): # output = backbone_with_fpn(test_input) # # print("Output feature maps:") # for k, v in output.items(): # print(f"{k}: {v.shape}") # model.eval() # output=model(test_input)