62 lines
2.4 KiB
Python
62 lines
2.4 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
|
|
from .hr_base import HRNET_base
|
|
from .tokenpose_base import TokenPose_TB_base
|
|
|
|
|
|
class TokenPose_B(nn.Module):
|
|
def __init__(self, cfg, **kwargs):
|
|
super(TokenPose_B, self).__init__()
|
|
|
|
extra = cfg["MODEL"]["EXTRA"]
|
|
self.cfg = cfg
|
|
print(cfg["MODEL"])
|
|
##################################################
|
|
self.pre_feature = HRNET_base(cfg, **kwargs)
|
|
self.transformer = TokenPose_TB_base(
|
|
feature_size=[cfg["MODEL"]["IMAGE_SIZE"][1] // 4, cfg["MODEL"]["IMAGE_SIZE"][0] // 4],
|
|
patch_size=[cfg["MODEL"]["PATCH_SIZE"][1], cfg["MODEL"]["PATCH_SIZE"][0]],
|
|
num_keypoints=cfg["MODEL"]["NUM_JOINTS"],
|
|
dim=cfg["MODEL"]["DIM"],
|
|
channels=cfg["MODEL"]["BASE_CHANNEL"],
|
|
depth=cfg["MODEL"]["TRANSFORMER_DEPTH"],
|
|
heads=cfg["MODEL"]["TRANSFORMER_HEADS"],
|
|
mlp_dim=cfg["MODEL"]["DIM"] * cfg["MODEL"]["TRANSFORMER_MLP_RATIO"],
|
|
apply_init=cfg["MODEL"]["INIT"],
|
|
hidden_heatmap_dim=cfg["MODEL"]["HEATMAP_SIZE"][1]
|
|
* cfg["MODEL"]["HEATMAP_SIZE"][0]
|
|
// 8,
|
|
heatmap_dim=cfg["MODEL"]["HEATMAP_SIZE"][1] * cfg["MODEL"]["HEATMAP_SIZE"][0],
|
|
heatmap_size=[cfg["MODEL"]["HEATMAP_SIZE"][1], cfg["MODEL"]["HEATMAP_SIZE"][0]],
|
|
pos_embedding_type=cfg["MODEL"]["POS_EMBEDDING_TYPE"],
|
|
estimate_jacobian=cfg["MODEL"]["ESTIMATE_JACOBIAN"],
|
|
temperature=cfg["MODEL"]["TEMPERATURE"],
|
|
fix_img2motion_attention=cfg["MODEL"]["FIX_IMG2MOTION_ATTENTION"],
|
|
)
|
|
###################################################
|
|
|
|
def forward(self, x):
|
|
if self.cfg["MODEL"]["DATA_PREPROCESS"]:
|
|
mean = torch.tensor([0.485, 0.456, 0.406]).to(x.device)
|
|
std = torch.tensor([0.229, 0.224, 0.225]).to(x.device)
|
|
mean = mean.view(1, 3, 1, 1)
|
|
std = std.view(1, 3, 1, 1)
|
|
x = x - mean
|
|
x = x / std
|
|
x = self.pre_feature(x)
|
|
x = self.transformer(x)
|
|
return x
|
|
|
|
def init_weights(self, pretrained=""):
|
|
self.pre_feature.init_weights(pretrained)
|
|
|
|
|
|
def get_pose_net(cfg, is_train, **kwargs):
|
|
# cfg = convert_dict_to_attrit_dict(cfg)
|
|
model = TokenPose_B(cfg, **kwargs)
|
|
if is_train and cfg["MODEL"]["INIT_WEIGHTS"]:
|
|
model.init_weights(cfg["MODEL"]["PRETRAINED"])
|
|
|
|
return model
|