盘算机视觉算法实现——浮水衣穿着状态智能识别 - IT评测·应用市场-qidao123.com技术社区

import torchvision
from torchvision.models.detection import FasterRCNN
from torchvision.models.detection.rpn import AnchorGenerator
def get_faster_rcnn_model(num_classes):
# 加载预训练的骨干网络
backbone = torchvision.models.mobilenet_v2(pretrained=True).features
backbone.out_channels = 1280
# 定义锚点生成器
anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512),),
aspect_ratios=((0.5, 1.0, 2.0),))
# 定义ROI pooling
roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=['0'],
output_size=7,
sampling_ratio=2)
# 构建Faster R-CNN模型
model = FasterRCNN(backbone,
num_classes=num_classes,
rpn_anchor_generator=anchor_generator,
box_roi_pool=roi_pooler)
return model

复制代码

import torch
from models.experimental import attempt_load
def load_yolov5_model(weights_path):
# 加载预训练的YOLOv5模型
model = attempt_load(weights_path, map_location=torch.device('cpu'))
return model

复制代码

import torch
import torch.nn as nn
class LifeJacketKeypointModel(nn.Module):
def __init__(self, num_keypoints):
super().__init__()
self.backbone = torchvision.models.resnet18(pretrained=True)
self.deconv_layers = self._make_deconv_layer()
self.final_layer = nn.Conv2d(256, num_keypoints, kernel_size=1)
def _make_deconv_layer(self):
layers = []
layers.append(nn.ConvTranspose2d(512, 256, kernel_size=4, stride=2, padding=1))
layers.append(nn.ReLU(inplace=True))
return nn.Sequential(*layers)
def forward(self, x):
x = self.backbone.conv1(x)
x = self.backbone.bn1(x)
x = self.backbone.relu(x)
x = self.backbone.maxpool(x)
x = self.backbone.layer1(x)
x = self.backbone.layer2(x)
x = self.backbone.layer3(x)
x = self.backbone.layer4(x)
x = self.deconv_layers(x)
x = self.final_layer(x)
return x

复制代码

from albumentations import (
Compose, RandomBrightnessContrast, HueSaturationValue,
MotionBlur, RandomRain, HorizontalFlip, Rotate
)
train_transform = Compose([
HorizontalFlip(p=0.5),
RandomBrightnessContrast(p=0.3),
HueSaturationValue(hue_shift_limit=20, sat_shift_limit=30, val_shift_limit=20, p=0.5),
MotionBlur(blur_limit=7, p=0.3),
RandomRain(drop_length=5, blur_value=3, p=0.1),
Rotate(limit=15, p=0.5)
])

复制代码

import torch
import cv2
import numpy as np
from pathlib import Path
import argparse
class LifeJacketDetector:
def __init__(self, weights_path, conf_thres=0.5, iou_thres=0.45):
self.model = self.load_model(weights_path)
self.conf_thres = conf_thres
self.iou_thres = iou_thres
self.classes = ['no_lifejacket', 'proper_worn', 'improper_worn']
self.colors = [(0, 0, 255), (0, 255, 0), (255, 0, 0)] # 红:未穿, 绿:正确, 蓝:错误
def load_model(self, weights_path):
model = torch.hub.load('ultralytics/yolov5', 'custom', path=weights_path)
return model
def detect(self, img):
# 推理
results = self.model(img)
pred = results.pred[0]
# 后处理
detections = []
for *xyxy, conf, cls in pred:
if conf > self.conf_thres:
x1, y1, x2, y2 = map(int, xyxy)
class_id = int(cls)
label = f"{self.classes[class_id]}: {conf:.2f}"
detections.append({
'bbox': [x1, y1, x2, y2],
'class_id': class_id,
'confidence': float(conf),
'label': label
})
return detections
def draw_detections(self, img, detections):
for det in detections:
x1, y1, x2, y2 = det['bbox']
class_id = det['class_id']
label = det['label']
# 绘制边界框
color = self.colors[class_id]
cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)
# 绘制标签背景
(w, h), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 1)
cv2.rectangle(img, (x1, y1 - 20), (x1 + w, y1), color, -1)
# 绘制标签文本
cv2.putText(img, label, (x1, y1 - 5),
cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255,255,255), 1)
return img
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--source', type=str, default='0', help='视频源，0为摄像头')
parser.add_argument('--weights', type=str, required=True, help='模型权重路径')
args = parser.parse_args()
# 初始化检测器
detector = LifeJacketDetector(args.weights)
# 初始化视频源
cap = cv2.VideoCapture(0 if args.source == '0' else args.source)
while True:
ret, frame = cap.read()
if not ret:
break
# 检测救生衣
detections = detector.detect(frame)
# 绘制结果
frame = detector.draw_detections(frame, detections)
# 显示
cv2.imshow('Life Jacket Detection', frame)
if cv2.waitKey(1) & 0xFF == ord('q'):
break
cap.release()
cv2.destroyAllWindows()
if __name__ == '__main__':
main()

复制代码