反光衣识别 AI算法：多模态融合，夜间/强光精准检测

import torch
import torch.nn as nn
import torch.nn.functional as F
import cv2
import numpy as np

# 1. 光照预处理模块：解决夜间/强光下的图像归一化
class LightNormalization(nn.Module):
    def forward(self, x):
        """
        输入：x - RGB图像张量，shape [B, 3, H, W]，取值范围0-1
        输出：光照归一化后的图像张量
        """
        # 全局亮度归一化
        mean = torch.mean(x, dim=[2, 3], keepdim=True)
        std = torch.std(x, dim=[2, 3], keepdim=True)
        x_normalized = (x - mean) / (std + 1e-8)
        
        # 自适应伽马校正：强光下调低伽马，夜间上调伽马
        gamma = torch.clamp(1.0 / (mean.mean() + 1e-8), 0.5, 2.0)
        x_gamma = torch.pow(torch.clamp(x, 0, 1), gamma)
        
        # 融合归一化和伽马校正结果
        return 0.7 * x_normalized + 0.3 * x_gamma

# 2. 反光区域特征提取模块（光反射模态）
class ReflectiveFeatureExtractor(nn.Module):
    def __init__(self, in_channels=3, out_channels=64):
        super().__init__()
        # 提取高光区域：基于亮度阈值+边缘检测
        self.conv1 = nn.Conv2d(in_channels, 32, 3, padding=1)
        self.conv2 = nn.Conv2d(32, out_channels, 3, padding=1)
        self.relu = nn.ReLU()

    def forward(self, x):
        """
        输入：x - RGB图像张量 [B, 3, H, W]
        输出：反光区域特征 [B, 64, H, W]
        """
        # 转换为灰度图，提取亮度特征
        gray = torch.mean(x, dim=1, keepdim=True)  # [B, 1, H, W]
        # 高亮区域掩码（反光衣的核心特征）
        highlight_mask = (gray > torch.quantile(gray, 0.95, dim=[2,3], keepdim=True)).float()
        
        # 卷积提取反光区域的纹理/边缘特征
        x = torch.cat([x, gray, highlight_mask], dim=1)  # 融合亮度+高亮掩码
        x = self.relu(self.conv1(x))
        reflective_feat = self.relu(self.conv2(x))
        return reflective_feat

# 3. 多模态融合模块（特征级融合）
class MultimodalFusion(nn.Module):
    def __init__(self, rgb_channels=256, reflective_channels=64, fusion_channels=256):
        super().__init__()
        # 对齐通道数
        self.reflective_proj = nn.Conv2d(reflective_channels, fusion_channels, 1)
        self.rgb_proj = nn.Conv2d(rgb_channels, fusion_channels, 1)
        # 注意力融合：自适应权重分配
        self.attention = nn.Sequential(
            nn.Conv2d(fusion_channels * 2, fusion_channels, 1),
            nn.ReLU(),
            nn.Conv2d(fusion_channels, 2, 1),  # 输出RGB和反光特征的权重
            nn.Softmax(dim=1)
        )

    def forward(self, rgb_feat, reflective_feat):
        """
        输入：
            rgb_feat: RGB分支特征 [B, 256, H, W]
            reflective_feat: 反光分支特征 [B, 64, H, W]
        输出：融合后的特征 [B, 256, H, W]
        """
        # 通道对齐
        rgb_proj = self.rgb_proj(rgb_feat)
        reflective_proj = self.reflective_proj(reflective_feat)
        
        # 注意力权重计算
        concat_feat = torch.cat([rgb_proj, reflective_proj], dim=1)
        weights = self.attention(concat_feat)  # [B, 2, H, W]
        
        # 加权融合
        fusion_feat = weights[:, 0:1, :, :] * rgb_proj + weights[:, 1:2, :, :] * reflective_proj
        return fusion_feat

# 4. 完整的反光衣检测网络
class ReflectiveVestDetector(nn.Module):
    def __init__(self):
        super().__init__()
        # 1. 预处理模块
        self.light_norm = LightNormalization()
        # 2. RGB分支（基于ResNet18的特征提取）
        from torchvision.models import resnet18
        self.rgb_backbone = nn.Sequential(*list(resnet18(pretrained=True).children())[:-3])
        # 3. 反光特征分支
        self.reflective_backbone = ReflectiveFeatureExtractor()
        # 4. 多模态融合
        self.fusion = MultimodalFusion(rgb_channels=256, reflective_channels=64, fusion_channels=256)
        # 5. 检测头（简化版，实际可替换为YOLO/RetinaNet等检测头）
        self.detection_head = nn.Sequential(
            nn.Conv2d(256, 128, 3, padding=1),
            nn.ReLU(),
            nn.Conv2d(128, 2, 1)  # 输出：类别（反光衣/背景）+ 置信度
        )

    def forward(self, x):
        """
        输入：x - RGB图像张量 [B, 3, H, W]，取值范围0-1
        输出：检测结果 [B, 2, H/8, W/8]（下采样8倍）
        """
        # 步骤1：光照预处理
        x_norm = self.light_norm(x)
        # 步骤2：双分支特征提取
        rgb_feat = self.rgb_backbone(x_norm)  # [B, 256, H/8, W/8]
        reflective_feat = self.reflective_backbone(x_norm)
        # 下采样对齐尺寸
        reflective_feat = F.interpolate(reflective_feat, size=rgb_feat.shape[2:], mode='bilinear')
        # 步骤3：多模态融合
        fusion_feat = self.fusion(rgb_feat, reflective_feat)
        # 步骤4：检测
        output = self.detection_head(fusion_feat)
        return output

# ------------------- 测试代码 -------------------
if __name__ == "__main__":
    # 1. 构建模型
    model = ReflectiveVestDetector()
    model.ｅｖａｌ()  # 推理模式
    
    # 2. 模拟输入（批量1，3通道，640x640）
    # 实际使用时可替换为真实图像：cv2.imread -> 转RGB -> 归一化 -> 转张量
    fake_input = torch.randn(1, 3, 640, 640) / 255.0  # 模拟0-1范围的图像
    
    # 3. 前向推理
    with torch.no_grad():  # 关闭梯度，加速推理
        output = model(fake_input)
    
    print(f"输入尺寸：{fake_input.shape}")
    print(f"输出尺寸：{output.shape}")  # [1, 2, 80, 80]（640/8=80）
    print("多模态反光衣检测模型推理完成！")