"""
单目深度估计用于 OOP 检测

流程：
1. DMS 摄像头 → RGB 图像
2. 单目深度模型 → 深度图
3. 深度图 + 关键点 → 3D 姿态估计
4. 3D 姿态 → OOP 分类

关键：零额外硬件，仅用现有 DMS 摄像头
"""

import torch
import torch.nn as nn
import numpy as np

class MonoDepthOOPDetector(nn.Module):
    """
    单目深度估计 → OOP 检测
    
    使用轻量化深度估计模型
    输出：深度图 + OOP 分类
    """
    
    def __init__(self):
        super().__init__()
        # 轻量深度估计 backbone
        self.depth_encoder = nn.Sequential(
            nn.Conv2d(3, 32, 7, stride=2, padding=3),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.Conv2d(32, 64, 5, stride=2, padding=2),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.Conv2d(64, 128, 3, stride=2, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(),
            nn.Conv2d(128, 256, 3, stride=1, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(),
        )
        
        # 深度解码器
        self.depth_decoder = nn.Sequential(
            nn.ConvTranspose2d(256, 128, 3, stride=2, padding=1, output_padding=1),
            nn.ReLU(),
            nn.ConvTranspose2d(128, 64, 3, stride=2, padding=1, output_padding=1),
            nn.ReLU(),
            nn.ConvTranspose2d(64, 32, 3, stride=2, padding=1, output_padding=1),
            nn.ReLU(),
            nn.Conv2d(32, 1, 1),  # 深度图
        )
        
        # OOP 分类头
        self.oop_head = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Flatten(),
            nn.Linear(256, 128),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(128, 6),  # 6类 OOP
        )
    
    def forward(self, x):
        feat = self.depth_encoder(x)
        depth = self.depth_decoder(feat)
        oop = self.oop_head(feat)
        return depth, oop

# OOP 分类
OOP_CLASSES = [
    'normal_sitting',    # 正常坐姿
    'leaning_forward',  # 前倾
    'leaning_side',     # 侧倾
    'leaning_back',     # 后仰
    'kneeling',         # 跪姿
    'empty_seat',       # 空座
]

# 测试
if __name__ == "__main__":
    model = MonoDepthOOPDetector()
    model.eval()
    
    x = torch.randn(1, 3, 224, 224)
    depth, oop = model(x)
    
    print(f"输入: {x.shape}")
    print(f"深度图: {depth.shape}")
    print(f"OOP 预测: {oop.shape} → {OOP_CLASSES[oop.argmax(1).item()]}")
    print(f"参数量: {sum(p.numel() for p in model.parameters()):,}")