#这是一个RCNN骨架的搭建 我的水平真的很拉 但是也得逼自己写一下看能不能写出来
import torch,torchvision
from torch import nn
from torchvision.models.detection.image_list import ImageList
from torchvision.models.detection.anchor_utils import AnchorGenerator
#需要传入的参数有 图片 X（batch，3,800,600）
class FasteRCNN(nn.Module):
    def __init__(self,anchor_sizes,aspect_ratios,ntype):
        super(FasteRCNN, self).__init__()
        self.ntype=ntype
        resnet = torchvision.models.resnet50(pretrained=True)
        self.backbone = nn.Sequential(*list(resnet.children())[:-2])  # 去掉avgpool和fc
        self.anchors_size=9
        self.anchor_generator = AnchorGenerator(anchor_sizes, aspect_ratios)
        self.RPN = RegionProposalNetwork(self.anchors_size)
        self.roihead = ROIhead(self.ntype)
    def forward(self, X):
        features = self.backbone(X)
        # 现在的大小就是(batch 2048 19 25)了
        image_list = ImageList(X, [(X.shape[2], X.shape[3])] * X.shape[0])
        anchors = self.anchor_generator(image_list, [features])
        predbias,predclassify=self.RPN(features)
        proposals = generate_proposals(anchors, predclassify, predbias, image_list.image_sizes)
        return self.roihead(features,proposals)

class RegionProposalNetwork(nn.Module):
    #anchors:anchorbox个数
    def __init__(self,anchors):
        super(RegionProposalNetwork, self).__init__()
        #先做一个3x3的conv
        self.conv1 = nn.Conv2d(2048, 2048, kernel_size=3, padding=1)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = nn.Conv2d(2048,4*anchors , kernel_size=1)
        self.conv3 = nn.Conv2d(2048, 2*anchors, kernel_size=1)
    def forward(self,features):
        features = self.relu(self.conv1(features))
        #第一个分支生成偏移量
        predbias = self.conv2(features)
        #第二个分支判断是否为背景
        predclassify = self.conv3(features)
        B, _, H, W = predbias.shape
        predbias = predbias.permute(0, 2, 3, 1).reshape(B, -1, 4)
        predclassify = predclassify.permute(0, 2, 3, 1).reshape(B, -1, 2)
        return predbias, predclassify

from torchvision.ops import nms


def generate_proposals(anchors, predclassify, predbias, image_sizes):
    """
    anchors: List[Tensor]，每个特征图对应一个 Tensor，形状 [A_i*H_i*W_i, 4]
    predclassify: Tensor [B, total_anchors, 2] （尚未 softmax）
    predbias: Tensor [B, total_anchors, 4]
    image_sizes: List[Tuple[int, int]]，每张图片的 (H, W)
    """
    # 1. 将所有层级的 anchor 合并成一个大的 Tensor

    # 2. 对 predclassify 做 softmax，取前景分数（索引1）
    scores = torch.softmax(predclassify, dim=-1)[:, :, 1]  # [B, total_anchors]

    batch_size = predclassify.shape[0]
    proposals = []

    for i in range(batch_size):
        # 当前图片的 anchor 和预测值
        cur_anchors = anchors[0]
        deltas = predbias[i]  # [total_anchors, 4]
        score_i = scores[i]  # [total_anchors]

        # 3. 将偏移量应用到 anchor 上（中心点格式转换）
        # anchor 格式: (x1, y1, x2, y2)
        width = cur_anchors[:, 2] - cur_anchors[:, 0]
        height = cur_anchors[:, 3] - cur_anchors[:, 1]
        ctr_x = cur_anchors[:, 0] + 0.5 * width
        ctr_y = cur_anchors[:, 1] + 0.5 * height

        dx = deltas[:, 0]
        dy = deltas[:, 1]
        dw = deltas[:, 2]
        dh = deltas[:, 3]

        pred_ctr_x = dx * width + ctr_x
        pred_ctr_y = dy * height + ctr_y
        pred_w = torch.exp(dw) * width
        pred_h = torch.exp(dh) * height

        # 转回 (x1, y1, x2, y2)
        pred_x1 = pred_ctr_x - 0.5 * pred_w
        pred_y1 = pred_ctr_y - 0.5 * pred_h
        pred_x2 = pred_ctr_x + 0.5 * pred_w
        pred_y2 = pred_ctr_y + 0.5 * pred_h

        boxes = torch.stack([pred_x1, pred_y1, pred_x2, pred_y2], dim=1)

        # 4. 裁剪到图像边界
        H, W = image_sizes[i]
        boxes[:, 0] = torch.clamp(boxes[:, 0], min=0, max=W)
        boxes[:, 1] = torch.clamp(boxes[:, 1], min=0, max=H)
        boxes[:, 2] = torch.clamp(boxes[:, 2], min=0, max=W)
        boxes[:, 3] = torch.clamp(boxes[:, 3], min=0, max=H)

        # 5. 剔除宽高 <= 0 的框
        keep = (boxes[:, 2] > boxes[:, 0]) & (boxes[:, 3] > boxes[:, 1])
        boxes = boxes[keep]
        score_i = score_i[keep]

        # 6. 按分数排序取 top-N（训练时取 12000，测试取 6000）
        pre_nms_top_n = 12000  # 可以根据 self.training 调整，但函数里无法获取，可以传入参数
        if len(score_i) > pre_nms_top_n:
            topk = torch.topk(score_i, pre_nms_top_n)
            boxes = boxes[topk.indices]
            score_i = score_i[topk.indices]

        # 7. NMS（阈值 0.7）
        keep = nms(boxes, score_i, 0.7)
        boxes = boxes[keep]
        score_i = score_i[keep]

        # 8. 最终取 top-N（训练取 2000，测试取 1000）
        post_nms_top_n = 2000
        if len(score_i) > post_nms_top_n:
            topk = torch.topk(score_i, post_nms_top_n)
            boxes = boxes[topk.indices]
            # score_i 可以丢弃，proposal 只需框
        proposals.append(boxes)

    return proposals  # List[Tensor]，每个 Tensor 形状 [N_i, 4]
import torchvision.ops.roi_align as roi_align
class ROIhead(nn.Module):
    def __init__(self,ntype):
        super(ROIhead, self).__init__()
        self.conv1 = nn.Conv2d(2048, 512, kernel_size=3,padding=1)
        self.fc1 = nn.Linear(512*7*7,4096)
        self.fc2 = nn.Linear(4096,4096)
        self.fc3_1 = nn.Linear(4096,ntype+1)
        self.fc3_2 = nn.Linear(4096,(ntype+1)*4)
        self.relu = nn.ReLU()
    def forward(self,features ,proposals):
        X=roi_align(features,proposals,output_size=7,spatial_scale=1.0/32)
        X=self.conv1(X)
        #(B,512,7,7)->(B,512*7*7)
        X=X.flatten(start_dim=1,end_dim=3)
        X = self.relu(self.fc2(self.relu(self.fc1(X))))
        return self.fc3_1(X),self.fc3_2(X).reshape(features.shape[0],-1,4)
model = FasteRCNN(anchor_sizes = ((32, 64, 128),),aspect_ratios = ((0.5, 1.0, 2.0),),ntype=20)