diff --git a/.gitignore b/.gitignore index 1a2d0fd..96c0f1a 100644 --- a/.gitignore +++ b/.gitignore @@ -2,6 +2,7 @@ module_articulate/models/*.pth module_spline/models/*.tar +module_mrfa/models/*.pth checkpoints/* !checkpoints/place_checkpoints_here.txt diff --git a/README.md b/README.md index 2595756..f7ce99b 100644 --- a/README.md +++ b/README.md @@ -40,6 +40,11 @@ https://github.com/user-attachments/assets/b2948efb-3b44-440b-bff2-dde7b95a9946 ![Workflow Spline](workflows/workflow_spline.png) +### MRFA + +[MRFA.json](workflows/MRFA.json) + +![Workflow MRFA](workflows/workflow_mrfa.png) ## Arguments @@ -73,6 +78,10 @@ Doesn't need any ### MRFA +* `model_name`: `celebvhq` or `vox`, which is trained on different datasets +* `use_relative`: Whether to use relative mode or not (absolute mode). Absolute mode is similar to FOMM's `adapt_movement_scale` set to False +* `relative_movement`, `relative_jacobian`, `adapt_movement_scale`: Same as FOMM + ## Installation 1. Clone the repo to `ComfyUI/custom_nodes/` @@ -134,4 +143,9 @@ For **Articulate**, download the model from [Pre-trained checkpoints](https://gi For **Spline**, download the model from [Pre-trained models](https://github.com/yoyo-nb/thin-plate-spline-motion-model?tab=readme-ov-file#pre-trained-models) section and place it here: `spline_module/models/vox.pth.tar`. To use `find_best_frame`, install `face-alignment`. -For **MRFA**, download the model from [Pretrained models](https://github.com/JialeTao/MRFA/?tab=readme-ov-file#pretrained-models) section and place it here: `mrfa_module/models/` \ No newline at end of file +For **MRFA**, download `MRFACheckpoints.tar.gz` from [Pretrained models](https://github.com/JialeTao/MRFA/?tab=readme-ov-file#pretrained-models) section and decompress it. Inside there should be two files. Place them here: + +``` +mrfa_module/models/celebvhq.pth +mrfa_module/models/vox.pth +``` \ No newline at end of file diff --git a/__init__.py b/__init__.py index 9edeec0..f59a68f 100644 --- a/__init__.py +++ b/__init__.py @@ -5,7 +5,8 @@ from .run import ( FOMM_Seg10Chooser, FOMM_Seg15Chooser, Articulate_Runner, - Spline_Runner + Spline_Runner, + MRFA_Runner ) NODE_CLASS_MAPPINGS = { @@ -16,6 +17,7 @@ NODE_CLASS_MAPPINGS = { "FOMM_Seg15Chooser": FOMM_Seg15Chooser, "Articulate_Runner": Articulate_Runner, "Spline_Runner": Spline_Runner, + "MRFA_Runner": MRFA_Runner, } NODE_DISPLAY_NAME_MAPPINGS = { @@ -26,6 +28,7 @@ NODE_DISPLAY_NAME_MAPPINGS = { "FOMM_Seg15Chooser": "FOMM Seg15 Chooser", "Articulate_Runner": "Articulate Runner", "Spline_Runner": "Spline Runner", + "MRFA_Runner": "MRFA Runner", } diff --git a/constants.py b/constants.py index 93b184b..ad76d08 100644 --- a/constants.py +++ b/constants.py @@ -106,4 +106,16 @@ SPLINE_MODES = ['relative', 'standard', 'avd'] SPLINE_DEFAULT = 'relative' SPLINE_MODEL_PATH = "module_spline/models/vox.pth.tar" -SPLINE_CFG_PATH = "module_spline/config/vox-256.yaml" \ No newline at end of file +SPLINE_CFG_PATH = "module_spline/config/vox-256.yaml" + +MRFA_MODEL_NAMES = ["celebvhq", "vox"] +MRFA_MODEL_PATHS = { + "celebvhq": "module_mrfa/models/celebvhq.pth", + "vox": "module_mrfa/models/vox.pth", +} +MRFA_DEFAULT_MODEL = "celebvhq" + +MRFA_CFG_PATHS = { + "celebvhq": "module_mrfa/configs/celebvhq.yaml", + "vox": "module_mrfa/configs/vox1.yaml", +} \ No newline at end of file diff --git a/inference_mrfa.py b/inference_mrfa.py new file mode 100644 index 0000000..6215144 --- /dev/null +++ b/inference_mrfa.py @@ -0,0 +1,131 @@ +import numpy as np +import torch +import tqdm +import yaml +from comfy.utils import ProgressBar +from scipy.spatial import ConvexHull + +from .module_mrfa.model import MRFA +from .module_mrfa.util import AntiAliasInterpolation2d + + +def mrfa_inference( + source_image: torch.Tensor, + driving_video: torch.Tensor, + config_path: str, + checkpoint_path: str, + use_relative=True, + relative_movement=True, + relative_jacobian=True, + adapt_movement_scale=False, +): + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + with open(config_path) as f: + config = yaml.full_load(f) + + model = MRFA(config) + model = torch.nn.DataParallel(model) + + checkpoint = torch.load(checkpoint_path, map_location=device) + model.load_state_dict(checkpoint["model"], strict=False) + model = model.module + model.eval().to(device) + + kp_detector = model.encoder + dense_motion_network = model.dense_motion + decoder = model.decoder + down = AntiAliasInterpolation2d(3, 0.25).to(device) + + source = source_image.to(device) # 4D + driving = driving_video.to(device) # + + num_frames = driving.shape[2] + pbar = ProgressBar(num_frames) + + predictions = [] + with torch.no_grad(): + # visualizations = [] + + kp_source = kp_detector.forward(source) + kp_driving_initial = kp_detector.forward(driving[:, :, 0]) + + for frame_idx in tqdm.tqdm( + range(num_frames), + desc=f"Generating {'Relative' if use_relative else 'Absolute'} mode", + ): + driving_frame = driving[:, :, frame_idx] + if use_relative: + kp_driving = kp_detector.forward(driving_frame) + kp_norm = normalize_kp( + kp_source=kp_source, + kp_driving=kp_driving, + kp_driving_initial=kp_driving_initial, + use_relative_movement=relative_movement, + use_relative_jacobian=relative_jacobian, + adapt_movement_scale=adapt_movement_scale, + ) + dense_motion = dense_motion_network.forward( + source, kp_norm, kp_source, bg_param=None + ) + + if config["train_params"]["prior_model"] == "tpsm": + kp_s_value = kp_source["kp"].view(source.shape[0], -1, 5, 2).mean(2) + kp_d_value = kp_norm["kp"].view(driving.shape[0], -1, 5, 2).mean(2) + else: + kp_s_value = kp_source["kp"] + kp_d_value = kp_norm["kp"] + + out, _, _ = decoder.forward( + kp_s_value, + kp_d_value, + dense_motion, + img=down(source), + img_full=source, + ) + + # visualization = Visualizer(**config['visualizer_params']).visualize(source=source, driving=driving_frame, out=out) + # visualizations.append(visualization) + + else: + abs_inp = {"source": source, "driving": driving_frame} + out, _, _, _, _ = model.forward(abs_inp, is_train=False) + + prediction = out.data.cpu().numpy() + prediction = np.transpose(prediction, [0, 2, 3, 1]).squeeze(0) + predictions.append(prediction) + + pbar.update_absolute(frame_idx, num_frames) + + return predictions + + +def normalize_kp( + kp_source, + kp_driving, + kp_driving_initial, + adapt_movement_scale=False, + use_relative_movement=False, + use_relative_jacobian=False, +): + if adapt_movement_scale: + source_area = ConvexHull(kp_source["kp"][0].data.cpu().numpy()).volume + driving_area = ConvexHull(kp_driving_initial["kp"][0].data.cpu().numpy()).volume + adapt_movement_scale = np.sqrt(source_area) / np.sqrt(driving_area) + else: + adapt_movement_scale = 1 + + kp_new = {k: v for k, v in kp_driving.items()} + + if use_relative_movement: + kp_value_diff = kp_driving["kp"] - kp_driving_initial["kp"] + kp_value_diff *= adapt_movement_scale + kp_new["kp"] = kp_value_diff + kp_source["kp"] + + if use_relative_jacobian: + jacobian_diff = torch.matmul( + kp_driving["jacobian"], torch.inverse(kp_driving_initial["jacobian"]) + ) + kp_new["jacobian"] = torch.matmul(jacobian_diff, kp_source["jacobian"]) + + return kp_new diff --git a/module_mrfa/bg_motion_predictor.py b/module_mrfa/bg_motion_predictor.py new file mode 100644 index 0000000..272683c --- /dev/null +++ b/module_mrfa/bg_motion_predictor.py @@ -0,0 +1,29 @@ +import torch +from torch import nn +from torchvision import models + + +class BGMotionPredictor(nn.Module): + """ + Module for background estimation, return single transformation, parametrized as 3x3 matrix. The third row is [0 0 1] + """ + + def __init__(self): + super(BGMotionPredictor, self).__init__() + self.bg_encoder = models.resnet18(pretrained=False) + self.bg_encoder.conv1 = nn.Conv2d( + 6, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False + ) + num_features = self.bg_encoder.fc.in_features + self.bg_encoder.fc = nn.Linear(num_features, 6) + self.bg_encoder.fc.weight.data.zero_() + self.bg_encoder.fc.bias.data.copy_( + torch.tensor([1, 0, 0, 0, 1, 0], dtype=torch.float) + ) + + def forward(self, source_image, driving_image): + bs = source_image.shape[0] + out = torch.eye(3).unsqueeze(0).repeat(bs, 1, 1).type(source_image.type()) + prediction = self.bg_encoder(torch.cat([source_image, driving_image], dim=1)) + out[:, :2, :] = prediction.view(bs, 2, 3) + return out diff --git a/module_mrfa/configs/celebvhq.yaml b/module_mrfa/configs/celebvhq.yaml new file mode 100644 index 0000000..02cfa17 --- /dev/null +++ b/module_mrfa/configs/celebvhq.yaml @@ -0,0 +1,186 @@ +dataset_params: + root_dir: ./data/celebvhq256-png/ + split_len: 11 + frame_shape: [256, 256, 3] + id_sampling: True + pairs_list: + augmentation_params: + flip_param: + horizontal_flip: True + time_flip: True + jitter_param: + brightness: 0.1 + contrast: 0.1 + saturation: 0.1 + hue: 0.1 + +dense_motion: + block_expansion: 64 + max_features: 1024 + num_blocks: 5 + scale_factor: 0.25 + num_kp: 10 + num_channels: 3 + estimate_occlusion_map: True +fomm_kp_detector: + block_expansion: 32 + num_kp: 10 + num_channels: 3 + max_features: 1024 + num_blocks: 5 + temperature: 0.1 + scale_factor: 0.25 + estimate_jacobian: True + estimate_occlusion: False +tpsm_kp_detector: + num_tps: 10 +tpsm_dense_motion: + block_expansion: 64 + max_features: 1024 + num_blocks: 5 + scale_factor: 0.25 + num_tps: 10 + num_channels: 3 + +raft_flow: + prior_only: False + num_kp: 10 + dim: 256 + size: 256 + generator: + num_channels: 3 + block_expansion: 64 + max_features: 512 + num_up_blocks: 5 + driving_encoder: + in_features: 10 + block_expansion: 32 + max_features: 512 + num_blocks: 5 + source_encoder: + in_features: 13 + block_expansion: 32 + max_features: 512 + num_blocks: 5 + +train_params: + ## [mtia, tpsm, fomm] + prior_model: mtia + num_epochs: 100 + num_repeats: 5 + epoch_milestones: [60, 90] + # ### enable this for tpsm + # epoch_milestones: [70, 90] + # ### enable this for tpsm + + lr: 2.0e-4 + batch_size: 80 + scales: [1, 0.5, 0.25, 0.125] + clip_grad: True + clip: 10 + bg_start: 0 + checkpoint_freq: 100 + ### for tpsm + dropout_epoch: 35 + dropout_maxp: 0.3 + dropout_startp: 0.1 + dropout_inc_epoch: 10 + ### for tpsm + + transform_params: + sigma_affine: 0.05 + sigma_tps: 0.005 + points_tps: 5 + loss_weights: + perceptual: [10, 10, 10, 10, 10] + equivariance: 10 + equivariance_jacobian: 10 + # ### enable this for tpsm + # equivariance_jacobian: 0 + # ### enable this for tpsm + +reconstruction_params: + num_videos: 1000 + format: '.mp4' + +animate_params: + num_pairs: 50 + format: '.mp4' + normalization_params: + adapt_movement_scale: False + use_relative_movement: True + use_relative_jacobian: True + +visualizer_params: + draw_border: True + colormap: 'gist_rainbow' + +mtia_kp_detector: + MODEL: + ESTIMATE_JACOBIAN: True + DATA_PREPROCESS: False + FIX_IMG2MOTION_ATTENTION: False + TRANSFORMER_DEPTH: 12 + TRANSFORMER_HEADS: 8 + DIM: 192 # 4*4*3 + INIT_WEIGHTS: False + NAME: pose_tokenpose_b + NUM_JOINTS: 10 + PRETRAINED: '' + PATCH_SIZE: [4, 4] + IMAGE_SIZE: [256, 256] + HEATMAP_SIZE: [64, 64] + + ## default + TAG_PER_JOINT: True + HIDDEN_HEATMAP_DIM: -1 + MULTI_TRANSFORMER_DEPTH: [12, 12] + MULTI_TRANSFORMER_HEADS: [16, 16] + MULTI_DIM: [48, 48] + NUM_BRANCHES: 1 + BASE_CHANNEL: 32 + TRANSFORMER_MLP_RATIO: 3 + POS_EMBEDDING_TYPE: 'sine-full' + TEMPERATURE: 0.1 + TARGET_TYPE: gaussian + INIT: True + + SIGMA: 2 + EXTRA: + PRETRAINED_LAYERS: + - 'conv1' + - 'bn1' + - 'conv2' + - 'bn2' + - 'layer1' + - 'transition1' + - 'stage2' + - 'transition2' + - 'stage3' + FINAL_CONV_KERNEL: 1 + STAGE2: + NUM_MODULES: 1 + NUM_BRANCHES: 2 + BLOCK: BASIC + NUM_BLOCKS: + - 4 + - 4 + NUM_CHANNELS: + - 32 + - 64 + FUSE_METHOD: SUM + STAGE3: + NUM_MODULES: 4 + NUM_BRANCHES: 3 + BLOCK: BASIC + NUM_BLOCKS: + - 4 + - 4 + - 4 + NUM_CHANNELS: + - 32 + - 64 + - 128 + FUSE_METHOD: SUM + ## default + \ No newline at end of file diff --git a/module_mrfa/configs/vox1.yaml b/module_mrfa/configs/vox1.yaml new file mode 100644 index 0000000..e81f029 --- /dev/null +++ b/module_mrfa/configs/vox1.yaml @@ -0,0 +1,186 @@ +dataset_params: + root_dir: ./data/voxceleb1-png/ + split_len: + frame_shape: [256, 256, 3] + id_sampling: True + pairs_list: + augmentation_params: + flip_param: + horizontal_flip: True + time_flip: True + jitter_param: + brightness: 0.1 + contrast: 0.1 + saturation: 0.1 + hue: 0.1 + +dense_motion: + block_expansion: 64 + max_features: 1024 + num_blocks: 5 + scale_factor: 0.25 + num_kp: 10 + num_channels: 3 + estimate_occlusion_map: True +fomm_kp_detector: + block_expansion: 32 + num_kp: 10 + num_channels: 3 + max_features: 1024 + num_blocks: 5 + temperature: 0.1 + scale_factor: 0.25 + estimate_jacobian: True + estimate_occlusion: False +tpsm_kp_detector: + num_tps: 10 +tpsm_dense_motion: + block_expansion: 64 + max_features: 1024 + num_blocks: 5 + scale_factor: 0.25 + num_tps: 10 + num_channels: 3 + +raft_flow: + prior_only: False + num_kp: 10 + dim: 256 + size: 256 + generator: + num_channels: 3 + block_expansion: 64 + max_features: 512 + num_up_blocks: 5 + driving_encoder: + in_features: 10 + block_expansion: 32 + max_features: 512 + num_blocks: 5 + source_encoder: + in_features: 13 + block_expansion: 32 + max_features: 512 + num_blocks: 5 + +train_params: + ## [mtia, tpsm, fomm] + prior_model: mtia + num_epochs: 100 + num_repeats: 150 + epoch_milestones: [60, 90] + # ### enable this for tpsm + # epoch_milestones: [70, 90] + # ### enable this for tpsm + + lr: 2.0e-4 + batch_size: 80 + scales: [1, 0.5, 0.25, 0.125] + clip_grad: True + clip: 10 + bg_start: 1000 + checkpoint_freq: 100 + ### for tpsm + dropout_epoch: 35 + dropout_maxp: 0.3 + dropout_startp: 0.1 + dropout_inc_epoch: 10 + ### for tpsm + + transform_params: + sigma_affine: 0.05 + sigma_tps: 0.005 + points_tps: 5 + loss_weights: + perceptual: [10, 10, 10, 10, 10] + equivariance: 10 + equivariance_jacobian: 10 + # ### enable this for tpsm + # equivariance_jacobian: 0 + # ### enable this for tpsm + +reconstruction_params: + num_videos: 1000 + format: '.mp4' + +animate_params: + num_pairs: 50 + format: '.mp4' + normalization_params: + adapt_movement_scale: False + use_relative_movement: True + use_relative_jacobian: True + +visualizer_params: + draw_border: True + colormap: 'gist_rainbow' + +mtia_kp_detector: + MODEL: + ESTIMATE_JACOBIAN: True + DATA_PREPROCESS: False + FIX_IMG2MOTION_ATTENTION: False + TRANSFORMER_DEPTH: 12 + TRANSFORMER_HEADS: 8 + DIM: 192 # 4*4*3 + INIT_WEIGHTS: False + NAME: pose_tokenpose_b + NUM_JOINTS: 10 + PRETRAINED: '' + PATCH_SIZE: [4, 4] + IMAGE_SIZE: [256, 256] + HEATMAP_SIZE: [64, 64] + + ## default + TAG_PER_JOINT: True + HIDDEN_HEATMAP_DIM: -1 + MULTI_TRANSFORMER_DEPTH: [12, 12] + MULTI_TRANSFORMER_HEADS: [16, 16] + MULTI_DIM: [48, 48] + NUM_BRANCHES: 1 + BASE_CHANNEL: 32 + TRANSFORMER_MLP_RATIO: 3 + POS_EMBEDDING_TYPE: 'sine-full' + TEMPERATURE: 0.1 + TARGET_TYPE: gaussian + INIT: True + + SIGMA: 2 + EXTRA: + PRETRAINED_LAYERS: + - 'conv1' + - 'bn1' + - 'conv2' + - 'bn2' + - 'layer1' + - 'transition1' + - 'stage2' + - 'transition2' + - 'stage3' + FINAL_CONV_KERNEL: 1 + STAGE2: + NUM_MODULES: 1 + NUM_BRANCHES: 2 + BLOCK: BASIC + NUM_BLOCKS: + - 4 + - 4 + NUM_CHANNELS: + - 32 + - 64 + FUSE_METHOD: SUM + STAGE3: + NUM_MODULES: 4 + NUM_BRANCHES: 3 + BLOCK: BASIC + NUM_BLOCKS: + - 4 + - 4 + - 4 + NUM_CHANNELS: + - 32 + - 64 + - 128 + FUSE_METHOD: SUM + ## default + \ No newline at end of file diff --git a/module_mrfa/dense_motion.py b/module_mrfa/dense_motion.py new file mode 100644 index 0000000..5cf758d --- /dev/null +++ b/module_mrfa/dense_motion.py @@ -0,0 +1,435 @@ +import math + +import torch +import torch.nn.functional as F +from torch import nn + +from .util import ( + TPS, + AntiAliasInterpolation2d, + Hourglass, + UpBlock2d, + from_homogeneous, + kp2gaussian, + make_coordinate_grid, + to_homogeneous, +) + + +class DenseMotionNetwork(nn.Module): + """ + Module that predicting a dense motion from sparse motion representation given by kp_source and kp_driving + """ + + def __init__( + self, + block_expansion, + num_blocks, + max_features, + num_kp, + num_channels, + estimate_occlusion_map=True, + scale_factor=1, + kp_variance=0.01, + ): + super(DenseMotionNetwork, self).__init__() + + infeatures = num_kp + 1 + self.infeatures = infeatures + + self.hourglass = Hourglass( + block_expansion=block_expansion, + in_features=infeatures * (num_channels + 1), + max_features=max_features, + num_blocks=num_blocks, + ) + self.mask = nn.Conv2d( + self.hourglass.out_filters, infeatures, kernel_size=(7, 7), padding=(3, 3) + ) + + if estimate_occlusion_map: + self.occlusion = nn.Conv2d( + self.hourglass.out_filters, 1, kernel_size=(7, 7), padding=(3, 3) + ) + else: + self.occlusion = None + + self.num_kp = num_kp + self.scale_factor = scale_factor + self.kp_variance = kp_variance + + if self.scale_factor != 1: + self.down = AntiAliasInterpolation2d(num_channels, self.scale_factor) + + def create_heatmap_representations(self, source_image, kp_driving, kp_source): + spatial_size = source_image.shape[2:] + gaussian_driving = kp2gaussian( + kp_driving["kp"], spatial_size=spatial_size, kp_variance=self.kp_variance + ) + gaussian_source = kp2gaussian( + kp_source["kp"], spatial_size=spatial_size, kp_variance=self.kp_variance + ) + heatmap = gaussian_driving - gaussian_source + + # adding background feature + zeros = torch.zeros(heatmap.shape[0], 1, spatial_size[0], spatial_size[1]).type( + heatmap.type() + ) + heatmap = torch.cat([zeros, heatmap], dim=1) + heatmap = heatmap.unsqueeze(2) + return heatmap + + def create_sparse_motions(self, source_image, kp_driving, kp_source, bg_param=None): + bs, _, h, w = source_image.shape + identity_grid = make_coordinate_grid((h, w), type=kp_source["kp"].type()) + identity_grid = identity_grid.view(1, 1, h, w, 2) + coordinate_grid = identity_grid - kp_driving["kp"].view( + bs, self.num_kp, 1, 1, 2 + ) + if "jacobian" in kp_driving: + jacobian = torch.matmul( + kp_source["jacobian"], torch.inverse(kp_driving["jacobian"]) + ) + jacobian = jacobian.unsqueeze(-3).unsqueeze(-3) + jacobian = jacobian.repeat(1, 1, h, w, 1, 1) + coordinate_grid = torch.matmul(jacobian, coordinate_grid.unsqueeze(-1)) + coordinate_grid = coordinate_grid.squeeze(-1) + # n * 10 * h * w * 2 + + driving_to_source = coordinate_grid + kp_source["kp"].view( + bs, self.num_kp, 1, 1, 2 + ) + + # #adding background feature + # identity_grid = identity_grid.repeat(bs, 1, 1, 1, 1) + + # adding background feature + if bg_param is None: + bg_grid = identity_grid.repeat(bs, 1, 1, 1, 1) + else: + bg_grid = identity_grid.repeat(bs, 1, 1, 1, 1) + bg_grid = to_homogeneous(bg_grid) + bg_grid = torch.matmul( + bg_param.view(bs, 1, 1, 1, 3, 3), bg_grid.unsqueeze(-1) + ).squeeze(-1) + bg_grid = from_homogeneous(bg_grid) + + sparse_motions = torch.cat([bg_grid, driving_to_source], dim=1) + return sparse_motions + + def create_deformed_source_image(self, source_image, sparse_motions): + bs, _, h, w = source_image.shape + source_repeat = ( + source_image.unsqueeze(1) + .unsqueeze(1) + .repeat(1, self.num_kp + 1, 1, 1, 1, 1) + ) + source_repeat = source_repeat.view(bs * (self.num_kp + 1), -1, h, w) + sparse_motions = sparse_motions.view((bs * (self.num_kp + 1), h, w, -1)) + sparse_deformed = F.grid_sample(source_repeat, sparse_motions) + sparse_deformed = sparse_deformed.view((bs, self.num_kp + 1, -1, h, w)) + return sparse_deformed + + def dropout_softmax(self, X, P): + """ + Dropout for TPS transformations. Eq(7) and Eq(8) in the paper. + """ + drop = ( + (torch.rand(X.shape[0], X.shape[1]) < (1 - P)).type(X.type()).to(X.device) + ) + drop[..., 0] = 1 + drop = drop.repeat(X.shape[2], X.shape[3], 1, 1).permute(2, 3, 0, 1) + + maxx = X.max(1).values.unsqueeze_(1) + X = X - maxx + X_exp = X.exp() + X[:, 1:, ...] /= 1 - P + mask_bool = drop == 0 + X_exp = X_exp.masked_fill(mask_bool, 0) + partition = X_exp.sum(dim=1, keepdim=True) + 1e-6 + return X_exp / partition + + def forward( + self, + source_image, + kp_driving, + kp_source, + bg_param=None, + dropout_flag=False, + dropout_p=0, + ): + if self.scale_factor != 1: + source_image = self.down(source_image) + + bs, _, h, w = source_image.shape + + out_dict = dict() + heatmap_representation = self.create_heatmap_representations( + source_image, kp_driving, kp_source + ) + sparse_motion = self.create_sparse_motions( + source_image, kp_driving, kp_source, bg_param=bg_param + ) + # n * 11 * h * w * 2 + deformed_source = self.create_deformed_source_image(source_image, sparse_motion) + # n * 11 * 3 * h * w + out_dict["sparse_deformed"] = deformed_source + + input = torch.cat([heatmap_representation, deformed_source], dim=2) + # n * 11(101,22) * 4 * h * w + input = input.view(bs, -1, h, w) + + prediction = self.hourglass(input) + + mask = self.mask(prediction) + out_dict["logit_mask"] = mask + if dropout_flag: + mask = self.dropout_softmax(mask, dropout_p) + else: + mask = F.softmax(mask, dim=1) + # mask = F.softmax(mask, dim=1) + out_dict["mask"] = mask + mask = mask.unsqueeze(2) + sparse_motion = sparse_motion.permute(0, 1, 4, 2, 3) + # n * 11 * 2 * h * w + deformation = (sparse_motion * mask).sum(dim=1) + deformation = deformation.permute(0, 2, 3, 1) + # n * h * w * 2 + + out_dict["deformation"] = deformation + + if self.occlusion is not None: + # occlusion_map = torch.sigmoid(self.occlusion(prediction)) + occlusion_map = self.occlusion(prediction) + out_dict["occlusion"] = occlusion_map + + return out_dict + + +class TPSDenseMotionNetwork(nn.Module): + """ + Module that estimating an optical flow and multi-resolution occlusion masks + from K TPS transformations and an affine transformation. + """ + + def __init__( + self, + block_expansion, + num_blocks, + max_features, + num_tps, + num_channels, + scale_factor=0.25, + bg=False, + multi_mask=False, + kp_variance=0.01, + ): + super(TPSDenseMotionNetwork, self).__init__() + + if scale_factor != 1: + self.down = AntiAliasInterpolation2d(num_channels, scale_factor) + self.scale_factor = scale_factor + self.multi_mask = multi_mask + + self.hourglass = Hourglass( + block_expansion=block_expansion, + in_features=(num_channels * (num_tps + 1) + num_tps * 5 + 1), + max_features=max_features, + num_blocks=num_blocks, + ) + + hourglass_output_size = self.hourglass.out_filters + # self.maps = nn.Conv2d(hourglass_output_size[-1], num_tps + 1, kernel_size=(7, 7), padding=(3, 3)) + self.maps = nn.Conv2d( + hourglass_output_size, num_tps + 1, kernel_size=(7, 7), padding=(3, 3) + ) + + if multi_mask: + up = [] + self.up_nums = int(math.log(1 / scale_factor, 2)) + self.occlusion_num = 4 + + channel = [hourglass_output_size[-1] // (2**i) for i in range(self.up_nums)] + for i in range(self.up_nums): + up.append( + UpBlock2d(channel[i], channel[i] // 2, kernel_size=3, padding=1) + ) + self.up = nn.ModuleList(up) + + channel = [ + hourglass_output_size[-i - 1] + for i in range(self.occlusion_num - self.up_nums)[::-1] + ] + for i in range(self.up_nums): + channel.append(hourglass_output_size[-1] // (2 ** (i + 1))) + occlusion = [] + + for i in range(self.occlusion_num): + occlusion.append( + nn.Conv2d(channel[i], 1, kernel_size=(7, 7), padding=(3, 3)) + ) + self.occlusion = nn.ModuleList(occlusion) + else: + # occlusion = [nn.Conv2d(hourglass_output_size[-1], 1, kernel_size=(7, 7), padding=(3, 3))] + occlusion = [ + nn.Conv2d(hourglass_output_size, 1, kernel_size=(7, 7), padding=(3, 3)) + ] + self.occlusion = nn.ModuleList(occlusion) + + self.num_tps = num_tps + self.bg = bg + self.kp_variance = kp_variance + + def create_heatmap_representations(self, source_image, kp_driving, kp_source): + spatial_size = source_image.shape[2:] + gaussian_driving = kp2gaussian( + kp_driving["kp"], spatial_size=spatial_size, kp_variance=self.kp_variance + ) + gaussian_source = kp2gaussian( + kp_source["kp"], spatial_size=spatial_size, kp_variance=self.kp_variance + ) + heatmap = gaussian_driving - gaussian_source + + zeros = ( + torch.zeros(heatmap.shape[0], 1, spatial_size[0], spatial_size[1]) + .type(heatmap.type()) + .to(heatmap.device) + ) + heatmap = torch.cat([zeros, heatmap], dim=1) + + return heatmap + + def create_transformations(self, source_image, kp_driving, kp_source, bg_param): + # K TPS transformaions + bs, _, h, w = source_image.shape + kp_1 = kp_driving["kp"] + kp_2 = kp_source["kp"] + kp_1 = kp_1.view(bs, -1, 5, 2) + kp_2 = kp_2.view(bs, -1, 5, 2) + trans = TPS(mode="kp", bs=bs, kp_1=kp_1, kp_2=kp_2) + driving_to_source = trans.transform_frame(source_image) + + identity_grid = make_coordinate_grid((h, w), type=kp_1.type()).to(kp_1.device) + identity_grid = identity_grid.view(1, 1, h, w, 2) + identity_grid = identity_grid.repeat(bs, 1, 1, 1, 1) + + # affine background transformation + if bg_param is not None: + identity_grid = to_homogeneous(identity_grid) + identity_grid = torch.matmul( + bg_param.view(bs, 1, 1, 1, 3, 3), identity_grid.unsqueeze(-1) + ).squeeze(-1) + identity_grid = from_homogeneous(identity_grid) + + transformations = torch.cat([identity_grid, driving_to_source], dim=1) + return transformations + + def create_deformed_source_image(self, source_image, transformations): + bs, _, h, w = source_image.shape + source_repeat = ( + source_image.unsqueeze(1) + .unsqueeze(1) + .repeat(1, self.num_tps + 1, 1, 1, 1, 1) + ) + source_repeat = source_repeat.view(bs * (self.num_tps + 1), -1, h, w) + transformations = transformations.view((bs * (self.num_tps + 1), h, w, -1)) + deformed = F.grid_sample(source_repeat, transformations, align_corners=True) + deformed = deformed.view((bs, self.num_tps + 1, -1, h, w)) + return deformed + + def dropout_softmax(self, X, P): + """ + Dropout for TPS transformations. Eq(7) and Eq(8) in the paper. + """ + drop = ( + (torch.rand(X.shape[0], X.shape[1]) < (1 - P)).type(X.type()).to(X.device) + ) + drop[..., 0] = 1 + drop = drop.repeat(X.shape[2], X.shape[3], 1, 1).permute(2, 3, 0, 1) + + maxx = X.max(1).values.unsqueeze_(1) + X = X - maxx + X_exp = X.exp() + X[:, 1:, ...] /= 1 - P + mask_bool = drop == 0 + X_exp = X_exp.masked_fill(mask_bool, 0) + partition = X_exp.sum(dim=1, keepdim=True) + 1e-6 + return X_exp / partition + + def forward( + self, + source_image, + kp_driving, + kp_source, + bg_param=None, + dropout_flag=False, + dropout_p=0, + ): + if self.scale_factor != 1: + source_image = self.down(source_image) + + bs, _, h, w = source_image.shape + + out_dict = dict() + heatmap_representation = self.create_heatmap_representations( + source_image, kp_driving, kp_source + ) + transformations = self.create_transformations( + source_image, kp_driving, kp_source, bg_param + ) + deformed_source = self.create_deformed_source_image( + source_image, transformations + ) + out_dict["deformed_source"] = deformed_source + # out_dict['transformations'] = transformations + deformed_source = deformed_source.view(bs, -1, h, w) + input = torch.cat([heatmap_representation, deformed_source], dim=1) + input = input.view(bs, -1, h, w) + + # prediction = self.hourglass(input, mode = 1) + prediction = self.hourglass(input) + + # contribution_maps = self.maps(prediction[-1]) + contribution_maps = self.maps(prediction) + if dropout_flag: + contribution_maps = self.dropout_softmax(contribution_maps, dropout_p) + else: + contribution_maps = F.softmax(contribution_maps, dim=1) + out_dict["contribution_maps"] = contribution_maps + out_dict["mask"] = contribution_maps + + # Combine the K+1 transformations + # Eq(6) in the paper + contribution_maps = contribution_maps.unsqueeze(2) + transformations = transformations.permute(0, 1, 4, 2, 3) + deformation = (transformations * contribution_maps).sum(dim=1) + deformation = deformation.permute(0, 2, 3, 1) + + out_dict["deformation"] = deformation # Optical Flow + + occlusion_map = [] + if self.multi_mask: + for i in range(self.occlusion_num - self.up_nums): + occlusion_map.append( + torch.sigmoid( + self.occlusion[i]( + prediction[self.up_nums - self.occlusion_num + i] + ) + ) + ) + prediction = prediction[-1] + for i in range(self.up_nums): + prediction = self.up[i](prediction) + occlusion_map.append( + torch.sigmoid( + self.occlusion[i + self.occlusion_num - self.up_nums]( + prediction + ) + ) + ) + else: + # occlusion_map.append(torch.sigmoid(self.occlusion[0](prediction[-1]))) + occlusion_map = self.occlusion[0](prediction) + + out_dict["occlusion"] = occlusion_map # Multi-resolution Occlusion Masks + return out_dict diff --git a/module_mrfa/generator.py b/module_mrfa/generator.py new file mode 100644 index 0000000..f632cc4 --- /dev/null +++ b/module_mrfa/generator.py @@ -0,0 +1,88 @@ +# import math + +import torch +# import torch.nn.functional as F +from torch import nn + +from .util import ChannelBlock2d, DownBlock2d, ResBlock2d, SameBlock2d, UpBlock2d + + +class OcclusionAwareGenerator(nn.Module): + def __init__(self, num_channels, block_expansion, max_features, num_up_blocks): + super(OcclusionAwareGenerator, self).__init__() + + self.num_up_blocks = num_up_blocks + self.first = SameBlock2d( + num_channels, block_expansion, kernel_size=(7, 7), padding=(3, 3) + ) + down_blocks = [] + up_blocks = [] + resblock = [] + channel_block = [] + for i in range(num_up_blocks): + in_features = min(max_features, block_expansion * (2**i)) + out_features = min(max_features, block_expansion * (2 ** (i + 1))) + down_blocks.append( + DownBlock2d( + in_features, out_features, kernel_size=(3, 3), padding=(1, 1) + ) + ) + decoder_in_feature = out_features + if i == num_up_blocks - 1: + decoder_in_feature = out_features + up_blocks.append( + UpBlock2d( + decoder_in_feature, in_features, kernel_size=(3, 3), padding=(1, 1) + ) + ) + resblock.append( + ResBlock2d(decoder_in_feature, kernel_size=(3, 3), padding=(1, 1)) + ) + channel_block.append( + ChannelBlock2d( + decoder_in_feature * 2, kernel_size=(3, 3), padding=(1, 1) + ) + ) + self.down_blocks = nn.ModuleList(down_blocks) + self.up_blocks = nn.ModuleList(up_blocks[::-1]) + self.resblock = nn.ModuleList(resblock[::-1]) + self.channel_block = nn.ModuleList(channel_block[::-1]) + self.final = nn.Conv2d( + block_expansion, num_channels, kernel_size=(7, 7), padding=(3, 3) + ) + + def encode(self, x): + f = [] + out = self.first(x) + f.append(out) + for i in range(len(self.down_blocks)): + out = self.down_blocks[i](out) + f.append(out) + # 256 128 64 32,...,16,8 + return f[::-1] + + def decode(self, warp_f, warp_img, occlusion, warp_f_c=None, occlusion_c=None): + # warp_f 32 64 128 256 + out = warp_f[0] * occlusion[0] + # out = warp_f[0] + if warp_f_c is not None: + out_c = warp_f_c[0] + out = torch.cat([out, out_c], dim=1) + for i in range(self.num_up_blocks): + if warp_f_c is not None: + out = self.channel_block[i](out) + out = self.resblock[i](out) + out = self.up_blocks[i](out) + out = warp_f[i + 1] * occlusion[i + 1] + out * (1 - occlusion[i + 1]) + if warp_f_c is not None and i != self.num_up_blocks - 1: + out_c = warp_f_c[i + 1] + out = torch.cat([out, out_c], dim=1) + out = self.final(out) + out = torch.sigmoid(out) + out = out * (1 - occlusion[-1]) + warp_img * occlusion[-1] + return out + + def forward(self, x): + x = self.encode(x) + x = self.decode(x) + return x diff --git a/module_mrfa/kp_detector.py b/module_mrfa/kp_detector.py new file mode 100644 index 0000000..95f69bb --- /dev/null +++ b/module_mrfa/kp_detector.py @@ -0,0 +1,213 @@ +# from collections import OrderedDict + +import torch +import torch.nn.functional as F +from torch import nn + +# from torch.jit.annotations import Dict +# from torch.nn import BatchNorm2d +from torchvision import models + +# from torchvision.models._utils import IntermediateLayerGetter +from .util import AntiAliasInterpolation2d, DownBlock2d, Hourglass, make_coordinate_grid + + +class KPDetector(nn.Module): + """ + Detecting a keypoints. Return keypoint position and jacobian near each keypoint. + """ + + def __init__( + self, + block_expansion=32, + num_kp=15, + num_channels=3, + max_features=1024, + num_blocks=5, + temperature=0.1, + scale_factor=0.25, + estimate_jacobian=False, + estimate_occlusion=False, + ): + super(KPDetector, self).__init__() + + self.predictor = Hourglass( + block_expansion, + in_features=num_channels, + max_features=max_features, + num_blocks=num_blocks, + ) + + self.kp = nn.Conv2d( + in_channels=self.predictor.out_filters, + out_channels=num_kp, + kernel_size=(7, 7), + padding=0, + ) + self.num_kp = num_kp + + self.estimate_jacobian = estimate_jacobian + if estimate_jacobian: + self.num_jacobian_maps = 1 + self.jacobian = nn.Conv2d( + in_channels=self.predictor.out_filters, + out_channels=4 * self.num_jacobian_maps, + kernel_size=(7, 7), + padding=0, + ) + self.jacobian.weight.data.zero_() + self.jacobian.bias.data.copy_( + torch.tensor([1, 0, 0, 1] * self.num_jacobian_maps, dtype=torch.float) + ) + + self.estimate_occlusion = estimate_occlusion + if estimate_occlusion: + kp_occlusion = [ + DownBlock2d( + self.predictor.out_filters, + block_expansion, + kernel_size=3, + padding=1, + ) + ] + kp_occlusion.append( + DownBlock2d( + block_expansion, block_expansion * 2, kernel_size=3, padding=1 + ) + ) + kp_occlusion.append( + DownBlock2d( + block_expansion * 2, block_expansion * 3, kernel_size=3, padding=1 + ) + ) + kp_occlusion.append( + DownBlock2d( + block_expansion * 3, block_expansion * 4, kernel_size=3, padding=1 + ) + ) + kp_occlusion.append( + nn.Conv2d( + block_expansion * 4, num_kp, kernel_size=(4, 4), padding=0, stride=4 + ) + ) + self.kp_occlusion = nn.Sequential(*kp_occlusion) + + self.temperature = temperature + self.scale_factor = scale_factor + if self.scale_factor != 1: + self.down = AntiAliasInterpolation2d(num_channels, self.scale_factor) + + def create_part_flows(self, source_image, kp_driving, kp_source, bg_params=None): + """ + Eq 4. in the paper T_{s<-d}(z) + """ + bs, _, h, w = source_image.shape + num_kp = kp_driving["kp"].shape[1] + identity_grid = make_coordinate_grid((h, w), type=kp_source["kp"].type()) + identity_grid = identity_grid.view(1, 1, h, w, 2) + coordinate_grid = identity_grid - kp_driving["kp"].view(bs, num_kp, 1, 1, 2) + if "jacobian" in kp_driving: + jacobian = torch.matmul( + kp_source["jacobian"], torch.inverse(kp_driving["jacobian"]) + ) + jacobian = jacobian.unsqueeze(-3).unsqueeze(-3) + jacobian = jacobian.repeat(1, 1, h, w, 1, 1) + coordinate_grid = torch.matmul(jacobian, coordinate_grid.unsqueeze(-1)) + coordinate_grid = coordinate_grid.squeeze(-1) + # n * 10 * h * w * 2 + driving_to_source = coordinate_grid + kp_source["kp"].view(bs, num_kp, 1, 1, 2) + flow = driving_to_source - identity_grid + flow = (h - 1) * flow / 2.0 + return flow.view(bs, num_kp * 2, h, w) + + # # #adding background feature + # # identity_grid = identity_grid.repeat(bs, 1, 1, 1, 1) + + # # adding background feature + # if bg_params is None: + # bg_grid = identity_grid.repeat(bs, 1, 1, 1, 1) + # else: + # bg_grid = identity_grid.repeat(bs, 1, 1, 1, 1) + # bg_grid = to_homogeneous(bg_grid) + # bg_grid = torch.matmul(bg_params.view(bs, 1, 1, 1, 3, 3), bg_grid.unsqueeze(-1)).squeeze(-1) + # bg_grid = from_homogeneous(bg_grid) + # sparse_motions = torch.cat([bg_grid, driving_to_source], dim=1) + # return sparse_motions + + def gaussian2kp(self, heatmap): + """ + Extract the mean and from a heatmap + """ + shape = heatmap.shape + heatmap = heatmap.unsqueeze(-1) + grid = ( + make_coordinate_grid(shape[2:], heatmap.type()).unsqueeze_(0).unsqueeze_(0) + ) + kp = (heatmap * grid).sum(dim=(2, 3)) # N * 10 * 2 + # kp = {'value': value} + + return kp + + def forward(self, x): + if self.scale_factor != 1: + x = self.down(x) + feature_map = self.predictor(x) + prediction = self.kp(feature_map) + final_shape = prediction.shape + heatmap = prediction.view(final_shape[0], final_shape[1], -1) + heatmap = F.softmax(heatmap / self.temperature, dim=2) + heatmap = heatmap.view(*final_shape) + out = self.gaussian2kp(heatmap) + if self.estimate_jacobian: + jacobian_map = self.jacobian(feature_map) + jacobian_map = jacobian_map.reshape( + final_shape[0], + self.num_jacobian_maps, + 4, + final_shape[2], + final_shape[3], + ) + heatmap = heatmap.unsqueeze(2) + jacobian = heatmap * jacobian_map + jacobian = jacobian.view(final_shape[0], final_shape[1], 4, -1) + jacobian = jacobian.sum(dim=-1) + jacobian = jacobian.view(jacobian.shape[0], jacobian.shape[1], 2, 2) + # N * 10 *2 *2 + out = {"kp": out, "jacobian": jacobian} + else: + out = {"kp": out} + + if self.estimate_occlusion: + # N*10*1*1 + kp_occlusion = self.kp_occlusion(feature_map) + kp_occlusion = torch.sigmoid(kp_occlusion) + out["kp_occlusion"] = kp_occlusion + # out['heatmap'] = heatmap + # out['feature_map'] = feature_map + return out + + +class TPSKPDetector(nn.Module): + """ + Predict K*5 keypoints. + """ + + def __init__(self, num_tps, **kwargs): + super(TPSKPDetector, self).__init__() + self.num_tps = num_tps + + self.fg_encoder = models.resnet18(pretrained=False) + num_features = self.fg_encoder.fc.in_features + self.fg_encoder.fc = nn.Linear(num_features, num_tps * 5 * 2) + + def forward(self, image): + fg_kp = self.fg_encoder(image) + ( + bs, + _, + ) = fg_kp.shape + fg_kp = torch.sigmoid(fg_kp) + fg_kp = fg_kp * 2 - 1 + out = {"kp": fg_kp.view(bs, self.num_tps * 5, -1)} + + return out diff --git a/module_mrfa/model.py b/module_mrfa/model.py new file mode 100644 index 0000000..3c5a1a1 --- /dev/null +++ b/module_mrfa/model.py @@ -0,0 +1,294 @@ +# from __future__ import absolute_import, division, print_function + +# import logging + +# import timm +# import math +# import os + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.autograd import grad +from torchvision import models + +from .bg_motion_predictor import BGMotionPredictor +from .dense_motion import DenseMotionNetwork, TPSDenseMotionNetwork + +# from .generator import OcclusionAwareGenerator +from .kp_detector import KPDetector, TPSKPDetector +from .raft import RaftFlow +from .transformer.pose_tokenpose_b import get_pose_net +from .util import AntiAliasInterpolation2d, make_coordinate_grid + + +class Transform: + """ + Random tps transformation for equivariance constraints. See Sec 3.3 + """ + + def __init__(self, bs, **kwargs): + noise = torch.normal( + mean=0, std=kwargs["sigma_affine"] * torch.ones([bs, 2, 3]) + ) + self.theta = noise + torch.eye(2, 3).view(1, 2, 3) + self.bs = bs + + if ("sigma_tps" in kwargs) and ("points_tps" in kwargs): + self.tps = True + self.control_points = make_coordinate_grid( + (kwargs["points_tps"], kwargs["points_tps"]), type=noise.type() + ) + self.control_points = self.control_points.unsqueeze(0) + self.control_params = torch.normal( + mean=0, + std=kwargs["sigma_tps"] + * torch.ones([bs, 1, kwargs["points_tps"] ** 2]), + ) + else: + self.tps = False + + def transform_frame(self, frame): + grid = make_coordinate_grid(frame.shape[2:], type=frame.type()).unsqueeze(0) + grid = grid.view(1, frame.shape[2] * frame.shape[3], 2) + grid = self.warp_coordinates(grid).view( + self.bs, frame.shape[2], frame.shape[3], 2 + ) + return F.grid_sample(frame, grid, padding_mode="reflection") + + def warp_coordinates(self, coordinates): + theta = self.theta.type(coordinates.type()) + theta = theta.unsqueeze(1) + # N * (HW) * 2 * 1 + transformed = ( + torch.matmul(theta[:, :, :, :2], coordinates.unsqueeze(-1)) + + theta[:, :, :, 2:] + ) + transformed = transformed.squeeze(-1) + + if self.tps: + control_points = self.control_points.type(coordinates.type()) + control_params = self.control_params.type(coordinates.type()) + # 1 * (HW) * 1 * 2 - 1 * 1 * 25 *2 + distances = coordinates.view( + coordinates.shape[0], -1, 1, 2 + ) - control_points.view(1, 1, -1, 2) + distances = torch.abs(distances).sum(-1) + + result = distances**2 + result = result * torch.log(distances + 1e-6) + result = result * control_params + result = result.sum(dim=2).view(self.bs, coordinates.shape[1], 1) + transformed = transformed + result + + return transformed + + def jacobian(self, coordinates): + new_coordinates = self.warp_coordinates(coordinates) + grad_x = grad(new_coordinates[..., 0].sum(), coordinates, create_graph=True) + grad_y = grad(new_coordinates[..., 1].sum(), coordinates, create_graph=True) + jacobian = torch.cat([grad_x[0].unsqueeze(-2), grad_y[0].unsqueeze(-2)], dim=-2) + return jacobian + + +class Vgg19(torch.nn.Module): + """ + Vgg19 network for perceptual loss. See Sec 3.3. + """ + + def __init__(self, requires_grad=False): + super(Vgg19, self).__init__() + vgg_pretrained_features = models.vgg19(pretrained=True).features + self.slice1 = torch.nn.Sequential() + self.slice2 = torch.nn.Sequential() + self.slice3 = torch.nn.Sequential() + self.slice4 = torch.nn.Sequential() + self.slice5 = torch.nn.Sequential() + for x in range(2): + self.slice1.add_module(str(x), vgg_pretrained_features[x]) + for x in range(2, 7): + self.slice2.add_module(str(x), vgg_pretrained_features[x]) + for x in range(7, 12): + self.slice3.add_module(str(x), vgg_pretrained_features[x]) + for x in range(12, 21): + self.slice4.add_module(str(x), vgg_pretrained_features[x]) + for x in range(21, 30): + self.slice5.add_module(str(x), vgg_pretrained_features[x]) + + self.mean = torch.nn.Parameter( + data=torch.Tensor(np.array([0.485, 0.456, 0.406]).reshape((1, 3, 1, 1))), + requires_grad=False, + ) + self.std = torch.nn.Parameter( + data=torch.Tensor(np.array([0.229, 0.224, 0.225]).reshape((1, 3, 1, 1))), + requires_grad=False, + ) + + if not requires_grad: + for param in self.parameters(): + param.requires_grad = False + + def forward(self, X): + # X = X.clamp(-1, 1) + # X = X / 2 + 0.5 + X = (X - self.mean) / self.std + h_relu1 = self.slice1(X) + h_relu2 = self.slice2(h_relu1) + h_relu3 = self.slice3(h_relu2) + h_relu4 = self.slice4(h_relu3) + h_relu5 = self.slice5(h_relu4) + out = [h_relu1, h_relu2, h_relu3, h_relu4, h_relu5] + return out + + +class ImagePyramide(torch.nn.Module): + """ + Create image pyramide for computing pyramide perceptual loss. See Sec 3.3 + """ + + def __init__(self, scales, num_channels=3): + super(ImagePyramide, self).__init__() + downs = {} + for scale in scales: + downs[str(scale).replace(".", "-")] = AntiAliasInterpolation2d( + num_channels, scale + ) + self.downs = nn.ModuleDict(downs) + + def forward(self, x): + out_dict = {} + for scale, down_module in self.downs.items(): + out_dict["prediction_" + str(scale).replace("-", ".")] = down_module(x) + return out_dict + + +class MRFA(nn.Module): + def __init__(self, cfg, **kwargs): + super(MRFA, self).__init__() + + self.cfg = cfg + train_params = cfg["train_params"] + self.train_params = train_params + self.scales = train_params["scales"] + self.loss_weights = train_params["loss_weights"] + self.pyramid = ImagePyramide(self.scales).cuda() + if sum(self.loss_weights["perceptual"]) != 0: + self.vgg = Vgg19().cuda() + + ################################################## + if train_params["prior_model"] == "fomm": + self.encoder = KPDetector(**cfg["fomm_kp_detector"]) + self.dense_motion = DenseMotionNetwork(**cfg["dense_motion"]) + elif train_params["prior_model"] == "tpsm": + self.encoder = TPSKPDetector(**cfg["tpsm_kp_detector"]) + self.dense_motion = TPSDenseMotionNetwork(**cfg["tpsm_dense_motion"]) + self.dropout_epoch = train_params["dropout_epoch"] + self.dropout_maxp = train_params["dropout_maxp"] + self.dropout_startp = train_params["dropout_startp"] + self.dropout_inc_epoch = train_params["dropout_inc_epoch"] + elif train_params["prior_model"] == "mtia": + self.encoder = get_pose_net(cfg["mtia_kp_detector"], is_train=True) + self.dense_motion = DenseMotionNetwork(**cfg["dense_motion"]) + else: + print("please specify the prior motion model") + + if train_params["bg_start"] < train_params["num_epochs"]: + self.bg_predictor = BGMotionPredictor() + self.bg_start = train_params["bg_start"] + + self.decoder = RaftFlow(**cfg["raft_flow"]) + self.down = AntiAliasInterpolation2d(3, 0.25) + + def forward(self, x, epoch=100, is_train=True): + kp_s = self.encoder(x["source"]) + kp_d = self.encoder(x["driving"]) + + img_down = self.down(x["source"]) + if epoch >= self.bg_start: + bg_param = self.bg_predictor(x["source"], x["driving"]) + else: + bg_param = None + + if self.train_params["prior_model"] == "tpsm" and epoch < self.dropout_epoch: + # dropout_p will linearly increase from dropout_startp to dropout_maxp + dropout_flag = True + dropout_p = min( + epoch / self.dropout_inc_epoch * self.dropout_maxp + + self.dropout_startp, + self.dropout_maxp, + ) + else: + dropout_flag = False + dropout_p = 0 + dense_motion = self.dense_motion( + x["source"], + kp_d, + kp_s, + bg_param=bg_param, + dropout_flag=dropout_flag, + dropout_p=dropout_p, + ) + + if self.train_params["prior_model"] == "tpsm": + kp_s_value = kp_s["kp"].view(x["source"].shape[0], -1, 5, 2).mean(2) + kp_d_value = kp_d["kp"].view(x["driving"].shape[0], -1, 5, 2).mean(2) + else: + kp_s_value = kp_s["kp"] + kp_d_value = kp_d["kp"] + + gen, warp_img, occlusion = self.decoder( + kp_s_value, kp_d_value, dense_motion, img=img_down, img_full=x["source"] + ) + warp_img = torch.cat([warp_img, occlusion.repeat(1, 3, 1, 1)], dim=3) + + loss_values = {} + + if not is_train: + return gen, warp_img, loss_values, kp_s["kp"], kp_d["kp"] + + pyramide_real = self.pyramid(x["driving"]) + pyramide_generated = self.pyramid(gen) + value_total = 0 + for scale in self.scales: + x_vgg = self.vgg(pyramide_generated["prediction_" + str(scale)]) + y_vgg = self.vgg(pyramide_real["prediction_" + str(scale)]) + + for i, weight in enumerate(self.loss_weights["perceptual"]): + value = torch.abs(x_vgg[i] - y_vgg[i].detach()).mean() + value_total += self.loss_weights["perceptual"][i] * value + loss_values["perceptual"] = value_total + + if self.loss_weights["equivariance"] != 0: + transform = Transform( + x["driving"].shape[0], **self.train_params["transform_params"] + ) + transformed_frame = transform.transform_frame(x["driving"]) + transformed_kp = self.encoder(transformed_frame) + value = torch.abs( + kp_d["kp"] - transform.warp_coordinates(transformed_kp["kp"]) + ) + value = value.mean() + loss_values["equivariance"] = self.loss_weights["equivariance"] * value + + if self.loss_weights["equivariance_jacobian"] != 0: + jacobian_transformed = torch.matmul( + transform.jacobian(transformed_kp["kp"]), transformed_kp["jacobian"] + ) + normed_driving = torch.inverse(kp_d["jacobian"]) + normed_transformed = jacobian_transformed + value = torch.matmul(normed_driving, normed_transformed) + eye = torch.eye(2).view(1, 1, 2, 2).type(value.type()) + value = torch.abs(eye - value) + loss_values["equivariance_jacobian"] = ( + self.loss_weights["equivariance_jacobian"] * value + ) + + if epoch >= self.bg_start: + bg_param_reverse = self.bg_predictor(x["driving"], x["source"]) + value = torch.matmul(bg_param, bg_param_reverse) + eye = torch.eye(3).view(1, 1, 3, 3).type(value.type()) + value = torch.abs(eye - value).mean() + loss_values["bg"] = 10 * value + + return gen, warp_img, loss_values, kp_s["kp"], kp_d["kp"] diff --git a/module_mrfa/models/.gitkeep b/module_mrfa/models/.gitkeep new file mode 100644 index 0000000..e69de29 diff --git a/module_mrfa/raft.py b/module_mrfa/raft.py new file mode 100644 index 0000000..6ee3b6b --- /dev/null +++ b/module_mrfa/raft.py @@ -0,0 +1,456 @@ +import math + +import torch +import torch.nn.functional as F +from einops import rearrange +from timm.models.layers.weight_init import trunc_normal_ +from torch import nn + +from .generator import OcclusionAwareGenerator +from .util import ( + # AntiAliasInterpolation2d, + Hourglass, + batch_bilinear_sampler, + bilinear_sampler, + coords_grid, + kp2gaussian, + # make_coordinate_grid, +) + + +class CorrBlock: + def __init__(self, corr, num_levels=2, radius=3): + self.num_levels = num_levels + self.radius = radius + self.corr_pyramid = [] + + self.corr_pyramid.append(corr) + for i in range(self.num_levels - 1): + corr = F.avg_pool2d(corr, 2, stride=2) + self.corr_pyramid.append(corr) + + def __call__(self, coords): + r = self.radius + coords = coords.permute(0, 2, 3, 1) + batch, h1, w1, _ = coords.shape + + out_pyramid = [] + for i in range(self.num_levels): + corr = self.corr_pyramid[i] + dx = torch.linspace(-r, r, 2 * r + 1, device=coords.device) + dy = torch.linspace(-r, r, 2 * r + 1, device=coords.device) + delta = torch.stack(torch.meshgrid(dy, dx), axis=-1) + + centroid_lvl = coords.reshape(batch * h1 * w1, 1, 1, 2) / 2**i + delta_lvl = delta.view(1, 2 * r + 1, 2 * r + 1, 2) + coords_lvl = centroid_lvl + delta_lvl + + if batch > 1 and h1 >= 128: + corr = batch_bilinear_sampler( + corr, coords_lvl, h=h1, w=w1, mini_batch=1 + ) + else: + corr = bilinear_sampler(corr, coords_lvl) + # corr = bilinear_sampler(corr, coords_lvl) + corr = corr.view(batch, h1, w1, -1) + out_pyramid.append(corr) + + out = torch.cat(out_pyramid, dim=-1) + return out.permute(0, 3, 1, 2).contiguous().float() + + +class BasicMotionEncoder(nn.Module): + def __init__(self, num_levels=2, radius=3): + super(BasicMotionEncoder, self).__init__() + cor_planes = num_levels * (2 * radius + 1) ** 2 + self.convc1 = nn.Conv2d(cor_planes, 128, 1, padding=0) + self.convc2 = nn.Conv2d(128, 96, 3, padding=1) + self.convf1 = nn.Conv2d(2, 128, 7, padding=3) + self.convf2 = nn.Conv2d(128, 64, 3, padding=1) + self.conv = nn.Conv2d(64 + 96, 128 - 2, 3, padding=1) + + def forward(self, delta_flow, corr): + cor = F.relu(self.convc1(corr)) + cor = F.relu(self.convc2(cor)) + flo = F.relu(self.convf1(delta_flow)) + flo = F.relu(self.convf2(flo)) + + cor_flo = torch.cat([cor, flo], dim=1) + out = F.relu(self.conv(cor_flo)) + return torch.cat([out, delta_flow], dim=1) + + +class RefineFlow(nn.Module): + def __init__(self): + super(RefineFlow, self).__init__() + self.convc1 = nn.Conv2d(192, 128, 3, padding=1) + # self.convc2 = nn.Conv2d(128, 64, 3, padding=1) + self.conv1 = nn.Conv2d(256, 128, 3, padding=1) + self.conv2 = nn.Conv2d(128, 2, 3, padding=1) + self.convo1 = nn.Conv2d(256, 128, 3, padding=1) + self.convo2 = nn.Conv2d(128, 1, 3, padding=1) + + def forward(self, m_f, warp_f): + c = F.relu(self.convc1(warp_f)) + ## c = F.relu(self.convc2(c)) + inp = torch.cat([m_f, c], dim=1) + flow = self.conv2(F.relu(self.conv1(inp))) + occ = self.convo2(F.relu(self.convo1(inp))) + out = torch.cat([flow, occ], dim=1) + return out, inp + # return flow, inp + + +class RaftFlow(nn.Module): + def __init__( + self, + prior_only=False, + num_kp=10, + dim=256, + size=256, + generator=None, + driving_encoder=None, + source_encoder=None, + ): + super(RaftFlow, self).__init__() + + self.scale = dim**-0.5 + self.size = size + ## basic flow resolution of the prior motion model + self.h = size // 4 + self.w = size // 4 + self.prior_only = prior_only + self.generator = OcclusionAwareGenerator(**generator) + # self.down = AntiAliasInterpolation2d(3, 0.25) + + ## start iteration from resolution of size//2**5 + channels = { + size // 32: 512, + size // 16: 512, + size // 8: 512, + size // 4: 256, + size // 2: 128, + size: 64, + # 512: 32, + } + self.total_iter = int(math.log(2**5, 2)) + 1 + self.num_iter = int(math.log(2**5, 2)) + 1 + self.basic_res_index = int(math.log(self.h // (size // 32), 2)) + + if not self.prior_only: + ### driving and source structure encoder + self.kp = Hourglass(**driving_encoder) + self.kp_img = Hourglass(**source_encoder) + self.kp_head = nn.Conv2d( + in_channels=self.kp.out_filters, + out_channels=dim, + kernel_size=1, + padding=0, + ) + self.kp_img_head = nn.Conv2d( + in_channels=self.kp_img.out_filters, + out_channels=dim, + kernel_size=1, + padding=0, + ) + # self._make_position_embedding(64, 64, num_kp, 'sine-full') + # self.pos_embedding = rearrange(self.pos_embedding, 'b (h w) c -> b c h w', h=self.h, w=self.w) + self.pos_embedding = nn.Parameter(torch.zeros(1, num_kp, self.h, self.w)) + trunc_normal_(self.pos_embedding, std=0.02) + ### driving and source structure encoder + + ### flow updater + self.corr_enc = BasicMotionEncoder() + self.refine = RefineFlow() + ## project different-scale warped features to the same dimension + self.to_context = nn.ModuleList() + for i in range(self.num_iter): + f_channel = channels[(size // 32) * (2**i)] + self.to_context.append(nn.Conv2d(f_channel, 192, 1, padding=0)) + ### flow updater + + def forward(self, kp_s, kp_d, dense_motion, img, img_full): + feature = self.generator.encode(img_full) + if img is None: + img = self.down(img_full) + b, _, h, w = img.shape + + out_warp_f = [] + out_occlusion = [] + out_flow = [] + + out_warp_f_c = [] + out_occlusion_c = [] + + ### prior only + if self.prior_only: + flow = dense_motion["deformation"] + occlusion = dense_motion["occlusion"] + for i in range(self.total_iter): + if flow.shape[2] != feature[i].shape[2]: + flow_res = F.interpolate( + flow.permute(0, 3, 1, 2), + size=feature[i].shape[2:], + mode="bilinear", + align_corners=True, + ) + occlusion_res = F.interpolate( + occlusion, + size=feature[i].shape[2:], + mode="bilinear", + align_corners=True, + ) + else: + flow_res = flow.permute(0, 3, 1, 2) + occlusion_res = occlusion + out_warp_f.append( + F.grid_sample(feature[i], flow_res.permute(0, 2, 3, 1)) + ) + out_occlusion.append(F.sigmoid(occlusion_res)) + warp_img = F.grid_sample(img_full, flow_res.permute(0, 2, 3, 1)) + out = self.generator.decode(out_warp_f, warp_img, out_occlusion) + for i in range(len(out_occlusion)): + out_occlusion[i] = F.interpolate( + out_occlusion[i], + size=self.size, + mode="bilinear", + align_corners=True, + ) + occlusion = torch.cat(out_occlusion, dim=3) + return out, warp_img, occlusion + ### prior only + + ### compute correlation volume at basic resolution + kp_s = kp2gaussian(kp_s, (h, w), 0.1) + self.pos_embedding + kp_d = kp2gaussian(kp_d, (h, w), 0.1) + self.pos_embedding + fe_s = self.kp_img(torch.cat([kp_s, img], dim=1)) + fe_d = self.kp(kp_d) + k_s = self.kp_img_head(fe_s) + q_d = self.kp_head(fe_d) + f_s = rearrange(k_s, "b c h w -> b (h w) c", h=self.h, w=self.w) + f_d = rearrange(q_d, "b c h w -> b (h w) c", h=self.h, w=self.w) + corr_volume = torch.einsum("bic,bjc->bij", f_d, f_s) * self.scale + ### compute correlation volume at basic resolution + + ### prior motion initialization + id_grid = coords_grid(b, self.h, self.w, corr_volume.device) + init_flow = (self.h - 1) * ( + dense_motion["deformation"].permute(0, 3, 1, 2) + 1 + ) / 2.0 - id_grid + init_occlusion = dense_motion["occlusion"] + ### prior motion initialization + + # ### NPMR only initialization + # id_grid = coords_grid(b, self.h, self.w, corr_volume.device) + # id_grid = rearrange(id_grid, 'b c h w -> b (h w) c', h=self.h, w=self.w) + # init_flow = torch.einsum('bij,bjc->bic', corr_volume.softmax(-1), id_grid) + # init_flow = rearrange(init_flow, 'b (h w) c-> b c h w', h=self.h, w=self.w) + # id_grid = rearrange(id_grid, 'b (h w) c -> b c h w', h=self.h, w=self.w) + # init_flow = init_flow - id_grid + # init_occlusion = None + # ### NPMR only initialization + + ## start iter resolution: size//32 = (size//4) // 8, where size//4 defines the basic flow resolution + flow = ( + F.interpolate( + init_flow, scale_factor=1.0 / 8.0, mode="bilinear", align_corners=True + ) + / 8.0 + ) + occlusion = ( + F.interpolate( + init_occlusion, + scale_factor=1.0 / 8.0, + mode="bilinear", + align_corners=True, + ) + if init_occlusion is not None + else init_occlusion + ) + ## reshape for pooling the driving image dimension, fitting different resolutions + corr_volume = rearrange( + corr_volume, "b (h w) n -> (b n) h w", h=self.h, w=self.w + ).unsqueeze(1) + + for i in range(self.total_iter): + ## identity grid with current-iter resolution + id_grid = coords_grid( + b, + self.size // 32 * (2**i), + self.size // 32 * (2**i), + corr_volume.device, + ) + flow_sample = flow + id_grid_sample = id_grid + + ### computing upsampled or downsampled structure correlation volume + if i < self.basic_res_index: + corr_volume_res = F.avg_pool2d( + corr_volume, + 2 ** (self.basic_res_index - i), + stride=2 ** (self.basic_res_index - i), + ) + scale = 2 ** (self.basic_res_index - i) + elif i == self.basic_res_index: + corr_volume_res = corr_volume + scale = 1 + elif i > self.basic_res_index: + ## for resolution larger than basic flow resolution, sampling with the basic flow resolution + corr_volume_res = corr_volume + scale = 0.5 ** (i - self.basic_res_index) + flow_sample = ( + F.interpolate( + flow, size=self.h, mode="bilinear", align_corners=True + ) + * scale + ) + id_grid_sample = coords_grid(b, self.h, self.w, corr_volume.device) + scale = 1 + ### computing upsampled or downsampled structure correlation volume + + if i < self.num_iter: + ## reshape for pooling the source image dimension, thus to obtain correlation pyramids + corr_volume_res = rearrange( + corr_volume_res, "(b n) c h w -> (b h w) c n", n=self.h * self.w + ) + corr_volume_res = rearrange( + corr_volume_res, "b c (p q) -> b c p q", p=self.h, q=self.w + ) + ## building correlation pyramids + corr_fn = CorrBlock(corr_volume_res) + ## sampling correlation features from the pyramids + corr = corr_fn((flow_sample + id_grid_sample) * scale) + if i > self.basic_res_index: + ## reshape from basic resolution to current resolution + corr = F.interpolate( + corr, size=flow.shape[2], mode="bilinear", align_corners=True + ) + ## non-prior-based motion feature + m_f = self.corr_enc(flow, corr) + + warp_f = bilinear_sampler( + feature[i], (flow + id_grid).permute(0, 2, 3, 1) + ) + warp_f = F.relu(self.to_context[i](warp_f)) + + d_flow, _ = self.refine(m_f, warp_f) + flow_w = flow + d_flow[:, 0:2, :, :] + d_occ = d_flow[:, 2:, :, :] + if occlusion is not None: + occlusion = occlusion + d_occ + else: + occlusion = d_occ + else: + flow_w = flow + + out = bilinear_sampler(feature[i], (flow_w + id_grid).permute(0, 2, 3, 1)) + out_occlusion.append(F.sigmoid(occlusion)) + out_warp_f.append(out) + + ## coarse warping (prior motion flow) + if i != self.basic_res_index: + flow_res = F.interpolate( + dense_motion["deformation"].permute(0, 3, 1, 2), + size=feature[i].shape[2:], + mode="bilinear", + align_corners=True, + ) + occlusion_res = F.interpolate( + dense_motion["occlusion"], + size=feature[i].shape[2:], + mode="bilinear", + align_corners=True, + ) + else: + flow_res = dense_motion["deformation"].permute(0, 3, 1, 2) + occlusion_res = dense_motion["occlusion"] + out_warp_f_c.append(F.grid_sample(feature[i], flow_res.permute(0, 2, 3, 1))) + out_occlusion_c.append(F.sigmoid(occlusion_res)) + ## coarse warping (prior motion flow) + + ### updating flow by directly scaling initial flow + if i < self.num_iter - 1: + ## scale between basic resolution and updated resolution + scale = 2 ** (self.basic_res_index - i) / 2.0 + d_f = ( + F.interpolate( + d_flow[:, 0:2, :, :], + scale_factor=2, + mode="bilinear", + align_corners=True, + ) + * 2 + ) + ## flow refinement of current iteration + flow = ( + d_f + + F.interpolate( + init_flow, + size=self.size // 32 * (2 ** (i + 1)), + mode="bilinear", + align_corners=True, + ) + / scale + ) + ## flow refinement of previous iterations + if i == 0: + d_f_pre = d_f + else: + flow = ( + flow + + F.interpolate( + d_f_pre, scale_factor=2, mode="bilinear", align_corners=True + ) + * 2 + ) + d_f_pre = ( + d_f + + F.interpolate( + d_f_pre, scale_factor=2, mode="bilinear", align_corners=True + ) + * 2 + ) + ## occlusion refinement + d_occ = F.interpolate( + d_occ, scale_factor=2, mode="bilinear", align_corners=True + ) + occlusion = d_occ + F.interpolate( + init_occlusion, + size=self.size // 32 * (2 ** (i + 1)), + mode="bilinear", + align_corners=True, + ) + if i == 0: + d_occ_pre = d_occ + else: + occlusion = occlusion + F.interpolate( + d_occ_pre, scale_factor=2, mode="bilinear", align_corners=True + ) + d_occ_pre = d_occ + F.interpolate( + d_occ_pre, scale_factor=2, mode="bilinear", align_corners=True + ) + ## for num_iter < total_iter + # elif i < self.total_iter-1: + # ## without refinement + # flow = F.interpolate(flow, scale_factor=2, mode='bilinear',align_corners=True) * 2 + # occlusion = F.interpolate(occlusion, scale_factor=2, mode='bilinear',align_corners=True) + + warp_img = bilinear_sampler(img_full, (flow + id_grid).permute(0, 2, 3, 1)) + out = self.generator.decode( + out_warp_f, warp_img, out_occlusion, out_warp_f_c, out_occlusion_c + ) + + out_occlusion.append(F.sigmoid(init_occlusion)) + out_occlusion_vis = [] + for i in range(len(out_occlusion)): + out_occlusion_vis.append( + F.interpolate( + out_occlusion[i], + size=self.size, + mode="bilinear", + align_corners=True, + ) + ) + occlusion = torch.cat(out_occlusion_vis, dim=3) + + return out, warp_img, occlusion diff --git a/module_mrfa/transformer/hr_base.py b/module_mrfa/transformer/hr_base.py new file mode 100644 index 0000000..39f0500 --- /dev/null +++ b/module_mrfa/transformer/hr_base.py @@ -0,0 +1,727 @@ +# ------------------------------------------------------------------------------ +# Copyright (c) Microsoft +# Licensed under the MIT License. +# Written by Bin Xiao (Bin.Xiao@microsoft.com) +# Modified by Sen Yang (yangsenius@seu.edu.cn) +# ------------------------------------------------------------------------------ + +import copy +import logging +import os + +import torch +import torch.nn.functional as F +from torch import nn + +BN_MOMENTUM = 0.1 +logger = logging.getLogger(__name__) + + +def conv3x3(in_planes, out_planes, stride=1): + """3x3 convolution with padding""" + return nn.Conv2d( + in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False + ) + + +class BasicBlock(nn.Module): + expansion = 1 + + def __init__(self, inplanes, planes, stride=1, downsample=None): + super(BasicBlock, self).__init__() + self.conv1 = conv3x3(inplanes, planes, stride) + self.bn1 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM) + self.relu = nn.ReLU(inplace=True) + self.conv2 = conv3x3(planes, planes) + self.bn2 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM) + self.downsample = downsample + self.stride = stride + + def forward(self, x): + residual = x + + out = self.conv1(x) + out = self.bn1(out) + out = self.relu(out) + + out = self.conv2(out) + out = self.bn2(out) + + if self.downsample is not None: + residual = self.downsample(x) + + out += residual + out = self.relu(out) + + return out + + +class Bottleneck(nn.Module): + expansion = 4 + + def __init__(self, inplanes, planes, stride=1, downsample=None): + super(Bottleneck, self).__init__() + self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False) + self.bn1 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM) + self.conv2 = nn.Conv2d( + planes, planes, kernel_size=3, stride=stride, padding=1, bias=False + ) + self.bn2 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM) + self.conv3 = nn.Conv2d( + planes, planes * self.expansion, kernel_size=1, bias=False + ) + self.bn3 = nn.BatchNorm2d(planes * self.expansion, momentum=BN_MOMENTUM) + self.relu = nn.ReLU(inplace=True) + self.downsample = downsample + self.stride = stride + + def forward(self, x): + residual = x + + out = self.conv1(x) + out = self.bn1(out) + out = self.relu(out) + + out = self.conv2(out) + out = self.bn2(out) + out = self.relu(out) + + out = self.conv3(out) + out = self.bn3(out) + + if self.downsample is not None: + residual = self.downsample(x) + + out += residual + out = self.relu(out) + + return out + + +def _get_clones(module, N): + return nn.ModuleList([copy.deepcopy(module) for i in range(N)]) + + +def _get_activation_fn(activation): + """Return an activation function given a string""" + if activation == "relu": + return F.relu + if activation == "gelu": + return F.gelu + if activation == "glu": + return F.glu + raise RuntimeError(f"activation should be relu/gelu, not {activation}.") + + +class HighResolutionModule(nn.Module): + def __init__( + self, + num_branches, + blocks, + num_blocks, + num_inchannels, + num_channels, + fuse_method, + multi_scale_output=True, + ): + super(HighResolutionModule, self).__init__() + self._check_branches( + num_branches, blocks, num_blocks, num_inchannels, num_channels + ) + + self.num_inchannels = num_inchannels + self.fuse_method = fuse_method + self.num_branches = num_branches + + self.multi_scale_output = multi_scale_output + + self.branches = self._make_branches( + num_branches, blocks, num_blocks, num_channels + ) + self.fuse_layers = self._make_fuse_layers() + self.relu = nn.ReLU(True) + + def _check_branches( + self, num_branches, blocks, num_blocks, num_inchannels, num_channels + ): + if num_branches != len(num_blocks): + error_msg = "NUM_BRANCHES({}) <> NUM_BLOCKS({})".format( + num_branches, len(num_blocks) + ) + logger.error(error_msg) + raise ValueError(error_msg) + + if num_branches != len(num_channels): + error_msg = "NUM_BRANCHES({}) <> NUM_CHANNELS({})".format( + num_branches, len(num_channels) + ) + logger.error(error_msg) + raise ValueError(error_msg) + + if num_branches != len(num_inchannels): + error_msg = "NUM_BRANCHES({}) <> NUM_INCHANNELS({})".format( + num_branches, len(num_inchannels) + ) + logger.error(error_msg) + raise ValueError(error_msg) + + def _make_one_branch(self, branch_index, block, num_blocks, num_channels, stride=1): + downsample = None + if ( + stride != 1 + or self.num_inchannels[branch_index] + != num_channels[branch_index] * block.expansion + ): + downsample = nn.Sequential( + nn.Conv2d( + self.num_inchannels[branch_index], + num_channels[branch_index] * block.expansion, + kernel_size=1, + stride=stride, + bias=False, + ), + nn.BatchNorm2d( + num_channels[branch_index] * block.expansion, momentum=BN_MOMENTUM + ), + ) + + layers = [] + layers.append( + block( + self.num_inchannels[branch_index], + num_channels[branch_index], + stride, + downsample, + ) + ) + self.num_inchannels[branch_index] = num_channels[branch_index] * block.expansion + for i in range(1, num_blocks[branch_index]): + layers.append( + block(self.num_inchannels[branch_index], num_channels[branch_index]) + ) + + return nn.Sequential(*layers) + + def _make_branches(self, num_branches, block, num_blocks, num_channels): + branches = [] + + for i in range(num_branches): + branches.append(self._make_one_branch(i, block, num_blocks, num_channels)) + + return nn.ModuleList(branches) + + def _make_fuse_layers(self): + if self.num_branches == 1: + return None + + num_branches = self.num_branches + num_inchannels = self.num_inchannels + fuse_layers = [] + for i in range(num_branches if self.multi_scale_output else 1): + fuse_layer = [] + for j in range(num_branches): + if j > i: + fuse_layer.append( + nn.Sequential( + nn.Conv2d( + num_inchannels[j], + num_inchannels[i], + 1, + 1, + 0, + bias=False, + ), + nn.BatchNorm2d(num_inchannels[i]), + nn.Upsample(scale_factor=2 ** (j - i), mode="nearest"), + ) + ) + elif j == i: + fuse_layer.append(None) + else: + conv3x3s = [] + for k in range(i - j): + if k == i - j - 1: + num_outchannels_conv3x3 = num_inchannels[i] + conv3x3s.append( + nn.Sequential( + nn.Conv2d( + num_inchannels[j], + num_outchannels_conv3x3, + 3, + 2, + 1, + bias=False, + ), + nn.BatchNorm2d(num_outchannels_conv3x3), + ) + ) + else: + num_outchannels_conv3x3 = num_inchannels[j] + conv3x3s.append( + nn.Sequential( + nn.Conv2d( + num_inchannels[j], + num_outchannels_conv3x3, + 3, + 2, + 1, + bias=False, + ), + nn.BatchNorm2d(num_outchannels_conv3x3), + nn.ReLU(True), + ) + ) + fuse_layer.append(nn.Sequential(*conv3x3s)) + fuse_layers.append(nn.ModuleList(fuse_layer)) + + return nn.ModuleList(fuse_layers) + + def get_num_inchannels(self): + return self.num_inchannels + + def forward(self, x): + if self.num_branches == 1: + return [self.branches[0](x[0])] + + for i in range(self.num_branches): + x[i] = self.branches[i](x[i]) + + x_fuse = [] + + for i in range(len(self.fuse_layers)): + y = x[0] if i == 0 else self.fuse_layers[i][0](x[0]) + for j in range(1, self.num_branches): + if i == j: + y = y + x[j] + else: + y = y + self.fuse_layers[i][j](x[j]) + x_fuse.append(self.relu(y)) + + return x_fuse + + +blocks_dict = {"BASIC": BasicBlock, "BOTTLENECK": Bottleneck} + + +class HRNET_base(nn.Module): + def __init__(self, cfg, **kwargs): + self.inplanes = 64 + extra = cfg["MODEL"]["EXTRA"] + super(HRNET_base, self).__init__() + + # stem net + self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=2, padding=1, bias=False) + self.bn1 = nn.BatchNorm2d(64, momentum=BN_MOMENTUM) + self.conv2 = nn.Conv2d(64, 64, kernel_size=3, stride=2, padding=1, bias=False) + self.bn2 = nn.BatchNorm2d(64, momentum=BN_MOMENTUM) + self.relu = nn.ReLU(inplace=True) + self.layer1 = self._make_layer(Bottleneck, 64, 4) + + self.stage2_cfg = extra["STAGE2"] + num_channels = self.stage2_cfg["NUM_CHANNELS"] + block = blocks_dict[self.stage2_cfg["BLOCK"]] + num_channels = [ + num_channels[i] * block.expansion for i in range(len(num_channels)) + ] + self.transition1 = self._make_transition_layer([256], num_channels) + self.stage2, pre_stage_channels = self._make_stage( + self.stage2_cfg, num_channels + ) + + self.stage3_cfg = extra["STAGE3"] + num_channels = self.stage3_cfg["NUM_CHANNELS"] + block = blocks_dict[self.stage3_cfg["BLOCK"]] + num_channels = [ + num_channels[i] * block.expansion for i in range(len(num_channels)) + ] + self.transition2 = self._make_transition_layer(pre_stage_channels, num_channels) + self.stage3, pre_stage_channels = self._make_stage( + self.stage3_cfg, num_channels, multi_scale_output=False + ) + + self.pretrained_layers = extra["PRETRAINED_LAYERS"] + + def _make_transition_layer(self, num_channels_pre_layer, num_channels_cur_layer): + num_branches_cur = len(num_channels_cur_layer) + num_branches_pre = len(num_channels_pre_layer) + + transition_layers = [] + for i in range(num_branches_cur): + if i < num_branches_pre: + if num_channels_cur_layer[i] != num_channels_pre_layer[i]: + transition_layers.append( + nn.Sequential( + nn.Conv2d( + num_channels_pre_layer[i], + num_channels_cur_layer[i], + 3, + 1, + 1, + bias=False, + ), + nn.BatchNorm2d(num_channels_cur_layer[i]), + nn.ReLU(inplace=True), + ) + ) + else: + transition_layers.append(None) + else: + conv3x3s = [] + for j in range(i + 1 - num_branches_pre): + inchannels = num_channels_pre_layer[-1] + outchannels = ( + num_channels_cur_layer[i] + if j == i - num_branches_pre + else inchannels + ) + conv3x3s.append( + nn.Sequential( + nn.Conv2d(inchannels, outchannels, 3, 2, 1, bias=False), + nn.BatchNorm2d(outchannels), + nn.ReLU(inplace=True), + ) + ) + transition_layers.append(nn.Sequential(*conv3x3s)) + + return nn.ModuleList(transition_layers) + + def _make_layer(self, block, planes, blocks, stride=1): + downsample = None + if stride != 1 or self.inplanes != planes * block.expansion: + downsample = nn.Sequential( + nn.Conv2d( + self.inplanes, + planes * block.expansion, + kernel_size=1, + stride=stride, + bias=False, + ), + nn.BatchNorm2d(planes * block.expansion, momentum=BN_MOMENTUM), + ) + + layers = [] + layers.append(block(self.inplanes, planes, stride, downsample)) + self.inplanes = planes * block.expansion + for i in range(1, blocks): + layers.append(block(self.inplanes, planes)) + + return nn.Sequential(*layers) + + def _make_stage(self, layer_config, num_inchannels, multi_scale_output=True): + num_modules = layer_config["NUM_MODULES"] + num_branches = layer_config["NUM_BRANCHES"] + num_blocks = layer_config["NUM_BLOCKS"] + num_channels = layer_config["NUM_CHANNELS"] + block = blocks_dict[layer_config["BLOCK"]] + fuse_method = layer_config["FUSE_METHOD"] + + modules = [] + for i in range(num_modules): + # multi_scale_output is only used last module + if not multi_scale_output and i == num_modules - 1: + reset_multi_scale_output = False + else: + reset_multi_scale_output = True + + modules.append( + HighResolutionModule( + num_branches, + block, + num_blocks, + num_inchannels, + num_channels, + fuse_method, + reset_multi_scale_output, + ) + ) + num_inchannels = modules[-1].get_num_inchannels() + + return nn.Sequential(*modules), num_inchannels + + def forward(self, x): + x = self.conv1(x) + x = self.bn1(x) + x = self.relu(x) + x = self.conv2(x) + x = self.bn2(x) + x = self.relu(x) + x = self.layer1(x) + + x_list = [] + for i in range(self.stage2_cfg["NUM_BRANCHES"]): + if self.transition1[i] is not None: + x_list.append(self.transition1[i](x)) + else: + x_list.append(x) + y_list = self.stage2(x_list) + + x_list = [] + for i in range(self.stage3_cfg["NUM_BRANCHES"]): + if self.transition2[i] is not None: + x_list.append(self.transition2[i](y_list[-1])) + else: + x_list.append(y_list[i]) + y_list = self.stage3(x_list) + return y_list[0] + + def init_weights(self, pretrained="", print_load_info=False): + logger.info("=> init weights from normal distribution") + for m in self.modules(): + if isinstance(m, nn.Conv2d): + nn.init.normal_(m.weight, std=0.001) + for name, _ in m.named_parameters(): + if name in ["bias"]: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.BatchNorm2d): + nn.init.constant_(m.weight, 1) + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.ConvTranspose2d): + nn.init.normal_(m.weight, std=0.001) + for name, _ in m.named_parameters(): + if name in ["bias"]: + nn.init.constant_(m.bias, 0) + + if os.path.isfile(pretrained): + pretrained_state_dict = torch.load( + pretrained, map_location=torch.device("cpu") + ) + logger.info("=> loading pretrained model {}".format(pretrained)) + + existing_state_dict = {} + for name, m in pretrained_state_dict.items(): + if ( + name.split(".")[0] in self.pretrained_layers + and name in self.state_dict() + or self.pretrained_layers[0] is "*" + ): + existing_state_dict[name] = m + if print_load_info: + print(":: {} is loaded from {}".format(name, pretrained)) + self.load_state_dict(existing_state_dict, strict=False) + elif pretrained: + logger.error("=> please download pre-trained models first!") + raise ValueError("{} is not exist!".format(pretrained)) + + +class HRNET4_base(nn.Module): + def __init__(self, cfg, **kwargs): + self.inplanes = 64 + extra = cfg["MODEL"]["EXTRA"] + super(HRNET4_base, self).__init__() + + # stem net + self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=2, padding=1, bias=False) + self.bn1 = nn.BatchNorm2d(64, momentum=BN_MOMENTUM) + self.conv2 = nn.Conv2d(64, 64, kernel_size=3, stride=2, padding=1, bias=False) + self.bn2 = nn.BatchNorm2d(64, momentum=BN_MOMENTUM) + self.relu = nn.ReLU(inplace=True) + self.layer1 = self._make_layer(Bottleneck, 64, 4) + + # stage 2 + self.stage2_cfg = extra["STAGE2"] + num_channels = self.stage2_cfg["NUM_CHANNELS"] + block = blocks_dict[self.stage2_cfg["BLOCK"]] + num_channels = [ + num_channels[i] * block.expansion for i in range(len(num_channels)) + ] + self.transition1 = self._make_transition_layer([256], num_channels) + self.stage2, pre_stage_channels = self._make_stage( + self.stage2_cfg, num_channels + ) + + # stage 3 + self.stage3_cfg = extra["STAGE3"] + num_channels = self.stage3_cfg["NUM_CHANNELS"] + block = blocks_dict[self.stage3_cfg["BLOCK"]] + num_channels = [ + num_channels[i] * block.expansion for i in range(len(num_channels)) + ] + self.transition2 = self._make_transition_layer(pre_stage_channels, num_channels) + self.stage3, pre_stage_channels = self._make_stage( + self.stage3_cfg, num_channels + ) + + # stage 4 + self.stage4_cfg = extra["STAGE4"] + num_channels = self.stage4_cfg["NUM_CHANNELS"] + block = blocks_dict[self.stage4_cfg["BLOCK"]] + num_channels = [ + num_channels[i] * block.expansion for i in range(len(num_channels)) + ] + self.transition3 = self._make_transition_layer(pre_stage_channels, num_channels) + self.stage4, pre_stage_channels = self._make_stage( + self.stage4_cfg, num_channels, multi_scale_output=False + ) + + self.pretrained_layers = extra["PRETRAINED_LAYERS"] + + def _make_transition_layer(self, num_channels_pre_layer, num_channels_cur_layer): + num_branches_cur = len(num_channels_cur_layer) + num_branches_pre = len(num_channels_pre_layer) + + transition_layers = [] + for i in range(num_branches_cur): + if i < num_branches_pre: + if num_channels_cur_layer[i] != num_channels_pre_layer[i]: + transition_layers.append( + nn.Sequential( + nn.Conv2d( + num_channels_pre_layer[i], + num_channels_cur_layer[i], + 3, + 1, + 1, + bias=False, + ), + nn.BatchNorm2d(num_channels_cur_layer[i]), + nn.ReLU(inplace=True), + ) + ) + else: + transition_layers.append(None) + else: + conv3x3s = [] + for j in range(i + 1 - num_branches_pre): + inchannels = num_channels_pre_layer[-1] + outchannels = ( + num_channels_cur_layer[i] + if j == i - num_branches_pre + else inchannels + ) + conv3x3s.append( + nn.Sequential( + nn.Conv2d(inchannels, outchannels, 3, 2, 1, bias=False), + nn.BatchNorm2d(outchannels), + nn.ReLU(inplace=True), + ) + ) + transition_layers.append(nn.Sequential(*conv3x3s)) + + return nn.ModuleList(transition_layers) + + def _make_layer(self, block, planes, blocks, stride=1): + downsample = None + if stride != 1 or self.inplanes != planes * block.expansion: + downsample = nn.Sequential( + nn.Conv2d( + self.inplanes, + planes * block.expansion, + kernel_size=1, + stride=stride, + bias=False, + ), + nn.BatchNorm2d(planes * block.expansion, momentum=BN_MOMENTUM), + ) + + layers = [] + layers.append(block(self.inplanes, planes, stride, downsample)) + self.inplanes = planes * block.expansion + for i in range(1, blocks): + layers.append(block(self.inplanes, planes)) + + return nn.Sequential(*layers) + + def _make_stage(self, layer_config, num_inchannels, multi_scale_output=True): + num_modules = layer_config["NUM_MODULES"] + num_branches = layer_config["NUM_BRANCHES"] + num_blocks = layer_config["NUM_BLOCKS"] + num_channels = layer_config["NUM_CHANNELS"] + block = blocks_dict[layer_config["BLOCK"]] + fuse_method = layer_config["FUSE_METHOD"] + + modules = [] + for i in range(num_modules): + # multi_scale_output is only used last module + if not multi_scale_output and i == num_modules - 1: + reset_multi_scale_output = False + else: + reset_multi_scale_output = True + + modules.append( + HighResolutionModule( + num_branches, + block, + num_blocks, + num_inchannels, + num_channels, + fuse_method, + reset_multi_scale_output, + ) + ) + num_inchannels = modules[-1].get_num_inchannels() + + return nn.Sequential(*modules), num_inchannels + + def forward(self, x): + x = self.conv1(x) + x = self.bn1(x) + x = self.relu(x) + x = self.conv2(x) + x = self.bn2(x) + x = self.relu(x) + x = self.layer1(x) + + # stage 2 + x_list = [] + for i in range(self.stage2_cfg["NUM_BRANCHES"]): + if self.transition1[i] is not None: + x_list.append(self.transition1[i](x)) + else: + x_list.append(x) + y_list = self.stage2(x_list) + + # stage 3 + x_list = [] + for i in range(self.stage3_cfg["NUM_BRANCHES"]): + if self.transition2[i] is not None: + x_list.append(self.transition2[i](y_list[-1])) + else: + x_list.append(y_list[i]) + y_list = self.stage3(x_list) + + # stage 4 + x_list = [] + for i in range(self.stage4_cfg["NUM_BRANCHES"]): + if self.transition3[i] is not None: + x_list.append(self.transition3[i](y_list[-1])) + else: + x_list.append(y_list[i]) + y_list = self.stage4(x_list) + + return y_list[0] + + def init_weights(self, pretrained="", print_load_info=False): + logger.info("=> init weights from normal distribution") + for m in self.modules(): + if isinstance(m, nn.Conv2d): + nn.init.normal_(m.weight, std=0.001) + for name, _ in m.named_parameters(): + if name in ["bias"]: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.BatchNorm2d): + nn.init.constant_(m.weight, 1) + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.ConvTranspose2d): + nn.init.normal_(m.weight, std=0.001) + for name, _ in m.named_parameters(): + if name in ["bias"]: + nn.init.constant_(m.bias, 0) + + if os.path.isfile(pretrained): + pretrained_state_dict = torch.load(pretrained) + logger.info("=> loading pretrained model {}".format(pretrained)) + + existing_state_dict = {} + for name, m in pretrained_state_dict.items(): + if ( + name.split(".")[0] in self.pretrained_layers + and name in self.state_dict() + or self.pretrained_layers[0] is "*" + ): + existing_state_dict[name] = m + if print_load_info: + print(":: {} is loaded from {}".format(name, pretrained)) + self.load_state_dict(existing_state_dict, strict=False) + elif pretrained: + logger.error("=> please download pre-trained models first!") + raise ValueError("{} is not exist!".format(pretrained)) diff --git a/module_mrfa/transformer/pose_tokenpose_b.py b/module_mrfa/transformer/pose_tokenpose_b.py new file mode 100644 index 0000000..cc554fd --- /dev/null +++ b/module_mrfa/transformer/pose_tokenpose_b.py @@ -0,0 +1,61 @@ +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 diff --git a/module_mrfa/transformer/tokenpose_base.py b/module_mrfa/transformer/tokenpose_base.py new file mode 100644 index 0000000..5531d7a --- /dev/null +++ b/module_mrfa/transformer/tokenpose_base.py @@ -0,0 +1,1236 @@ +import math + +import torch +import torch.nn.functional as F +from einops import rearrange, repeat +from timm.models.layers.weight_init import trunc_normal_ +from torch import nn + +from ..util import make_coordinate_grid + +MIN_NUM_PATCHES = 16 +BN_MOMENTUM = 0.1 + + +class Residual(nn.Module): + def __init__(self, fn, num_keypoints=10): + super().__init__() + self.fn = fn + self.num_keypoints = num_keypoints + + def forward(self, x, **kwargs): + return self.fn(x, **kwargs) + x + + +class PreNorm(nn.Module): + def __init__(self, dim, fn, fusion_factor=1): + super().__init__() + self.norm = nn.LayerNorm(dim * fusion_factor) + self.fn = fn + + def forward(self, x, **kwargs): + return self.fn(self.norm(x), **kwargs) + + +class QKVPreNorm(nn.Module): + def __init__(self, dim, fn, fusion_factor=1): + super().__init__() + self.norm = nn.LayerNorm(dim * fusion_factor) + self.fn = fn + + def forward(self, q, k, v, **kwargs): + return self.fn(self.norm(q), self.norm(k), self.norm(v), **kwargs) + + +class FeedForward(nn.Module): + def __init__(self, dim, hidden_dim, dropout=0.0): + super().__init__() + self.net = nn.Sequential( + nn.Linear(dim, hidden_dim), + nn.GELU(), + nn.Dropout(dropout), + nn.Linear(hidden_dim, dim), + nn.Dropout(dropout), + ) + + def forward(self, x): + return self.net(x) + + +class Attention(nn.Module): + def __init__( + self, + dim, + heads=8, + dropout=0.0, + num_keypoints=None, + scale_with_head=False, + fix_img2motion_attention=False, + num_img_tokens=None, + ): + super().__init__() + self.heads = heads + self.scale = (dim // heads) ** -0.5 if scale_with_head else dim**-0.5 + + self.to_qkv = nn.Linear(dim, dim * 3, bias=False) + self.to_out = nn.Sequential(nn.Linear(dim, dim), nn.Dropout(dropout)) + self.num_keypoints = num_keypoints + self.fix_img2motion_attention = fix_img2motion_attention + if fix_img2motion_attention: + num_tokens = num_keypoints + num_img_tokens + mask = torch.zeros(num_tokens, num_tokens) + mask[self.num_keypoints :, 0 : self.num_keypoints] += -1000.0 + mask = mask.unsqueeze(0).unsqueeze(0) + self.register_buffer("attn_mask", mask) + + # @get_local('dots') + # @get_local('attn') + def forward(self, x, mask=None): + b, n, _, h = *x.shape, self.heads + qkv = self.to_qkv(x).chunk(3, dim=-1) + q, k, v = map(lambda t: rearrange(t, "b n (h d) -> b h n d", h=h), qkv) + + dots = torch.einsum("bhid,bhjd->bhij", q, k) * self.scale + mask_value = -torch.finfo(dots.dtype).max + + if mask is not None: + mask = F.pad(mask.flatten(1), (1, 0), value=True) + assert mask.shape[-1] == dots.shape[-1], "mask has incorrect dimensions" + mask = mask[:, None, :] * mask[:, :, None] + dots.masked_fill_(~mask, mask_value) + del mask + + if self.fix_img2motion_attention: + dots += self.attn_mask + attn = dots.softmax(dim=-1) + # print(attn[0,0,-1,-1]) + out = torch.einsum("bhij,bhjd->bhid", attn, v) + out = rearrange(out, "b h n d -> b n (h d)") + out = self.to_out(out) + return out + + +class QKVAttention(nn.Module): + def __init__( + self, dim, heads=8, dropout=0.0, num_keypoints=None, scale_with_head=False + ): + super().__init__() + self.heads = heads + self.scale = (dim // heads) ** -0.5 if scale_with_head else dim**-0.5 + + self.to_k = nn.Linear(dim, dim, bias=False) + self.to_v = nn.Linear(dim, dim, bias=False) + self.to_q = nn.Linear(dim, dim, bias=False) + self.to_out = nn.Sequential(nn.Linear(dim, dim), nn.Dropout(dropout)) + self.num_keypoints = num_keypoints + + # @get_local('dots') + # @get_local('attn') + def forward(self, x, k=None, v=None, mask=None): + b, n, _, h = *x.shape, self.heads + q = self.to_q(x) + k = self.to_k(k) + v = self.to_v(v) + q = rearrange(q, "b n (h d) -> b h n d", h=h) + k = rearrange(k, "b n (h d) -> b h n d", h=h) + v = rearrange(v, "b n (h d) -> b h n d", h=h) + + dots = torch.einsum("bhid,bhjd->bhij", q, k) * self.scale + mask_value = -torch.finfo(dots.dtype).max + + if mask is not None: + mask = F.pad(mask.flatten(1), (1, 0), value=True) + assert mask.shape[-1] == dots.shape[-1], "mask has incorrect dimensions" + mask = mask[:, None, :] * mask[:, :, None] + dots.masked_fill_(~mask, mask_value) + del mask + attn = dots.softmax(dim=-1) + # print(attn[0,0,-1,-1]) + out = torch.einsum("bhij,bhjd->bhid", attn, v) + out = rearrange(out, "b h n d -> b n (h d)") + out = self.to_out(out) + return out + + +class Transformer(nn.Module): + def __init__( + self, + dim, + depth, + heads, + mlp_dim, + dropout, + num_keypoints=None, + all_attn=False, + scale_with_head=False, + fix_img2motion_attention=False, + num_patches=256, + ): + super().__init__() + self.layers = nn.ModuleList([]) + self.all_attn = all_attn + self.num_keypoints = num_keypoints + self.fix_img2motion_attention = fix_img2motion_attention + for d in range(depth): + # if not d < depth // 2: + # fix_img2motion_attention = False + self.layers.append( + nn.ModuleList( + [ + Residual( + PreNorm( + dim, + Attention( + dim, + heads=heads, + dropout=dropout, + num_keypoints=num_keypoints, + scale_with_head=scale_with_head, + fix_img2motion_attention=fix_img2motion_attention, + num_img_tokens=num_patches, + ), + ), + num_keypoints=self.num_keypoints, + ), + Residual( + PreNorm(dim, FeedForward(dim, mlp_dim, dropout=dropout)) + ), + ] + ) + ) + + def forward(self, x, mask=None, pos=None): + img_token = x[:, self.num_keypoints :] + for idx, (attn, ff) in enumerate(self.layers): + if idx > 0 and self.all_attn: + x[:, self.num_keypoints :] += pos + x = attn(x, mask=mask) + x = ff(x) + return x + + +class TransformerEncoder(nn.Module): + def __init__( + self, + dim, + depth, + heads, + mlp_dim, + dropout, + num_keypoints=None, + all_attn=False, + scale_with_head=False, + num_patches=256, + v_pos=False, + ): + super().__init__() + self.encoder = nn.ModuleList([]) + self.all_attn = all_attn + self.num_keypoints = num_keypoints + self.v_pos = v_pos + for d in range(depth): + self.encoder.append( + nn.ModuleList( + [ + Residual( + QKVPreNorm( + dim, + QKVAttention( + dim, + heads=heads, + dropout=dropout, + num_keypoints=num_keypoints, + scale_with_head=scale_with_head, + ), + ), + num_keypoints=self.num_keypoints, + ), + Residual( + PreNorm(dim, FeedForward(dim, mlp_dim, dropout=dropout)) + ), + ] + ) + ) + + def forward(self, x, mask=None, pos=None): + for idx, (attn, ff) in enumerate(self.encoder): + if self.all_attn: + if self.v_pos: + x[:, self.num_keypoints :] += pos + q = x + k = x + v = x + else: + v = x.clone() + x[:, self.num_keypoints :] += pos + q = x + k = x + x = attn(q, k=k, v=v, mask=mask) + x = ff(x) + return x + + +class Bottleneck(nn.Module): + expansion = 4 + + def __init__(self, inplanes, planes, stride=1, downsample=None): + super(Bottleneck, self).__init__() + self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False) + self.bn1 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM) + self.conv2 = nn.Conv2d( + planes, planes, kernel_size=3, stride=stride, padding=1, bias=False + ) + self.bn2 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM) + self.conv3 = nn.Conv2d( + planes, planes * self.expansion, kernel_size=1, bias=False + ) + self.bn3 = nn.BatchNorm2d(planes * self.expansion, momentum=BN_MOMENTUM) + self.relu = nn.ReLU(inplace=True) + self.downsample = downsample + self.stride = stride + + def forward(self, x): + residual = x + + out = self.conv1(x) + out = self.bn1(out) + out = self.relu(out) + + out = self.conv2(out) + out = self.bn2(out) + out = self.relu(out) + + out = self.conv3(out) + out = self.bn3(out) + + if self.downsample is not None: + residual = self.downsample(x) + + out += residual + out = self.relu(out) + + return out + + +class TokenPose_TB_base(nn.Module): + def __init__( + self, + *, + feature_size, + patch_size, + num_keypoints, + dim, + depth, + heads, + mlp_dim, + apply_init=False, + apply_multi=True, + hidden_heatmap_dim=64 * 6, + heatmap_dim=64 * 64, + heatmap_size=[64, 64], + channels=3, + dropout=0.0, + emb_dropout=0.0, + pos_embedding_type="sine-full", + estimate_jacobian=True, + temperature=0.1, + spatial_kp_head=False, + jacobian_token=True, + hidden_dim=False, + affine_jacobian=False, + fix_img2motion_attention=False, + ): + super().__init__() + assert isinstance(feature_size, list) and isinstance( + patch_size, list + ), "image_size and patch_size should be list" + assert ( + feature_size[0] % patch_size[0] == 0 + and feature_size[1] % patch_size[1] == 0 + ), "Image dimensions must be divisible by the patch size." + num_patches = (feature_size[0] // (patch_size[0])) * ( + feature_size[1] // (patch_size[1]) + ) + patch_dim = channels * patch_size[0] * patch_size[1] + # assert pos_embedding_type in ['sine','learnable','sine-full'] + + self.inplanes = 64 + self.patch_size = patch_size + self.heatmap_size = heatmap_size + self.num_keypoints = num_keypoints + self.num_patches = num_patches + self.pos_embedding_type = pos_embedding_type + self.all_attn = self.pos_embedding_type == "sine-full" + + self.jacobian_token = jacobian_token + if jacobian_token: + ## additional token for jacobian + num_keypoints = 2 * num_keypoints + self.num_keypoints = num_keypoints + self.keypoint_token = nn.Parameter(torch.zeros(1, self.num_keypoints, dim)) + h, w = ( + feature_size[0] // (self.patch_size[0]), + feature_size[1] // (self.patch_size[1]), + ) + self._make_position_embedding(w, h, dim, pos_embedding_type) + + self.patch_to_embedding = nn.Linear(patch_dim, dim) + self.dropout = nn.Dropout(emb_dropout) + + # transformer + self.transformer = Transformer( + dim, + depth, + heads, + mlp_dim, + dropout, + num_keypoints=num_keypoints, + all_attn=self.all_attn, + scale_with_head=True, + fix_img2motion_attention=fix_img2motion_attention, + num_patches=num_patches, + ) + + self.to_keypoint_token = nn.Identity() + ## original mlphead for keypoint heatmap prediction + self.spatial_kp_head = spatial_kp_head + if spatial_kp_head: + self.mlp_head = ( + nn.Sequential( + nn.LayerNorm(dim), + nn.Linear(dim, hidden_heatmap_dim), + nn.LayerNorm(hidden_heatmap_dim), + nn.Linear(hidden_heatmap_dim, heatmap_dim), + ) + if (dim <= hidden_heatmap_dim * 0.5 and apply_multi) + else nn.Sequential(nn.LayerNorm(dim), nn.Linear(dim, heatmap_dim)) + ) + self.temperature = temperature + ## original mlphead for keypoint heatmap prediction + else: + if hidden_dim: + self.mlp_head = nn.Sequential( + nn.LayerNorm(dim), + nn.Linear(dim, dim), + nn.LayerNorm(dim), + nn.Linear(dim, 2), + ) + else: + self.mlp_head = nn.Sequential( + nn.LayerNorm(dim), + nn.Linear(dim, 2), + ) + trunc_normal_(self.keypoint_token, std=0.02) + if apply_init: + self.apply(self._init_weights) + + ## light-weight head for direct jacobian regression + if estimate_jacobian: + if hidden_dim: + self.mlp_head_jacobian = nn.Sequential( + nn.LayerNorm(dim), + nn.Linear(dim, dim), + nn.LayerNorm(dim), + nn.Linear(dim, 4), + ) + self._init_weights(self.mlp_head_jacobian[0]) + self._init_weights(self.mlp_head_jacobian[1]) + self._init_weights(self.mlp_head_jacobian[2]) + else: + self.mlp_head_jacobian = nn.Sequential( + nn.LayerNorm(dim), + nn.Linear(dim, 4), + ) + self._init_weights(self.mlp_head_jacobian[0]) + self.mlp_head_jacobian[-1].weight.data.zero_() + self.mlp_head_jacobian[-1].bias.data.copy_( + torch.tensor([1, 0, 0, 1], dtype=torch.float) + ) + self.affine_jacobian = affine_jacobian + ## light-weight head for direct keypoint and jacobian regression + + def _make_position_embedding(self, w, h, d_model, pe_type="sine"): + """ + d_model: embedding size in transformer encoder + """ + assert pe_type in ["none", "learnable", "sine", "sine-full"] + if pe_type == "none": + self.pos_embedding = None + print("==> Without any PositionEmbedding~") + else: + with torch.no_grad(): + self.pe_h = h + self.pe_w = w + length = self.pe_h * self.pe_w + if pe_type == "learnable": + self.pos_embedding = nn.Parameter( + torch.zeros(1, self.num_patches + self.num_keypoints, d_model) + ) + trunc_normal_(self.pos_embedding, std=0.02) + print("==> Add Learnable PositionEmbedding~") + else: + self.pos_embedding = nn.Parameter( + self._make_sine_position_embedding(d_model), requires_grad=False + ) + print("==> Add Sine PositionEmbedding~") + + def _make_sine_position_embedding( + self, d_model, temperature=10000, scale=2 * math.pi + ): + h, w = self.pe_h, self.pe_w + area = torch.ones(1, h, w) # [b, h, w] + y_embed = area.cumsum(1, dtype=torch.float32) + x_embed = area.cumsum(2, dtype=torch.float32) + + one_direction_feats = d_model // 2 + + eps = 1e-6 + y_embed = y_embed / (y_embed[:, -1:, :] + eps) * scale + x_embed = x_embed / (x_embed[:, :, -1:] + eps) * scale + + dim_t = torch.arange(one_direction_feats, dtype=torch.float32) + dim_t = temperature ** (2 * (dim_t // 2) / one_direction_feats) + + pos_x = x_embed[:, :, :, None] / dim_t + pos_y = y_embed[:, :, :, None] / dim_t + pos_x = torch.stack( + (pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4 + ).flatten(3) + pos_y = torch.stack( + (pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4 + ).flatten(3) + pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2) + pos = pos.flatten(2).permute(0, 2, 1) + return pos + + def _make_layer(self, block, planes, blocks, stride=1): + downsample = None + if stride != 1 or self.inplanes != planes * block.expansion: + downsample = nn.Sequential( + nn.Conv2d( + self.inplanes, + planes * block.expansion, + kernel_size=1, + stride=stride, + bias=False, + ), + nn.BatchNorm2d(planes * block.expansion, momentum=BN_MOMENTUM), + ) + + layers = [] + layers.append(block(self.inplanes, planes, stride, downsample)) + self.inplanes = planes * block.expansion + for i in range(1, blocks): + layers.append(block(self.inplanes, planes)) + + return nn.Sequential(*layers) + + def _init_weights(self, m): + # print("Initialization...") + if isinstance(m, nn.Linear): + trunc_normal_(m.weight, std=0.02) + if isinstance(m, nn.Linear) and m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.LayerNorm): + nn.init.constant_(m.bias, 0) + nn.init.constant_(m.weight, 1.0) + + def gaussian2kp(self, heatmap): + """ + Extract the mean and from a heatmap + """ + shape = heatmap.shape + heatmap = heatmap.unsqueeze(-1) + grid = ( + make_coordinate_grid(shape[2:], heatmap.type()).unsqueeze_(0).unsqueeze_(0) + ) + value = (heatmap * grid).sum(dim=(2, 3)) # N * 10 * 2 + kp = {"value": value} + + return kp + + def forward(self, feature, mask=None): + p = self.patch_size + # transformer + x = rearrange( + feature, "b c (h p1) (w p2) -> b (h w) (p1 p2 c)", p1=p[0], p2=p[1] + ) + # print(x.shape) + x = self.patch_to_embedding(x) + + b, n, _ = x.shape + + keypoint_tokens = repeat(self.keypoint_token, "() n d -> b n d", b=b) + if self.pos_embedding_type in ["sine", "sine-full"]: + x += self.pos_embedding[:, :n] + x = torch.cat((keypoint_tokens, x), dim=1) + elif self.pos_embedding_type == "none": + x = torch.cat((keypoint_tokens, x), dim=1) + else: + x = torch.cat((keypoint_tokens, x), dim=1) + x += self.pos_embedding[:, : (n + self.num_keypoints)] + x = self.dropout(x) + + x = self.transformer(x, mask, self.pos_embedding) + if self.jacobian_token: + x_keypoint = self.to_keypoint_token(x[:, 0 : self.num_keypoints // 2]) + else: + x_keypoint = self.to_keypoint_token(x[:, 0 : self.num_keypoints]) + if self.spatial_kp_head: + heatmap = self.mlp_head(x_keypoint) + heatmap = F.softmax(heatmap / self.temperature, dim=2) + heatmap = rearrange( + heatmap, + "b c (p1 p2) -> b c p1 p2", + p1=self.heatmap_size[0], + p2=self.heatmap_size[1], + ) + out = self.gaussian2kp(heatmap) + out["heatmap"] = heatmap + else: + # keypoint = torch.tanh(self.mlp_head(x_keypoint)) + keypoint = 2 * F.sigmoid(self.mlp_head(x_keypoint)) - 1 + out = {"kp": keypoint} + if self.mlp_head_jacobian is not None: + if self.jacobian_token: + x_jacobian = self.to_keypoint_token( + x[:, self.num_keypoints // 2 : self.num_keypoints] + ) + else: + x_jacobian = x_keypoint + jacobian = self.mlp_head_jacobian(x_jacobian) + if self.affine_jacobian: + theta = jacobian[:, :, 0:2] + theta = theta / (torch.norm(theta, p=2, dim=-1, keepdim=True) + 1e-10) + cos_theta = theta[:, :, 0:1] + sin_theta = theta[:, :, 1:2] + rotate = torch.cat( + (cos_theta, -sin_theta, sin_theta, cos_theta), dim=-1 + ) + rotate = rearrange(rotate, "b c (p1 p2) -> b c p1 p2", p1=2, p2=2) + + scale = jacobian[:, :, 2:] + # scale = torch.abs(jacobian[:,:,2:]) + 0.1 + scale = torch.tanh(scale) * 0.9 + 1 + scale = 1 / scale + scale_x = scale[:, :, 0:1] + scale_y = scale[:, :, 1:2] + scale = torch.cat( + ( + scale_x, + torch.zeros(scale_x.shape).to(x.device), + torch.zeros(scale_x.shape).to(x.device), + scale_y, + ), + dim=-1, + ) + scale = rearrange(scale, "b c (p1 p2) -> b c p1 p2", p1=2, p2=2) + + jacobian = torch.matmul(rotate, scale) + else: + jacobian = rearrange(jacobian, "b c (p1 p2) -> b c p1 p2", p1=2, p2=2) + out["jacobian"] = jacobian + return out + + +class TokenPose_S_base(nn.Module): + def __init__( + self, + *, + image_size, + patch_size, + num_keypoints, + dim, + depth, + heads, + mlp_dim, + apply_init=False, + apply_multi=True, + hidden_heatmap_dim=64 * 6, + heatmap_dim=64 * 48, + heatmap_size=[64, 48], + channels=3, + dropout=0.0, + emb_dropout=0.0, + pos_embedding_type="learnable", + estimate_jacobian=False, + temperature=0.1, + jacobian_token=True, + hidden_dim=False, + affine_jacobian=False, + ): + super().__init__() + assert isinstance(image_size, list) and isinstance( + patch_size, list + ), "image_size and patch_size should be list" + assert ( + image_size[0] % patch_size[0] == 0 and image_size[1] % patch_size[1] == 0 + ), "Image dimensions must be divisible by the patch size." + num_patches = (image_size[0] // (4 * patch_size[0])) * ( + image_size[1] // (4 * patch_size[1]) + ) + patch_dim = channels * patch_size[0] * patch_size[1] + assert ( + num_patches > MIN_NUM_PATCHES + ), f"your number of patches ({num_patches}) is way too small for attention to be effective (at least 16). Try decreasing your patch size" + assert pos_embedding_type in ["sine", "none", "learnable", "sine-full"] + + self.inplanes = 64 + self.patch_size = patch_size + self.heatmap_size = heatmap_size + self.num_keypoints = num_keypoints + self.num_patches = num_patches + self.pos_embedding_type = pos_embedding_type + self.all_attn = self.pos_embedding_type == "sine-full" + + self.jacobian_token = jacobian_token + if jacobian_token: + ## additional token for jacobian + num_keypoints = 2 * num_keypoints + self.num_keypoints = num_keypoints + + self.keypoint_token = nn.Parameter(torch.zeros(1, self.num_keypoints, dim)) + h, w = ( + image_size[0] // (4 * self.patch_size[0]), + image_size[1] // (4 * self.patch_size[1]), + ) + self._make_position_embedding(w, h, dim, pos_embedding_type) + + self.patch_to_embedding = nn.Linear(patch_dim, dim) + self.dropout = nn.Dropout(emb_dropout) + + # stem net + self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=2, padding=1, bias=False) + self.bn1 = nn.BatchNorm2d(64, momentum=BN_MOMENTUM) + self.conv2 = nn.Conv2d(64, 64, kernel_size=3, stride=2, padding=1, bias=False) + self.bn2 = nn.BatchNorm2d(64, momentum=BN_MOMENTUM) + self.relu = nn.ReLU(inplace=True) + self.layer1 = self._make_layer(Bottleneck, 64, 4) + + # transformer + self.transformer = Transformer( + dim, + depth, + heads, + mlp_dim, + dropout, + num_keypoints=num_keypoints, + all_attn=self.all_attn, + ) + + self.to_keypoint_token = nn.Identity() + + if hidden_dim: + self.mlp_head = nn.Sequential( + nn.LayerNorm(dim), + nn.Linear(dim, dim), + nn.LayerNorm(dim), + nn.Linear(dim, 2), + ) + else: + self.mlp_head = nn.Sequential( + nn.LayerNorm(dim), + nn.Linear(dim, 2), + ) + trunc_normal_(self.keypoint_token, std=0.02) + if apply_init: + self.apply(self._init_weights) + + ## light-weight head for direct jacobian regression + if estimate_jacobian: + if hidden_dim: + self.mlp_head_jacobian = nn.Sequential( + nn.LayerNorm(dim), + nn.Linear(dim, dim), + nn.LayerNorm(dim), + nn.Linear(dim, 4), + ) + self._init_weights(self.mlp_head_jacobian[0]) + self._init_weights(self.mlp_head_jacobian[1]) + self._init_weights(self.mlp_head_jacobian[2]) + else: + self.mlp_head_jacobian = nn.Sequential( + nn.LayerNorm(dim), + nn.Linear(dim, 4), + ) + self._init_weights(self.mlp_head_jacobian[0]) + self.mlp_head_jacobian[-1].weight.data.zero_() + self.mlp_head_jacobian[-1].bias.data.copy_( + torch.tensor([1, 0, 0, 1], dtype=torch.float) + ) + self.temperature = temperature + self.affine_jacobian = affine_jacobian + ## light-weight head for direct keypoint and jacobian regression + + def _make_position_embedding(self, w, h, d_model, pe_type="sine"): + """ + d_model: embedding size in transformer encoder + """ + assert pe_type in ["none", "learnable", "sine", "sine-full"] + if pe_type == "none": + self.pos_embedding = None + print("==> Without any PositionEmbedding~") + else: + with torch.no_grad(): + self.pe_h = h + self.pe_w = w + length = self.pe_h * self.pe_w + if pe_type == "learnable": + self.pos_embedding = nn.Parameter( + torch.zeros(1, self.num_patches + self.num_keypoints, d_model) + ) + trunc_normal_(self.pos_embedding, std=0.02) + print("==> Add Learnable PositionEmbedding~") + else: + self.pos_embedding = nn.Parameter( + self._make_sine_position_embedding(d_model), requires_grad=False + ) + print("==> Add Sine PositionEmbedding~") + + def _make_sine_position_embedding( + self, d_model, temperature=10000, scale=2 * math.pi + ): + h, w = self.pe_h, self.pe_w + area = torch.ones(1, h, w) # [b, h, w] + y_embed = area.cumsum(1, dtype=torch.float32) + x_embed = area.cumsum(2, dtype=torch.float32) + + one_direction_feats = d_model // 2 + + eps = 1e-6 + y_embed = y_embed / (y_embed[:, -1:, :] + eps) * scale + x_embed = x_embed / (x_embed[:, :, -1:] + eps) * scale + + dim_t = torch.arange(one_direction_feats, dtype=torch.float32) + dim_t = temperature ** (2 * (dim_t // 2) / one_direction_feats) + + pos_x = x_embed[:, :, :, None] / dim_t + pos_y = y_embed[:, :, :, None] / dim_t + pos_x = torch.stack( + (pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4 + ).flatten(3) + pos_y = torch.stack( + (pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4 + ).flatten(3) + pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2) + pos = pos.flatten(2).permute(0, 2, 1) + return pos + + def _make_layer(self, block, planes, blocks, stride=1): + downsample = None + if stride != 1 or self.inplanes != planes * block.expansion: + downsample = nn.Sequential( + nn.Conv2d( + self.inplanes, + planes * block.expansion, + kernel_size=1, + stride=stride, + bias=False, + ), + nn.BatchNorm2d(planes * block.expansion, momentum=BN_MOMENTUM), + ) + + layers = [] + layers.append(block(self.inplanes, planes, stride, downsample)) + self.inplanes = planes * block.expansion + for i in range(1, blocks): + layers.append(block(self.inplanes, planes)) + + return nn.Sequential(*layers) + + def _init_weights(self, m): + # print("Initialization...") + if isinstance(m, nn.Linear): + trunc_normal_(m.weight, std=0.02) + if isinstance(m, nn.Linear) and m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.LayerNorm): + nn.init.constant_(m.bias, 0) + nn.init.constant_(m.weight, 1.0) + + def forward(self, img, mask=None): + p = self.patch_size + # stem net + x = self.conv1(img) + x = self.bn1(x) + x = self.relu(x) + x = self.conv2(x) + x = self.bn2(x) + x = self.relu(x) + x = self.layer1(x) + + # transformer + x = rearrange(x, "b c (h p1) (w p2) -> b (h w) (p1 p2 c)", p1=p[0], p2=p[1]) + x = self.patch_to_embedding(x) + b, n, _ = x.shape + + keypoint_tokens = repeat(self.keypoint_token, "() n d -> b n d", b=b) + if self.pos_embedding_type in ["sine", "sine-full"]: # + x += self.pos_embedding[:, :n] + x = torch.cat((keypoint_tokens, x), dim=1) + elif self.pos_embedding_type == "learnable": + x = torch.cat((keypoint_tokens, x), dim=1) + x += self.pos_embedding[:, : (n + self.num_keypoints)] + x = self.dropout(x) + + x = self.transformer(x, mask, self.pos_embedding) + x_keypoint = self.to_keypoint_token(x[:, 0 : self.num_keypoints // 2]) + keypoint = 2 * F.sigmoid(self.mlp_head(x_keypoint)) - 1 + out = {"value": keypoint} + if self.mlp_head_jacobian is not None: + x_jacobian = self.to_keypoint_token( + x[:, self.num_keypoints // 2 : self.num_keypoints] + ) + # x_jacobian = x_keypoint + jacobian = self.mlp_head_jacobian(x_jacobian) + if self.affine_jacobian: + theta = jacobian[:, :, 0:2] + theta = theta / (torch.norm(theta, p=2, dim=-1, keepdim=True) + 1e-10) + cos_theta = theta[:, :, 0:1] + sin_theta = theta[:, :, 1:2] + rotate = torch.cat( + (cos_theta, -sin_theta, sin_theta, cos_theta), dim=-1 + ) + rotate = rearrange(rotate, "b c (p1 p2) -> b c p1 p2", p1=2, p2=2) + + scale = torch.abs(jacobian[:, :, 2:]) + 0.1 + # scale = torch.tanh(scale)*0.9+1 + # scale = 1/scale + scale_x = scale[:, :, 0:1] + scale_y = scale[:, :, 1:2] + scale = torch.cat( + ( + scale_x, + torch.zeros(scale_x.shape).to(x.device), + torch.zeros(scale_x.shape).to(x.device), + scale_y, + ), + dim=-1, + ) + scale = rearrange(scale, "b c (p1 p2) -> b c p1 p2", p1=2, p2=2) + + jacobian = torch.matmul(rotate, scale) + else: + # jacobian = rearrange(jacobian,'b c (p0 p1 p2) -> b c p0 p1 p2',p0=4,p1=self.heatmap_size[0],p2=self.heatmap_size[1]) + # jacobian = heatmap.unsqueeze(2) * jacobian + # jacobian = jacobian.sum(-1).sum(-1) + jacobian = rearrange(jacobian, "b c (p1 p2) -> b c p1 p2", p1=2, p2=2) + out["jacobian"] = jacobian + return out + + +class TokenPose_L_base(nn.Module): + def __init__( + self, + *, + feature_size, + patch_size, + num_keypoints, + dim, + depth, + heads, + mlp_dim, + apply_init=False, + hidden_heatmap_dim=64 * 6, + heatmap_dim=64 * 48, + heatmap_size=[64, 48], + channels=3, + dropout=0.0, + emb_dropout=0.0, + pos_embedding_type="learnable", + estimate_jacobian=False, + temperature=0.1, + jacobian_token=True, + hidden_dim=False, + affine_jacobian=False, + ): + super().__init__() + assert isinstance(feature_size, list) and isinstance( + patch_size, list + ), "image_size and patch_size should be list" + assert ( + feature_size[0] % patch_size[0] == 0 + and feature_size[1] % patch_size[1] == 0 + ), "Image dimensions must be divisible by the patch size." + num_patches = (feature_size[0] // (patch_size[0])) * ( + feature_size[1] // (patch_size[1]) + ) + patch_dim = channels * patch_size[0] * patch_size[1] + assert pos_embedding_type in ["sine", "learnable", "sine-full"] + + self.inplanes = 64 + self.patch_size = patch_size + self.heatmap_size = heatmap_size + self.num_keypoints = num_keypoints + self.num_patches = num_patches + self.pos_embedding_type = pos_embedding_type + self.all_attn = self.pos_embedding_type == "sine-full" + + self.jacobian_token = jacobian_token + if jacobian_token: + ## additional token for jacobian + num_keypoints = 2 * num_keypoints + self.num_keypoints = num_keypoints + self.keypoint_token = nn.Parameter(torch.zeros(1, self.num_keypoints, dim)) + h, w = ( + feature_size[0] // (self.patch_size[0]), + feature_size[1] // (self.patch_size[1]), + ) + + # for normal + self._make_position_embedding(w, h, dim, pos_embedding_type) + + self.patch_to_embedding = nn.Linear(patch_dim, dim) + self.dropout = nn.Dropout(emb_dropout) + + # transformer + self.transformer1 = Transformer( + dim, + depth, + heads, + mlp_dim, + dropout, + num_keypoints=num_keypoints, + all_attn=self.all_attn, + scale_with_head=True, + ) + self.transformer2 = Transformer( + dim, + depth, + heads, + mlp_dim, + dropout, + num_keypoints=num_keypoints, + all_attn=self.all_attn, + scale_with_head=True, + ) + self.transformer3 = Transformer( + dim, + depth, + heads, + mlp_dim, + dropout, + num_keypoints=num_keypoints, + all_attn=self.all_attn, + scale_with_head=True, + ) + + self.to_keypoint_token = nn.Identity() + if hidden_dim: + self.mlp_head = nn.Sequential( + # nn.LayerNorm(dim), + # nn.Linear(dim, dim), + # nn.LayerNorm(dim), + # nn.Linear(dim, 2), + nn.LayerNorm(dim * 3), + nn.Linear(dim * 3, dim * 3), + nn.LayerNorm(dim * 3), + nn.Linear(dim * 3, 2), + ) + else: + self.mlp_head = nn.Sequential( + # nn.LayerNorm(dim), + # nn.Linear(dim, 2), + nn.LayerNorm(dim * 3), + nn.Linear(dim * 3, 2), + ) + trunc_normal_(self.keypoint_token, std=0.02) + if apply_init: + self.apply(self._init_weights) + + ## light-weight head for direct jacobian regression + if estimate_jacobian: + if hidden_dim: + self.mlp_head_jacobian = nn.Sequential( + # nn.LayerNorm(dim), + # nn.Linear(dim, dim), + # nn.LayerNorm(dim), + # nn.Linear(dim, 4), + nn.LayerNorm(dim * 3), + nn.Linear(dim * 3, dim * 3), + nn.LayerNorm(dim * 3), + nn.Linear(dim * 3, 4), + ) + self._init_weights(self.mlp_head_jacobian[0]) + self._init_weights(self.mlp_head_jacobian[1]) + self._init_weights(self.mlp_head_jacobian[2]) + else: + self.mlp_head_jacobian = nn.Sequential( + # nn.LayerNorm(dim), + # nn.Linear(dim, 4), + nn.LayerNorm(dim * 3), + nn.Linear(dim * 3, 4), + ) + self._init_weights(self.mlp_head_jacobian[0]) + self.mlp_head_jacobian[-1].weight.data.zero_() + self.mlp_head_jacobian[-1].bias.data.copy_( + torch.tensor([1, 0, 0, 1], dtype=torch.float) + ) + self.temperature = temperature + self.affine_jacobian = affine_jacobian + ## light-weight head for direct keypoint and jacobian regression + + def _make_position_embedding(self, w, h, d_model, pe_type="sine"): + """ + d_model: embedding size in transformer encoder + """ + assert pe_type in ["none", "learnable", "sine", "sine-full"] + if pe_type == "none": + self.pos_embedding = None + print("==> Without any PositionEmbedding~") + else: + with torch.no_grad(): + self.pe_h = h + self.pe_w = w + length = self.pe_h * self.pe_w + if pe_type == "learnable": + self.pos_embedding = nn.Parameter( + torch.zeros(1, self.num_patches + self.num_keypoints, d_model) + ) + trunc_normal_(self.pos_embedding, std=0.02) + print("==> Add Learnable PositionEmbedding~") + else: + self.pos_embedding = nn.Parameter( + self._make_sine_position_embedding(d_model), requires_grad=False + ) + print("==> Add Sine PositionEmbedding~") + + def _make_sine_position_embedding( + self, d_model, temperature=10000, scale=2 * math.pi + ): + h, w = self.pe_h, self.pe_w + area = torch.ones(1, h, w) # [b, h, w] + y_embed = area.cumsum(1, dtype=torch.float32) + x_embed = area.cumsum(2, dtype=torch.float32) + + one_direction_feats = d_model // 2 + + eps = 1e-6 + y_embed = y_embed / (y_embed[:, -1:, :] + eps) * scale + x_embed = x_embed / (x_embed[:, :, -1:] + eps) * scale + + dim_t = torch.arange(one_direction_feats, dtype=torch.float32) + dim_t = temperature ** (2 * (dim_t // 2) / one_direction_feats) + + pos_x = x_embed[:, :, :, None] / dim_t + pos_y = y_embed[:, :, :, None] / dim_t + pos_x = torch.stack( + (pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4 + ).flatten(3) + pos_y = torch.stack( + (pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4 + ).flatten(3) + pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2) + pos = pos.flatten(2).permute(0, 2, 1) + return pos + + def _make_layer(self, block, planes, blocks, stride=1): + downsample = None + if stride != 1 or self.inplanes != planes * block.expansion: + downsample = nn.Sequential( + nn.Conv2d( + self.inplanes, + planes * block.expansion, + kernel_size=1, + stride=stride, + bias=False, + ), + nn.BatchNorm2d(planes * block.expansion, momentum=BN_MOMENTUM), + ) + + layers = [] + layers.append(block(self.inplanes, planes, stride, downsample)) + self.inplanes = planes * block.expansion + for i in range(1, blocks): + layers.append(block(self.inplanes, planes)) + + return nn.Sequential(*layers) + + def _init_weights(self, m): + # print("Initialization...") + if isinstance(m, nn.Linear): + trunc_normal_(m.weight, std=0.02) + if isinstance(m, nn.Linear) and m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.LayerNorm): + nn.init.constant_(m.bias, 0) + nn.init.constant_(m.weight, 1.0) + + def forward(self, feature, mask=None): + p = self.patch_size + # transformer + x = rearrange( + feature, "b c (h p1) (w p2) -> b (h w) (p1 p2 c)", p1=p[0], p2=p[1] + ) + x = self.patch_to_embedding(x) + + b, n, _ = x.shape + + keypoint_tokens = repeat(self.keypoint_token, "() n d -> b n d", b=b) + if self.pos_embedding_type in ["sine", "sine-full"]: + x += self.pos_embedding[:, :n] + x = torch.cat((keypoint_tokens, x), dim=1) + else: + x = torch.cat((keypoint_tokens, x), dim=1) + x += self.pos_embedding[:, : (n + self.num_keypoints)] + x = self.dropout(x) + + x1 = self.transformer1(x, mask, self.pos_embedding) + x2 = self.transformer2(x1, mask, self.pos_embedding) + x3 = self.transformer3(x2, mask, self.pos_embedding) + + x1_out = self.to_keypoint_token(x1[:, 0 : self.num_keypoints]) + x2_out = self.to_keypoint_token(x2[:, 0 : self.num_keypoints]) + x3_out = self.to_keypoint_token(x3[:, 0 : self.num_keypoints]) + + x = torch.cat((x1_out, x2_out, x3_out), dim=2) + # x = x3 + + if self.jacobian_token: + x_keypoint = self.to_keypoint_token(x[:, 0 : self.num_keypoints // 2]) + else: + x_keypoint = self.to_keypoint_token(x[:, 0 : self.num_keypoints]) + keypoint = 2 * F.sigmoid(self.mlp_head(x_keypoint)) - 1 + out = {"value": keypoint} + if self.mlp_head_jacobian is not None: + if self.jacobian_token: + x_jacobian = self.to_keypoint_token( + x[:, self.num_keypoints // 2 : self.num_keypoints] + ) + else: + x_jacobian = x_keypoint + jacobian = self.mlp_head_jacobian(x_jacobian) + if self.affine_jacobian: + theta = jacobian[:, :, 0:2] + theta = theta / (torch.norm(theta, p=2, dim=-1, keepdim=True) + 1e-10) + cos_theta = theta[:, :, 0:1] + sin_theta = theta[:, :, 1:2] + rotate = torch.cat( + (cos_theta, -sin_theta, sin_theta, cos_theta), dim=-1 + ) + rotate = rearrange(rotate, "b c (p1 p2) -> b c p1 p2", p1=2, p2=2) + + scale = torch.abs(jacobian[:, :, 2:]) + 0.1 + # scale = torch.tanh(scale)*0.9+1 + # scale = 1/scale + scale_x = scale[:, :, 0:1] + scale_y = scale[:, :, 1:2] + scale = torch.cat( + ( + scale_x, + torch.zeros(scale_x.shape).to(x.device), + torch.zeros(scale_x.shape).to(x.device), + scale_y, + ), + dim=-1, + ) + scale = rearrange(scale, "b c (p1 p2) -> b c p1 p2", p1=2, p2=2) + + jacobian = torch.matmul(rotate, scale) + else: + # jacobian = rearrange(jacobian,'b c (p0 p1 p2) -> b c p0 p1 p2',p0=4,p1=self.heatmap_size[0],p2=self.heatmap_size[1]) + # jacobian = heatmap.unsqueeze(2) * jacobian + # jacobian = jacobian.sum(-1).sum(-1) + jacobian = rearrange(jacobian, "b c (p1 p2) -> b c p1 p2", p1=2, p2=2) + out["jacobian"] = jacobian + return out diff --git a/module_mrfa/util.py b/module_mrfa/util.py new file mode 100644 index 0000000..f999e96 --- /dev/null +++ b/module_mrfa/util.py @@ -0,0 +1,503 @@ +import torch +import torch.nn.functional as F +from torch import nn +from torch.nn import BatchNorm2d + + +class AttributeDict(dict): + def __getattr__(self, attr): + try: + return self[attr] + except AttributeError as e: + raise AttributeError(f"{e}") + + def __setattr__(self, attr, value): + self[attr] = value + + +def convert_dict_to_attrit_dict(dict_instance): + att_dict = AttributeDict() + for key, value in dict_instance.items(): + if isinstance(value, dict): + value = convert_dict_to_attrit_dict(value) + att_dict.__setattr__(key, value) + return att_dict + + +def bilinear_sampler(img, coords, mode="bilinear", mask=False): + """Wrapper for grid_sample, uses pixel coordinates""" + H, W = img.shape[-2:] + xgrid, ygrid = coords.split([1, 1], dim=-1) + xgrid = 2 * xgrid / (W - 1) - 1 + ygrid = 2 * ygrid / (H - 1) - 1 + + grid = torch.cat([xgrid, ygrid], dim=-1) + img = F.grid_sample(img, grid, align_corners=True) + if mask: + mask = (xgrid > -1) & (ygrid > -1) & (xgrid < 1) & (ygrid < 1) + return img, mask.float() + return img + + +def batch_bilinear_sampler( + img, coords, mode="bilinear", mask=False, h=256, w=256, mini_batch=4 +): + """Wrapper for grid_sample, uses pixel coordinates""" + batch = img.shape[0] // (h * w) + # img = img.view(batch,h,w,img.shape[1],img.shape[2],img.shape[3]) + # coords = coords.view(batch,h,w,coords.shape[1],coords.shape[2],coords.shape[3]) + num_iter = batch // mini_batch + sampled = [] + for i in range(num_iter): + img_i = img[i * mini_batch * h * w : (i + 1) * mini_batch * h * w, :, :, :] + coords_i = coords[ + i * mini_batch * h * w : (i + 1) * mini_batch * h * w, :, :, : + ] + sampled.append(bilinear_sampler(img_i, coords_i)) + return torch.cat(sampled, dim=0) + + +def coords_grid(batch, ht, wd, device): + coords = torch.meshgrid( + torch.arange(ht, device=device), torch.arange(wd, device=device) + ) + coords = torch.stack(coords[::-1], dim=0).float() + return coords[None].repeat(batch, 1, 1, 1) + + +def kp2gaussian(kp, spatial_size, kp_variance): + """ + Transform a keypoint into gaussian like representation + """ + # mean = kp['value'] + mean = kp + # N * 10 * 2 + + coordinate_grid = make_coordinate_grid(spatial_size, mean.type()) + number_of_leading_dimensions = len(mean.shape) - 1 + shape = (1,) * number_of_leading_dimensions + coordinate_grid.shape + # 1,1,h,w,2 + coordinate_grid = coordinate_grid.view(*shape) + repeats = mean.shape[:number_of_leading_dimensions] + (1, 1, 1) + # N * 10 * 1 *1 *1 + coordinate_grid = coordinate_grid.repeat(*repeats) + # N,10,h,w,2 + + # Preprocess kp shape + shape = mean.shape[:number_of_leading_dimensions] + (1, 1, 2) + # N * 10 * 1 * 1 * 2 + mean = mean.view(*shape) + + mean_sub = coordinate_grid - mean + + out = torch.exp(-0.5 * (mean_sub**2).sum(-1) / kp_variance) + # N * 10 * h * w + + return out + + +def make_coordinate_grid(spatial_size, type): + """ + Create a meshgrid [-1,1] x [-1,1] of given spatial_size. + """ + h, w = spatial_size + x = torch.arange(w).type(type) + y = torch.arange(h).type(type) + + x = 2 * (x / (w - 1)) - 1 + y = 2 * (y / (h - 1)) - 1 + + yy = y.view(-1, 1).repeat(1, w) + xx = x.view(1, -1).repeat(h, 1) + # xx [[-1,-1/3,1/3,1],[-1,-1/3,1/3,1],[-1,-1/3,1/3,1],[-1,-1/3,1/3,1]] + # yy [[-1,-1,-1,-1],[-1/3,-1/3,-1/3,-1/3],[1/3,1/3,1/3,1/3],[1,1,1,1]] + + meshed = torch.cat([xx.unsqueeze_(2), yy.unsqueeze_(2)], 2) + + return meshed + + +class ChannelBlock2d(nn.Module): + """ + Res block, preserve spatial resolution. + """ + + def __init__(self, in_features, kernel_size, padding): + super(ChannelBlock2d, self).__init__() + self.conv1 = nn.Conv2d( + in_channels=in_features, + out_channels=in_features // 2, + kernel_size=kernel_size, + padding=padding, + ) + # self.conv2 = nn.Conv2d(in_channels=in_features//2, out_channels=in_features//2, kernel_size=kernel_size, + # padding=padding) + self.norm1 = BatchNorm2d(in_features, affine=True) + # self.norm2 = BatchNorm2d(in_features, affine=True) + + def forward(self, x): + out = self.norm1(x) + out = F.relu(out) + out = self.conv1(out) + # out = self.norm2(out) + # out = F.relu(out) + # out = self.conv2(out) + # out += x + return out + + +class ResBlock2d(nn.Module): + """ + Res block, preserve spatial resolution. + """ + + def __init__(self, in_features, kernel_size, padding): + super(ResBlock2d, self).__init__() + self.conv1 = nn.Conv2d( + in_channels=in_features, + out_channels=in_features, + kernel_size=kernel_size, + padding=padding, + ) + self.conv2 = nn.Conv2d( + in_channels=in_features, + out_channels=in_features, + kernel_size=kernel_size, + padding=padding, + ) + self.norm1 = BatchNorm2d(in_features, affine=True) + self.norm2 = BatchNorm2d(in_features, affine=True) + + def forward(self, x): + out = self.norm1(x) + out = F.relu(out) + out = self.conv1(out) + out = self.norm2(out) + out = F.relu(out) + out = self.conv2(out) + out += x + return out + + +class UpBlock2d(nn.Module): + """ + Upsampling block for use in decoder. + """ + + def __init__(self, in_features, out_features, kernel_size=3, padding=1, groups=1): + super(UpBlock2d, self).__init__() + + self.conv = nn.Conv2d( + in_channels=in_features, + out_channels=out_features, + kernel_size=kernel_size, + padding=padding, + groups=groups, + ) + self.norm = BatchNorm2d(out_features, affine=True) + + def forward(self, x): + out = F.interpolate(x, scale_factor=2) + out = self.conv(out) + out = self.norm(out) + out = F.relu(out) + return out + + +class DownBlock2d(nn.Module): + """ + Downsampling block for use in encoder. + """ + + def __init__(self, in_features, out_features, kernel_size=3, padding=1, groups=1): + super(DownBlock2d, self).__init__() + self.conv = nn.Conv2d( + in_channels=in_features, + out_channels=out_features, + kernel_size=kernel_size, + padding=padding, + groups=groups, + ) + self.norm = BatchNorm2d(out_features, affine=True) + self.pool = nn.AvgPool2d(kernel_size=(2, 2)) + + def forward(self, x): + out = self.conv(x) + out = self.norm(out) + out = F.relu(out) + out = self.pool(out) + return out + + +class SameBlock2d(nn.Module): + """ + Simple block, preserve spatial resolution. + """ + + def __init__(self, in_features, out_features, groups=1, kernel_size=3, padding=1): + super(SameBlock2d, self).__init__() + self.conv = nn.Conv2d( + in_channels=in_features, + out_channels=out_features, + kernel_size=kernel_size, + padding=padding, + groups=groups, + ) + self.norm = BatchNorm2d(out_features, affine=True) + + def forward(self, x): + out = self.conv(x) + out = self.norm(out) + out = F.relu(out) + return out + + +class Encoder(nn.Module): + """ + Hourglass Encoder + """ + + def __init__(self, block_expansion, in_features, num_blocks=3, max_features=256): + super(Encoder, self).__init__() + + down_blocks = [] + for i in range(num_blocks): + down_blocks.append( + DownBlock2d( + in_features + if i == 0 + else min(max_features, block_expansion * (2**i)), + min(max_features, block_expansion * (2 ** (i + 1))), + kernel_size=3, + padding=1, + ) + ) + self.down_blocks = nn.ModuleList(down_blocks) + + def forward(self, x): + outs = [x] + for down_block in self.down_blocks: + outs.append(down_block(outs[-1])) + return outs + + +class Decoder(nn.Module): + """ + Hourglass Decoder + """ + + def __init__(self, block_expansion, in_features, num_blocks=3, max_features=256): + super(Decoder, self).__init__() + + up_blocks = [] + + for i in range(num_blocks)[::-1]: + in_filters = (1 if i == num_blocks - 1 else 2) * min( + max_features, block_expansion * (2 ** (i + 1)) + ) + out_filters = min(max_features, block_expansion * (2**i)) + up_blocks.append( + UpBlock2d(in_filters, out_filters, kernel_size=3, padding=1) + ) + + self.up_blocks = nn.ModuleList(up_blocks) + self.out_filters = block_expansion + in_features + + def forward(self, x): + out = x.pop() + for up_block in self.up_blocks: + out = up_block(out) + skip = x.pop() + out = torch.cat([out, skip], dim=1) + return out + + +class Hourglass(nn.Module): + """ + Hourglass architecture. + """ + + def __init__(self, block_expansion, in_features, num_blocks=3, max_features=256): + super(Hourglass, self).__init__() + self.encoder = Encoder(block_expansion, in_features, num_blocks, max_features) + self.decoder = Decoder(block_expansion, in_features, num_blocks, max_features) + self.out_filters = self.decoder.out_filters + + def forward(self, x): + return self.decoder(self.encoder(x)) + + +class AntiAliasInterpolation2d(nn.Module): + """ + Band-limited downsampling, for better preservation of the input signal. + """ + + def __init__(self, channels, scale): + super(AntiAliasInterpolation2d, self).__init__() + sigma = (1 / scale - 1) / 2 + kernel_size = 2 * round(sigma * 4) + 1 + self.ka = kernel_size // 2 + self.kb = self.ka - 1 if kernel_size % 2 == 0 else self.ka + + kernel_size = [kernel_size, kernel_size] + sigma = [sigma, sigma] + # The gaussian kernel is the product of the + # gaussian function of each dimension. + kernel = 1 + meshgrids = torch.meshgrid( + [torch.arange(size, dtype=torch.float32) for size in kernel_size] + ) + for size, std, mgrid in zip(kernel_size, sigma, meshgrids): + mean = (size - 1) / 2 + kernel *= torch.exp(-((mgrid - mean) ** 2) / (2 * std**2)) + + # Make sure sum of values in gaussian kernel equals 1. + kernel = kernel / torch.sum(kernel) + # Reshape to depthwise convolutional weight + kernel = kernel.view(1, 1, *kernel.size()) + kernel = kernel.repeat(channels, *[1] * (kernel.dim() - 1)) + + self.register_buffer("weight", kernel) + self.groups = channels + self.scale = scale + + def forward(self, input): + if self.scale == 1.0: + return input + + out = F.pad(input, (self.ka, self.kb, self.ka, self.kb)) + out = F.conv2d(out, weight=self.weight, groups=self.groups) + out = F.interpolate(out, scale_factor=(self.scale, self.scale)) + + return out + + +def to_homogeneous(coordinates): + ones_shape = list(coordinates.shape) + ones_shape[-1] = 1 + ones = torch.ones(ones_shape).type(coordinates.type()) + + return torch.cat([coordinates, ones], dim=-1) + + +def from_homogeneous(coordinates): + return coordinates[..., :2] / coordinates[..., 2:3] + + +class TPS: + """ + TPS transformation, mode 'kp' for Eq(2) in the paper, mode 'random' for equivariance loss. + """ + + def __init__(self, mode, bs, **kwargs): + self.bs = bs + self.mode = mode + if mode == "random": + noise = torch.normal( + mean=0, std=kwargs["sigma_affine"] * torch.ones([bs, 2, 3]) + ) + self.theta = noise + torch.eye(2, 3).view(1, 2, 3) + self.control_points = make_coordinate_grid( + (kwargs["points_tps"], kwargs["points_tps"]), type=noise.type() + ) + self.control_points = self.control_points.unsqueeze(0) + self.control_params = torch.normal( + mean=0, + std=kwargs["sigma_tps"] + * torch.ones([bs, 1, kwargs["points_tps"] ** 2]), + ) + elif mode == "kp": + kp_1 = kwargs["kp_1"] + kp_2 = kwargs["kp_2"] + device = kp_1.device + kp_type = kp_1.type() + self.gs = kp_1.shape[1] + n = kp_1.shape[2] + K = torch.norm(kp_1[:, :, :, None] - kp_1[:, :, None, :], dim=4, p=2) + K = K**2 + K = K * torch.log(K + 1e-9) + + one1 = ( + torch.ones(self.bs, kp_1.shape[1], kp_1.shape[2], 1) + .to(device) + .type(kp_type) + ) + kp_1p = torch.cat([kp_1, one1], 3) + + zero = torch.zeros(self.bs, kp_1.shape[1], 3, 3).to(device).type(kp_type) + P = torch.cat([kp_1p, zero], 2) + L = torch.cat([K, kp_1p.permute(0, 1, 3, 2)], 2) + L = torch.cat([L, P], 3) + + zero = torch.zeros(self.bs, kp_1.shape[1], 3, 2).to(device).type(kp_type) + Y = torch.cat([kp_2, zero], 2) + one = torch.eye(L.shape[2]).expand(L.shape).to(device).type(kp_type) * 0.01 + L = L + one + + param = torch.matmul(torch.inverse(L), Y) + self.theta = param[:, :, n:, :].permute(0, 1, 3, 2) + + self.control_points = kp_1 + self.control_params = param[:, :, :n, :] + else: + raise Exception("Error TPS mode") + + def transform_frame(self, frame): + grid = ( + make_coordinate_grid(frame.shape[2:], type=frame.type()) + .unsqueeze(0) + .to(frame.device) + ) + grid = grid.view(1, frame.shape[2] * frame.shape[3], 2) + shape = [self.bs, frame.shape[2], frame.shape[3], 2] + if self.mode == "kp": + shape.insert(1, self.gs) + grid = self.warp_coordinates(grid).view(*shape) + return grid + + def warp_coordinates(self, coordinates): + theta = self.theta.type(coordinates.type()).to(coordinates.device) + control_points = self.control_points.type(coordinates.type()).to( + coordinates.device + ) + control_params = self.control_params.type(coordinates.type()).to( + coordinates.device + ) + + if self.mode == "kp": + transformed = ( + torch.matmul(theta[:, :, :, :2], coordinates.permute(0, 2, 1)) + + theta[:, :, :, 2:] + ) + + distances = coordinates.view( + coordinates.shape[0], 1, 1, -1, 2 + ) - control_points.view(self.bs, control_points.shape[1], -1, 1, 2) + + distances = distances**2 + result = distances.sum(-1) + result = result * torch.log(result + 1e-9) + result = torch.matmul(result.permute(0, 1, 3, 2), control_params) + transformed = transformed.permute(0, 1, 3, 2) + result + + elif self.mode == "random": + theta = theta.unsqueeze(1) + transformed = ( + torch.matmul(theta[:, :, :, :2], coordinates.unsqueeze(-1)) + + theta[:, :, :, 2:] + ) + transformed = transformed.squeeze(-1) + ances = coordinates.view( + coordinates.shape[0], -1, 1, 2 + ) - control_points.view(1, 1, -1, 2) + distances = ances**2 + + result = distances.sum(-1) + result = result * torch.log(result + 1e-9) + result = result * control_params + result = result.sum(dim=2).view(self.bs, coordinates.shape[1], 1) + transformed = transformed + result + else: + raise Exception("Error TPS mode") + + return transformed diff --git a/pyproject.toml b/pyproject.toml index cfa910a..89191c8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,9 +1,9 @@ [project] name = "comfyui-firstordermm" description = "ComfyUI-native nodes to run First Order Motion Model for Image Animation and its non-diffusion-based successors. [a/https://github.com/AliaksandrSiarohin/first-order-model](https://github.com/AliaksandrSiarohin/first-order-model)" -version = "1.0.3" +version = "1.0.4" license = { file = "LICENSE" } -dependencies = ["numpy", "torch", "scipy", "yaml", "matplotlib"] +dependencies = ["numpy", "torch", "scipy", "pyyaml", "matplotlib", "einops", "timm"] [project.urls] Repository = "https://github.com/FuouM/ComfyUI-FirstOrderMM" diff --git a/requirements.txt b/requirements.txt index 5eecb23..7443ff8 100644 --- a/requirements.txt +++ b/requirements.txt @@ -2,4 +2,6 @@ numpy torch scipy pyyaml -matplotlib \ No newline at end of file +matplotlib +einops +timm \ No newline at end of file diff --git a/run.py b/run.py index 340dae2..9fc0491 100644 --- a/run.py +++ b/run.py @@ -12,6 +12,10 @@ import torch from .constants import ( ARTICULATE_CFG_PATH, ARTICULATE_MODEL_PATH, + MRFA_CFG_PATHS, + MRFA_DEFAULT_MODEL, + MRFA_MODEL_NAMES, + MRFA_MODEL_PATHS, SPLINE_CFG_PATH, SPLINE_DEFAULT, SPLINE_MODEL_PATH, @@ -29,6 +33,7 @@ from .constants import ( from .face_parsing.face_parsing_loader import load_face_parser_model from .inference_articulate import articulate_inference from .inference_fomm import inference, inference_best_frame, load_checkpoint +from .inference_mrfa import mrfa_inference from .inference_partswap import load_partswap_checkpoint, partswap_inference from .inference_spline import spline_inference from .seg_viz import visualize_frame @@ -449,7 +454,7 @@ class Spline_Runner: find_best_frame: bool, audio=None, ): - print(f"{type(source_image)=}") + print(f"{type(source_image)=}") # [B, H, W, C] print(f"{type(driving_video_input)=}") print(f"{source_image.shape=}") print(f"{driving_video_input.shape=}") @@ -486,6 +491,99 @@ class Spline_Runner: ) +class MRFA_Runner: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "source_image": ("IMAGE",), + "driving_video_input": ("IMAGE",), + "model_name": (MRFA_MODEL_NAMES, {"default": MRFA_DEFAULT_MODEL}), + "frame_rate": ("FLOAT", {"default": 30.0}), + "use_relative": ( + "BOOLEAN", + {"default": True}, + ), + "relative_movement": ( + "BOOLEAN", + {"default": True}, + ), + "relative_jacobian": ( + "BOOLEAN", + {"default": True}, + ), + "adapt_movement_scale": ( + "BOOLEAN", + {"default": False}, + ), + }, + "optional": {"audio": ("AUDIO",)}, + } + + RETURN_TYPES = ( + "IMAGE", + "AUDIO", + "FLOAT", + ) + RETURN_NAMES = ( + "images", + "audio", + "frame_rate", + ) + FUNCTION = "todo" + CATEGORY = "FirstOrderMM" + + def todo( + self, + source_image, + driving_video_input, + model_name: str, + frame_rate: float, + use_relative: bool, + relative_movement: bool, + relative_jacobian: bool, + adapt_movement_scale: bool, + audio=None, + ): + print(f"{type(source_image)=}") # [B, H, W, C] + print(f"{type(driving_video_input)=}") + print(f"{source_image.shape=}") + print(f"{driving_video_input.shape=}") + print(f"{type(audio)=}") + print(base_dir) + + config_path = f"{base_dir}/{MRFA_CFG_PATHS[model_name]}" + checkpoint_path = f"{base_dir}/{MRFA_MODEL_PATHS[model_name]}" + + source_image = reshape_image(source_image, (256, 256)) + driving_video = reshape_image(driving_video_input, (256, 256)).unsqueeze(0) + driving_video = driving_video.permute(0, 2, 1, 3, 4) + + print("After reshaping") + print(f"{source_image.shape=}") + print(f"{driving_video.shape=}") + params = { + "source_image": source_image, + "driving_video": driving_video, + "config_path": config_path, + "checkpoint_path": checkpoint_path, + "use_relative": use_relative, + "relative_movement": relative_movement, + "relative_jacobian": relative_jacobian, + "adapt_movement_scale": adapt_movement_scale, + } + + predictions = mrfa_inference(**params) + + output_images = out_video(predictions) + + return ( + output_images, + audio, + frame_rate, + ) + + def serialize_integers(int_list): return "_".join(map(str, int_list)) diff --git a/workflows/MRFA.json b/workflows/MRFA.json new file mode 100644 index 0000000..7b12dbf --- /dev/null +++ b/workflows/MRFA.json @@ -0,0 +1,423 @@ +{ + "last_node_id": 7, + "last_link_id": 18, + "nodes": [ + { + "id": 2, + "type": "LoadImage", + "pos": [ + 83, + 105 + ], + "size": { + "0": 315, + "1": 314 + }, + "flags": {}, + "order": 0, + "mode": 0, + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 13 + ], + "shape": 3, + "slot_index": 0 + }, + { + "name": "MASK", + "type": "MASK", + "links": null, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "LoadImage" + }, + "widgets_values": [ + "Rick1.png", + "image" + ] + }, + { + "id": 5, + "type": "VHS_VideoInfoLoaded", + "pos": [ + 753, + 383 + ], + "size": { + "0": 304.79998779296875, + "1": 106 + }, + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [ + { + "name": "video_info", + "type": "VHS_VIDEOINFO", + "link": 7 + } + ], + "outputs": [ + { + "name": "fps🟦", + "type": "FLOAT", + "links": [ + 18 + ], + "shape": 3, + "slot_index": 0 + }, + { + "name": "frame_count🟦", + "type": "INT", + "links": null, + "shape": 3 + }, + { + "name": "duration🟦", + "type": "FLOAT", + "links": null, + "shape": 3 + }, + { + "name": "width🟦", + "type": "INT", + "links": null, + "shape": 3 + }, + { + "name": "height🟦", + "type": "INT", + "links": null, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "VHS_VideoInfoLoaded" + }, + "widgets_values": {} + }, + { + "id": 6, + "type": "VHS_VideoCombine", + "pos": [ + 1236, + 163 + ], + "size": [ + 315, + 615 + ], + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 15 + }, + { + "name": "audio", + "type": "AUDIO", + "link": 16 + }, + { + "name": "meta_batch", + "type": "VHS_BatchManager", + "link": null + }, + { + "name": "vae", + "type": "VAE", + "link": null + }, + { + "name": "frame_rate", + "type": "FLOAT", + "link": 17, + "widget": { + "name": "frame_rate" + } + } + ], + "outputs": [ + { + "name": "Filenames", + "type": "VHS_FILENAMES", + "links": null, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "VHS_VideoCombine" + }, + "widgets_values": { + "frame_rate": 8, + "loop_count": 0, + "filename_prefix": "MRFA", + "format": "video/h264-mp4", + "pix_fmt": "yuv420p", + "crf": 0, + "save_metadata": true, + "pingpong": false, + "save_output": true, + "videopreview": { + "hidden": false, + "paused": false, + "params": { + "filename": "MRFA_00003.mp4", + "subfolder": "", + "type": "output", + "format": "video/h264-mp4", + "frame_rate": 29.970029970029966 + } + } + } + }, + { + "id": 3, + "type": "VHS_LoadVideo", + "pos": [ + 430, + 196 + ], + "size": [ + 235.1999969482422, + 491.1999969482422 + ], + "flags": { + "collapsed": false + }, + "order": 1, + "mode": 0, + "inputs": [ + { + "name": "meta_batch", + "type": "VHS_BatchManager", + "link": null + }, + { + "name": "vae", + "type": "VAE", + "link": null + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 14 + ], + "shape": 3, + "slot_index": 0 + }, + { + "name": "frame_count", + "type": "INT", + "links": null, + "shape": 3 + }, + { + "name": "audio", + "type": "AUDIO", + "links": [], + "shape": 3, + "slot_index": 2 + }, + { + "name": "video_info", + "type": "VHS_VIDEOINFO", + "links": [ + 7 + ], + "shape": 3, + "slot_index": 3 + } + ], + "properties": { + "Node name for S&R": "VHS_LoadVideo" + }, + "widgets_values": { + "video": "damedane.mp4", + "force_rate": 0, + "force_size": "Disabled", + "custom_width": 512, + "custom_height": 512, + "frame_load_cap": 0, + "skip_first_frames": 0, + "select_every_nth": 1, + "choose video to upload": "image", + "videopreview": { + "hidden": false, + "paused": false, + "params": { + "frame_load_cap": 0, + "skip_first_frames": 0, + "force_rate": 0, + "filename": "damedane.mp4", + "type": "input", + "format": "video/mp4", + "select_every_nth": 1 + } + } + } + }, + { + "id": 7, + "type": "MRFA_Runner", + "pos": [ + 779, + 107 + ], + "size": { + "0": 380.4000244140625, + "1": 218 + }, + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [ + { + "name": "source_image", + "type": "IMAGE", + "link": 13 + }, + { + "name": "driving_video_input", + "type": "IMAGE", + "link": 14 + }, + { + "name": "audio", + "type": "AUDIO", + "link": null + }, + { + "name": "frame_rate", + "type": "FLOAT", + "link": 18, + "widget": { + "name": "frame_rate" + } + } + ], + "outputs": [ + { + "name": "images", + "type": "IMAGE", + "links": [ + 15 + ], + "shape": 3, + "slot_index": 0 + }, + { + "name": "audio", + "type": "AUDIO", + "links": [ + 16 + ], + "shape": 3, + "slot_index": 1 + }, + { + "name": "frame_rate", + "type": "FLOAT", + "links": [ + 17 + ], + "shape": 3, + "slot_index": 2 + } + ], + "properties": { + "Node name for S&R": "MRFA_Runner" + }, + "widgets_values": [ + "celebvhq", + 30, + true, + true, + true, + true + ] + } + ], + "links": [ + [ + 7, + 3, + 3, + 5, + 0, + "VHS_VIDEOINFO" + ], + [ + 13, + 2, + 0, + 7, + 0, + "IMAGE" + ], + [ + 14, + 3, + 0, + 7, + 1, + "IMAGE" + ], + [ + 15, + 7, + 0, + 6, + 0, + "IMAGE" + ], + [ + 16, + 7, + 1, + 6, + 1, + "AUDIO" + ], + [ + 17, + 7, + 2, + 6, + 4, + "FLOAT" + ], + [ + 18, + 5, + 0, + 7, + 3, + "FLOAT" + ] + ], + "groups": [], + "config": {}, + "extra": { + "ds": { + "scale": 1, + "offset": [ + -17, + 60 + ] + } + }, + "version": 0.4 +} \ No newline at end of file diff --git a/workflows/workflow_mrfa.png b/workflows/workflow_mrfa.png new file mode 100644 index 0000000..6d075e6 Binary files /dev/null and b/workflows/workflow_mrfa.png differ