260 lines
9.7 KiB
Python
260 lines
9.7 KiB
Python
import os
|
|
import sys
|
|
|
|
# os.environ["CUDA_VISIBLE_DEVICES"] = '3'
|
|
# 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)
|
|
|
|
stage: List[nn.Module] = []
|
|
for _ in range(3):
|
|
stage.append(block(block_setting[0].input_channels, layer_scale, 0))
|
|
self.stage_smooth = nn.Sequential(*stage).to(device)
|
|
|
|
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.stage_smooth(input)
|
|
x = self.first(x)
|
|
f = self.layers(x)
|
|
x = self.stage_pose(f)
|
|
pose = torch.sigmoid(self.last(x))
|
|
|
|
# x = self.stage_mask(f)
|
|
# mask = torch.sigmoid(self.last_mask(x))
|
|
|
|
# return pose, mask
|
|
return pose, torch.zeros((x.shape[0], 1, x.shape[2], x.shape[3]), device=x.device)
|
|
|
|
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=7, 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_in_back, pose_out, mask_out = input
|
|
features, _ = self.encoder(torch.cat([mask, pose_in, pose_in_back], 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 = torch.einsum('bij,bjk->bik', x, rotation)
|
|
x = x.view(x.size(0), x.size(1), 16, 16)
|
|
pose, mask_out_pred = self.decoder(x)
|
|
pose = pose * mask_out
|
|
# mask_out = torch.round(mask_out)
|
|
# mask_out = (mask_out > 0.5).to(torch.int64)
|
|
# pose = pose * (mask_out > 0.5).to(torch.float32)
|
|
|
|
return pose, mask_out_pred, rotation
|
|
# x = pose
|
|
# return x, torch.zeros((x.shape[0], 1, x.shape[2], x.shape[3]), device=x.device), torch.zeros((x.shape[0], 256, 256), device=x.device)
|
|
|
|
class DumbPoseEncoder(nn.Module):
|
|
def __init__(self, in_channels=7, out_channels=3, device='cuda'):
|
|
super(DumbPoseEncoder, self).__init__()
|
|
|
|
# Encoder
|
|
self.encoder_conv1 = nn.Conv2d(in_channels, 64, kernel_size=3, padding=1)
|
|
self.encoder_conv2 = nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1)
|
|
self.encoder_conv3 = nn.Conv2d(128, 256, kernel_size=3, stride=2, padding=1)
|
|
self.encoder_conv4 = nn.Conv2d(256, 512, kernel_size=3, padding=1)
|
|
self.encoder_pool = nn.MaxPool2d(kernel_size=2, stride=2)
|
|
|
|
# Decoder
|
|
self.decoder_upsample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
|
|
self.decoder_conv1 = nn.Conv2d(512, 256, kernel_size=3, padding=1)
|
|
self.decoder_upsample_2 = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
|
|
self.decoder_conv2 = nn.Conv2d(256, 128, kernel_size=3, padding=1)
|
|
self.decoder_upsample_3 = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
|
|
self.decoder_conv3 = nn.Conv2d(128, 64, kernel_size=3, padding=1)
|
|
self.decoder_conv4 = nn.Conv2d(64, out_channels, kernel_size=3, padding=1)
|
|
|
|
def forward(self, x):
|
|
# Encoder
|
|
mask, pose_in, pose_out = x
|
|
x = torch.cat([mask, pose_in, pose_out], dim=1)
|
|
x1 = nn.functional.relu(self.encoder_conv1(x))
|
|
x = self.encoder_pool(x1)
|
|
x = nn.functional.relu(self.encoder_conv2(x))
|
|
x = nn.functional.relu(self.encoder_conv3(x))
|
|
x = nn.functional.relu(self.encoder_conv4(x))
|
|
|
|
# Decoder
|
|
x = self.decoder_upsample(x)
|
|
x = nn.functional.relu(self.decoder_conv1(x))
|
|
x = self.decoder_upsample_2(x)
|
|
x = nn.functional.relu(self.decoder_conv2(x))
|
|
x = self.decoder_upsample_3(x)
|
|
x = nn.functional.relu(self.decoder_conv3(x))
|
|
x = nn.functional.sigmoid(self.decoder_conv4(x))
|
|
|
|
return x, torch.zeros((x.shape[0], 1, x.shape[2], x.shape[3]), device=x.device), torch.zeros((x.shape[0], 256, 256), device=x.device)
|
|
|
|
|
|
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 |