100 lines
3.6 KiB
Python
100 lines
3.6 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
|
|
class Discriminator(nn.Module):
|
|
|
|
def __init__(self):
|
|
"""
|
|
Pose + Shape discriminator proposed in HMR
|
|
"""
|
|
super(Discriminator, self).__init__()
|
|
|
|
self.num_joints = 23
|
|
# poses_alone
|
|
self.D_conv1 = nn.Conv2d(9, 32, kernel_size=1)
|
|
nn.init.xavier_uniform_(self.D_conv1.weight)
|
|
nn.init.zeros_(self.D_conv1.bias)
|
|
self.relu = nn.ReLU(inplace=True)
|
|
self.D_conv2 = nn.Conv2d(32, 32, kernel_size=1)
|
|
nn.init.xavier_uniform_(self.D_conv2.weight)
|
|
nn.init.zeros_(self.D_conv2.bias)
|
|
pose_out = []
|
|
for i in range(self.num_joints):
|
|
pose_out_temp = nn.Linear(32, 1)
|
|
nn.init.xavier_uniform_(pose_out_temp.weight)
|
|
nn.init.zeros_(pose_out_temp.bias)
|
|
pose_out.append(pose_out_temp)
|
|
self.pose_out = nn.ModuleList(pose_out)
|
|
|
|
# betas
|
|
self.betas_fc1 = nn.Linear(10, 10)
|
|
nn.init.xavier_uniform_(self.betas_fc1.weight)
|
|
nn.init.zeros_(self.betas_fc1.bias)
|
|
self.betas_fc2 = nn.Linear(10, 5)
|
|
nn.init.xavier_uniform_(self.betas_fc2.weight)
|
|
nn.init.zeros_(self.betas_fc2.bias)
|
|
self.betas_out = nn.Linear(5, 1)
|
|
nn.init.xavier_uniform_(self.betas_out.weight)
|
|
nn.init.zeros_(self.betas_out.bias)
|
|
|
|
# poses_joint
|
|
self.D_alljoints_fc1 = nn.Linear(32*self.num_joints, 1024)
|
|
nn.init.xavier_uniform_(self.D_alljoints_fc1.weight)
|
|
nn.init.zeros_(self.D_alljoints_fc1.bias)
|
|
self.D_alljoints_fc2 = nn.Linear(1024, 1024)
|
|
nn.init.xavier_uniform_(self.D_alljoints_fc2.weight)
|
|
nn.init.zeros_(self.D_alljoints_fc2.bias)
|
|
self.D_alljoints_out = nn.Linear(1024, 1)
|
|
nn.init.xavier_uniform_(self.D_alljoints_out.weight)
|
|
nn.init.zeros_(self.D_alljoints_out.bias)
|
|
|
|
|
|
def forward(self, poses: torch.Tensor, betas: torch.Tensor) -> torch.Tensor:
|
|
"""
|
|
Forward pass of the discriminator.
|
|
Args:
|
|
poses (torch.Tensor): Tensor of shape (B, 23, 3, 3) containing a batch of SMPL body poses (excluding the global orientation).
|
|
betas (torch.Tensor): Tensor of shape (B, 10) containign a batch of SMPL beta coefficients.
|
|
Returns:
|
|
torch.Tensor: Discriminator output with shape (B, 25)
|
|
"""
|
|
#import ipdb; ipdb.set_trace()
|
|
#bn = poses.shape[0]
|
|
# poses B x 207
|
|
#poses = poses.reshape(bn, -1)
|
|
# poses B x num_joints x 1 x 9
|
|
poses = poses.reshape(-1, self.num_joints, 1, 9)
|
|
bn = poses.shape[0]
|
|
# poses B x 9 x num_joints x 1
|
|
poses = poses.permute(0, 3, 1, 2).contiguous()
|
|
|
|
# poses_alone
|
|
poses = self.D_conv1(poses)
|
|
poses = self.relu(poses)
|
|
poses = self.D_conv2(poses)
|
|
poses = self.relu(poses)
|
|
|
|
poses_out = []
|
|
for i in range(self.num_joints):
|
|
poses_out_ = self.pose_out[i](poses[:, :, i, 0])
|
|
poses_out.append(poses_out_)
|
|
poses_out = torch.cat(poses_out, dim=1)
|
|
|
|
# betas
|
|
betas = self.betas_fc1(betas)
|
|
betas = self.relu(betas)
|
|
betas = self.betas_fc2(betas)
|
|
betas = self.relu(betas)
|
|
betas_out = self.betas_out(betas)
|
|
|
|
# poses_joint
|
|
poses = poses.reshape(bn,-1)
|
|
poses_all = self.D_alljoints_fc1(poses)
|
|
poses_all = self.relu(poses_all)
|
|
poses_all = self.D_alljoints_fc2(poses_all)
|
|
poses_all = self.relu(poses_all)
|
|
poses_all_out = self.D_alljoints_out(poses_all)
|
|
|
|
disc_out = torch.cat((poses_out, betas_out, poses_all_out), 1)
|
|
return disc_out
|