luni/reconstruction
1
0
Fork 0
Code Issues 6 Pull Requests Packages Projects 1 Releases Wiki Activity

pose encoder model

Browse Source
This commit is contained in:
luni
2023-05-09 23:05:07 +00:00
parent 3f60403167
commit a8647951aa
5 changed files with 288 additions and 93 deletions
Download Patch File Download Diff File Expand all files Collapse all files
data/pose_renderer.py
+3 -3
View File
@@ -32,14 +32,14 @@ class PoseRenderer(Module):
super().__init__()
colors = np.load('custom_data/pose_colors.npy')
self.textures = TexturesVertex(verts_features=torch.from_numpy(colors).to(device)[None, :, :])
self.cam = torch.tensor([5.0, 0.0, 0.0], device=device).cpu().numpy()
self.cam = torch.tensor([3.0, 0.0, 0.0], device=device).cpu().numpy()
def forward(self, verts: torch.Tensor, faces: torch.Tensor):
cameras, K = get_camera(self.cam, [480, 640], verts.get_device())
cameras, K = get_camera(self.cam, [512, 512], verts.get_device(), res_x=512, res_y=512)
raster_settings = RasterizationSettings(
image_size=(480, 640),
image_size=(512, 512),
blur_radius=0.0,
faces_per_pixel=1,
bin_size=0
models/convnext.py
+74 -84
View File
@@ -94,6 +94,7 @@ class ConvNeXtBase(nn.Module):
layer_scale: float = 1e-6,
block: Optional[Callable[..., nn.Module]] = None,
norm_layer: Optional[Callable[..., nn.Module]] = None,
input_features: int = 3,
device: str = 'cuda',
) -> None:
super().__init__()
@@ -115,7 +116,7 @@ class ConvNeXtBase(nn.Module):
# Stem
firstconv_output_channels = block_setting[0].input_channels
self.stem = Conv2dNormActivation(
3,
input_features,
firstconv_output_channels,
kernel_size=4,
stride=4,
@@ -125,6 +126,7 @@ class ConvNeXtBase(nn.Module):
bias=True,
).to(device)
layers_list: List[nn.Module] = []
total_stage_blocks = sum(cnf.num_layers for cnf in block_setting)
stage_block_id = 0
@@ -136,16 +138,17 @@ class ConvNeXtBase(nn.Module):
sd_prob = stochastic_depth_prob * stage_block_id / (total_stage_blocks - 1.0)
stage.append(block(cnf.input_channels, layer_scale, sd_prob))
stage_block_id += 1
self.stages.append(nn.Sequential(*stage).to(device))
layers_list.append(nn.Sequential(*stage))
if cnf.out_channels is not None:
# Downsampling
self.downsamples.append(
layers_list.append(
nn.Sequential(
norm_layer(cnf.input_channels),
nn.Conv2d(cnf.input_channels, cnf.out_channels, kernel_size=2, stride=2),
).to(device)
)
self.layers = nn.Sequential(*layers_list).to(device)
self.act = nn.Tanh()
for m in self.modules():
@@ -157,11 +160,7 @@ class ConvNeXtBase(nn.Module):
def _forward_impl(self, x: Tensor) -> Tensor:
x = self.stem(x)
features: List[Tensor] = []
for i in range(4 - 1):
x = self.stages[i](x)
features.append(x)
x = self.downsamples[i](x)
x = self.stages[-1](x)
x = self.layers(x)
x = self.act(x)
return x, features
@@ -171,7 +170,7 @@ class ConvNeXtBase(nn.Module):
class ConvNeXtReverseBase2(nn.Module):
class ConvNeXtReverseDepthNormals(nn.Module):
def __init__(
self,
block_setting: List[CNBlockConfig],
@@ -217,8 +216,8 @@ class ConvNeXtReverseBase2(nn.Module):
total_stage_blocks = sum(cnf.num_layers for cnf in block_setting)
stage_block_id = 0
self.stages: List[nn.Module] = []
self.upsamples: List[nn.Module] = []
layers_list: List[nn.Module] = []
for cnf in block_setting[1:]:
if cnf.input_channels is not None:
# Bottlenecks
@@ -228,19 +227,21 @@ class ConvNeXtReverseBase2(nn.Module):
sd_prob = stochastic_depth_prob * stage_block_id / (total_stage_blocks - 1.0)
stage.append(block(cnf.input_channels, layer_scale, sd_prob))
stage_block_id += 1
self.stages.append(nn.Sequential(*stage).to(device))
layers_list.append(nn.Sequential(*stage).to(device))
else:
self.stages.append(nn.Identity().to(device))
layers_list.append(nn.Identity().to(device))
if cnf.out_channels is not None:
# Upnsampling
self.upsamples.append(
# Upsampling
layers_list.append(
nn.Sequential(
norm_layer(cnf.input_channels),
nn.ConvTranspose2d(cnf.input_channels, cnf.out_channels, kernel_size=2, stride=2),
).to(device)
)
self.layers = nn.Sequential(*layers_list).to(device)
for cnf in block_setting[-1:]:
# Bottlenecks
stage: List[nn.Module] = []
@@ -275,7 +276,10 @@ class ConvNeXtReverseBase2(nn.Module):
def _forward_impl(self, x) -> Tensor:
input, features = x
if self.use_skip:
input, features = x
else:
input = x
x = self.first(input)
for i in range(5 - 2):
if self.use_skip:
@@ -306,10 +310,14 @@ class ConvNeXtReverseBase(nn.Module):
stochastic_depth_prob: float = 0.0,
layer_scale: float = 1e-6,
block: Optional[Callable[..., nn.Module]] = None,
norm_layer: Optional[Callable[..., nn.Module]] = None
norm_layer: Optional[Callable[..., nn.Module]] = None,
device: str = 'cuda',
use_skip=False,
) -> None:
super().__init__()
self.use_skip = use_skip
if not block_setting:
raise ValueError("The block_setting should not be empty")
elif not (isinstance(block_setting, Sequence) and all([isinstance(s, CNBlockConfig) for s in block_setting])):
@@ -321,71 +329,44 @@ class ConvNeXtReverseBase(nn.Module):
if norm_layer is None:
norm_layer = partial(LayerNorm2d, eps=1e-6)
layers: List[nn.Module] = []
# add bottleneck which will be executed at the end, to process after the final larger upsampling
# Stem
firstconv_output_channels = block_setting[0].input_channels
layers.append(
nn.Sequential(
nn.Conv2d(16, 1, kernel_size=7, padding=3, bias=True),
)
)
stage: List[nn.Module] = []
for _ in range(3):
stage.append(block(16, layer_scale, 0.0))
layers.append(nn.Sequential(*stage))
layers.append(
# Convtr NormActivation(
# firstconv_output_channels,
# 3,
# kernel_size=4,
# stride=4,
# padding=0,
# norm_layer=None,
# activation_layer=None,
# bias=True,
# )
nn.Sequential(
nn.ConvTranspose2d(
firstconv_output_channels,
16,
kernel_size=4,
stride=4,
padding=0,
),
)
)
self.first = nn.Sequential(
nn.ConvTranspose2d(block_setting[0].input_channels, block_setting[0].out_channels, kernel_size=2, stride=2),
norm_layer(block_setting[0].out_channels),
).to(device)
total_stage_blocks = sum(cnf.num_layers for cnf in block_setting)
stage_block_id = 0
for cnf in block_setting:
# Bottlenecks
stage: List[nn.Module] = []
for _ in range(cnf.num_layers):
# adjust stochastic depth probability based on the depth of the stage block
sd_prob = stochastic_depth_prob * stage_block_id / (total_stage_blocks - 1.0)
stage.append(block(cnf.input_channels, layer_scale, sd_prob))
stage_block_id += 1
layers.append(nn.Sequential(*stage))
layers_list: List[nn.Module] = []
for cnf in block_setting[1:]:
if cnf.input_channels is not None:
# Bottlenecks
stage: List[nn.Module] = []
for _ in range(cnf.num_layers):
# adjust stochastic depth probability based on the depth of the stage block
sd_prob = stochastic_depth_prob * stage_block_id / (total_stage_blocks - 1.0)
stage.append(block(cnf.input_channels, layer_scale, sd_prob))
stage_block_id += 1
layers_list.append(nn.Sequential(*stage).to(device))
else:
layers_list.append(nn.Identity().to(device))
if cnf.out_channels is not None:
# Upsampling
layers.append(
# Upnsampling
layers_list.append(
nn.Sequential(
norm_layer(cnf.out_channels),
nn.ConvTranspose2d(cnf.out_channels, cnf.input_channels, kernel_size=2, stride=2),
)
norm_layer(cnf.input_channels),
nn.ConvTranspose2d(cnf.input_channels, cnf.out_channels, kernel_size=2, stride=2),
).to(device)
)
self.features = nn.Sequential(*reversed(layers))
self.layers = nn.Sequential(*layers_list).to(device)
self.last = nn.Conv2d(
block_setting[-1].input_channels,
1,
kernel_size=3,
padding='same',
).to(device)
for m in self.modules():
if isinstance(m, (nn.Conv2d, nn.Linear)):
@@ -393,12 +374,20 @@ class ConvNeXtReverseBase(nn.Module):
if m.bias is not None:
nn.init.zeros_(m.bias)
def _forward_impl(self, x: Tensor) -> Tensor:
x = self.features(x)
# x = torch.tanh(x)
return x
def forward(self, x: Tensor) -> Tensor:
def _forward_impl(self, x) -> Tensor:
if self.use_skip:
raise NotImplementedError
input, features = x
else:
input = x
x = self.first(input)
x = self.layers(x)
out = self.last(x)
return out
def forward(self, x) -> Tensor:
return self._forward_impl(x)
@@ -407,13 +396,14 @@ def _convnext(
stochastic_depth_prob: float,
weights: Optional[WeightsEnum],
reversed: bool,
input_features: int,
**kwargs: Any,
) -> ConvNeXtBase:
if reversed:
model = ConvNeXtReverseBase2(block_setting, stochastic_depth_prob=stochastic_depth_prob, **kwargs)
model = ConvNeXtReverseBase(block_setting, stochastic_depth_prob=stochastic_depth_prob,**kwargs)
else:
model = ConvNeXtBase(block_setting, stochastic_depth_prob=stochastic_depth_prob, **kwargs)
model = ConvNeXtBase(block_setting, stochastic_depth_prob=stochastic_depth_prob, input_features=input_features, **kwargs)
if weights is not None:
model.load_state_dict(weights.get_state_dict())
@@ -422,10 +412,10 @@ def _convnext(
class ConvNeXt(nn.Module):
def __init__(self, block_setting, stochastic_depth_prob=0.0, reversed=False, **kwargs) -> None:
def __init__(self, block_setting, stochastic_depth_prob=0.0, reversed=False, input_features=3, **kwargs) -> None:
super().__init__()
self.model = _convnext(block_setting, stochastic_depth_prob, None, reversed, **kwargs)
self.model = _convnext(block_setting, stochastic_depth_prob, None, reversed, input_features, **kwargs)
def forward(self, x: Tensor) -> Tensor:
x = self.model(x)
models/pose_encoder.py
+205
View File
@@ -0,0 +1,205 @@
import os
import sys
os.environ["CUDA_VISIBLE_DEVICES"] = '7'
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
import torch
from torch import nn, Tensor
from torch.nn import functional as F
from functools import partial
from typing import Any, Callable, List, Optional, Sequence
torch.set_float32_matmul_precision('high')
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
from models.convnext import ConvNeXt, CNBlockConfig, CNBlock, LayerNorm2d
class ConvNeXtReverseMaskPose(nn.Module):
def __init__(
self,
block_setting: List[CNBlockConfig],
stochastic_depth_prob: float = 0.0,
layer_scale: float = 1e-6,
block: Optional[Callable[..., nn.Module]] = None,
norm_layer: Optional[Callable[..., nn.Module]] = None,
device: str = 'cuda',
use_skip=False,
) -> None:
super().__init__()
self.use_skip = use_skip
if not block_setting:
raise ValueError("The block_setting should not be empty")
elif not (isinstance(block_setting, Sequence) and all([isinstance(s, CNBlockConfig) for s in block_setting])):
raise TypeError("The block_setting should be List[CNBlockConfig]")
if block is None:
block = CNBlock
if norm_layer is None:
norm_layer = partial(LayerNorm2d, eps=1e-6)
self.first = nn.Sequential(
nn.ConvTranspose2d(block_setting[0].input_channels, block_setting[0].out_channels, kernel_size=2, stride=2),
norm_layer(block_setting[0].out_channels),
).to(device)
total_stage_blocks = sum(cnf.num_layers for cnf in block_setting)
stage_block_id = 0
layers_list: List[nn.Module] = []
for cnf in block_setting[1:]:
if cnf.input_channels is not None:
# Bottlenecks
stage: List[nn.Module] = []
for _ in range(cnf.num_layers):
# adjust stochastic depth probability based on the depth of the stage block
sd_prob = stochastic_depth_prob * stage_block_id / (total_stage_blocks - 1.0)
stage.append(block(cnf.input_channels, layer_scale, sd_prob))
stage_block_id += 1
layers_list.append(nn.Sequential(*stage).to(device))
else:
layers_list.append(nn.Identity().to(device))
if cnf.out_channels is not None:
# Upsampling
layers_list.append(
nn.Sequential(
norm_layer(cnf.input_channels),
nn.ConvTranspose2d(cnf.input_channels, cnf.out_channels, kernel_size=2, stride=2),
).to(device)
)
self.layers = nn.Sequential(*layers_list[:-1]).to(device)
self.stage_pose = nn.Sequential(*layers_list[-1:]).to(device)
for cnf in block_setting[-1:]:
# Bottlenecks
stage: List[nn.Module] = []
for _ in range(cnf.num_layers):
stage.append(block(cnf.input_channels, layer_scale, 0))
self.stage_mask = nn.Sequential(*stage).to(device)
self.last = nn.Conv2d(
block_setting[-1].input_channels,
3,
kernel_size=3,
padding='same',
).to(device)
self.last_mask = nn.Conv2d(
block_setting[-1].input_channels,
1,
kernel_size=3,
padding='same',
).to(device)
for m in self.modules():
if isinstance(m, (nn.Conv2d, nn.Linear)):
nn.init.trunc_normal_(m.weight, std=0.02)
if m.bias is not None:
nn.init.zeros_(m.bias)
def _forward_impl(self, x) -> Tensor:
if self.use_skip:
raise NotImplementedError
input, features = x
else:
input = x
x = self.first(input)
f = self.layers(x)
x = self.stage_pose(f)
pose = self.last(x)
x = self.stage_mask(f)
mask = self.last_mask(x)
return pose, mask
def forward(self, x) -> Tensor:
return self._forward_impl(x)
class PoseEncoder(nn.Module):
def __init__(self, features=128, device='cuda'):
super().__init__()
block_setting = [
CNBlockConfig(96, 192, 3),
CNBlockConfig(192, 384, 3),
CNBlockConfig(384, features, 9),
CNBlockConfig(features, None, 3),
]
block_setting_reverse = [
CNBlockConfig(features, 384, 3),
CNBlockConfig(384, 192, 9),
CNBlockConfig(192, 96, 3),
CNBlockConfig(96, 27, 3),
CNBlockConfig(27, 9, 3),
CNBlockConfig(9, None, 3),
]
rotation_block_setting = [
CNBlockConfig(64, 128, 3),
CNBlockConfig(128, 128, 3),
CNBlockConfig(128, 256, 9),
CNBlockConfig(256, None, 3),
]
rotation_block_setting_reverse = [
CNBlockConfig(256, 128, 3),
CNBlockConfig(128, 64, 9),
CNBlockConfig(64, 32, 3),
CNBlockConfig(32, 16, 3),
CNBlockConfig(16, None, 3),
]
self.encoder = ConvNeXt(block_setting, input_features=4, device=device)
self.decoder = ConvNeXtReverseMaskPose(block_setting_reverse, device=device)
self.rotation_encoder = ConvNeXt(rotation_block_setting, input_features=7, device=device)
self.rotation_decoder = ConvNeXt(rotation_block_setting_reverse, reversed=True, device=device)
self.flatten = nn.Flatten()
def forward(self, input) -> torch.Tensor:
mask, pose_in, pose_out = input
features, _ = self.encoder(torch.cat([mask, pose_in], dim=1))
x = features.view(features.size(0), features.size(1), -1)
y, _ = self.rotation_encoder(torch.cat([mask, pose_in, pose_out], dim=1))
rotation = self.rotation_decoder(y)
rotation = torch.squeeze(rotation, dim=1)
rotation = torch.softmax(rotation, dim=1)
x = torch.einsum('bij,bkj->bik', x, rotation)
x = x.view(x.size(0), x.size(1), 16, 16)
pose, mask = self.decoder(x)
return pose, mask, rotation
if __name__ == '__main__':
# test the model
model = PoseEncoder()
# model = torch.compile(model)
print(model)
print(sum(p.numel() for p in model.parameters() if p.requires_grad))
print(sum(p.numel() for p in model.parameters()))
x = torch.rand((4, 1, 512, 512)).to('cuda')
y = torch.rand((4, 3, 512, 512)).to('cuda')
z = torch.rand((4, 3, 512, 512)).to('cuda')
output = model((x, y, z))
from torchviz import make_dot
graph = make_dot(output, params=dict(model.named_parameters()))
graph.render("model_graph", format="pdf")
pass
utils/data_utils.py
+2 -2
View File
@@ -113,10 +113,10 @@ def proj(camLoc):
return intrinsic().dot(extrinsic(camLoc))
def get_camera(cam, img_size, device):
def get_camera(cam, img_size, device, res_x=640, res_y=480):
size = np.array(img_size, dtype=np.float32)
ext = torch.tensor(extrinsic(cam).astype(np.float32), device=device)
K = torch.tensor(intrinsic()[None, :, :], device=device)
K = torch.tensor(intrinsic(res_x, res_y)[None, :, :], device=device)
cameras = cameras_from_opencv_projection(R=ext[None, :, :3], tvec=ext[None, :, 3], camera_matrix=K, image_size=torch.tensor(size[None, :], device=device))
return cameras, K
utils/pose_coloring.py
+4 -4
View File
@@ -35,14 +35,14 @@ class PoseRenderer(Module):
super().__init__()
colors = np.load('custom_data/pose_colors.npy')
self.textures = TexturesVertex(verts_features=torch.from_numpy(colors).to(device)[None, :, :])
self.cam = torch.tensor([5.0, 0.0, 0.0], device=device).cpu().numpy()
self.cam = torch.tensor([3.0, 0.0, 0.0], device=device).cpu().numpy()
def forward(self, verts: torch.Tensor, faces: torch.Tensor):
cameras, K = get_camera(self.cam, [480, 640], verts.get_device())
cameras, K = get_camera(self.cam, [512, 512], verts.get_device(), res_x=512, res_y=512)
raster_settings = RasterizationSettings(
image_size=(480, 640),
image_size=(512, 512),
blur_radius=0.0,
faces_per_pixel=1,
bin_size=0
@@ -63,7 +63,7 @@ if __name__ == '__main__':
smpl = SMPL(path='superSMPL', device='cuda')
pose = torch.zeros((1,72), device='cuda')
# generate random pose
pose[:,3:] = torch.randn((1,69), device='cuda') * 0.1
# pose[:,3:] = torch.randn((1,69), device='cuda') * 0.1
body, _ = smpl(pose, torch.zeros((1,11), device='cuda'), remove=False)
colors = np.zeros((body.shape[1], 3))