# import torch
# import torch.nn as nn
# import torchvision.models as models
#
#
# class ResNet50Backbone(nn.Module):
#     def __init__(self, num_classes=5, num_stacks=1, pretrained=True):
#         super(ResNet50Backbone, self).__init__()
#
#         # 加载预训练的ResNet50
#         if pretrained:
#             resnet = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V1)
#         else:
#             resnet = models.resnet50(weights=None)
#
#         # 移除最后的全连接层
#         self.backbone = nn.Sequential(
#             resnet.conv1,#特征图分辨率降低为1/2，通道数从3升为64
#             resnet.bn1,
#             resnet.relu,
#             resnet.maxpool,#特征图分辨率降低为1/4，通道数仍然为64
#             resnet.layer1,#stride为1，不改变分辨率，依然为1/4，通道数从64升为64*4=256
#             # resnet.layer2,#stride为2，特征图分辨率降低为1/8，通道数从256升为128*4=512
#             # resnet.layer3,#stride为2，特征图分辨率降低为1/16，通道数从512升为256*4=1024
#             # resnet.layer4,#stride为2，特征图分辨率降低为1/32，通道数从512升为256*4=2048
#         )
#
#         # 多任务输出层
#         self.score_layers = nn.ModuleList([
#             nn.Sequential(
#                 nn.Conv2d(256, 128, kernel_size=3, padding=1),
#                 nn.BatchNorm2d(128),
#                 nn.ReLU(inplace=True),
#                 nn.Conv2d(128, num_classes, kernel_size=1)
#             )
#             for _ in range(num_stacks)
#         ])
#
#         # 上采样层，确保输出大小为128x128
#         self.upsample = nn.Upsample(
#             scale_factor=0.25,
#             mode='bilinear',
#             align_corners=True
#         )
#
#     def forward(self, x):
#         # 主干网络特征提取
#         x = self.backbone(x)
#
#         # # 调整通道数
#         # x = self.channel_adjust(x)
#         #
#         # # 上采样到128x128
#         # x = self.upsample(x)
#
#         # 多堆栈输出
#         outputs = []
#         for score_layer in self.score_layers:
#             output = score_layer(x)
#             outputs.append(output)
#
#         # 返回第一个输出（如果有多个堆栈）
#         return outputs, x
#
#
# def resnet50(**kwargs):
#     model = ResNet50Backbone(
#         num_classes=kwargs.get("num_classes", 5),
#         num_stacks=kwargs.get("num_stacks", 1),
#         pretrained=kwargs.get("pretrained", True)
#     )
#     return model
#
#
# __all__ = ["ResNet50Backbone", "resnet50"]
#
#
# # 测试网络输出
# model = resnet50(num_classes=5, num_stacks=1)
#
#
# # 方法1：直接传入图像张量
# x = torch.randn(2, 3, 512, 512)
# outputs, feature = model(x)
# print("Outputs length:", len(outputs))
# print("Output[0] shape:", outputs[0].shape)
# print("Feature shape:", feature.shape)

import torch
import torch.nn as nn
import torchvision
import torchvision.models as models
from torchvision.models.detection.backbone_utils import _validate_trainable_layers, _resnet_fpn_extractor


class ResNet50Backbone1(nn.Module):
    def __init__(self, num_classes=5, num_stacks=1, pretrained=True):
        super(ResNet50Backbone1, self).__init__()

        # 加载预训练的ResNet50
        if pretrained:
            # self.resnet = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V1)
            trainable_backbone_layers = _validate_trainable_layers(True, None, 5, 3)
            backbone = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V1)
            self.resnet = _resnet_fpn_extractor(backbone, trainable_backbone_layers)
        else:
            self.resnet = models.resnet50(weights=None)

    def forward(self, x):
        features = self.resnet(x)

        # if self.training:
        #     proposals, matched_idxs, labels, regression_targets = self.select_training_samples20(proposals, targets)
        # else:
        #     labels = None
        #     regression_targets = None
        #     matched_idxs = None
        # box_features = self.box_roi_pool(features, proposals, image_shapes)  # ROI 映射到固定尺寸的特征表示上
        # # print(f"box_features:{box_features.shape}")  # [2048, 256, 7, 7] 建议框统一大小
        # box_features = self.box_head(box_features)  #
        # # print(f"box_features:{box_features.shape}")  # [N, 1024] 经过头部网络处理后的特征向量
        # class_logits, box_regression = self.box_predictor(box_features)
        #
        # result: List[Dict[str, torch.Tensor]] = []
        # losses = {}
        # if self.training:
        #     if labels is None:
        #         raise ValueError("labels cannot be None")
        #     if regression_targets is None:
        #         raise ValueError("regression_targets cannot be None")
        #     loss_classifier, loss_box_reg = fastrcnn_loss(class_logits, box_regression, labels, regression_targets)
        #     losses = {"loss_classifier": loss_classifier, "loss_box_reg": loss_box_reg}
        # else:
        #     boxes, scores, labels = self.postprocess_detections(class_logits, box_regression, proposals, image_shapes)
        #     num_images = len(boxes)
        #     for i in range(num_images):
        #         result.append(
        #             {
        #                 "boxes": boxes[i],
        #                 "labels": labels[i],
        #                 "scores": scores[i],
        #             }
        #         )
        #     print(f"boxes:{boxes[0].shape}")

        return x['0']

        # # 多堆栈输出
        # outputs = []
        # for score_layer in self.score_layers:
        #     output = score_layer(x)
        #     outputs.append(output)
        #
        # # 返回第一个输出（如果有多个堆栈）
        # return outputs, x


def resnet501(**kwargs):
    model = ResNet50Backbone1(
        num_classes=kwargs.get("num_classes", 5),
        num_stacks=kwargs.get("num_stacks", 1),
        pretrained=kwargs.get("pretrained", True)
    )
    return model


# __all__ = ["ResNet50Backbone1", "resnet501"]
#
# # 测试网络输出
# model = resnet501(num_classes=5, num_stacks=1)
#
# # 方法1：直接传入图像张量
# x = torch.randn(2, 3, 512, 512)
# # outputs, feature = model(x)
# # print("Outputs length:", len(outputs))
# # print("Output[0] shape:", outputs[0].shape)
# # print("Feature shape:", feature.shape)
# feature = model(x)
# print("Feature shape:", feature.keys())