muse pose

MusePose

@@ -0,0 +1 @@
+Subproject commit 124543e3ff347b508a2c489c4344f5f40190c5d3

extract_dwpose_from_vid.py

@@ -0,0 +1,101 @@
+import concurrent.futures
+import os
+import random
+from pathlib import Path
+
+import numpy as np
+
+from src.dwpose import DWposeDetector
+from src.utils.util import get_fps, read_frames, save_videos_from_pil
+
+# Extract dwpose mp4 videos from raw videos
+# /path/to/video_dataset/*/*.mp4 -> /path/to/video_dataset_dwpose/*/*.mp4
+
+
+def process_single_video(video_path, detector, root_dir, save_dir):
+    relative_path = os.path.relpath(video_path, root_dir)
+    print(relative_path, video_path, root_dir)
+    out_path = os.path.join(save_dir, relative_path)
+    if os.path.exists(out_path):
+        return
+
+    output_dir = Path(os.path.dirname(os.path.join(save_dir, relative_path)))
+    if not output_dir.exists():
+        output_dir.mkdir(parents=True, exist_ok=True)
+
+    fps = get_fps(video_path)
+    frames = read_frames(video_path)
+    kps_results = []
+    for i, frame_pil in enumerate(frames):
+        result, score = detector(frame_pil)
+        score = np.mean(score, axis=-1)
+
+        kps_results.append(result)
+
+    save_videos_from_pil(kps_results, out_path, fps=fps)
+
+
+def process_batch_videos(video_list, detector, root_dir, save_dir):
+    for i, video_path in enumerate(video_list):
+        print(f"Process {i}/{len(video_list)} video")
+        process_single_video(video_path, detector, root_dir, save_dir)
+
+
+if __name__ == "__main__":
+    # -----
+    # NOTE:
+    # python tools/extract_dwpose_from_vid.py --video_root /path/to/video_dir
+    # -----
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--video_root", type=str)
+    parser.add_argument(
+        "--save_dir", type=str, help="Path to save extracted pose videos"
+    )
+    parser.add_argument("-j", type=int, default=4, help="Num workers")
+    args = parser.parse_args()
+    num_workers = args.j
+    if args.save_dir is None:
+        save_dir = args.video_root + "_dwpose"
+    else:
+        save_dir = args.save_dir
+    if not os.path.exists(save_dir):
+        os.makedirs(save_dir)
+    cuda_visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES", "0")
+    gpu_ids = [int(id) for id in range(len(cuda_visible_devices.split(",")))]
+    print(f"avaliable gpu ids: {gpu_ids}")
+
+    # collect all video_folder paths
+    video_mp4_paths = set()
+    for root, dirs, files in os.walk(args.video_root):
+        for name in files:
+            if name.endswith(".mp4"):
+                video_mp4_paths.add(os.path.join(root, name))
+    video_mp4_paths = list(video_mp4_paths)
+    random.shuffle(video_mp4_paths)
+
+    # split into chunks,
+    batch_size = (len(video_mp4_paths) + num_workers - 1) // num_workers
+    print(f"Num videos: {len(video_mp4_paths)} {batch_size = }")
+    video_chunks = [
+        video_mp4_paths[i : i + batch_size]
+        for i in range(0, len(video_mp4_paths), batch_size)
+    ]
+
+    with concurrent.futures.ThreadPoolExecutor() as executor:
+        futures = []
+        for i, chunk in enumerate(video_chunks):
+            # init detector
+            gpu_id = gpu_ids[i % len(gpu_ids)]
+            detector = DWposeDetector()
+            # torch.cuda.set_device(gpu_id)
+            detector = detector.to(f"cuda:{gpu_id}")
+
+            futures.append(
+                executor.submit(
+                    process_batch_videos, chunk, detector, args.video_root, save_dir
+                )
+            )
+        for future in concurrent.futures.as_completed(futures):
+            future.result()

handler.py

@@ -4,28 +4,32 @@ from PIL import Image
 import base64
 from io import BytesIO
 import numpy as np
-from diffusers import AutoencoderKL, DDIMScheduler
+# from diffusers import AutoencoderKL, DDIMScheduler
 from einops import repeat
 from omegaconf import OmegaConf
-from transformers import CLIPVisionModelWithProjection
+# from transformers import CLIPVisionModelWithProjection
 import cv2
 import os
 import sys
 import skvideo.io
-from src.models.pose_guider import PoseGuider
-from src.models.unet_2d_condition import UNet2DConditionModel
-from src.models.unet_3d import UNet3DConditionModel
-from src.pipelines.pipeline_pose2vid_long import Pose2VideoPipeline
-from src.utils.util import read_frames, get_fps, save_videos_grid
+# from src.models.pose_guider import PoseGuider
+# from src.models.unet_2d_condition import UNet2DConditionModel
+# from src.models.unet_3d import UNet3DConditionModel
+# from src.pipelines.pipeline_pose2vid_long import Pose2VideoPipeline
+# from src.utils.util import read_frames, get_fps, save_videos_grid
 import roop.globals
 from roop.core import start, decode_execution_providers, suggest_max_memory, suggest_execution_threads
 from roop.utilities import normalize_output_path
 from roop.processors.frame.core import get_frame_processors_modules
 
-import onnxruntime as ort
+# import onnxruntime as ort
 import gc
 import subprocess
 
+import requests
+import tempfile
+
+
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
 if device.type != 'cuda':
@@ -39,11 +43,21 @@ class EndpointHandler():
         if not os.path.exists(config_path):
             raise FileNotFoundError(f"The configuration file was not found at: {config_path}")
 
+        self.run_post_install()
         self.config = OmegaConf.load(config_path)
         self.weight_dtype = torch.float16
         self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
         self.pipeline = None
-        self._initialize_pipeline()
+        # self._initialize_pipeline()
+
+    def run_post_install(self):
+        try:
+            result = subprocess.run(['bash', 'post_install.sh'], check=True, capture_output=True, text=True)
+            print("Post-install script ran successfully.")
+            print(result.stdout)
+        except subprocess.CalledProcessError as e:
+            print("Error running post-install script.")
+            print(e.stderr)
 
     def _initialize_pipeline(self):
         base_dir = os.path.dirname(os.path.abspath(__file__))
@@ -127,9 +141,9 @@ class EndpointHandler():
 
         return cropped_face
 
-    def _swap_face(self, source_image, target_video_path):
-        source_path = "input.jpg"
-        source_image.save(source_path, format="JPEG", quality=95)
+    def _swap_face(self, source_path, target_video_path):
+        # source_path = "input.jpg"
+        # source_image.save(source_path, format="JPEG", quality=95)
         output_path = "output.mp4"
 
         roop.globals.source_path = source_path
@@ -141,8 +155,8 @@ class EndpointHandler():
         roop.globals.keep_audio = True
         roop.globals.keep_frames = False
         roop.globals.many_faces = False
-        roop.globals.video_encoder = "libx264"
-        roop.globals.video_quality = 100
+        # roop.globals.video_encoder = "libx264"
+        roop.globals.video_quality = 50
         roop.globals.max_memory = suggest_max_memory()
 
         # Set GPU execution provider
@@ -250,83 +264,119 @@ class EndpointHandler():
         if result.returncode != 0:
             raise RuntimeError(f"FFmpeg slow down failed with exit code {result.returncode}")
 
+    def download_file(self, url: str, save_path: str):
+        response = requests.get(url, stream=True)
+        if response.status_code == 200:
+            with open(save_path, 'wb') as f:
+                for chunk in response.iter_content(chunk_size=8192):
+                    f.write(chunk)
+        else:
+            raise ValueError(f"Failed to download file from {url}")
+
+    def print_directory_contents(self, directory):
+        for root, dirs, files in os.walk(directory):
+            level = root.replace(directory, '').count(os.sep)
+            indent = ' ' * 4 * (level)
+            print(f"{indent}{os.path.basename(root)}/")
+            subindent = ' ' * 4 * (level + 1)
+            for f in files:
+                print(f"{subindent}{f}")
+
     def __call__(self, data: Any) -> Dict[str, str]:
         inputs = data.get("inputs", {})
-        ref_image_base64 = inputs.get("ref_image", "")
-        pose_video_path = inputs.get("pose_video_path", "")
-        width = inputs.get("width", 512)
-        height = inputs.get("height", 768)
-        length = inputs.get("length", 24)
-        num_inference_steps = inputs.get("num_inference_steps", 25)
-        cfg = inputs.get("cfg", 3.5)
-        seed = inputs.get("seed", 123)
-
-        ref_image = Image.open(BytesIO(base64.b64decode(ref_image_base64)))
-
-        # Get the base directory of the current file
-        base_dir = os.path.dirname(os.path.abspath(__file__))
-        
-        # Update pose_video_path to use the base directory
-        pose_video_path = os.path.join(base_dir, pose_video_path)
-        
-        if not os.path.exists(pose_video_path):
-            raise FileNotFoundError(f"The pose video was not found at: {pose_video_path}")
-
-        # Speed up the pose video by 4x
-        sped_up_pose_video_path = os.path.join(base_dir, "sped_up_pose_video.mp4")
-        self.speed_up_video(pose_video_path, sped_up_pose_video_path, factor=4)
-
-        torch.manual_seed(seed)
-        pose_images = read_frames(sped_up_pose_video_path)
-        src_fps = get_fps(sped_up_pose_video_path)
-        
-        pose_list = []
-        total_length = min(length, len(pose_images))
-        for pose_image_pil in pose_images[:total_length]:
-            pose_list.append(pose_image_pil)
-
-        video = self.pipeline(
-            ref_image,
-            pose_list,
-            width=width,
-            height=height,
-            video_length=total_length,
-            num_inference_steps=num_inference_steps,
-            guidance_scale=cfg
-        ).videos
-
-        save_dir = os.path.join(base_dir, "output", "gradio")
-        if not os.path.exists(save_dir):
-            os.makedirs(save_dir, exist_ok=True)
-        animation_path = os.path.join(save_dir, "animation_output.mp4")
-        save_videos_grid(video, animation_path, n_rows=1, fps=src_fps)
-
-        # Crop the face from the reference image and save it
-        cropped_face_path = os.path.join(save_dir, "cropped_face.jpg")
-        cropped_face = self._crop_face(ref_image, save_path=cropped_face_path)
-
-        # Delete the pipeline and clear CUDA cache to free up memory
-        del self.pipeline
-        torch.cuda.empty_cache()
-
-        # Perform face swapping
-        swapped_face_video_path = self._swap_face(cropped_face, animation_path)
-
-        # Slow down the produced video by 4x
-        slowed_down_animation_path = os.path.join(save_dir, "slowed_down_animation_output.mp4")
-        self.slow_down_video(swapped_face_video_path, slowed_down_animation_path, factor=4)
-
-        # Clear CUDA cache before RIFE interpolation
-        torch.cuda.empty_cache()
-
-        # Perform RIFE interpolation
-        rife_output_path = os.path.join(save_dir, "completed_result.mp4")
-        self.run_rife_interpolation(slowed_down_animation_path, rife_output_path, multi=2, scale=0.5)
-
-        # Encode the final video in base64
-        with open(rife_output_path, "rb") as video_file:
-            video_base64 = base64.b64encode(video_file.read()).decode("utf-8")
-
-        torch.cuda.empty_cache()
-
-        return {"video": video_base64}
+        ref_image_url = inputs.get("ref_image_url", "")
+        video_url = inputs.get("video_url", "")
+
+        # Create a unique temporary directory for this request
+        with tempfile.TemporaryDirectory() as temp_dir:
+            print(f"Temporary directory created at {temp_dir}")  # Debug statement
+            video_root = os.path.join(temp_dir, "dw_poses_videos")
+            os.makedirs(video_root, exist_ok=True)
+            downloaded_video_path = os.path.join(video_root, "downloaded_video.mp4")
+            downloaded_image_path = os.path.join(video_root, "downloaded_image.jpg")
+
+            # Download the video from the URL
+            self.download_file(video_url, downloaded_video_path)
+
+            # Download the reference image from the URL
+            self.download_file(ref_image_url, downloaded_image_path)
+            ref_image = Image.open(downloaded_image_path)
+
+            pose_output_path = os.path.join(temp_dir, "pose_videos")
+
+            # Run the extract_dwpose_from_vid.py script
+            command = [
+                "python", "./MusePose/pose_align.py",
+                "--imgfn_refer", downloaded_image_path,
+                "--vidfn", './pose_video.mp4',
+                "--output_dir", pose_output_path
+            ]
+            result = subprocess.run(command, capture_output=True, text=True)
+            if result.returncode != 0:
+                raise RuntimeError(f"Error running extract_dwpose_from_vid.py: {result.stderr}")
+
+            # Locate the extracted pose video
+            pose_video_path = os.path.join(pose_output_path, "pose_video.mp4")
+
+            if not os.path.exists(pose_video_path):
+                print(f"Error running extract_dwpose_from_vid.py: {result.stderr}")
+                print("Contents of the temporary directory:")
+                self.print_directory_contents(temp_dir)
+                raise FileNotFoundError(f"The pose video was not found at: {pose_video_path}")
+
+            # Speed up the pose video by 4x
+            sped_up_pose_video_path = os.path.join(temp_dir, "sped_up_pose_video.mp4")
+            self.speed_up_video(pose_video_path, sped_up_pose_video_path, factor=1)
+
+            dancing_video_dir = os.path.join(temp_dir, "dancing_video")
+            dancing_video_path_final = os.path.join(temp_dir, "dancing_video", "dance.mp4") #This is in create_video, can change there
+
+            command = [
+                "python", "./MusePose/create_video.py",
+                "--ref_image_path", downloaded_image_path,
+                "--pose_video_path", sped_up_pose_video_path,
+                "-W", "512",
+                "-H", "512",
+                "--output_dir", dancing_video_dir
+            ]
+            result = subprocess.run(command, capture_output=True, text=True)
+            if result.returncode != 0:
+                raise RuntimeError(f"Error running extract_dwpose_from_vid.py: {result.stderr}")
+
+            # save_dir = os.path.join(temp_dir, "output")
+            # if not os.path.exists(save_dir):
+            #     os.makedirs(save_dir, exist_ok=True)
+            # animation_path = os.path.join(save_dir, "animation_output.mp4")
+            # save_videos_grid(video, animation_path, n_rows=1, fps=src_fps)
+
+            # Crop the face from the reference image and save it
+            cropped_face_path = os.path.join(temp_dir, "cropped_face.jpg")
+            cropped_face = self._crop_face(ref_image, save_path=cropped_face_path)
+
+            # Delete the pipeline and clear CUDA cache to free up memory
+            del self.pipeline
+            torch.cuda.empty_cache()
+
+            # Perform face swapping
+            # self.print_directory_contents(temp_dir)
+            # swapped_face_video_path = self._swap_face(cropped_face_path, animation_path)
+
+            # Slow down the produced video by 4x
+            self.print_directory_contents(temp_dir)
+            slowed_down_animation_path = os.path.join(temp_dir, "slowed_down_animation_output.mp4")
+            self.slow_down_video(dancing_video_path_final, slowed_down_animation_path, factor=1)
+
+            # Clear CUDA cache before RIFE interpolation
+            torch.cuda.empty_cache()
+
+            # Perform RIFE interpolation
+            # rife_output_path = os.path.join(temp_dir, "completed_result.mp4")
+            # self.run_rife_interpolation(slowed_down_animation_path, rife_output_path, multi=2, scale=0.5)
+
+            # Encode the final video in base64
+            with open(slowed_down_animation_path, "rb") as video_file:
+                video_base64 = base64.b64encode(video_file.read()).decode("utf-8")
+
+            torch.cuda.empty_cache()
+
+            return {"video": video_base64}

post_install.sh

@@ -0,0 +1,14 @@
+#!/bin/bash
+
+pip install --no-cache-dir -U openmim 
+mim install mmengine 
+mim install "mmcv>=2.0.1" 
+mim install "mmdet>=3.1.0" 
+mim install "mmpose>=1.1.0" 
+
+# onnxruntime==1.16.3; sys_platform == 'darwin' and platform_machine != 'arm64'
+# onnxruntime-silicon==1.13.1; sys_platform == 'darwin' and platform_machine == 'arm64'
+# onnxruntime-gpu==1.16.3; sys_platform != 'darwin'
+# onnxruntime-coreml==1.13.1; python_version == '3.9' and sys_platform == 'darwin' and platform_machine != 'arm64'
+#onnx==1.14.0
+#protobuf==4.23.2

pretrained_weights/DWPose/dw-ll_ucoco_384.onnx

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:724f4ff2439ed61afb86fb8a1951ec39c6220682803b4a8bd4f598cd913b1843
-size 134399116

pretrained_weights/DWPose/yolox_l.onnx

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:7860ae79de6c89a3c1eb72ae9a2756c0ccfbe04b7791bb5880afabd97855a411
-size 216746733

pretrained_weights/denoising_unet.pth

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:b9e5a2c34fac369e8a922972ca2210916c6af175a0dad907deccf6235816ad52
-size 3438374293

pretrained_weights/image_encoder/config.json

@@ -1,23 +0,0 @@
-{
-  "_name_or_path": "/home/jpinkney/.cache/huggingface/diffusers/models--lambdalabs--sd-image-variations-diffusers/snapshots/ca6f97f838ae1b5bf764f31363a21f388f4d8f3e/image_encoder",
-  "architectures": [
-    "CLIPVisionModelWithProjection"
-  ],
-  "attention_dropout": 0.0,
-  "dropout": 0.0,
-  "hidden_act": "quick_gelu",
-  "hidden_size": 1024,
-  "image_size": 224,
-  "initializer_factor": 1.0,
-  "initializer_range": 0.02,
-  "intermediate_size": 4096,
-  "layer_norm_eps": 1e-05,
-  "model_type": "clip_vision_model",
-  "num_attention_heads": 16,
-  "num_channels": 3,
-  "num_hidden_layers": 24,
-  "patch_size": 14,
-  "projection_dim": 768,
-  "torch_dtype": "float32",
-  "transformers_version": "4.25.1"
-}

pretrained_weights/image_encoder/pytorch_model.bin

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:89d2aa29b5fdf64f3ad4f45fb4227ea98bc45156bbae673b85be1af7783dbabb
-size 1215993967

pretrained_weights/motion_module.pth

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:0d11e01a281b39880da2efeea892215c1313e5713fca3d100a7fbb72ee312ef9
-size 1817900227

pretrained_weights/pose_guider.pth

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:1a8b7c1b4db92980fd977b4fd003c1396bbae9a9cdea00c35d452136d5e4f488
-size 4351337

pretrained_weights/reference_unet.pth

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:beddccb08d49a8b29b0f4d6d456c6521d4382a8d8d48884fa60ba8802509c214
-size 3438323817

pretrained_weights/sd-vae-ft-mse/config.json

@@ -1,29 +0,0 @@
-{
-  "_class_name": "AutoencoderKL",
-  "_diffusers_version": "0.4.2",
-  "act_fn": "silu",
-  "block_out_channels": [
-    128,
-    256,
-    512,
-    512
-  ],
-  "down_block_types": [
-    "DownEncoderBlock2D",
-    "DownEncoderBlock2D",
-    "DownEncoderBlock2D",
-    "DownEncoderBlock2D"
-  ],
-  "in_channels": 3,
-  "latent_channels": 4,
-  "layers_per_block": 2,
-  "norm_num_groups": 32,
-  "out_channels": 3,
-  "sample_size": 256,
-  "up_block_types": [
-    "UpDecoderBlock2D",
-    "UpDecoderBlock2D",
-    "UpDecoderBlock2D",
-    "UpDecoderBlock2D"
-  ]
-}

pretrained_weights/sd-vae-ft-mse/diffusion_pytorch_model.bin

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:1b4889b6b1d4ce7ae320a02dedaeff1780ad77d415ea0d744b476155c6377ddc
-size 334707217

pretrained_weights/stable-diffusion-v1-5/unet/config.json

@@ -1,36 +0,0 @@
-{
-  "_class_name": "UNet2DConditionModel",
-  "_diffusers_version": "0.6.0",
-  "act_fn": "silu",
-  "attention_head_dim": 8,
-  "block_out_channels": [
-    320,
-    640,
-    1280,
-    1280
-  ],
-  "center_input_sample": false,
-  "cross_attention_dim": 768,
-  "down_block_types": [
-    "CrossAttnDownBlock2D",
-    "CrossAttnDownBlock2D",
-    "CrossAttnDownBlock2D",
-    "DownBlock2D"
-  ],
-  "downsample_padding": 1,
-  "flip_sin_to_cos": true,
-  "freq_shift": 0,
-  "in_channels": 4,
-  "layers_per_block": 2,
-  "mid_block_scale_factor": 1,
-  "norm_eps": 1e-05,
-  "norm_num_groups": 32,
-  "out_channels": 4,
-  "sample_size": 64,
-  "up_block_types": [
-    "UpBlock2D",
-    "CrossAttnUpBlock2D",
-    "CrossAttnUpBlock2D",
-    "CrossAttnUpBlock2D"
-  ]
-}

pretrained_weights/stable-diffusion-v1-5/unet/diffusion_pytorch_model.bin

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:c7da0e21ba7ea50637bee26e81c220844defdf01aafca02b2c42ecdadb813de4
-size 3438354725

requirements.txt

@@ -23,6 +23,7 @@ gfpgan==1.3.8
 gradio==3.41.2
 onnxruntime-coreml==1.13.1; python_version == '3.9' and sys_platform == 'darwin' and platform_machine != 'arm64'
 transformers==4.41.1
+controlnet-aux==0.0.7
 
 # Add additional dependencies
 diffusers==0.24.0
@@ -53,4 +54,6 @@ torchsde==0.2.5
 
 # Additional dependencies for RIFE
 sk-video==1.1.10
-moviepy==1.0.3
+moviepy==1.0.3
+
+requests==2.32.3

roop-unleashed

@@ -0,0 +1 @@
+Subproject commit ed6e3dbcf875213051dbc3b095e570afd3557463

sampler.py

@@ -7,15 +7,11 @@ import io
 # Initialize the handler
 handler = EndpointHandler()
 
-# Read a sample reference image and encode it in base64
-with open("rithwik.png", "rb") as image_file:
-    ref_image_base64 = base64.b64encode(image_file.read()).decode("utf-8")
-
 # Define sample inputs
 inputs = {
     "inputs": {
-        "ref_image": ref_image_base64,
-        "pose_video_path": "pose_video.mp4",
+        "ref_image_url": "https://media.discordapp.net/attachments/1183633414612594708/1245882096116043887/image.jpg?ex=665a5d9f&is=66590c1f&hm=3065fed7b8f5bd13aa2c8ad7d97e625dd4c2977589dbe7d8c13d024b782ab25a&=&format=webp&width=672&height=1194",
+        "video_url": "https://cdn.discordapp.com/attachments/1237667074210267217/1245971599660679208/pose.mov?ex=665ab0fa&is=66595f7a&hm=63691e23a23ebd8657a10ec708d63a06046a124c3940aa133de22a94aa1fd6c5&",
         "width": 378,
         "height": 504,
         "length": 24,
@@ -29,11 +25,11 @@ inputs = {
 output = handler(inputs)
 
 # # Decode the base64 video output
-# video_base64 = output.get("video", "")
-# video_bytes = base64.b64decode(video_base64)
+video_base64 = output.get("video", "")
+video_bytes = base64.b64decode(video_base64)
 
-# # Save the video to a file
-# with open("output_video.mp4", "wb") as video_file:
-#     video_file.write(video_bytes)
+# Save the video to a file
+with open("output_video.mp4", "wb") as video_file:
+    video_file.write(video_bytes)
 
 print("Inference completed. Output video saved as output_video.mp4")

src/dataset/dance_image.py

@@ -0,0 +1,124 @@
+import json
+import random
+
+import torch
+import torchvision.transforms as transforms
+from decord import VideoReader
+from PIL import Image
+from torch.utils.data import Dataset
+from transformers import CLIPImageProcessor
+
+
+class HumanDanceDataset(Dataset):
+    def __init__(
+        self,
+        img_size,
+        img_scale=(1.0, 1.0),
+        img_ratio=(0.9, 1.0),
+        drop_ratio=0.1,
+        data_meta_paths=["./data/fahsion_meta.json"],
+        sample_margin=30,
+    ):
+        super().__init__()
+
+        self.img_size = img_size
+        self.img_scale = img_scale
+        self.img_ratio = img_ratio
+        self.sample_margin = sample_margin
+
+        # -----
+        # vid_meta format:
+        # [{'video_path': , 'kps_path': , 'other':},
+        #  {'video_path': , 'kps_path': , 'other':}]
+        # -----
+        vid_meta = []
+        for data_meta_path in data_meta_paths:
+            vid_meta.extend(json.load(open(data_meta_path, "r")))
+        self.vid_meta = vid_meta
+
+        self.clip_image_processor = CLIPImageProcessor()
+
+        self.transform = transforms.Compose(
+            [
+                transforms.RandomResizedCrop(
+                    self.img_size,
+                    scale=self.img_scale,
+                    ratio=self.img_ratio,
+                    interpolation=transforms.InterpolationMode.BILINEAR,
+                ),
+                transforms.ToTensor(),
+                transforms.Normalize([0.5], [0.5]),
+            ]
+        )
+
+        self.cond_transform = transforms.Compose(
+            [
+                transforms.RandomResizedCrop(
+                    self.img_size,
+                    scale=self.img_scale,
+                    ratio=self.img_ratio,
+                    interpolation=transforms.InterpolationMode.BILINEAR,
+                ),
+                transforms.ToTensor(),
+            ]
+        )
+
+        self.drop_ratio = drop_ratio
+
+    def augmentation(self, image, transform, state=None):
+        if state is not None:
+            torch.set_rng_state(state)
+        return transform(image)
+
+    def __getitem__(self, index):
+        video_meta = self.vid_meta[index]
+        video_path = video_meta["video_path"]
+        kps_path = video_meta["kps_path"]
+
+        video_reader = VideoReader(video_path)
+        kps_reader = VideoReader(kps_path)
+
+        assert len(video_reader) == len(
+            kps_reader
+        ), f"{len(video_reader) = } != {len(kps_reader) = } in {video_path}"
+
+        video_length = len(video_reader)
+
+        margin = min(self.sample_margin, video_length)
+
+        ref_img_idx = random.randint(0, video_length - 1)
+        if ref_img_idx + margin < video_length:
+            tgt_img_idx = random.randint(ref_img_idx + margin, video_length - 1)
+        elif ref_img_idx - margin > 0:
+            tgt_img_idx = random.randint(0, ref_img_idx - margin)
+        else:
+            tgt_img_idx = random.randint(0, video_length - 1)
+
+        ref_img = video_reader[ref_img_idx]
+        ref_img_pil = Image.fromarray(ref_img.asnumpy())
+        tgt_img = video_reader[tgt_img_idx]
+        tgt_img_pil = Image.fromarray(tgt_img.asnumpy())
+
+        tgt_pose = kps_reader[tgt_img_idx]
+        tgt_pose_pil = Image.fromarray(tgt_pose.asnumpy())
+
+        state = torch.get_rng_state()
+        tgt_img = self.augmentation(tgt_img_pil, self.transform, state)
+        tgt_pose_img = self.augmentation(tgt_pose_pil, self.cond_transform, state)
+        ref_img_vae = self.augmentation(ref_img_pil, self.transform, state)
+        clip_image = self.clip_image_processor(
+            images=ref_img_pil, return_tensors="pt"
+        ).pixel_values[0]
+
+        sample = dict(
+            video_dir=video_path,
+            img=tgt_img,
+            tgt_pose=tgt_pose_img,
+            ref_img=ref_img_vae,
+            clip_images=clip_image,
+        )
+
+        return sample
+
+    def __len__(self):
+        return len(self.vid_meta)

src/dataset/dance_video.py

@@ -0,0 +1,137 @@
+import json
+import random
+from typing import List
+
+import numpy as np
+import pandas as pd
+import torch
+import torchvision.transforms as transforms
+from decord import VideoReader
+from PIL import Image
+from torch.utils.data import Dataset
+from transformers import CLIPImageProcessor
+
+
+class HumanDanceVideoDataset(Dataset):
+    def __init__(
+        self,
+        sample_rate,
+        n_sample_frames,
+        width,
+        height,
+        img_scale=(1.0, 1.0),
+        img_ratio=(0.9, 1.0),
+        drop_ratio=0.1,
+        data_meta_paths=["./data/fashion_meta.json"],
+    ):
+        super().__init__()
+        self.sample_rate = sample_rate
+        self.n_sample_frames = n_sample_frames
+        self.width = width
+        self.height = height
+        self.img_scale = img_scale
+        self.img_ratio = img_ratio
+
+        vid_meta = []
+        for data_meta_path in data_meta_paths:
+            vid_meta.extend(json.load(open(data_meta_path, "r")))
+        self.vid_meta = vid_meta
+
+        self.clip_image_processor = CLIPImageProcessor()
+
+        self.pixel_transform = transforms.Compose(
+            [
+                transforms.RandomResizedCrop(
+                    (height, width),
+                    scale=self.img_scale,
+                    ratio=self.img_ratio,
+                    interpolation=transforms.InterpolationMode.BILINEAR,
+                ),
+                transforms.ToTensor(),
+                transforms.Normalize([0.5], [0.5]),
+            ]
+        )
+
+        self.cond_transform = transforms.Compose(
+            [
+                transforms.RandomResizedCrop(
+                    (height, width),
+                    scale=self.img_scale,
+                    ratio=self.img_ratio,
+                    interpolation=transforms.InterpolationMode.BILINEAR,
+                ),
+                transforms.ToTensor(),
+            ]
+        )
+
+        self.drop_ratio = drop_ratio
+
+    def augmentation(self, images, transform, state=None):
+        if state is not None:
+            torch.set_rng_state(state)
+        if isinstance(images, List):
+            transformed_images = [transform(img) for img in images]
+            ret_tensor = torch.stack(transformed_images, dim=0)  # (f, c, h, w)
+        else:
+            ret_tensor = transform(images)  # (c, h, w)
+        return ret_tensor
+
+    def __getitem__(self, index):
+        video_meta = self.vid_meta[index]
+        video_path = video_meta["video_path"]
+        kps_path = video_meta["kps_path"]
+
+        video_reader = VideoReader(video_path)
+        kps_reader = VideoReader(kps_path)
+
+        assert len(video_reader) == len(
+            kps_reader
+        ), f"{len(video_reader) = } != {len(kps_reader) = } in {video_path}"
+
+        video_length = len(video_reader)
+
+        clip_length = min(
+            video_length, (self.n_sample_frames - 1) * self.sample_rate + 1
+        )
+        start_idx = random.randint(0, video_length - clip_length)
+        batch_index = np.linspace(
+            start_idx, start_idx + clip_length - 1, self.n_sample_frames, dtype=int
+        ).tolist()
+
+        # read frames and kps
+        vid_pil_image_list = []
+        pose_pil_image_list = []
+        for index in batch_index:
+            img = video_reader[index]
+            vid_pil_image_list.append(Image.fromarray(img.asnumpy()))
+            img = kps_reader[index]
+            pose_pil_image_list.append(Image.fromarray(img.asnumpy()))
+
+        ref_img_idx = random.randint(0, video_length - 1)
+        ref_img = Image.fromarray(video_reader[ref_img_idx].asnumpy())
+
+        # transform
+        state = torch.get_rng_state()
+        pixel_values_vid = self.augmentation(
+            vid_pil_image_list, self.pixel_transform, state
+        )
+        pixel_values_pose = self.augmentation(
+            pose_pil_image_list, self.cond_transform, state
+        )
+        pixel_values_ref_img = self.augmentation(ref_img, self.pixel_transform, state)
+        clip_ref_img = self.clip_image_processor(
+            images=ref_img, return_tensors="pt"
+        ).pixel_values[0]
+
+        sample = dict(
+            video_dir=video_path,
+            pixel_values_vid=pixel_values_vid,
+            pixel_values_pose=pixel_values_pose,
+            pixel_values_ref_img=pixel_values_ref_img,
+            clip_ref_img=clip_ref_img,
+        )
+
+        return sample
+
+    def __len__(self):
+        return len(self.vid_meta)

src/dwpose/__init__.py

@@ -0,0 +1,123 @@
+# https://github.com/IDEA-Research/DWPose
+# Openpose
+# Original from CMU https://github.com/CMU-Perceptual-Computing-Lab/openpose
+# 2nd Edited by https://github.com/Hzzone/pytorch-openpose
+# 3rd Edited by ControlNet
+# 4th Edited by ControlNet (added face and correct hands)
+
+import copy
+import os
+
+os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
+import cv2
+import numpy as np
+import torch
+from controlnet_aux.util import HWC3, resize_image
+from PIL import Image
+
+from . import util
+from .wholebody import Wholebody
+
+
+def draw_pose(pose, H, W):
+    bodies = pose["bodies"]
+    faces = pose["faces"]
+    hands = pose["hands"]
+    candidate = bodies["candidate"]
+    subset = bodies["subset"]
+    canvas = np.zeros(shape=(H, W, 3), dtype=np.uint8)
+
+    canvas = util.draw_bodypose(canvas, candidate, subset)
+
+    canvas = util.draw_handpose(canvas, hands)
+
+    canvas = util.draw_facepose(canvas, faces)
+
+    return canvas
+
+
+class DWposeDetector:
+    def __init__(self):
+        pass
+
+    def to(self, device):
+        self.pose_estimation = Wholebody(device)
+        return self
+
+    def cal_height(self, input_image):
+        input_image = cv2.cvtColor(
+            np.array(input_image, dtype=np.uint8), cv2.COLOR_RGB2BGR
+        )
+
+        input_image = HWC3(input_image)
+        H, W, C = input_image.shape
+        with torch.no_grad():
+            candidate, subset = self.pose_estimation(input_image)
+            nums, keys, locs = candidate.shape
+            # candidate[..., 0] /= float(W)
+            # candidate[..., 1] /= float(H)
+            body = candidate
+        return body[0, ..., 1].min(), body[..., 1].max() - body[..., 1].min()
+
+    def __call__(
+        self,
+        input_image,
+        detect_resolution=512,
+        image_resolution=512,
+        output_type="pil",
+        **kwargs,
+    ):
+        input_image = cv2.cvtColor(
+            np.array(input_image, dtype=np.uint8), cv2.COLOR_RGB2BGR
+        )
+
+        input_image = HWC3(input_image)
+        input_image = resize_image(input_image, detect_resolution)
+        H, W, C = input_image.shape
+        with torch.no_grad():
+            candidate, subset = self.pose_estimation(input_image)
+            nums, keys, locs = candidate.shape
+            candidate[..., 0] /= float(W)
+            candidate[..., 1] /= float(H)
+            score = subset[:, :18]
+            max_ind = np.mean(score, axis=-1).argmax(axis=0)
+            score = score[[max_ind]]
+            body = candidate[:, :18].copy()
+            body = body[[max_ind]]
+            nums = 1
+            body = body.reshape(nums * 18, locs)
+            body_score = copy.deepcopy(score)
+            for i in range(len(score)):
+                for j in range(len(score[i])):
+                    if score[i][j] > 0.3:
+                        score[i][j] = int(18 * i + j)
+                    else:
+                        score[i][j] = -1
+
+            un_visible = subset < 0.3
+            candidate[un_visible] = -1
+
+            foot = candidate[:, 18:24]
+
+            faces = candidate[[max_ind], 24:92]
+
+            hands = candidate[[max_ind], 92:113]
+            hands = np.vstack([hands, candidate[[max_ind], 113:]])
+
+            bodies = dict(candidate=body, subset=score)
+            pose = dict(bodies=bodies, hands=hands, faces=faces)
+
+            detected_map = draw_pose(pose, H, W)
+            detected_map = HWC3(detected_map)
+
+            img = resize_image(input_image, image_resolution)
+            H, W, C = img.shape
+
+            detected_map = cv2.resize(
+                detected_map, (W, H), interpolation=cv2.INTER_LINEAR
+            )
+
+            if output_type == "pil":
+                detected_map = Image.fromarray(detected_map)
+
+            return detected_map, body_score

src/dwpose/onnxdet.py

@@ -0,0 +1,130 @@
+# https://github.com/IDEA-Research/DWPose
+import cv2
+import numpy as np
+import onnxruntime
+
+
+def nms(boxes, scores, nms_thr):
+    """Single class NMS implemented in Numpy."""
+    x1 = boxes[:, 0]
+    y1 = boxes[:, 1]
+    x2 = boxes[:, 2]
+    y2 = boxes[:, 3]
+
+    areas = (x2 - x1 + 1) * (y2 - y1 + 1)
+    order = scores.argsort()[::-1]
+
+    keep = []
+    while order.size > 0:
+        i = order[0]
+        keep.append(i)
+        xx1 = np.maximum(x1[i], x1[order[1:]])
+        yy1 = np.maximum(y1[i], y1[order[1:]])
+        xx2 = np.minimum(x2[i], x2[order[1:]])
+        yy2 = np.minimum(y2[i], y2[order[1:]])
+
+        w = np.maximum(0.0, xx2 - xx1 + 1)
+        h = np.maximum(0.0, yy2 - yy1 + 1)
+        inter = w * h
+        ovr = inter / (areas[i] + areas[order[1:]] - inter)
+
+        inds = np.where(ovr <= nms_thr)[0]
+        order = order[inds + 1]
+
+    return keep
+
+
+def multiclass_nms(boxes, scores, nms_thr, score_thr):
+    """Multiclass NMS implemented in Numpy. Class-aware version."""
+    final_dets = []
+    num_classes = scores.shape[1]
+    for cls_ind in range(num_classes):
+        cls_scores = scores[:, cls_ind]
+        valid_score_mask = cls_scores > score_thr
+        if valid_score_mask.sum() == 0:
+            continue
+        else:
+            valid_scores = cls_scores[valid_score_mask]
+            valid_boxes = boxes[valid_score_mask]
+            keep = nms(valid_boxes, valid_scores, nms_thr)
+            if len(keep) > 0:
+                cls_inds = np.ones((len(keep), 1)) * cls_ind
+                dets = np.concatenate(
+                    [valid_boxes[keep], valid_scores[keep, None], cls_inds], 1
+                )
+                final_dets.append(dets)
+    if len(final_dets) == 0:
+        return None
+    return np.concatenate(final_dets, 0)
+
+
+def demo_postprocess(outputs, img_size, p6=False):
+    grids = []
+    expanded_strides = []
+    strides = [8, 16, 32] if not p6 else [8, 16, 32, 64]
+
+    hsizes = [img_size[0] // stride for stride in strides]
+    wsizes = [img_size[1] // stride for stride in strides]
+
+    for hsize, wsize, stride in zip(hsizes, wsizes, strides):
+        xv, yv = np.meshgrid(np.arange(wsize), np.arange(hsize))
+        grid = np.stack((xv, yv), 2).reshape(1, -1, 2)
+        grids.append(grid)
+        shape = grid.shape[:2]
+        expanded_strides.append(np.full((*shape, 1), stride))
+
+    grids = np.concatenate(grids, 1)
+    expanded_strides = np.concatenate(expanded_strides, 1)
+    outputs[..., :2] = (outputs[..., :2] + grids) * expanded_strides
+    outputs[..., 2:4] = np.exp(outputs[..., 2:4]) * expanded_strides
+
+    return outputs
+
+
+def preprocess(img, input_size, swap=(2, 0, 1)):
+    if len(img.shape) == 3:
+        padded_img = np.ones((input_size[0], input_size[1], 3), dtype=np.uint8) * 114
+    else:
+        padded_img = np.ones(input_size, dtype=np.uint8) * 114
+
+    r = min(input_size[0] / img.shape[0], input_size[1] / img.shape[1])
+    resized_img = cv2.resize(
+        img,
+        (int(img.shape[1] * r), int(img.shape[0] * r)),
+        interpolation=cv2.INTER_LINEAR,
+    ).astype(np.uint8)
+    padded_img[: int(img.shape[0] * r), : int(img.shape[1] * r)] = resized_img
+
+    padded_img = padded_img.transpose(swap)
+    padded_img = np.ascontiguousarray(padded_img, dtype=np.float32)
+    return padded_img, r
+
+
+def inference_detector(session, oriImg):
+    input_shape = (640, 640)
+    img, ratio = preprocess(oriImg, input_shape)
+
+    ort_inputs = {session.get_inputs()[0].name: img[None, :, :, :]}
+    output = session.run(None, ort_inputs)
+    predictions = demo_postprocess(output[0], input_shape)[0]
+
+    boxes = predictions[:, :4]
+    scores = predictions[:, 4:5] * predictions[:, 5:]
+
+    boxes_xyxy = np.ones_like(boxes)
+    boxes_xyxy[:, 0] = boxes[:, 0] - boxes[:, 2] / 2.0
+    boxes_xyxy[:, 1] = boxes[:, 1] - boxes[:, 3] / 2.0
+    boxes_xyxy[:, 2] = boxes[:, 0] + boxes[:, 2] / 2.0
+    boxes_xyxy[:, 3] = boxes[:, 1] + boxes[:, 3] / 2.0
+    boxes_xyxy /= ratio
+    dets = multiclass_nms(boxes_xyxy, scores, nms_thr=0.45, score_thr=0.1)
+    if dets is not None:
+        final_boxes, final_scores, final_cls_inds = dets[:, :4], dets[:, 4], dets[:, 5]
+        isscore = final_scores > 0.3
+        iscat = final_cls_inds == 0
+        isbbox = [i and j for (i, j) in zip(isscore, iscat)]
+        final_boxes = final_boxes[isbbox]
+    else:
+        return []
+
+    return final_boxes

src/dwpose/onnxpose.py

@@ -0,0 +1,370 @@
+# https://github.com/IDEA-Research/DWPose
+from typing import List, Tuple
+
+import cv2
+import numpy as np
+import onnxruntime as ort
+
+
+def preprocess(
+    img: np.ndarray, out_bbox, input_size: Tuple[int, int] = (192, 256)
+) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """Do preprocessing for RTMPose model inference.
+
+    Args:
+        img (np.ndarray): Input image in shape.
+        input_size (tuple): Input image size in shape (w, h).
+
+    Returns:
+        tuple:
+        - resized_img (np.ndarray): Preprocessed image.
+        - center (np.ndarray): Center of image.
+        - scale (np.ndarray): Scale of image.
+    """
+    # get shape of image
+    img_shape = img.shape[:2]
+    out_img, out_center, out_scale = [], [], []
+    if len(out_bbox) == 0:
+        out_bbox = [[0, 0, img_shape[1], img_shape[0]]]
+    for i in range(len(out_bbox)):
+        x0 = out_bbox[i][0]
+        y0 = out_bbox[i][1]
+        x1 = out_bbox[i][2]
+        y1 = out_bbox[i][3]
+        bbox = np.array([x0, y0, x1, y1])
+
+        # get center and scale
+        center, scale = bbox_xyxy2cs(bbox, padding=1.25)
+
+        # do affine transformation
+        resized_img, scale = top_down_affine(input_size, scale, center, img)
+
+        # normalize image
+        mean = np.array([123.675, 116.28, 103.53])
+        std = np.array([58.395, 57.12, 57.375])
+        resized_img = (resized_img - mean) / std
+
+        out_img.append(resized_img)
+        out_center.append(center)
+        out_scale.append(scale)
+
+    return out_img, out_center, out_scale
+
+
+def inference(sess: ort.InferenceSession, img: np.ndarray) -> np.ndarray:
+    """Inference RTMPose model.
+
+    Args:
+        sess (ort.InferenceSession): ONNXRuntime session.
+        img (np.ndarray): Input image in shape.
+
+    Returns:
+        outputs (np.ndarray): Output of RTMPose model.
+    """
+    all_out = []
+    # build input
+    for i in range(len(img)):
+        input = [img[i].transpose(2, 0, 1)]
+
+        # build output
+        sess_input = {sess.get_inputs()[0].name: input}
+        sess_output = []
+        for out in sess.get_outputs():
+            sess_output.append(out.name)
+
+        # run model
+        outputs = sess.run(sess_output, sess_input)
+        all_out.append(outputs)
+
+    return all_out
+
+
+def postprocess(
+    outputs: List[np.ndarray],
+    model_input_size: Tuple[int, int],
+    center: Tuple[int, int],
+    scale: Tuple[int, int],
+    simcc_split_ratio: float = 2.0,
+) -> Tuple[np.ndarray, np.ndarray]:
+    """Postprocess for RTMPose model output.
+
+    Args:
+        outputs (np.ndarray): Output of RTMPose model.
+        model_input_size (tuple): RTMPose model Input image size.
+        center (tuple): Center of bbox in shape (x, y).
+        scale (tuple): Scale of bbox in shape (w, h).
+        simcc_split_ratio (float): Split ratio of simcc.
+
+    Returns:
+        tuple:
+        - keypoints (np.ndarray): Rescaled keypoints.
+        - scores (np.ndarray): Model predict scores.
+    """
+    all_key = []
+    all_score = []
+    for i in range(len(outputs)):
+        # use simcc to decode
+        simcc_x, simcc_y = outputs[i]
+        keypoints, scores = decode(simcc_x, simcc_y, simcc_split_ratio)
+
+        # rescale keypoints
+        keypoints = keypoints / model_input_size * scale[i] + center[i] - scale[i] / 2
+        all_key.append(keypoints[0])
+        all_score.append(scores[0])
+
+    return np.array(all_key), np.array(all_score)
+
+
+def bbox_xyxy2cs(
+    bbox: np.ndarray, padding: float = 1.0
+) -> Tuple[np.ndarray, np.ndarray]:
+    """Transform the bbox format from (x,y,w,h) into (center, scale)
+
+    Args:
+        bbox (ndarray): Bounding box(es) in shape (4,) or (n, 4), formatted
+            as (left, top, right, bottom)
+        padding (float): BBox padding factor that will be multilied to scale.
+            Default: 1.0
+
+    Returns:
+        tuple: A tuple containing center and scale.
+        - np.ndarray[float32]: Center (x, y) of the bbox in shape (2,) or
+            (n, 2)
+        - np.ndarray[float32]: Scale (w, h) of the bbox in shape (2,) or
+            (n, 2)
+    """
+    # convert single bbox from (4, ) to (1, 4)
+    dim = bbox.ndim
+    if dim == 1:
+        bbox = bbox[None, :]
+
+    # get bbox center and scale
+    x1, y1, x2, y2 = np.hsplit(bbox, [1, 2, 3])
+    center = np.hstack([x1 + x2, y1 + y2]) * 0.5
+    scale = np.hstack([x2 - x1, y2 - y1]) * padding
+
+    if dim == 1:
+        center = center[0]
+        scale = scale[0]
+
+    return center, scale
+
+
+def _fix_aspect_ratio(bbox_scale: np.ndarray, aspect_ratio: float) -> np.ndarray:
+    """Extend the scale to match the given aspect ratio.
+
+    Args:
+        scale (np.ndarray): The image scale (w, h) in shape (2, )
+        aspect_ratio (float): The ratio of ``w/h``
+
+    Returns:
+        np.ndarray: The reshaped image scale in (2, )
+    """
+    w, h = np.hsplit(bbox_scale, [1])
+    bbox_scale = np.where(
+        w > h * aspect_ratio,
+        np.hstack([w, w / aspect_ratio]),
+        np.hstack([h * aspect_ratio, h]),
+    )
+    return bbox_scale
+
+
+def _rotate_point(pt: np.ndarray, angle_rad: float) -> np.ndarray:
+    """Rotate a point by an angle.
+
+    Args:
+        pt (np.ndarray): 2D point coordinates (x, y) in shape (2, )
+        angle_rad (float): rotation angle in radian
+
+    Returns:
+        np.ndarray: Rotated point in shape (2, )
+    """
+    sn, cs = np.sin(angle_rad), np.cos(angle_rad)
+    rot_mat = np.array([[cs, -sn], [sn, cs]])
+    return rot_mat @ pt
+
+
+def _get_3rd_point(a: np.ndarray, b: np.ndarray) -> np.ndarray:
+    """To calculate the affine matrix, three pairs of points are required. This
+    function is used to get the 3rd point, given 2D points a & b.
+
+    The 3rd point is defined by rotating vector `a - b` by 90 degrees
+    anticlockwise, using b as the rotation center.
+
+    Args:
+        a (np.ndarray): The 1st point (x,y) in shape (2, )
+        b (np.ndarray): The 2nd point (x,y) in shape (2, )
+
+    Returns:
+        np.ndarray: The 3rd point.
+    """
+    direction = a - b
+    c = b + np.r_[-direction[1], direction[0]]
+    return c
+
+
+def get_warp_matrix(
+    center: np.ndarray,
+    scale: np.ndarray,
+    rot: float,
+    output_size: Tuple[int, int],
+    shift: Tuple[float, float] = (0.0, 0.0),
+    inv: bool = False,
+) -> np.ndarray:
+    """Calculate the affine transformation matrix that can warp the bbox area
+    in the input image to the output size.
+
+    Args:
+        center (np.ndarray[2, ]): Center of the bounding box (x, y).
+        scale (np.ndarray[2, ]): Scale of the bounding box
+            wrt [width, height].
+        rot (float): Rotation angle (degree).
+        output_size (np.ndarray[2, ] | list(2,)): Size of the
+            destination heatmaps.
+        shift (0-100%): Shift translation ratio wrt the width/height.
+            Default (0., 0.).
+        inv (bool): Option to inverse the affine transform direction.
+            (inv=False: src->dst or inv=True: dst->src)
+
+    Returns:
+        np.ndarray: A 2x3 transformation matrix
+    """
+    shift = np.array(shift)
+    src_w = scale[0]
+    dst_w = output_size[0]
+    dst_h = output_size[1]
+
+    # compute transformation matrix
+    rot_rad = np.deg2rad(rot)
+    src_dir = _rotate_point(np.array([0.0, src_w * -0.5]), rot_rad)
+    dst_dir = np.array([0.0, dst_w * -0.5])
+
+    # get four corners of the src rectangle in the original image
+    src = np.zeros((3, 2), dtype=np.float32)
+    src[0, :] = center + scale * shift
+    src[1, :] = center + src_dir + scale * shift
+    src[2, :] = _get_3rd_point(src[0, :], src[1, :])
+
+    # get four corners of the dst rectangle in the input image
+    dst = np.zeros((3, 2), dtype=np.float32)
+    dst[0, :] = [dst_w * 0.5, dst_h * 0.5]
+    dst[1, :] = np.array([dst_w * 0.5, dst_h * 0.5]) + dst_dir
+    dst[2, :] = _get_3rd_point(dst[0, :], dst[1, :])
+
+    if inv:
+        warp_mat = cv2.getAffineTransform(np.float32(dst), np.float32(src))
+    else:
+        warp_mat = cv2.getAffineTransform(np.float32(src), np.float32(dst))
+
+    return warp_mat
+
+
+def top_down_affine(
+    input_size: dict, bbox_scale: dict, bbox_center: dict, img: np.ndarray
+) -> Tuple[np.ndarray, np.ndarray]:
+    """Get the bbox image as the model input by affine transform.
+
+    Args:
+        input_size (dict): The input size of the model.
+        bbox_scale (dict): The bbox scale of the img.
+        bbox_center (dict): The bbox center of the img.
+        img (np.ndarray): The original image.
+
+    Returns:
+        tuple: A tuple containing center and scale.
+        - np.ndarray[float32]: img after affine transform.
+        - np.ndarray[float32]: bbox scale after affine transform.
+    """
+    w, h = input_size
+    warp_size = (int(w), int(h))
+
+    # reshape bbox to fixed aspect ratio
+    bbox_scale = _fix_aspect_ratio(bbox_scale, aspect_ratio=w / h)
+
+    # get the affine matrix
+    center = bbox_center
+    scale = bbox_scale
+    rot = 0
+    warp_mat = get_warp_matrix(center, scale, rot, output_size=(w, h))
+
+    # do affine transform
+    img = cv2.warpAffine(img, warp_mat, warp_size, flags=cv2.INTER_LINEAR)
+
+    return img, bbox_scale
+
+
+def get_simcc_maximum(
+    simcc_x: np.ndarray, simcc_y: np.ndarray
+) -> Tuple[np.ndarray, np.ndarray]:
+    """Get maximum response location and value from simcc representations.
+
+    Note:
+        instance number: N
+        num_keypoints: K
+        heatmap height: H
+        heatmap width: W
+
+    Args:
+        simcc_x (np.ndarray): x-axis SimCC in shape (K, Wx) or (N, K, Wx)
+        simcc_y (np.ndarray): y-axis SimCC in shape (K, Wy) or (N, K, Wy)
+
+    Returns:
+        tuple:
+        - locs (np.ndarray): locations of maximum heatmap responses in shape
+            (K, 2) or (N, K, 2)
+        - vals (np.ndarray): values of maximum heatmap responses in shape
+            (K,) or (N, K)
+    """
+    N, K, Wx = simcc_x.shape
+    simcc_x = simcc_x.reshape(N * K, -1)
+    simcc_y = simcc_y.reshape(N * K, -1)
+
+    # get maximum value locations
+    x_locs = np.argmax(simcc_x, axis=1)
+    y_locs = np.argmax(simcc_y, axis=1)
+    locs = np.stack((x_locs, y_locs), axis=-1).astype(np.float32)
+    max_val_x = np.amax(simcc_x, axis=1)
+    max_val_y = np.amax(simcc_y, axis=1)
+
+    # get maximum value across x and y axis
+    mask = max_val_x > max_val_y
+    max_val_x[mask] = max_val_y[mask]
+    vals = max_val_x
+    locs[vals <= 0.0] = -1
+
+    # reshape
+    locs = locs.reshape(N, K, 2)
+    vals = vals.reshape(N, K)
+
+    return locs, vals
+
+
+def decode(
+    simcc_x: np.ndarray, simcc_y: np.ndarray, simcc_split_ratio
+) -> Tuple[np.ndarray, np.ndarray]:
+    """Modulate simcc distribution with Gaussian.
+
+    Args:
+        simcc_x (np.ndarray[K, Wx]): model predicted simcc in x.
+        simcc_y (np.ndarray[K, Wy]): model predicted simcc in y.
+        simcc_split_ratio (int): The split ratio of simcc.
+
+    Returns:
+        tuple: A tuple containing center and scale.
+        - np.ndarray[float32]: keypoints in shape (K, 2) or (n, K, 2)
+        - np.ndarray[float32]: scores in shape (K,) or (n, K)
+    """
+    keypoints, scores = get_simcc_maximum(simcc_x, simcc_y)
+    keypoints /= simcc_split_ratio
+
+    return keypoints, scores
+
+
+def inference_pose(session, out_bbox, oriImg):
+    h, w = session.get_inputs()[0].shape[2:]
+    model_input_size = (w, h)
+    resized_img, center, scale = preprocess(oriImg, out_bbox, model_input_size)
+    outputs = inference(session, resized_img)
+    keypoints, scores = postprocess(outputs, model_input_size, center, scale)
+
+    return keypoints, scores

src/dwpose/util.py

@@ -0,0 +1,378 @@
+# https://github.com/IDEA-Research/DWPose
+import math
+import numpy as np
+import matplotlib
+import cv2
+
+
+eps = 0.01
+
+
+def smart_resize(x, s):
+    Ht, Wt = s
+    if x.ndim == 2:
+        Ho, Wo = x.shape
+        Co = 1
+    else:
+        Ho, Wo, Co = x.shape
+    if Co == 3 or Co == 1:
+        k = float(Ht + Wt) / float(Ho + Wo)
+        return cv2.resize(
+            x,
+            (int(Wt), int(Ht)),
+            interpolation=cv2.INTER_AREA if k < 1 else cv2.INTER_LANCZOS4,
+        )
+    else:
+        return np.stack([smart_resize(x[:, :, i], s) for i in range(Co)], axis=2)
+
+
+def smart_resize_k(x, fx, fy):
+    if x.ndim == 2:
+        Ho, Wo = x.shape
+        Co = 1
+    else:
+        Ho, Wo, Co = x.shape
+    Ht, Wt = Ho * fy, Wo * fx
+    if Co == 3 or Co == 1:
+        k = float(Ht + Wt) / float(Ho + Wo)
+        return cv2.resize(
+            x,
+            (int(Wt), int(Ht)),
+            interpolation=cv2.INTER_AREA if k < 1 else cv2.INTER_LANCZOS4,
+        )
+    else:
+        return np.stack([smart_resize_k(x[:, :, i], fx, fy) for i in range(Co)], axis=2)
+
+
+def padRightDownCorner(img, stride, padValue):
+    h = img.shape[0]
+    w = img.shape[1]
+
+    pad = 4 * [None]
+    pad[0] = 0  # up
+    pad[1] = 0  # left
+    pad[2] = 0 if (h % stride == 0) else stride - (h % stride)  # down
+    pad[3] = 0 if (w % stride == 0) else stride - (w % stride)  # right
+
+    img_padded = img
+    pad_up = np.tile(img_padded[0:1, :, :] * 0 + padValue, (pad[0], 1, 1))
+    img_padded = np.concatenate((pad_up, img_padded), axis=0)
+    pad_left = np.tile(img_padded[:, 0:1, :] * 0 + padValue, (1, pad[1], 1))
+    img_padded = np.concatenate((pad_left, img_padded), axis=1)
+    pad_down = np.tile(img_padded[-2:-1, :, :] * 0 + padValue, (pad[2], 1, 1))
+    img_padded = np.concatenate((img_padded, pad_down), axis=0)
+    pad_right = np.tile(img_padded[:, -2:-1, :] * 0 + padValue, (1, pad[3], 1))
+    img_padded = np.concatenate((img_padded, pad_right), axis=1)
+
+    return img_padded, pad
+
+
+def transfer(model, model_weights):
+    transfered_model_weights = {}
+    for weights_name in model.state_dict().keys():
+        transfered_model_weights[weights_name] = model_weights[
+            ".".join(weights_name.split(".")[1:])
+        ]
+    return transfered_model_weights
+
+
+def draw_bodypose(canvas, candidate, subset):
+    H, W, C = canvas.shape
+    candidate = np.array(candidate)
+    subset = np.array(subset)
+
+    stickwidth = 4
+
+    limbSeq = [
+        [2, 3],
+        [2, 6],
+        [3, 4],
+        [4, 5],
+        [6, 7],
+        [7, 8],
+        [2, 9],
+        [9, 10],
+        [10, 11],
+        [2, 12],
+        [12, 13],
+        [13, 14],
+        [2, 1],
+        [1, 15],
+        [15, 17],
+        [1, 16],
+        [16, 18],
+        [3, 17],
+        [6, 18],
+    ]
+
+    colors = [
+        [255, 0, 0],
+        [255, 85, 0],
+        [255, 170, 0],
+        [255, 255, 0],
+        [170, 255, 0],
+        [85, 255, 0],
+        [0, 255, 0],
+        [0, 255, 85],
+        [0, 255, 170],
+        [0, 255, 255],
+        [0, 170, 255],
+        [0, 85, 255],
+        [0, 0, 255],
+        [85, 0, 255],
+        [170, 0, 255],
+        [255, 0, 255],
+        [255, 0, 170],
+        [255, 0, 85],
+    ]
+
+    for i in range(17):
+        for n in range(len(subset)):
+            index = subset[n][np.array(limbSeq[i]) - 1]
+            if -1 in index:
+                continue
+            Y = candidate[index.astype(int), 0] * float(W)
+            X = candidate[index.astype(int), 1] * float(H)
+            mX = np.mean(X)
+            mY = np.mean(Y)
+            length = ((X[0] - X[1]) ** 2 + (Y[0] - Y[1]) ** 2) ** 0.5
+            angle = math.degrees(math.atan2(X[0] - X[1], Y[0] - Y[1]))
+            polygon = cv2.ellipse2Poly(
+                (int(mY), int(mX)), (int(length / 2), stickwidth), int(angle), 0, 360, 1
+            )
+            cv2.fillConvexPoly(canvas, polygon, colors[i])
+
+    canvas = (canvas * 0.6).astype(np.uint8)
+
+    for i in range(18):
+        for n in range(len(subset)):
+            index = int(subset[n][i])
+            if index == -1:
+                continue
+            x, y = candidate[index][0:2]
+            x = int(x * W)
+            y = int(y * H)
+            cv2.circle(canvas, (int(x), int(y)), 4, colors[i], thickness=-1)
+
+    return canvas
+
+
+def draw_handpose(canvas, all_hand_peaks):
+    H, W, C = canvas.shape
+
+    edges = [
+        [0, 1],
+        [1, 2],
+        [2, 3],
+        [3, 4],
+        [0, 5],
+        [5, 6],
+        [6, 7],
+        [7, 8],
+        [0, 9],
+        [9, 10],
+        [10, 11],
+        [11, 12],
+        [0, 13],
+        [13, 14],
+        [14, 15],
+        [15, 16],
+        [0, 17],
+        [17, 18],
+        [18, 19],
+        [19, 20],
+    ]
+
+    for peaks in all_hand_peaks:
+        peaks = np.array(peaks)
+
+        for ie, e in enumerate(edges):
+            x1, y1 = peaks[e[0]]
+            x2, y2 = peaks[e[1]]
+            x1 = int(x1 * W)
+            y1 = int(y1 * H)
+            x2 = int(x2 * W)
+            y2 = int(y2 * H)
+            if x1 > eps and y1 > eps and x2 > eps and y2 > eps:
+                cv2.line(
+                    canvas,
+                    (x1, y1),
+                    (x2, y2),
+                    matplotlib.colors.hsv_to_rgb([ie / float(len(edges)), 1.0, 1.0])
+                    * 255,
+                    thickness=2,
+                )
+
+        for i, keyponit in enumerate(peaks):
+            x, y = keyponit
+            x = int(x * W)
+            y = int(y * H)
+            if x > eps and y > eps:
+                cv2.circle(canvas, (x, y), 4, (0, 0, 255), thickness=-1)
+    return canvas
+
+
+def draw_facepose(canvas, all_lmks):
+    H, W, C = canvas.shape
+    for lmks in all_lmks:
+        lmks = np.array(lmks)
+        for lmk in lmks:
+            x, y = lmk
+            x = int(x * W)
+            y = int(y * H)
+            if x > eps and y > eps:
+                cv2.circle(canvas, (x, y), 3, (255, 255, 255), thickness=-1)
+    return canvas
+
+
+# detect hand according to body pose keypoints
+# please refer to https://github.com/CMU-Perceptual-Computing-Lab/openpose/blob/master/src/openpose/hand/handDetector.cpp
+def handDetect(candidate, subset, oriImg):
+    # right hand: wrist 4, elbow 3, shoulder 2
+    # left hand: wrist 7, elbow 6, shoulder 5
+    ratioWristElbow = 0.33
+    detect_result = []
+    image_height, image_width = oriImg.shape[0:2]
+    for person in subset.astype(int):
+        # if any of three not detected
+        has_left = np.sum(person[[5, 6, 7]] == -1) == 0
+        has_right = np.sum(person[[2, 3, 4]] == -1) == 0
+        if not (has_left or has_right):
+            continue
+        hands = []
+        # left hand
+        if has_left:
+            left_shoulder_index, left_elbow_index, left_wrist_index = person[[5, 6, 7]]
+            x1, y1 = candidate[left_shoulder_index][:2]
+            x2, y2 = candidate[left_elbow_index][:2]
+            x3, y3 = candidate[left_wrist_index][:2]
+            hands.append([x1, y1, x2, y2, x3, y3, True])
+        # right hand
+        if has_right:
+            right_shoulder_index, right_elbow_index, right_wrist_index = person[
+                [2, 3, 4]
+            ]
+            x1, y1 = candidate[right_shoulder_index][:2]
+            x2, y2 = candidate[right_elbow_index][:2]
+            x3, y3 = candidate[right_wrist_index][:2]
+            hands.append([x1, y1, x2, y2, x3, y3, False])
+
+        for x1, y1, x2, y2, x3, y3, is_left in hands:
+            # pos_hand = pos_wrist + ratio * (pos_wrist - pos_elbox) = (1 + ratio) * pos_wrist - ratio * pos_elbox
+            # handRectangle.x = posePtr[wrist*3] + ratioWristElbow * (posePtr[wrist*3] - posePtr[elbow*3]);
+            # handRectangle.y = posePtr[wrist*3+1] + ratioWristElbow * (posePtr[wrist*3+1] - posePtr[elbow*3+1]);
+            # const auto distanceWristElbow = getDistance(poseKeypoints, person, wrist, elbow);
+            # const auto distanceElbowShoulder = getDistance(poseKeypoints, person, elbow, shoulder);
+            # handRectangle.width = 1.5f * fastMax(distanceWristElbow, 0.9f * distanceElbowShoulder);
+            x = x3 + ratioWristElbow * (x3 - x2)
+            y = y3 + ratioWristElbow * (y3 - y2)
+            distanceWristElbow = math.sqrt((x3 - x2) ** 2 + (y3 - y2) ** 2)
+            distanceElbowShoulder = math.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2)
+            width = 1.5 * max(distanceWristElbow, 0.9 * distanceElbowShoulder)
+            # x-y refers to the center --> offset to topLeft point
+            # handRectangle.x -= handRectangle.width / 2.f;
+            # handRectangle.y -= handRectangle.height / 2.f;
+            x -= width / 2
+            y -= width / 2  # width = height
+            # overflow the image
+            if x < 0:
+                x = 0
+            if y < 0:
+                y = 0
+            width1 = width
+            width2 = width
+            if x + width > image_width:
+                width1 = image_width - x
+            if y + width > image_height:
+                width2 = image_height - y
+            width = min(width1, width2)
+            # the max hand box value is 20 pixels
+            if width >= 20:
+                detect_result.append([int(x), int(y), int(width), is_left])
+
+    """
+    return value: [[x, y, w, True if left hand else False]].
+    width=height since the network require squared input.
+    x, y is the coordinate of top left 
+    """
+    return detect_result
+
+
+# Written by Lvmin
+def faceDetect(candidate, subset, oriImg):
+    # left right eye ear 14 15 16 17
+    detect_result = []
+    image_height, image_width = oriImg.shape[0:2]
+    for person in subset.astype(int):
+        has_head = person[0] > -1
+        if not has_head:
+            continue
+
+        has_left_eye = person[14] > -1
+        has_right_eye = person[15] > -1
+        has_left_ear = person[16] > -1
+        has_right_ear = person[17] > -1
+
+        if not (has_left_eye or has_right_eye or has_left_ear or has_right_ear):
+            continue
+
+        head, left_eye, right_eye, left_ear, right_ear = person[[0, 14, 15, 16, 17]]
+
+        width = 0.0
+        x0, y0 = candidate[head][:2]
+
+        if has_left_eye:
+            x1, y1 = candidate[left_eye][:2]
+            d = max(abs(x0 - x1), abs(y0 - y1))
+            width = max(width, d * 3.0)
+
+        if has_right_eye:
+            x1, y1 = candidate[right_eye][:2]
+            d = max(abs(x0 - x1), abs(y0 - y1))
+            width = max(width, d * 3.0)
+
+        if has_left_ear:
+            x1, y1 = candidate[left_ear][:2]
+            d = max(abs(x0 - x1), abs(y0 - y1))
+            width = max(width, d * 1.5)
+
+        if has_right_ear:
+            x1, y1 = candidate[right_ear][:2]
+            d = max(abs(x0 - x1), abs(y0 - y1))
+            width = max(width, d * 1.5)
+
+        x, y = x0, y0
+
+        x -= width
+        y -= width
+
+        if x < 0:
+            x = 0
+
+        if y < 0:
+            y = 0
+
+        width1 = width * 2
+        width2 = width * 2
+
+        if x + width > image_width:
+            width1 = image_width - x
+
+        if y + width > image_height:
+            width2 = image_height - y
+
+        width = min(width1, width2)
+
+        if width >= 20:
+            detect_result.append([int(x), int(y), int(width)])
+
+    return detect_result
+
+
+# get max index of 2d array
+def npmax(array):
+    arrayindex = array.argmax(1)
+    arrayvalue = array.max(1)
+    i = arrayvalue.argmax()
+    j = arrayindex[i]
+    return i, j

src/dwpose/wholebody.py

@@ -0,0 +1,48 @@
+# https://github.com/IDEA-Research/DWPose
+from pathlib import Path
+
+import cv2
+import numpy as np
+import onnxruntime as ort
+
+from .onnxdet import inference_detector
+from .onnxpose import inference_pose
+
+ModelDataPathPrefix = Path("./pretrained_weights")
+
+
+class Wholebody:
+    def __init__(self, device="cuda:0"):
+        providers = (
+            ["CPUExecutionProvider"] if device == "cpu" else ["CUDAExecutionProvider"]
+        )
+        onnx_det = ModelDataPathPrefix.joinpath("DWPose/yolox_l.onnx")
+        onnx_pose = ModelDataPathPrefix.joinpath("DWPose/dw-ll_ucoco_384.onnx")
+
+        self.session_det = ort.InferenceSession(
+            path_or_bytes=onnx_det, providers=providers
+        )
+        self.session_pose = ort.InferenceSession(
+            path_or_bytes=onnx_pose, providers=providers
+        )
+
+    def __call__(self, oriImg):
+        det_result = inference_detector(self.session_det, oriImg)
+        keypoints, scores = inference_pose(self.session_pose, det_result, oriImg)
+
+        keypoints_info = np.concatenate((keypoints, scores[..., None]), axis=-1)
+        # compute neck joint
+        neck = np.mean(keypoints_info[:, [5, 6]], axis=1)
+        # neck score when visualizing pred
+        neck[:, 2:4] = np.logical_and(
+            keypoints_info[:, 5, 2:4] > 0.3, keypoints_info[:, 6, 2:4] > 0.3
+        ).astype(int)
+        new_keypoints_info = np.insert(keypoints_info, 17, neck, axis=1)
+        mmpose_idx = [17, 6, 8, 10, 7, 9, 12, 14, 16, 13, 15, 2, 1, 4, 3]
+        openpose_idx = [1, 2, 3, 4, 6, 7, 8, 9, 10, 12, 13, 14, 15, 16, 17]
+        new_keypoints_info[:, openpose_idx] = new_keypoints_info[:, mmpose_idx]
+        keypoints_info = new_keypoints_info
+
+        keypoints, scores = keypoints_info[..., :2], keypoints_info[..., 2]
+
+        return keypoints, scores

tools/download_weights.py

@@ -0,0 +1,111 @@
+import os
+from pathlib import Path, PurePosixPath
+
+from huggingface_hub import hf_hub_download
+
+
+def prepare_base_model():
+    print(f'Preparing base stable-diffusion-v1-5 weights...')
+    local_dir = "./pretrained_weights/stable-diffusion-v1-5"
+    os.makedirs(local_dir, exist_ok=True)
+    for hub_file in ["unet/config.json", "unet/diffusion_pytorch_model.bin"]:
+        path = Path(hub_file)
+        saved_path = local_dir / path
+        if os.path.exists(saved_path):
+            continue
+        hf_hub_download(
+            repo_id="runwayml/stable-diffusion-v1-5",
+            subfolder=PurePosixPath(path.parent),
+            filename=PurePosixPath(path.name),
+            local_dir=local_dir,
+        )
+
+
+def prepare_image_encoder():
+    print(f"Preparing image encoder weights...")
+    local_dir = "./pretrained_weights"
+    os.makedirs(local_dir, exist_ok=True)
+    for hub_file in ["image_encoder/config.json", "image_encoder/pytorch_model.bin"]:
+        path = Path(hub_file)
+        saved_path = local_dir / path
+        if os.path.exists(saved_path):
+            continue
+        hf_hub_download(
+            repo_id="lambdalabs/sd-image-variations-diffusers",
+            subfolder=PurePosixPath(path.parent),
+            filename=PurePosixPath(path.name),
+            local_dir=local_dir,
+        )
+
+
+def prepare_dwpose():
+    print(f"Preparing DWPose weights...")
+    local_dir = "./pretrained_weights/DWPose"
+    os.makedirs(local_dir, exist_ok=True)
+    for hub_file in [
+        "dw-ll_ucoco_384.onnx",
+        "yolox_l.onnx",
+    ]:
+        path = Path(hub_file)
+        saved_path = local_dir / path
+        if os.path.exists(saved_path):
+            continue
+
+        hf_hub_download(
+            repo_id="yzd-v/DWPose",
+            subfolder=PurePosixPath(path.parent),
+            filename=PurePosixPath(path.name),
+            local_dir=local_dir,
+        )
+
+
+def prepare_vae():
+    print(f"Preparing vae weights...")
+    local_dir = "./pretrained_weights/sd-vae-ft-mse"
+    os.makedirs(local_dir, exist_ok=True)
+    for hub_file in [
+        "config.json",
+        "diffusion_pytorch_model.bin",
+    ]:
+        path = Path(hub_file)
+        saved_path = local_dir / path
+        if os.path.exists(saved_path):
+            continue
+
+        hf_hub_download(
+            repo_id="stabilityai/sd-vae-ft-mse",
+            subfolder=PurePosixPath(path.parent),
+            filename=PurePosixPath(path.name),
+            local_dir=local_dir,
+        )
+
+
+def prepare_anyone():
+    print(f"Preparing AnimateAnyone weights...")
+    local_dir = "./pretrained_weights"
+    os.makedirs(local_dir, exist_ok=True)
+    for hub_file in [
+        "denoising_unet.pth",
+        "motion_module.pth",
+        "pose_guider.pth",
+        "reference_unet.pth",
+    ]:
+        path = Path(hub_file)
+        saved_path = local_dir / path
+        if os.path.exists(saved_path):
+            continue
+
+        hf_hub_download(
+            repo_id="patrolli/AnimateAnyone",
+            subfolder=PurePosixPath(path.parent),
+            filename=PurePosixPath(path.name),
+            local_dir=local_dir,
+        )
+
+if __name__ == '__main__':
+    prepare_base_model()
+    prepare_image_encoder()
+    prepare_dwpose()
+    prepare_vae()
+    prepare_anyone()
+    

tools/extract_meta_info.py

@@ -0,0 +1,37 @@
+import argparse
+import json
+import os
+
+# -----
+# [{'vid': , 'kps': , 'other':},
+#  {'vid': , 'kps': , 'other':}]
+# -----
+# python tools/extract_meta_info.py --root_path /path/to/video_dir --dataset_name fashion
+# -----
+parser = argparse.ArgumentParser()
+parser.add_argument("--root_path", type=str)
+parser.add_argument("--dataset_name", type=str)
+parser.add_argument("--meta_info_name", type=str)
+
+args = parser.parse_args()
+
+if args.meta_info_name is None:
+    args.meta_info_name = args.dataset_name
+
+pose_dir = args.root_path + "_dwpose"
+
+# collect all video_folder paths
+video_mp4_paths = set()
+for root, dirs, files in os.walk(args.root_path):
+    for name in files:
+        if name.endswith(".mp4"):
+            video_mp4_paths.add(os.path.join(root, name))
+video_mp4_paths = list(video_mp4_paths)
+
+meta_infos = []
+for video_mp4_path in video_mp4_paths:
+    relative_video_name = os.path.relpath(video_mp4_path, args.root_path)
+    kps_path = os.path.join(pose_dir, relative_video_name)
+    meta_infos.append({"video_path": video_mp4_path, "kps_path": kps_path})
+
+json.dump(meta_infos, open(f"./data/{args.meta_info_name}_meta.json", "w"))

tools/facetracker_api.py

@@ -0,0 +1,62 @@
+import copy
+import os, sys
+import math
+import numpy as np
+import cv2
+sys.path.append("OpenSeeFace/")
+from tracker import Tracker, get_model_base_path
+
+features = ["eye_l", "eye_r", "eyebrow_steepness_l", "eyebrow_updown_l", "eyebrow_quirk_l", "eyebrow_steepness_r", "eyebrow_updown_r", "eyebrow_quirk_r", "mouth_corner_updown_l", "mouth_corner_inout_l", "mouth_corner_updown_r", "mouth_corner_inout_r", "mouth_open", "mouth_wide"]
+
+
+def face_image(frame, save_path=None):
+    height, width, c = frame.shape
+    tracker = Tracker(width, height, threshold=None, max_threads=1, max_faces=1, discard_after=10, scan_every=3, silent=False, model_type=3, model_dir=None, 
+                      no_gaze=False, detection_threshold=0.4, use_retinaface=0, max_feature_updates=900, static_model=True, try_hard=False)
+    faces = tracker.predict(frame)
+    frame = np.zeros_like(frame)
+    detected = False
+    face_lms = None
+    for face_num, f in enumerate(faces):
+        f = copy.copy(f)
+        if f.eye_blink is None:
+            f.eye_blink = [1, 1]
+        right_state = "O" if f.eye_blink[0] > 0.30 else "-"
+        left_state = "O" if f.eye_blink[1] > 0.30 else "-"
+        detected = True
+        if not f.success:
+            pts_3d = np.zeros((70, 3), np.float32)
+        if face_num == 0:
+            face_lms = f.lms
+        for pt_num, (x,y,c) in enumerate(f.lms):
+            if pt_num == 66 and (f.eye_blink[0] < 0.30 or c < 0.20):
+                continue
+            if pt_num == 67 and (f.eye_blink[1] < 0.30 or c < 0.20):
+                continue
+            x = int(x + 0.5)
+            y = int(y + 0.5)
+
+            color = (0, 255, 0)
+            if pt_num >= 66:
+                color = (255, 255, 0)
+            if not (x < 0 or y < 0 or x >= height or y >= width):
+                cv2.circle(frame, (y, x), 1, color, -1)
+        if f.rotation is not None:
+            projected = cv2.projectPoints(f.contour, f.rotation, f.translation, tracker.camera, tracker.dist_coeffs)
+            for [(x,y)] in projected[0]:
+                x = int(x + 0.5)
+                y = int(y + 0.5)
+                if not (x < 0 or y < 0 or x >= height or y >= width):
+                    frame[int(x), int(y)] = (0, 255, 255)
+                x += 1
+                if not (x < 0 or y < 0 or x >= height or y >= width):
+                    frame[int(x), int(y)] = (0, 255, 255)
+                y += 1
+                if not (x < 0 or y < 0 or x >= height or y >= width):
+                    frame[int(x), int(y)] = (0, 255, 255)
+                x -= 1
+                if not (x < 0 or y < 0 or x >= height or y >= width):
+                    frame[int(x), int(y)] = (0, 255, 255)
+    if save_path is not None:
+        cv2.imwrite(save_path, frame)
+    return frame, face_lms

tools/vid2pose.py

@@ -0,0 +1,38 @@
+from src.dwpose import DWposeDetector
+import os
+from pathlib import Path
+
+from src.utils.util import get_fps, read_frames, save_videos_from_pil
+import numpy as np
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--video_path", type=str)
+    args = parser.parse_args()
+
+    if not os.path.exists(args.video_path):
+        raise ValueError(f"Path: {args.video_path} not exists")
+
+    dir_path, video_name = (
+        os.path.dirname(args.video_path),
+        os.path.splitext(os.path.basename(args.video_path))[0],
+    )
+    out_path = os.path.join(dir_path, video_name + "_kps.mp4")
+
+    detector = DWposeDetector()
+    detector = detector.to(f"cuda")
+
+    fps = get_fps(args.video_path)
+    frames = read_frames(args.video_path)
+    kps_results = []
+    for i, frame_pil in enumerate(frames):
+        result, score = detector(frame_pil)
+        score = np.mean(score, axis=-1)
+
+        kps_results.append(result)
+
+    print(out_path)
+    save_videos_from_pil(kps_results, out_path, fps=fps)