172 lines
6.5 KiB
Python
172 lines
6.5 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
from .gnn import pose_proj as pproj
|
|
from .cnn.cnn_multitask import MultitaskCNN
|
|
from .gnn.step_regressor import StepRegressor, RelativePositionEncoder
|
|
|
|
|
|
class SPIGA(nn.Module):
|
|
def __init__(self, num_landmarks=98, num_edges=15, steps=3, **kwargs):
|
|
|
|
super(SPIGA, self).__init__()
|
|
|
|
# Model parameters
|
|
self.steps = steps # Cascaded regressors
|
|
self.embedded_dim = 512 # GAT input channel
|
|
self.nstack = 4 # Number of stacked GATs per step
|
|
self.kwindow = 7 # Output cropped window dimension (kernel)
|
|
self.swindow = 0.25 # Scale of the cropped window at first step (Dft. 25% w.r.t the input featuremap)
|
|
self.offset_ratio = [self.swindow/(2**step)/2 for step in range(self.steps)]
|
|
|
|
# CNN parameters
|
|
self.num_landmarks = num_landmarks
|
|
self.num_edges = num_edges
|
|
|
|
# Initialize backbone
|
|
self.visual_cnn = MultitaskCNN(num_landmarks=self.num_landmarks, num_edges=self.num_edges)
|
|
# Features dimensions
|
|
self.img_res = self.visual_cnn.img_res
|
|
self.visual_res = self.visual_cnn.out_res
|
|
self.visual_dim = self.visual_cnn.ch_dim
|
|
|
|
# Initialize Pose head
|
|
self.channels_pose = 6
|
|
self.pose_fc = nn.Linear(self.visual_cnn.ch_dim, self.channels_pose)
|
|
|
|
# Initialize feature extractors:
|
|
# Relative positional encoder
|
|
shape_dim = 2 * (self.num_landmarks - 1)
|
|
shape_encoder = []
|
|
for step in range(self.steps):
|
|
shape_encoder.append(RelativePositionEncoder(shape_dim, self.embedded_dim, [256, 256]))
|
|
self.shape_encoder = nn.ModuleList(shape_encoder)
|
|
# Diagonal mask used to compute relative positions
|
|
diagonal_mask = (torch.ones(self.num_landmarks, self.num_landmarks) - torch.eye(self.num_landmarks)).type(torch.bool)
|
|
self.diagonal_mask = nn.parameter.Parameter(diagonal_mask, requires_grad=False)
|
|
|
|
# Visual feature extractor
|
|
conv_window = []
|
|
theta_S = []
|
|
for step in range(self.steps):
|
|
# S matrix per step
|
|
WH = self.visual_res # Width/height of ftmap
|
|
Wout = self.swindow / (2 ** step) * WH # Width/height of the window
|
|
K = self.kwindow # Kernel or resolution of the window
|
|
scale = K / WH * (Wout - 1) / (K - 1) # Scale of the affine transformation
|
|
# Rescale matrix S
|
|
theta_S_stp = torch.tensor([[scale, 0], [0, scale]])
|
|
theta_S.append(nn.parameter.Parameter(theta_S_stp, requires_grad=False))
|
|
|
|
# Convolutional to embedded to BxLxCx1x1
|
|
conv_window.append(nn.Conv2d(self.visual_dim, self.embedded_dim, self.kwindow))
|
|
|
|
self.theta_S = nn.ParameterList(theta_S)
|
|
self.conv_window = nn.ModuleList(conv_window)
|
|
|
|
# Initialize GAT modules
|
|
self.gcn = nn.ModuleList([StepRegressor(self.embedded_dim, 256, self.nstack) for i in range(self.steps)])
|
|
|
|
def forward(self, data):
|
|
# Inputs: Visual features and points projections
|
|
pts_proj, features = self.backbone_forward(data)
|
|
# Visual field
|
|
visual_field = features['VisualField'][-1]
|
|
|
|
# Params compute only once
|
|
gat_prob = []
|
|
features['Landmarks'] = []
|
|
for step in range(self.steps):
|
|
# Features generation
|
|
embedded_ft = self.extract_embedded(pts_proj, visual_field, step)
|
|
|
|
# GAT inference
|
|
offset, gat_prob = self.gcn[step](embedded_ft, gat_prob)
|
|
offset = F.hardtanh(offset)
|
|
|
|
# Update coordinates
|
|
pts_proj = pts_proj + self.offset_ratio[step] * offset
|
|
features['Landmarks'].append(pts_proj.clone())
|
|
|
|
features['GATProb'] = gat_prob
|
|
return features
|
|
|
|
def backbone_forward(self, data):
|
|
# Inputs: Image and model3D
|
|
imgs = data[0]
|
|
model3d = data[1]
|
|
cam_matrix = data[2]
|
|
|
|
# HourGlass Forward
|
|
features = self.visual_cnn(imgs)
|
|
|
|
# Head pose estimation
|
|
pose_raw = features['HGcore'][-1]
|
|
B, L, _, _ = pose_raw.shape
|
|
pose = pose_raw.reshape(B, L)
|
|
pose = self.pose_fc(pose)
|
|
features['Pose'] = pose.clone()
|
|
|
|
# Project model 3D
|
|
euler = pose[:, 0:3]
|
|
trl = pose[:, 3:]
|
|
rot = pproj.euler_to_rotation_matrix(euler)
|
|
pts_proj = pproj.projectPoints(model3d, rot, trl, cam_matrix)
|
|
pts_proj = pts_proj / self.visual_res
|
|
|
|
return pts_proj, features
|
|
|
|
def extract_embedded(self, pts_proj, receptive_field, step):
|
|
# Visual features
|
|
visual_ft = self.extract_visual_embedded(pts_proj, receptive_field, step)
|
|
# Shape features
|
|
shape_ft = self.calculate_distances(pts_proj)
|
|
shape_ft = self.shape_encoder[step](shape_ft)
|
|
# Addition
|
|
embedded_ft = visual_ft + shape_ft
|
|
return embedded_ft
|
|
|
|
def extract_visual_embedded(self, pts_proj, receptive_field, step):
|
|
# Affine matrix generation
|
|
B, L, _ = pts_proj.shape # Pts_proj range:[0,1]
|
|
centers = pts_proj + 0.5 / self.visual_res # BxLx2
|
|
centers = centers.reshape(B * L, 2) # B*Lx2
|
|
theta_trl = (-1 + centers * 2).unsqueeze(-1) # BxLx2x1
|
|
theta_s = self.theta_S[step] # 2x2
|
|
theta_s = theta_s.repeat(B * L, 1, 1) # B*Lx2x2
|
|
theta = torch.cat((theta_s, theta_trl), -1) # B*Lx2x3
|
|
|
|
# Generate crop grid
|
|
B, C, _, _ = receptive_field.shape
|
|
grid = torch.nn.functional.affine_grid(theta, (B * L, C, self.kwindow, self.kwindow))
|
|
grid = grid.reshape(B, L, self.kwindow, self.kwindow, 2)
|
|
grid = grid.reshape(B, L, self.kwindow * self.kwindow, 2)
|
|
|
|
# Crop windows
|
|
crops = torch.nn.functional.grid_sample(receptive_field, grid, padding_mode="border") # BxCxLxK*K
|
|
crops = crops.transpose(1, 2) # BxLxCxK*K
|
|
crops = crops.reshape(B * L, C, self.kwindow, self.kwindow)
|
|
|
|
# Flatten features
|
|
visual_ft = self.conv_window[step](crops)
|
|
_, Cout, _, _ = visual_ft.shape
|
|
visual_ft = visual_ft.reshape(B, L, Cout)
|
|
|
|
return visual_ft
|
|
|
|
def calculate_distances(self, pts_proj):
|
|
B, L, _ = pts_proj.shape # BxLx2
|
|
pts_a = pts_proj.unsqueeze(-2).repeat(1, 1, L, 1)
|
|
pts_b = pts_a.transpose(1, 2)
|
|
dist = pts_a - pts_b
|
|
dist_wo_self = dist[:, self.diagonal_mask, :].reshape(B, L, -1)
|
|
return dist_wo_self
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|