1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85
| import numpy as np import torch import torch.nn as nn
class OccupantPoseEstimator(nn.Module): """ 乘员3D姿态估计网络 """ def __init__(self): super().__init__() self.backbone = HRNetBackbone() self.lift_3d = nn.Sequential( nn.Linear(34, 128), nn.ReLU(), nn.Linear(128, 256), nn.ReLU(), nn.Linear(256, 51) ) self.depth_refine = DepthRefinement() def forward(self, rgb_image, depth_image=None): """ 估计3D姿态 """ keypoints_2d = self.backbone(rgb_image) keypoints_flat = keypoints_2d.view(-1, 51)[:, :34] keypoints_3d = self.lift_3d(keypoints_flat) keypoints_3d = keypoints_3d.view(-1, 17, 3) if depth_image is not None: keypoints_3d = self.depth_refine(keypoints_2d, keypoints_3d, depth_image) return { 'keypoints_2d': keypoints_2d, 'keypoints_3d': keypoints_3d } def detect_oop(self, keypoints_3d): """ 检测姿态异常 """ KEYPOINTS = { 'nose': 0, 'left_shoulder': 5, 'right_shoulder': 6, 'left_hip': 11, 'right_hip': 12 } shoulder_center = (keypoints_3d[KEYPOINTS['left_shoulder']] + keypoints_3d[KEYPOINTS['right_shoulder']]) / 2 hip_center = (keypoints_3d[KEYPOINTS['left_hip']] + keypoints_3d[KEYPOINTS['right_hip']]) / 2 torso_vector = shoulder_center - hip_center torso_angle = np.arctan2(torso_vector[2], torso_vector[1]) * 180 / np.pi is_forward = torso_angle > 30 head_pos = keypoints_3d[KEYPOINTS['nose']] is_head_close = head_pos[2] < 0.5 is_oop = is_forward or is_head_close return { 'is_oop': is_oop, 'torso_angle': torso_angle, 'is_forward': is_forward, 'is_head_close': is_head_close }
|