安全帽识别 AI算法：小目标、遮挡、暗光解决方案

from ultralytics import YOLO
import cv2
import albumentations as A

# 1. 小目标数据增强配置（Albumentations）
small_obj_aug = A.Compose([
    # 随机裁剪小目标区域并放大
    A.RandomSizedCrop(min_max_height=(200, 640), height=640, width=640, p=0.5),
    # 小目标复制粘贴增强
    A.CoarseDropout(max_holes=8, max_height=32, max_width=32, min_holes=2, 
                   min_height=8, min_width=8, fill_value=0, p=0.3),
    # 多尺度缩放，重点保留小目标
    A.MultiScaleResize(scale_range=(0.5, 2.0), p=1.0),
], bbox_params=A.BboxParams(format='yolo', label_fields=['class_labels']))

# 2. YOLOv8小目标优化配置
def train_yolov8_small_obj():
    # 加载预训练模型
    model = YOLO('yolov8s.pt')
    
    # 定制小目标锚框（安全帽常见尺寸：10-50像素）
    small_anchors = [[10,10], [20,20], [30,30], [40,40], [50,50]]
    model.model.args['anchors'] = small_anchors
    
    # 训练配置：高分辨率输入、增强小目标检测
    results = model.train(
        data='safety_hat.yaml',  # 数据集配置文件
        imgsz=1280,              # 高分辨率输入（提升小目标特征）
        epochs=100,
        batch=8,
        lr0=0.01,
        augment=True,            # 开启数据增强
        hsv_h=0.015,             # 色彩增强（辅助暗光场景）
        hsv_s=0.7,
        hsv_v=0.4,
        mixup=0.1,               # 混合增强
        copy_paste=0.3,          # 小目标复制粘贴增强
        device=0,                # GPU加速
    )
    return results

# 3. 推理时小目标后处理优化
def detect_small_obj(image_path, model):
    # 读取图像
    img = cv2.imread(image_path)
    # 多尺度推理（提升小目标召回率）
    results = model(img, imgsz=[640, 1280], conf=0.25, iou=0.45)
    
    # 小目标置信度修正（避免漏检）
    for r in results:
        boxes = r.boxes
        for box in boxes:
            # 小目标（面积<1000像素）提升置信度
            area = (box.xyxy[0][2] - box.xyxy[0][0]) * (box.xyxy[0][3] - box.xyxy[0][1])
            if area < 1000:
                box.conf = box.conf * 1.2  # 置信度提升20%
    return results

import torch
import torch.nn as nn

# 1. CBAM注意力模块（增强可见区域特征）
class CBAM(nn.Module):
    def __init__(self, channel, reduction=16):
        super(CBAM, self).__init__()
        # 通道注意力
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.max_pool = nn.AdaptiveMaxPool2d(1)
        self.fc = nn.Sequential(
            nn.Linear(channel, channel // reduction),
            nn.ReLU(inplace=True),
            nn.Linear(channel // reduction, channel)
        )
        # 空间注意力
        self.spatial = nn.Sequential(
            nn.Conv2d(2, 1, 7, padding=3, bias=False),
            nn.Sigmoid()
        )
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        # 通道注意力
        avg_out = self.fc(self.avg_pool(x).view(x.size(0), -1)).view(x.size(0), -1, 1, 1)
        max_out = self.fc(self.max_pool(x).view(x.size(0), -1)).view(x.size(0), -1, 1, 1)
        channel_att = self.sigmoid(avg_out + max_out)
        x = x * channel_att
        
        # 空间注意力
        avg_out = torch.mean(x, dim=1, keepdim=True)
        max_out, _ = torch.max(x, dim=1, keepdim=True)
        spatial_att = self.spatial(torch.cat([avg_out, max_out], dim=1))
        x = x * spatial_att
        return x

# 2. 结合人体检测的遮挡安全帽识别
def detect_occluded_hat(image_path, hat_model, person_model):
    # 第一步：检测人体
    person_results = person_model(image_path)
    person_boxes = person_results[0].boxes
    
    # 第二步：在人体头部区域内检测安全帽
    hat_detections = []
    for person_box in person_boxes:
        # 提取人体头部区域（上1/3部分）
        x1, y1, x2, y2 = person_box.xyxy[0]
        head_y1 = y1
        head_y2 = y1 + (y2 - y1) * 0.3
        head_region = img[int(head_y1):int(head_y2), int(x1):int(x2)]
        
        # 在头部区域内检测安全帽
        if head_region.size > 0:
            hat_results = hat_model(head_region, conf=0.2)
            # 映射回原图坐标
            for hat_box in hat_results[0].boxes:
                hx1, hy1, hx2, hy2 = hat_box.xyxy[0]
                hat_box.xyxy[0] = [x1+hx1, head_y1+hy1, x1+hx2, head_y1+hy2]
                hat_detections.append(hat_box)
    
    return hat_detections

import cv2
import numpy as np

# 1. 暗光图像预处理（Retinex算法）
def retinex_enhance(image, sigma_list=[15, 80, 250]):
    """Retinex算法：提升暗光图像的对比度和细节"""
    img = image.astype(np.float64) + 1e-6  # 避免除零
    
    # 多尺度高斯滤波
    retinex = np.zeros_like(img)
    for sigma in sigma_list:
        gaussian = cv2.GaussianBlur(img, (0, 0), sigma)
        retinex += np.log10(img) - np.log10(gaussian)
    
    retinex = retinex / len(sigma_list)
    # 归一化
    retinex = (retinex - np.min(retinex)) / (np.max(retinex) - np.min(retinex)) * 255
    retinex = np.clip(retinex, 0, 255).astype(np.uint8)
    
    # CLAHE增强对比度
    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
    if len(retinex.shape) == 3:
        retinex_yuv = cv2.cvtColor(retinex, cv2.COLOR_BGR2YUV)
        retinex_yuv[:,:,0] = clahe.apply(retinex_yuv[:,:,0])
        retinex = cv2.cvtColor(retinex_yuv, cv2.COLOR_YUV2BGR)
    else:
        retinex = clahe.apply(retinex)
    
    return retinex

# 2. 暗光场景推理流程
def detect_low_light_hat(image_path, model):
    # 读取暗光图像
    img = cv2.imread(image_path)
    # 暗光增强预处理
    enhanced_img = retinex_enhance(img)
    # 模型推理
    results = model(enhanced_img, conf=0.2, iou=0.45)
    return results, enhanced_img

# 3. 训练时的暗光数据增强
def low_light_augmentation(image):
    """模拟暗光场景，用于训练数据增强"""
    # 随机降低亮度
    brightness = np.random.uniform(0.1, 0.7)
    image = cv2.convertScaleAbs(image, alpha=1, beta=brightness*255 - 255)
    # 添加噪声（模拟暗光下的传感器噪声）
    noise = np.random.normal(0, np.random.uniform(5, 20), image.shape)
    image = np.clip(image + noise, 0, 255).astype(np.uint8)
    return image