From a8647951aa6684ba8e66462c78bbb5abd12d7514 Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Tue, 9 May 2023 23:05:07 +0000 Subject: [PATCH] pose encoder model --- data/pose_renderer.py | 6 +- models/convnext.py | 158 +++++++++++++++---------------- models/pose_encoder.py | 205 +++++++++++++++++++++++++++++++++++++++++ utils/data_utils.py | 4 +- utils/pose_coloring.py | 8 +- 5 files changed, 288 insertions(+), 93 deletions(-) create mode 100644 models/pose_encoder.py diff --git a/data/pose_renderer.py b/data/pose_renderer.py index fafb926..0fd054b 100644 --- a/data/pose_renderer.py +++ b/data/pose_renderer.py @@ -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 diff --git a/models/convnext.py b/models/convnext.py index a8b6f70..ef5986d 100644 --- a/models/convnext.py +++ b/models/convnext.py @@ -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) diff --git a/models/pose_encoder.py b/models/pose_encoder.py new file mode 100644 index 0000000..c095598 --- /dev/null +++ b/models/pose_encoder.py @@ -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 \ No newline at end of file diff --git a/utils/data_utils.py b/utils/data_utils.py index 8ab55fe..1560523 100644 --- a/utils/data_utils.py +++ b/utils/data_utils.py @@ -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 diff --git a/utils/pose_coloring.py b/utils/pose_coloring.py index 7dae36e..85f3777 100644 --- a/utils/pose_coloring.py +++ b/utils/pose_coloring.py @@ -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))