# 现有DMS注意力数据集的视野盲区
frontal_fov_blindspots = {
    "变道场景": "侧方车辆在60°视野外，无法建模注视转移",
    "转弯场景": "交叉口行人/骑车人在外围视野外",
    "停车场": "多方向交通参与者，需要360°感知",
    "城市复杂路口": "左转/右转时注视重点在侧方",
    "应急场景": "后方紧急车辆需要后视镜注视"
}

# DriverGaze360解决方案
gaze360_solution = {
    "摄像头": "全景360°摄像头（前后左右覆盖）",
    "注视标注": "全场景空间注意力图",
    "对象级引导": "辅助语义分割头联合学习",
    "驾驶场景": "变道、转弯、交叉口、停车场等全场景"
}

import torch
import torch.nn as nn

class DriverGaze360Net(nn.Module):
    """
    DriverGaze360-Net: 全景驾驶员注意力预测网络
    
    CVPR 2026
    核心创新：
    1. 全景图像输入（360°视野）
    2. 注意力图+语义分割联合学习
    3. 对象级引导提升空间感知
    """
    
    def __init__(self, num_classes_seg=20):
        super().__init__()
        
        # Swin Transformer编码器（全景图像）
        self.encoder = SwinTransformerEncoder(
            img_size=(512, 1024),  # 全景分辨率
            patch_size=4,
            embed_dim=128,
            depths=[2, 2, 18, 2],
            num_heads=[4, 8, 16, 32]
        )
        
        # 注意力图解码器
        self.attention_decoder = AttentionDecoder(
            in_channels=128,
            mid_channels=64,
            out_channels=1  # 注意力图
        )
        
        # 辅助语义分割头（对象级引导）
        self.segmentation_head = SegmentationHead(
            in_channels=128,
            num_classes=num_classes_seg
        )
    
    def forward(self, panoramic_image):
        """
        Args:
            panoramic_image: 全景图像 (B, 3, 512, 1024)
        
        Returns:
            attention_map: 驾驶员注意力图 (B, 1, H, W)
            seg_map: 语义分割图 (B, C, H, W) - 辅助任务
        """
        # 编码
        features = self.encoder(panoramic_image)
        
        # 注意力图解码
        attention_map = self.attention_decoder(features)
        
        # 语义分割（辅助任务）
        seg_map = self.segmentation_head(features)
        
        return attention_map, seg_map


class AttentionDecoder(nn.Module):
    """多尺度注意力图解码器"""
    def __init__(self, in_channels, mid_channels, out_channels):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Conv2d(in_channels, mid_channels, 3, 1, 1),
            nn.BatchNorm2d(mid_channels),
            nn.ReLU(),
            nn.Conv2d(mid_channels, mid_channels, 3, 1, 1),
            nn.BatchNorm2d(mid_channels),
            nn.ReLU(),
            nn.Conv2d(mid_channels, out_channels, 1),
            nn.Sigmoid()  # 归一化到[0,1]
        )
    
    def forward(self, x):
        return self.layers(x)


class SegmentationHead(nn.Module):
    """语义分割辅助头"""
    def __init__(self, in_channels, num_classes):
        super().__init__()
        self.head = nn.Sequential(
            nn.Conv2d(in_channels, 64, 3, 1, 1),
            nn.ReLU(),
            nn.Conv2d(64, num_classes, 1)
        )
    
    def forward(self, x):
        return self.head(x)