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
|
import torch
import torch.nn as nn
from models.auto_encoder import AutoEncoder
class RGBPartNet(nn.Module):
def __init__(
self,
ae_in_channels: int = 3,
ae_in_size: tuple[int, int] = (64, 48),
ae_feature_channels: int = 64,
f_a_c_p_dims: tuple[int, int, int] = (128, 128, 64),
image_log_on: bool = False
):
super().__init__()
(self.f_a_dim, self.f_c_dim, self.f_p_dim) = f_a_c_p_dims
self.image_log_on = image_log_on
self.ae = AutoEncoder(
ae_in_channels, ae_in_size, ae_feature_channels, f_a_c_p_dims
)
def forward(self, x_c1, x_c2=None):
# Step 1: Disentanglement
# n, t, c, h, w
((x_c, x_p), losses, images) = self._disentangle(x_c1, x_c2)
if self.training:
losses = torch.stack(losses)
return losses, images
else:
return x_c, x_p
def _disentangle(self, x_c1_t2, x_c2_t2=None):
n, t, c, h, w = x_c1_t2.size()
device = x_c1_t2.device
x_c1_t1 = x_c1_t2[:, torch.randperm(t), :, :, :]
if self.training:
((f_a_, f_c_, f_p_), losses) = self.ae(x_c1_t2, x_c1_t1, x_c2_t2)
# Decode features
with torch.no_grad():
x_c = self._decode_cano_feature(f_c_, n, t, device)
x_p = self._decode_pose_feature(f_p_, n, t, c, h, w, device)
i_a, i_c, i_p = None, None, None
if self.image_log_on:
i_a = self._decode_appr_feature(f_a_, n, t, device)
# Continue decoding canonical features
i_c = self.ae.decoder.trans_conv3(x_c)
i_c = torch.sigmoid(self.ae.decoder.trans_conv4(i_c))
i_p = x_p
return (x_c, x_p), losses, (i_a, i_c, i_p)
else: # evaluating
f_c_, f_p_ = self.ae(x_c1_t2)
x_c = self._decode_cano_feature(f_c_, n, t, device)
x_p = self._decode_pose_feature(f_p_, n, t, c, h, w, device)
return (x_c, x_p), None, None
def _decode_appr_feature(self, f_a_, n, t, device):
# Decode appearance features
f_a = f_a_.view(n, t, -1)
x_a = self.ae.decoder(
f_a.mean(1),
torch.zeros((n, self.f_c_dim), device=device),
torch.zeros((n, self.f_p_dim), device=device)
)
return x_a
def _decode_cano_feature(self, f_c_, n, t, device):
# Decode average canonical features to higher dimension
f_c = f_c_.view(n, t, -1)
x_c = self.ae.decoder(
torch.zeros((n, self.f_a_dim), device=device),
f_c.mean(1),
torch.zeros((n, self.f_p_dim), device=device),
cano_only=True
)
return x_c
def _decode_pose_feature(self, f_p_, n, t, c, h, w, device):
# Decode pose features to images
x_p_ = self.ae.decoder(
torch.zeros((n * t, self.f_a_dim), device=device),
torch.zeros((n * t, self.f_c_dim), device=device),
f_p_
)
x_p = x_p_.view(n, t, c, h, w)
return x_p
|
