146 lines
6.0 KiB
Python
146 lines
6.0 KiB
Python
#这是一个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)
|