基于深度学习的目标检测（Deep Learning-based Object Detection）

pythonCopy codeimport torch
import torchvision
from torchvision.models.detection import FasterRCNN
from torchvision.models.detection.rpn import AnchorGenerator
# 加载预训练的Faster R-CNN模型
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
# 替换模型的分类器
num_classes = 2  # 假设我们要检测的目标有两类
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = torchvision.models.detection.faster_rcnn.FastRCNNPredictor(in_features, num_classes)
# 定义数据预处理和数据加载器
transform = torchvision.transforms.Compose([
    torchvision.transforms.ToTensor()
])
dataset = torchvision.datasets.CocoDetection(root='data/', annFile='annotations.json', transform=transform)
data_loader = torch.utils.data.DataLoader(dataset, batch_size=4, shuffle=True, num_workers=2)
# 定义优化器和损失函数
params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params, lr=0.005, momentum=0.9, weight_decay=0.0005)
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
model.to(device)
# 训练模型
num_epochs = 10
for epoch in range(num_epochs):
    model.train()
    total_loss = 0
    for images, targets in data_loader:
        images = list(image.to(device) for image in images)
        targets = [{k: v.to(device) for k, v in t.items()} for t in targets]
        loss_dict = model(images, targets)
        losses = sum(loss for loss in loss_dict.values())
        optimizer.zero_grad()
        losses.backward()
        optimizer.step()
        total_loss += losses.item()
    lr_scheduler.step()
    
    print(f'Epoch {epoch+1}/{num_epochs}, Loss: {total_loss/len(data_loader)}')
# 使用模型进行目标检测
model.eval()
images, targets = next(iter(data_loader))
images = list(image.to(device) for image in images)
targets = [{k: v.to(device) for k, v in t.items()} for t in targets]
with torch.no_grad():
    predictions = model(images)
# 输出预测结果
print(predictions)

pythonCopy codeimport cv2
import numpy as np
# 加载YOLO模型和类别标签
net = cv2.dnn.readNetFromDarknet('yolov3.cfg', 'yolov3.weights')
classes = []
with open('coco.names', 'r') as f:
    classes = [line.strip() for line in f.readlines()]
# 加载图像
image = cv2.imread('image.jpg')
height, width, _ = image.shape
# 图像预处理
blob = cv2.dnn.blobFromImage(image, 1/255.0, (416, 416), swapRB=True, crop=False)
net.setInput(blob)
# 运行前向传播
outputs = net.forward(net.getUnconnectedOutLayersNames())
# 解析检测结果
boxes = []
confidences = []
class_ids = []
for output in outputs:
    for detection in output:
        scores = detection[5:]
        class_id = np.argmax(scores)
        confidence = scores[class_id]
        if confidence > 0.5:
            center_x = int(detection[0] * width)
            center_y = int(detection[1] * height)
            w = int(detection[2] * width)
            h = int(detection[3] * height)
            x = int(center_x - w/2)
            y = int(center_y - h/2)
            boxes.append([x, y, w, h])
            confidences.append(float(confidence))
            class_ids.append(class_id)
# 应用非极大值抑制
indices = cv2.dnn.NMSBoxes(boxes, confidences, 0.5, 0.4)
# 绘制边界框和标签
font = cv2.FONT_HERSHEY_SIMPLEX
for i in indices:
    i = i[0]
    box = boxes[i]
    x, y, w, h = box
    label = f'{classes[class_ids[i]]}: {confidences[i]:.2f}'
    color = (0, 255, 0)
    cv2.rectangle(image, (x, y), (x+w, y+h), color, 2)
    cv2.putText(image, label, (x, y-10), font, 0.5, color, 2)
# 显示结果图像
cv2.imshow('Image', image)
cv2.waitKey(0)
cv2.destroyAllWindows()

pythonCopy codeimport torch
import torchvision
from torchvision.models.detection import ssdlite320_mobilenet_v3_large
# 加载SSD模型和类别标签
model = ssdlite320_mobilenet_v3_large(pretrained=True)
classes = [
    'background', 'person', 'bicycle', 'car', 'motorcycle',
    'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light',
    'fire hydrant', 'N/A', 'stop sign', 'parking meter', 'bench',
    'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant',
    'bear', 'zebra', 'giraffe', 'N/A', 'backpack', 'umbrella',
    'N/A', 'N/A', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis',
    'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove',
    'skateboard', 'surfboard', 'tennis racket', 'bottle', 'N/A',
    'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl',
    'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot',
    'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
    'potted plant', 'bed', 'N/A', 'dining table', 'N/A', 'N/A',
    'toilet', 'N/A', 'tv', 'laptop', 'mouse', 'remote',
    'keyboard', 'cell phone', 'microwave', 'oven', 'toaster',
    'sink', 'refrigerator', 'N/A', 'book', 'clock', 'vase',
    'scissors', 'teddy bear', 'hair drier', 'toothbrush'
]
# 加载图像
image = Image.open('image.jpg')
# 图像预处理
transform = torchvision.transforms.Compose([
    torchvision.transforms.ToTensor()
])
image_tensor = transform(image)
image_tensor = torch.unsqueeze(image_tensor, 0)
# 运行模型
model.eval()
with torch.no_grad():
    predictions = model(image_tensor)
# 解析检测结果
boxes = predictions[0]['boxes'].tolist()
scores = predictions[0]['scores'].tolist()
class_ids = predictions[0]['labels'].tolist()
# 绘制边界框和标签
draw = ImageDraw.Draw(image)
font = ImageFont.truetype('arial.ttf', size=12)
for box, score, class_id in zip(boxes, scores, class_ids):
    if score > 0.5:
        x, y, w, h = box
        label = f'{classes[class_id]}: {score:.2f}'
        draw.rectangle([(x, y), (x + w, y + h)], outline='red')
        draw.text((x, y), label, fill='red', font=font)
# 显示结果图像
image.show()