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 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111
| import torch import torch.nn as nn
class DepthIRPoseNet(nn.Module): """ 深度+红外3D姿态估计网络 论文方法复现 """ def __init__(self, num_keypoints=17, pretrained=True): super().__init__() self.num_keypoints = num_keypoints self.depth_encoder = nn.Sequential( nn.Conv2d(1, 32, kernel_size=7, stride=2, padding=3), nn.BatchNorm2d(32), nn.ReLU(inplace=True), nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1), nn.BatchNorm2d(64), nn.ReLU(inplace=True), nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1), nn.BatchNorm2d(128), nn.ReLU(inplace=True), nn.Conv2d(128, 256, kernel_size=3, stride=2, padding=1), nn.BatchNorm2d(256), nn.ReLU(inplace=True), ) self.ir_encoder = nn.Sequential( nn.Conv2d(1, 32, kernel_size=7, stride=2, padding=3), nn.BatchNorm2d(32), nn.ReLU(inplace=True), nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1), nn.BatchNorm2d(64), nn.ReLU(inplace=True), nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1), nn.BatchNorm2d(128), nn.ReLU(inplace=True), nn.Conv2d(128, 256, kernel_size=3, stride=2, padding=1), nn.BatchNorm2d(256), nn.ReLU(inplace=True), ) self.fusion = nn.Sequential( nn.Conv2d(512, 512, kernel_size=3, padding=1), nn.BatchNorm2d(512), nn.ReLU(inplace=True), ) self.keypoint_head = nn.Sequential( nn.AdaptiveAvgPool2d((1, 1)), nn.Flatten(), nn.Linear(512, 512), nn.ReLU(inplace=True), nn.Linear(512, num_keypoints * 3), ) def forward(self, depth_img, ir_img): """ 前向传播 Args: depth_img: 深度图 (B, 1, H, W) ir_img: 红外图 (B, 1, H, W) Returns: keypoints_3d: 3D关键点坐标 (B, num_keypoints, 3) """ depth_feat = self.depth_encoder(depth_img) ir_feat = self.ir_encoder(ir_img) fused_feat = torch.cat([depth_feat, ir_feat], dim=1) fused_feat = self.fusion(fused_feat) keypoints_flat = self.keypoint_head(fused_feat) keypoints_3d = keypoints_flat.view(-1, self.num_keypoints, 3) return keypoints_3d
if __name__ == "__main__": model = DepthIRPoseNet(num_keypoints=17) model.eval() depth_input = torch.randn(1, 1, 480, 640) ir_input = torch.randn(1, 1, 480, 640) with torch.no_grad(): keypoints_3d = model(depth_input, ir_input) print(f"输出关键点形状: {keypoints_3d.shape}")
|