faster rcnn目标检测代码

以下是使用PyTorch实现Faster R-CNN目标检测的示例代码：

import torch
import torchvision
import torchvision.transforms as transforms
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torch.autograd import Variable

from collections import OrderedDict
import numpy as np
import cv2
import os

# 定义超参数
learning_rate = 0.001
num_epochs = 5
batch_size = 16

# 定义数据预处理
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
])

# 加载COCO数据集
train_data = torchvision.datasets.CocoDetection(root='./data', annFile='./data/annotations/instances_train2017.json', transform=transform)
train_loader = DataLoader(train_data, batch_size=batch_size, shuffle=True)

# 定义模型
class FasterRCNN(nn.Module):
    def __init__(self):
        super(FasterRCNN, self).__init__()
        self.features = nn.Sequential(OrderedDict([
            ('conv1', nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1)),
            ('relu1', nn.ReLU(inplace=True)),
            ('conv2', nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1)),
            ('relu2', nn.ReLU(inplace=True)),
            ('conv3', nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1)),
            ('relu3', nn.ReLU(inplace=True)),
            ('conv4', nn.Conv2d(256, 512, kernel_size=3, stride=1, padding=1)),
            ('relu4', nn.ReLU(inplace=True)),
            ('conv5', nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1)),
            ('relu5', nn.ReLU(inplace=True)),
            ('maxpool', nn.MaxPool2d(kernel_size=2, stride=2))
        ]))

        self.roi_pool = nn.AdaptiveMaxPool2d((7, 7))

        self.head = nn.Sequential(OrderedDict([
            ('fc6', nn.Linear(7 * 7 * 512, 4096)),
            ('relu6', nn.ReLU(inplace=True)),
            ('drop6', nn.Dropout()),
            ('fc7', nn.Linear(4096, 4096)),
            ('relu7', nn.ReLU(inplace=True)),
            ('drop7', nn.Dropout())
        ]))

        self.cls_score = nn.Linear(4096, 80)
        self.bbox_pred = nn.Linear(4096, 320)

    def forward(self, x, rois):
        x = self.features(x)
        x = self.roi_pool(x)
        x = self.head(x.view(x.size(0), -1))
        cls_score = self.cls_score(x)
        bbox_pred = self.bbox_pred(x)

        return cls_score, bbox_pred

# 定义损失函数和优化器
model = FasterRCNN()
criterion_cls = nn.CrossEntropyLoss()
criterion_bbox = nn.SmoothL1Loss()
optimizer = optim.SGD(model.parameters(), lr=learning_rate, momentum=0.9)

# 训练模型
for epoch in range(num_epochs):
    for i, (images, targets) in enumerate(train_loader):
        images = Variable(images)
        targets = [{k: Variable(v) for k, v in t.items()} for t in targets]

        # 前向传播
        cls_scores, bbox_preds = model(images, targets)
        cls_scores = torch.cat(cls_scores, dim=0)
        bbox_preds = torch.cat(bbox_preds, dim=0)

        cls_targets = torch.cat([t['labels'] for t in targets])
        bbox_targets = torch.cat([t['boxes'] for t in targets])

        # 计算损失
        loss_cls = criterion_cls(cls_scores, cls_targets)
        loss_bbox = criterion_bbox(bbox_preds, bbox_targets)
        loss = loss_cls + loss_bbox

        # 反向传播和优化
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if (i + 1) % 100 == 0:
            print("Epoch [{}/{}], Iteration [{}/{}], Loss: {:.4f}"
                  .format(epoch + 1, num_epochs, i + 1, len(train_loader), loss.item()))

# 保存模型
torch.save(model.state_dict(), 'faster_rcnn.pth')

注意：以上代码仅是一个示例，实际使用时可能需要进行适当的修改和调整。