Fuou Marinas
+3
-2
@@ -1,7 +1,8 @@
|
||||
# Stuff
|
||||
|
||||
articulate_module/models/*.pth
|
||||
spline_module/models/*.tar
|
||||
module_articulate/models/*.pth
|
||||
module_spline/models/*.tar
|
||||
module_mrfa/models/*.pth
|
||||
checkpoints/*
|
||||
!checkpoints/place_checkpoints_here.txt
|
||||
|
||||
|
||||
@@ -7,6 +7,7 @@ Now supports:
|
||||
1. Face Swapping using [Motion Supervised co-part Segmentation](https://github.com/AliaksandrSiarohin/motion-cosegmentation):
|
||||
2. [Motion Representations for Articulated Animation](https://github.com/snap-research/articulated-animation)
|
||||
3. [Thin-Plate Spline Motion Model for Image Animation](https://github.com/yoyo-nb/thin-plate-spline-motion-model)
|
||||
4. [Learning Motion Refinement for Unsupervised Face Animation](https://github.com/JialeTao/MRFA/)
|
||||
|
||||
More will come soon
|
||||
|
||||
@@ -39,6 +40,11 @@ https://github.com/user-attachments/assets/b2948efb-3b44-440b-bff2-dde7b95a9946
|
||||
|
||||
|
||||
|
||||
### MRFA
|
||||
|
||||
[MRFA.json](workflows/MRFA.json)
|
||||
|
||||
|
||||
|
||||
## Arguments
|
||||
|
||||
@@ -69,6 +75,13 @@ Doesn't need any
|
||||
* `avd`: similar to `relative`, may yield better but more "jittery/jumpy" result
|
||||
* `find_best_frame`: Same as FOMM
|
||||
|
||||
|
||||
### 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/`
|
||||
@@ -87,13 +100,16 @@ pip install -r requirements.txt
|
||||
pip install face-alignment
|
||||
```
|
||||
|
||||
## Model
|
||||
## Models
|
||||
|
||||
FOMM currently supporting `vox` and `vox-adv`. Models can and must be manually downloaded from:
|
||||
### FOMM and Part Swap
|
||||
|
||||
**FOMM**: `vox` and `vox-adv` from
|
||||
* [AliaksandrSiarohin/first-order-model](https://github.com/AliaksandrSiarohin/first-order-model)
|
||||
* [graphemecluster/first-order-model-demo](https://github.com/graphemecluster/first-order-model-demo)
|
||||
|
||||
Part Swap currently supporting Seg-based models and FOMM (`vox` and `vox-adv`) models.
|
||||
**Part Swap**
|
||||
|
||||
* `vox-5segments`
|
||||
* `vox-10segments`
|
||||
* `vox-15segments`
|
||||
@@ -126,6 +142,56 @@ resnet18-5c106cde.pth
|
||||
79999_iter.pth
|
||||
```
|
||||
|
||||
For **Articulate**, download the model from [Pre-trained checkpoints](https://github.com/snap-research/articulated-animation?tab=readme-ov-file#pre-trained-checkpoints) section and place it here: `articulate_module/models/vox256.pth`
|
||||
### Other
|
||||
|
||||
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`.
|
||||
| Model Arch | File Path | Source |
|
||||
|------------|-----------|--------|
|
||||
| **Articulate** | `articulate_module/models/vox256.pth` | [Articulated Animation (Pre-trained checkpoints)](https://github.com/snap-research/articulated-animation?tab=readme-ov-file#pre-trained-checkpoints) |
|
||||
| **Spline** | `spline_module/models/vox.pth.tar` | [Thin Plate Spline Motion Model (Pre-trained models)](https://github.com/yoyo-nb/thin-plate-spline-motion-model?tab=readme-ov-file#pre-trained-models) |
|
||||
| **MRFA** (celebvhq) | `mrfa_module/models/celebvhq.pth` | [MRFA (Pre-trained checkpoints)](https://github.com/JialeTao/MRFA/?tab=readme-ov-file#pretrained-models) |
|
||||
| **MRFA** (vox) | `mrfa_module/models/vox.pth` | [MRFA (Pre-trained checkpoints)](https://github.com/JialeTao/MRFA/?tab=readme-ov-file#pretrained-models) |
|
||||
|
||||
Notes:
|
||||
- For **Spline**, to use `find_best_frame`, follow above instructions to install `face-alignment` with its models.
|
||||
|
||||
## Credits
|
||||
|
||||
```
|
||||
@InProceedings{Siarohin_2019_NeurIPS,
|
||||
author={Siarohin, Aliaksandr and Lathuilière, Stéphane and Tulyakov, Sergey and Ricci, Elisa and Sebe, Nicu},
|
||||
title={First Order Motion Model for Image Animation},
|
||||
booktitle = {Conference on Neural Information Processing Systems (NeurIPS)},
|
||||
month = {December},
|
||||
year = {2019}
|
||||
}
|
||||
```
|
||||
|
||||
```
|
||||
@InProceedings{Siarohin_2019_NeurIPS,
|
||||
author={Siarohin, Aliaksandr and Lathuilière, Stéphane and Tulyakov, Sergey and Ricci, Elisa and Sebe, Nicu},
|
||||
title={First Order Motion Model for Image Animation},
|
||||
booktitle = {Conference on Neural Information Processing Systems (NeurIPS)},
|
||||
month = {December},
|
||||
year = {2019}
|
||||
}
|
||||
```
|
||||
|
||||
```
|
||||
@inproceedings{siarohin2021motion,
|
||||
author={Siarohin, Aliaksandr and Woodford, Oliver and Ren, Jian and Chai, Menglei and Tulyakov, Sergey},
|
||||
title={Motion Representations for Articulated Animation},
|
||||
booktitle = {CVPR},
|
||||
year = {2021}
|
||||
}
|
||||
```
|
||||
|
||||
```
|
||||
@inproceedings{
|
||||
tao2023learning,
|
||||
title={Learning Motion Refinement for Unsupervised Face Animation},
|
||||
author={Jiale Tao and Shuhang Gu and Wen Li and Lixin Duan},
|
||||
booktitle={Thirty-seventh Conference on Neural Information Processing Systems},
|
||||
year={2023},
|
||||
url={https://openreview.net/forum?id=m9uHv1Pxq7}
|
||||
}
|
||||
```
|
||||
+4
-1
@@ -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",
|
||||
}
|
||||
|
||||
|
||||
|
||||
+16
-4
@@ -99,11 +99,23 @@ seg_model_dict = {
|
||||
"vox-adv": seg10_segment_names,
|
||||
}
|
||||
|
||||
ARTICULATE_MODEL_PATH = "articulate_module/models/vox256.pth"
|
||||
ARTICULATE_CFG_PATH = "articulate_module/config/vox256.yaml"
|
||||
ARTICULATE_MODEL_PATH = "module_articulate/models/vox256.pth"
|
||||
ARTICULATE_CFG_PATH = "module_articulate/config/vox256.yaml"
|
||||
|
||||
SPLINE_MODES = ['relative', 'standard', 'avd']
|
||||
SPLINE_DEFAULT = 'relative'
|
||||
|
||||
SPLINE_MODEL_PATH = "spline_module/models/vox.pth.tar"
|
||||
SPLINE_CFG_PATH = "spline_module/config/vox-256.yaml"
|
||||
SPLINE_MODEL_PATH = "module_spline/models/vox.pth.tar"
|
||||
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",
|
||||
}
|
||||
@@ -1,35 +0,0 @@
|
||||
import yaml
|
||||
import torch
|
||||
|
||||
from .modules.generator import OcclusionAwareGenerator
|
||||
from .modules.keypoint_detector import KPDetector
|
||||
|
||||
def load_checkpoint(
|
||||
config_path: str, checkpoint_path: str, cpu=False
|
||||
) -> tuple[OcclusionAwareGenerator, KPDetector]:
|
||||
with open(config_path) as f:
|
||||
config = yaml.full_load(f)
|
||||
|
||||
generator = OcclusionAwareGenerator(
|
||||
**config["model_params"]["generator_params"],
|
||||
**config["model_params"]["common_params"],
|
||||
)
|
||||
kp_detector = KPDetector(
|
||||
**config["model_params"]["kp_detector_params"],
|
||||
**config["model_params"]["common_params"],
|
||||
)
|
||||
|
||||
if not cpu:
|
||||
generator.cuda()
|
||||
kp_detector.cuda()
|
||||
checkpoint = torch.load(checkpoint_path)
|
||||
else:
|
||||
checkpoint = torch.load(checkpoint_path, map_location=torch.device("cpu"))
|
||||
|
||||
generator.load_state_dict(checkpoint["generator"])
|
||||
kp_detector.load_state_dict(checkpoint["kp_detector"])
|
||||
|
||||
generator.eval()
|
||||
kp_detector.eval()
|
||||
|
||||
return generator, kp_detector
|
||||
@@ -5,10 +5,10 @@ import yaml
|
||||
from comfy.utils import ProgressBar
|
||||
from scipy.spatial import ConvexHull
|
||||
|
||||
from .articulate_module.avd_network import AVDNetwork
|
||||
from .articulate_module.bg_motion_predictor import BGMotionPredictor
|
||||
from .articulate_module.generator import Generator
|
||||
from .articulate_module.region_predictor import RegionPredictor
|
||||
from .module_articulate.avd_network import AVDNetwork
|
||||
from .module_articulate.bg_motion_predictor import BGMotionPredictor
|
||||
from .module_articulate.generator import Generator
|
||||
from .module_articulate.region_predictor import RegionPredictor
|
||||
from .sync_batchnorm.replicate import DataParallelWithCallback
|
||||
|
||||
|
||||
@@ -1,11 +1,12 @@
|
||||
import torch
|
||||
import numpy as np
|
||||
import torch
|
||||
import yaml
|
||||
from comfy.utils import ProgressBar
|
||||
from tqdm import tqdm
|
||||
from scipy.spatial import ConvexHull
|
||||
from tqdm import tqdm
|
||||
|
||||
from .modules.generator import OcclusionAwareGenerator
|
||||
from .modules.keypoint_detector import KPDetector
|
||||
|
||||
from .utils import tensor_to_numpy
|
||||
|
||||
try:
|
||||
@@ -24,6 +25,7 @@ def normalize_kp_best_frame(kp):
|
||||
kp[:, :2] = kp[:, :2] / area
|
||||
return kp
|
||||
|
||||
|
||||
def normalize_kp(
|
||||
kp_source,
|
||||
kp_driving,
|
||||
@@ -55,11 +57,12 @@ def normalize_kp(
|
||||
|
||||
return kp_new
|
||||
|
||||
|
||||
def find_best_frame(source_image, driving_video: list, cpu=False) -> int:
|
||||
if not face_alignment_available:
|
||||
print("Face Alignment module were not found. Best frame set to 0.")
|
||||
return 0
|
||||
|
||||
|
||||
print("Initializing Face Alignment...")
|
||||
fa = face_alignment.FaceAlignment(
|
||||
face_alignment.LandmarksType.TWO_D,
|
||||
@@ -93,6 +96,7 @@ def find_best_frame(source_image, driving_video: list, cpu=False) -> int:
|
||||
|
||||
return frame_num
|
||||
|
||||
|
||||
def inference(
|
||||
source_image,
|
||||
driving_video: list,
|
||||
@@ -182,4 +186,35 @@ def inference_best_frame(
|
||||
predictions_first = predictions_first[::-1] # Reverse the first half back
|
||||
predictions.extend(predictions_first)
|
||||
predictions.extend(predictions_second)
|
||||
return predictions
|
||||
return predictions
|
||||
|
||||
|
||||
def load_checkpoint(
|
||||
config_path: str, checkpoint_path: str, cpu=False
|
||||
) -> tuple[OcclusionAwareGenerator, KPDetector]:
|
||||
with open(config_path) as f:
|
||||
config = yaml.full_load(f)
|
||||
|
||||
generator = OcclusionAwareGenerator(
|
||||
**config["model_params"]["generator_params"],
|
||||
**config["model_params"]["common_params"],
|
||||
)
|
||||
kp_detector = KPDetector(
|
||||
**config["model_params"]["kp_detector_params"],
|
||||
**config["model_params"]["common_params"],
|
||||
)
|
||||
|
||||
if not cpu:
|
||||
generator.cuda()
|
||||
kp_detector.cuda()
|
||||
checkpoint = torch.load(checkpoint_path)
|
||||
else:
|
||||
checkpoint = torch.load(checkpoint_path, map_location=torch.device("cpu"))
|
||||
|
||||
generator.load_state_dict(checkpoint["generator"])
|
||||
kp_detector.load_state_dict(checkpoint["kp_detector"])
|
||||
|
||||
generator.eval()
|
||||
kp_detector.eval()
|
||||
|
||||
return generator, kp_detector
|
||||
@@ -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
|
||||
@@ -0,0 +1,274 @@
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import yaml
|
||||
from comfy.utils import ProgressBar
|
||||
from tqdm import tqdm
|
||||
|
||||
from .modules.dense_motion import DenseMotionNetwork
|
||||
from .modules.reconstruction_module import ReconstructionModule
|
||||
from .modules.segmentation_module import SegmentationModule
|
||||
from .modules.util import AntiAliasInterpolation2d
|
||||
from .sync_batchnorm.replicate import DataParallelWithCallback
|
||||
|
||||
|
||||
def face_parse_seg(
|
||||
source_image, face_parser_model, in_size=(512, 512), out_size=(64, 64)
|
||||
):
|
||||
seg = F.interpolate(source_image, size=in_size)
|
||||
seg = (seg - face_parser_model.mean) / face_parser_model.std
|
||||
seg = torch.softmax(face_parser_model(seg)[0], dim=1)
|
||||
seg = F.interpolate(seg, size=out_size)
|
||||
return seg
|
||||
|
||||
|
||||
class PartSwapGenerator(ReconstructionModule):
|
||||
def __init__(self, blend_scale=1, first_order_motion_model=False, **kwargs):
|
||||
super(PartSwapGenerator, self).__init__(**kwargs)
|
||||
if blend_scale == 1:
|
||||
self.blend_downsample = lambda x: x
|
||||
else:
|
||||
self.blend_downsample = AntiAliasInterpolation2d(1, blend_scale)
|
||||
|
||||
if first_order_motion_model:
|
||||
self.dense_motion_network = DenseMotionNetwork(
|
||||
block_expansion=64,
|
||||
num_blocks=5,
|
||||
max_features=1024,
|
||||
num_kp=10,
|
||||
num_channels=3,
|
||||
estimate_occlusion_map=True,
|
||||
scale_factor=0.25,
|
||||
)
|
||||
else:
|
||||
self.dense_motion_network = None
|
||||
|
||||
def forward(
|
||||
self,
|
||||
source_image,
|
||||
target_image,
|
||||
seg_target,
|
||||
seg_source,
|
||||
blend_mask,
|
||||
use_source_segmentation=False,
|
||||
):
|
||||
# Encoding of source image
|
||||
enc_source = self.first(source_image)
|
||||
for i in range(len(self.down_blocks)):
|
||||
enc_source = self.down_blocks[i](enc_source)
|
||||
|
||||
# Encoding of target image
|
||||
enc_target = self.first(target_image)
|
||||
for i in range(len(self.down_blocks)):
|
||||
enc_target = self.down_blocks[i](enc_target)
|
||||
|
||||
output_dict = {}
|
||||
# Compute flow field for source image
|
||||
if self.dense_motion_network is None:
|
||||
segment_motions = self.segment_motion(seg_target, seg_source)
|
||||
segment_motions = segment_motions.permute(0, 1, 4, 2, 3)
|
||||
mask = seg_target["segmentation"].unsqueeze(2)
|
||||
deformation = (segment_motions * mask).sum(dim=1)
|
||||
deformation = deformation.permute(0, 2, 3, 1)
|
||||
else:
|
||||
motion = self.dense_motion_network.forward_partswap(
|
||||
source_image=source_image, seg_target=seg_target, seg_source=seg_source
|
||||
)
|
||||
deformation = motion["deformation"]
|
||||
|
||||
# Deform source encoding according to the motion
|
||||
enc_source = self.deform_input(enc_source, deformation)
|
||||
|
||||
if self.estimate_visibility:
|
||||
if self.dense_motion_network is None:
|
||||
visibility = seg_source["segmentation"][:, 1:].sum(
|
||||
dim=1, keepdim=True
|
||||
) * (
|
||||
1
|
||||
- seg_target["segmentation"][:, 1:]
|
||||
.sum(dim=1, keepdim=True)
|
||||
.detach()
|
||||
)
|
||||
visibility = 1 - visibility
|
||||
else:
|
||||
visibility = motion["visibility"]
|
||||
|
||||
if (
|
||||
enc_source.shape[2] != visibility.shape[2]
|
||||
or enc_source.shape[3] != visibility.shape[3]
|
||||
):
|
||||
visibility = F.interpolate(
|
||||
visibility, size=enc_source.shape[2:], mode="bilinear"
|
||||
)
|
||||
enc_source = enc_source * visibility
|
||||
|
||||
blend_mask = self.blend_downsample(blend_mask)
|
||||
# If source segmentation is provided use it should be deformed before blending
|
||||
if use_source_segmentation:
|
||||
blend_mask = self.deform_input(blend_mask, deformation)
|
||||
|
||||
out = enc_target * (1 - blend_mask) + enc_source * blend_mask
|
||||
|
||||
out = self.bottleneck(out)
|
||||
|
||||
for i in range(len(self.up_blocks)):
|
||||
out = self.up_blocks[i](out)
|
||||
|
||||
out = self.final(out)
|
||||
out = F.sigmoid(out)
|
||||
|
||||
output_dict["prediction"] = out
|
||||
|
||||
return output_dict
|
||||
|
||||
|
||||
def partswap_inference(
|
||||
swap_indices: list[int],
|
||||
source_image,
|
||||
target_video,
|
||||
reconstruction_module: PartSwapGenerator,
|
||||
segmentation_module: SegmentationModule,
|
||||
use_source_seg: bool,
|
||||
face_parser_model=None,
|
||||
hard_edges=False,
|
||||
cpu=False,
|
||||
) -> list:
|
||||
with torch.no_grad():
|
||||
predictions = []
|
||||
seg_targets = []
|
||||
if not cpu:
|
||||
source_image = source_image.cuda()
|
||||
target_video = target_video.cuda()
|
||||
print("Using GPU")
|
||||
|
||||
seg_source = segmentation_module(source_image)
|
||||
# print(f"{seg_source['segmentation'].shape}")
|
||||
|
||||
num_frames = target_video.shape[2]
|
||||
pbar = ProgressBar(num_frames)
|
||||
for frame_idx in tqdm(range(num_frames)):
|
||||
target_frame = target_video[:, :, frame_idx]
|
||||
seg_target = segmentation_module(target_frame)
|
||||
seg_targets.append(seg_target)
|
||||
|
||||
if face_parser_model is not None:
|
||||
blend_mask = face_parse_seg(
|
||||
source_image if use_source_seg else target_frame,
|
||||
face_parser_model,
|
||||
(512, 512),
|
||||
(64, 64),
|
||||
)
|
||||
|
||||
else:
|
||||
blend_mask = (
|
||||
seg_source["segmentation"]
|
||||
if use_source_seg
|
||||
else seg_target["segmentation"]
|
||||
)
|
||||
|
||||
blend_mask = blend_mask[:, swap_indices].sum(dim=1, keepdim=True)
|
||||
if hard_edges:
|
||||
blend_mask = (blend_mask > 0.5).type(blend_mask.type())
|
||||
# print(f"{blend_mask.shape=}")
|
||||
out = reconstruction_module(
|
||||
source_image,
|
||||
target_frame,
|
||||
seg_source=seg_source,
|
||||
seg_target=seg_target,
|
||||
blend_mask=blend_mask,
|
||||
use_source_segmentation=use_source_seg,
|
||||
)
|
||||
|
||||
predictions.append(
|
||||
np.transpose(out["prediction"].data.cpu().numpy(), [0, 2, 3, 1])[0]
|
||||
)
|
||||
|
||||
pbar.update_absolute(frame_idx, num_frames)
|
||||
|
||||
return seg_source, seg_targets, predictions
|
||||
|
||||
|
||||
def partial_state_dict_load(module, state_dict):
|
||||
own_state = module.state_dict()
|
||||
for name, param in state_dict.items():
|
||||
if name not in own_state:
|
||||
continue
|
||||
|
||||
if isinstance(param, torch.nn.Parameter):
|
||||
# backwards compatibility for serialized parameters
|
||||
param = param.data
|
||||
own_state[name].copy_(param)
|
||||
|
||||
|
||||
def load_reconstruction_module(module, checkpoint):
|
||||
if "generator" in checkpoint:
|
||||
partial_state_dict_load(module, checkpoint["generator"])
|
||||
else:
|
||||
module.load_state_dict(checkpoint["reconstruction_module"])
|
||||
|
||||
|
||||
def load_segmentation_module(module, checkpoint):
|
||||
if "kp_detector" in checkpoint:
|
||||
partial_state_dict_load(module, checkpoint["kp_detector"])
|
||||
module.state_dict()["affine.weight"].copy_(
|
||||
checkpoint["kp_detector"]["jacobian.weight"]
|
||||
)
|
||||
module.state_dict()["affine.bias"].copy_(
|
||||
checkpoint["kp_detector"]["jacobian.bias"]
|
||||
)
|
||||
module.state_dict()["shift.weight"].copy_(
|
||||
checkpoint["kp_detector"]["kp.weight"]
|
||||
)
|
||||
module.state_dict()["shift.bias"].copy_(checkpoint["kp_detector"]["kp.bias"])
|
||||
if "semantic_seg.weight" in checkpoint["kp_detector"]:
|
||||
module.state_dict()["segmentation.weight"].copy_(
|
||||
checkpoint["kp_detector"]["semantic_seg.weight"]
|
||||
)
|
||||
module.state_dict()["segmentation.bias"].copy_(
|
||||
checkpoint["kp_detector"]["semantic_seg.bias"]
|
||||
)
|
||||
else:
|
||||
print("Segmentation part initialized at random.")
|
||||
else:
|
||||
module.load_state_dict(checkpoint["segmentation_module"])
|
||||
|
||||
|
||||
def load_partswap_checkpoint(
|
||||
config_path: str,
|
||||
checkpoint_path: str,
|
||||
blend_scale: float,
|
||||
use_fomm: bool,
|
||||
cpu=False,
|
||||
):
|
||||
with open(config_path) as f:
|
||||
config = yaml.full_load(f)
|
||||
|
||||
reconstruction_module = PartSwapGenerator(
|
||||
blend_scale=blend_scale,
|
||||
first_order_motion_model=use_fomm,
|
||||
**config["model_params"]["reconstruction_module_params"],
|
||||
**config["model_params"]["common_params"],
|
||||
)
|
||||
|
||||
segmentation_module = SegmentationModule(
|
||||
**config["model_params"]["segmentation_module_params"],
|
||||
**config["model_params"]["common_params"],
|
||||
)
|
||||
|
||||
if cpu:
|
||||
checkpoint = torch.load(checkpoint_path, map_location=torch.device("cpu"))
|
||||
else:
|
||||
segmentation_module.cuda()
|
||||
checkpoint = torch.load(checkpoint_path)
|
||||
|
||||
load_reconstruction_module(reconstruction_module, checkpoint)
|
||||
load_segmentation_module(segmentation_module, checkpoint)
|
||||
|
||||
if not cpu:
|
||||
reconstruction_module = DataParallelWithCallback(reconstruction_module)
|
||||
segmentation_module = DataParallelWithCallback(segmentation_module)
|
||||
|
||||
reconstruction_module.eval()
|
||||
segmentation_module.eval()
|
||||
|
||||
return reconstruction_module, segmentation_module
|
||||
@@ -5,12 +5,12 @@ import yaml
|
||||
import tqdm
|
||||
from comfy.utils import ProgressBar
|
||||
|
||||
from .fomm_inference import find_best_frame
|
||||
from .inference_fomm import find_best_frame
|
||||
|
||||
from .spline_module.avd_network import AVDNetwork
|
||||
from .spline_module.dense_motion import DenseMotionNetwork
|
||||
from .spline_module.keypoint_detector import KPDetector
|
||||
from .spline_module.inpainting_network import InpaintingNetwork
|
||||
from .module_spline.avd_network import AVDNetwork
|
||||
from .module_spline.dense_motion import DenseMotionNetwork
|
||||
from .module_spline.keypoint_detector import KPDetector
|
||||
from .module_spline.inpainting_network import InpaintingNetwork
|
||||
|
||||
|
||||
def spline_inference(
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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"]
|
||||
@@ -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
|
||||
@@ -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))
|
||||
@@ -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
|
||||
File diff suppressed because it is too large
Load Diff
@@ -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
|
||||
@@ -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
|
||||
@@ -1,84 +0,0 @@
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from tqdm import tqdm
|
||||
|
||||
from comfy.utils import ProgressBar
|
||||
from .modules.segmentation_module import SegmentationModule
|
||||
from .partswap_model import PartSwapGenerator
|
||||
|
||||
|
||||
def face_parse_seg(
|
||||
source_image, face_parser_model, in_size=(512, 512), out_size=(64, 64)
|
||||
):
|
||||
seg = F.interpolate(source_image, size=in_size)
|
||||
seg = (seg - face_parser_model.mean) / face_parser_model.std
|
||||
seg = torch.softmax(face_parser_model(seg)[0], dim=1)
|
||||
seg = F.interpolate(seg, size=out_size)
|
||||
return seg
|
||||
|
||||
|
||||
def partswap_inference(
|
||||
swap_indices: list[int],
|
||||
source_image,
|
||||
target_video,
|
||||
reconstruction_module: PartSwapGenerator,
|
||||
segmentation_module: SegmentationModule,
|
||||
use_source_seg: bool,
|
||||
face_parser_model=None,
|
||||
hard_edges=False,
|
||||
cpu=False,
|
||||
) -> list:
|
||||
with torch.no_grad():
|
||||
predictions = []
|
||||
seg_targets = []
|
||||
if not cpu:
|
||||
source_image = source_image.cuda()
|
||||
target_video = target_video.cuda()
|
||||
print("Using GPU")
|
||||
|
||||
seg_source = segmentation_module(source_image)
|
||||
# print(f"{seg_source['segmentation'].shape}")
|
||||
|
||||
num_frames = target_video.shape[2]
|
||||
pbar = ProgressBar(num_frames)
|
||||
for frame_idx in tqdm(range(num_frames)):
|
||||
target_frame = target_video[:, :, frame_idx]
|
||||
seg_target = segmentation_module(target_frame)
|
||||
seg_targets.append(seg_target)
|
||||
|
||||
if face_parser_model is not None:
|
||||
blend_mask = face_parse_seg(
|
||||
source_image if use_source_seg else target_frame,
|
||||
face_parser_model,
|
||||
(512, 512),
|
||||
(64, 64),
|
||||
)
|
||||
|
||||
else:
|
||||
blend_mask = (
|
||||
seg_source["segmentation"]
|
||||
if use_source_seg
|
||||
else seg_target["segmentation"]
|
||||
)
|
||||
|
||||
blend_mask = blend_mask[:, swap_indices].sum(dim=1, keepdim=True)
|
||||
if hard_edges:
|
||||
blend_mask = (blend_mask > 0.5).type(blend_mask.type())
|
||||
print(f"{blend_mask.shape=}")
|
||||
out = reconstruction_module(
|
||||
source_image,
|
||||
target_frame,
|
||||
seg_source=seg_source,
|
||||
seg_target=seg_target,
|
||||
blend_mask=blend_mask,
|
||||
use_source_segmentation=use_source_seg,
|
||||
)
|
||||
|
||||
predictions.append(
|
||||
np.transpose(out["prediction"].data.cpu().numpy(), [0, 2, 3, 1])[0]
|
||||
)
|
||||
|
||||
pbar.update_absolute(frame_idx, num_frames)
|
||||
|
||||
return seg_source, seg_targets, predictions
|
||||
@@ -1,92 +0,0 @@
|
||||
import torch
|
||||
import yaml
|
||||
|
||||
from .modules.segmentation_module import SegmentationModule
|
||||
from .partswap_model import PartSwapGenerator
|
||||
from .sync_batchnorm.replicate import DataParallelWithCallback
|
||||
|
||||
|
||||
def partial_state_dict_load(module, state_dict):
|
||||
own_state = module.state_dict()
|
||||
for name, param in state_dict.items():
|
||||
if name not in own_state:
|
||||
continue
|
||||
|
||||
if isinstance(param, torch.nn.Parameter):
|
||||
# backwards compatibility for serialized parameters
|
||||
param = param.data
|
||||
own_state[name].copy_(param)
|
||||
|
||||
|
||||
def load_reconstruction_module(module, checkpoint):
|
||||
if "generator" in checkpoint:
|
||||
partial_state_dict_load(module, checkpoint["generator"])
|
||||
else:
|
||||
module.load_state_dict(checkpoint["reconstruction_module"])
|
||||
|
||||
|
||||
def load_segmentation_module(module, checkpoint):
|
||||
if "kp_detector" in checkpoint:
|
||||
partial_state_dict_load(module, checkpoint["kp_detector"])
|
||||
module.state_dict()["affine.weight"].copy_(
|
||||
checkpoint["kp_detector"]["jacobian.weight"]
|
||||
)
|
||||
module.state_dict()["affine.bias"].copy_(
|
||||
checkpoint["kp_detector"]["jacobian.bias"]
|
||||
)
|
||||
module.state_dict()["shift.weight"].copy_(
|
||||
checkpoint["kp_detector"]["kp.weight"]
|
||||
)
|
||||
module.state_dict()["shift.bias"].copy_(checkpoint["kp_detector"]["kp.bias"])
|
||||
if "semantic_seg.weight" in checkpoint["kp_detector"]:
|
||||
module.state_dict()["segmentation.weight"].copy_(
|
||||
checkpoint["kp_detector"]["semantic_seg.weight"]
|
||||
)
|
||||
module.state_dict()["segmentation.bias"].copy_(
|
||||
checkpoint["kp_detector"]["semantic_seg.bias"]
|
||||
)
|
||||
else:
|
||||
print("Segmentation part initialized at random.")
|
||||
else:
|
||||
module.load_state_dict(checkpoint["segmentation_module"])
|
||||
|
||||
|
||||
def load_partswap_checkpoint(
|
||||
config_path: str,
|
||||
checkpoint_path: str,
|
||||
blend_scale: float,
|
||||
use_fomm: bool,
|
||||
cpu=False,
|
||||
):
|
||||
with open(config_path) as f:
|
||||
config = yaml.full_load(f)
|
||||
|
||||
reconstruction_module = PartSwapGenerator(
|
||||
blend_scale=blend_scale,
|
||||
first_order_motion_model=use_fomm,
|
||||
**config["model_params"]["reconstruction_module_params"],
|
||||
**config["model_params"]["common_params"],
|
||||
)
|
||||
|
||||
segmentation_module = SegmentationModule(
|
||||
**config["model_params"]["segmentation_module_params"],
|
||||
**config["model_params"]["common_params"],
|
||||
)
|
||||
|
||||
if cpu:
|
||||
checkpoint = torch.load(checkpoint_path, map_location=torch.device("cpu"))
|
||||
else:
|
||||
segmentation_module.cuda()
|
||||
checkpoint = torch.load(checkpoint_path)
|
||||
|
||||
load_reconstruction_module(reconstruction_module, checkpoint)
|
||||
load_segmentation_module(segmentation_module, checkpoint)
|
||||
|
||||
if not cpu:
|
||||
reconstruction_module = DataParallelWithCallback(reconstruction_module)
|
||||
segmentation_module = DataParallelWithCallback(segmentation_module)
|
||||
|
||||
reconstruction_module.eval()
|
||||
segmentation_module.eval()
|
||||
|
||||
return reconstruction_module, segmentation_module
|
||||
@@ -1,105 +0,0 @@
|
||||
import torch.nn.functional as F
|
||||
|
||||
from .modules.dense_motion import DenseMotionNetwork
|
||||
from .modules.reconstruction_module import ReconstructionModule
|
||||
from .modules.util import AntiAliasInterpolation2d
|
||||
|
||||
|
||||
class PartSwapGenerator(ReconstructionModule):
|
||||
def __init__(self, blend_scale=1, first_order_motion_model=False, **kwargs):
|
||||
super(PartSwapGenerator, self).__init__(**kwargs)
|
||||
if blend_scale == 1:
|
||||
self.blend_downsample = lambda x: x
|
||||
else:
|
||||
self.blend_downsample = AntiAliasInterpolation2d(1, blend_scale)
|
||||
|
||||
if first_order_motion_model:
|
||||
self.dense_motion_network = DenseMotionNetwork(
|
||||
block_expansion=64,
|
||||
num_blocks=5,
|
||||
max_features=1024,
|
||||
num_kp=10,
|
||||
num_channels=3,
|
||||
estimate_occlusion_map=True,
|
||||
scale_factor=0.25,
|
||||
)
|
||||
else:
|
||||
self.dense_motion_network = None
|
||||
|
||||
def forward(
|
||||
self,
|
||||
source_image,
|
||||
target_image,
|
||||
seg_target,
|
||||
seg_source,
|
||||
blend_mask,
|
||||
use_source_segmentation=False,
|
||||
):
|
||||
# Encoding of source image
|
||||
enc_source = self.first(source_image)
|
||||
for i in range(len(self.down_blocks)):
|
||||
enc_source = self.down_blocks[i](enc_source)
|
||||
|
||||
# Encoding of target image
|
||||
enc_target = self.first(target_image)
|
||||
for i in range(len(self.down_blocks)):
|
||||
enc_target = self.down_blocks[i](enc_target)
|
||||
|
||||
output_dict = {}
|
||||
# Compute flow field for source image
|
||||
if self.dense_motion_network is None:
|
||||
segment_motions = self.segment_motion(seg_target, seg_source)
|
||||
segment_motions = segment_motions.permute(0, 1, 4, 2, 3)
|
||||
mask = seg_target["segmentation"].unsqueeze(2)
|
||||
deformation = (segment_motions * mask).sum(dim=1)
|
||||
deformation = deformation.permute(0, 2, 3, 1)
|
||||
else:
|
||||
motion = self.dense_motion_network.forward_partswap(
|
||||
source_image=source_image, seg_target=seg_target, seg_source=seg_source
|
||||
)
|
||||
deformation = motion["deformation"]
|
||||
|
||||
# Deform source encoding according to the motion
|
||||
enc_source = self.deform_input(enc_source, deformation)
|
||||
|
||||
if self.estimate_visibility:
|
||||
if self.dense_motion_network is None:
|
||||
visibility = seg_source["segmentation"][:, 1:].sum(
|
||||
dim=1, keepdim=True
|
||||
) * (
|
||||
1
|
||||
- seg_target["segmentation"][:, 1:]
|
||||
.sum(dim=1, keepdim=True)
|
||||
.detach()
|
||||
)
|
||||
visibility = 1 - visibility
|
||||
else:
|
||||
visibility = motion["visibility"]
|
||||
|
||||
if (
|
||||
enc_source.shape[2] != visibility.shape[2]
|
||||
or enc_source.shape[3] != visibility.shape[3]
|
||||
):
|
||||
visibility = F.interpolate(
|
||||
visibility, size=enc_source.shape[2:], mode="bilinear"
|
||||
)
|
||||
enc_source = enc_source * visibility
|
||||
|
||||
blend_mask = self.blend_downsample(blend_mask)
|
||||
# If source segmentation is provided use it should be deformed before blending
|
||||
if use_source_segmentation:
|
||||
blend_mask = self.deform_input(blend_mask, deformation)
|
||||
|
||||
out = enc_target * (1 - blend_mask) + enc_source * blend_mask
|
||||
|
||||
out = self.bottleneck(out)
|
||||
|
||||
for i in range(len(self.up_blocks)):
|
||||
out = self.up_blocks[i](out)
|
||||
|
||||
out = self.final(out)
|
||||
out = F.sigmoid(out)
|
||||
|
||||
output_dict["prediction"] = out
|
||||
|
||||
return output_dict
|
||||
+3
-3
@@ -1,9 +1,9 @@
|
||||
[project]
|
||||
name = "comfyui-firstordermm"
|
||||
description = "A ComfyUI-native node to run First Order Motion Model for Image Animation. [a/https://github.com/AliaksandrSiarohin/first-order-model](https://github.com/AliaksandrSiarohin/first-order-model)"
|
||||
version = "1.0.4"
|
||||
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.5"
|
||||
license = { file = "LICENSE" }
|
||||
dependencies = ["numpy", "torch", "scipy", "pyyaml", "matplotlib"]
|
||||
dependencies = ["numpy", "torch", "scipy", "pyyaml", "matplotlib", "einops", "timm"]
|
||||
|
||||
[project.urls]
|
||||
Repository = "https://github.com/FuouM/ComfyUI-FirstOrderMM"
|
||||
|
||||
@@ -3,3 +3,5 @@ torch
|
||||
scipy
|
||||
pyyaml
|
||||
matplotlib
|
||||
einops
|
||||
timm
|
||||
@@ -2,17 +2,20 @@
|
||||
@author: Fuou Marinas
|
||||
@title: ComfyUI-FirstOrderMM
|
||||
@nickname: FOMM
|
||||
@description: Run First Order Motion Model for Image Animation in ComfyUI.
|
||||
@description: ComfyUI-native nodes to run First Order Motion Model for Image Animation and its non-diffusion-based successors.
|
||||
"""
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
|
||||
from .articulate_inference import articulate_inference
|
||||
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,
|
||||
@@ -28,12 +31,12 @@ from .constants import (
|
||||
supported_model_names,
|
||||
)
|
||||
from .face_parsing.face_parsing_loader import load_face_parser_model
|
||||
from .fomm_inference import inference, inference_best_frame
|
||||
from .fomm_loader import load_checkpoint
|
||||
from .partswap_inference import partswap_inference
|
||||
from .partswap_loader import load_partswap_checkpoint
|
||||
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
|
||||
from .spline_inference import spline_inference
|
||||
from .utils import (
|
||||
build_seg_arguments,
|
||||
get_config_path,
|
||||
@@ -451,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=}")
|
||||
@@ -488,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))
|
||||
|
||||
|
||||
@@ -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
|
||||
}
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 253 KiB |
Reference in New Issue
Block a user