File size: 3,901 Bytes
7ca9b42 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 |
import time
import shutil
import gradio as gr
import os
import json
import torch
import argparse
import numpy as np
from lib.data import get_meanpose
from lib.network import get_autoencoder
from lib.util.motion import preprocess_mixamo, preprocess_test, postprocess
from lib.util.general import get_config
from lib.operation import rotate_and_maybe_project_world
from itertools import combinations
from lib.util.visualization import motion2video
from PIL import Image
def load_and_preprocess(path, config, mean_pose, std_pose):
motion3d = np.load(path)
# length must be multiples of 8 due to the size of convolution
_, _, T = motion3d.shape
T = (T // 8) * 8
motion3d = motion3d[:, :, :T]
# project to 2d
motion_proj = motion3d[:, [0, 2], :]
# reformat for mixamo data
motion_proj = preprocess_mixamo(motion_proj, unit=1.0)
# preprocess for network input
motion_proj, start = preprocess_test(motion_proj, mean_pose, std_pose, config.data.unit)
motion_proj = motion_proj.reshape((-1, motion_proj.shape[-1]))
motion_proj = torch.from_numpy(motion_proj).float()
return motion_proj, start
def handle_motion_generation(npy1,npy2):
path1 = './data/a.npy'
path2 = './data/b.npy'
np.save(path1,npy1)
np.save(path2,npy2)
config_path = './configs/transmomo.yaml' # 替换为您的配置文件路径
description_path = "./data/mse_description.json"
checkpoint_path = './data/autoencoder_00200000.pt'
out_dir_path = './output' # 替换为输出目录的路径
config = get_config(config_path)
ae = get_autoencoder(config)
ae.load_state_dict(torch.load(checkpoint_path))
ae.cuda()
ae.eval()
mean_pose, std_pose = get_meanpose("test", config.data)
# print("loaded model")
description = json.load(open(description_path))
chars = list(description.keys())
os.makedirs(out_dir_path, exist_ok=True)
# path1 = '/home/fazhong/studio/transmomo.pytorch/data/mixamo/36_800_24/test/PUMPKINHULK_L/Back_Squat/motions/2.npy'
# path2 = '/home/fazhong/studio/transmomo.pytorch/data/mixamo/36_800_24/test/PUMPKINHULK_L/Golf_Post_Shot/motions/3.npy'
out_path1 = os.path.join(out_dir_path, "adv.npy")
x_a, x_a_start = load_and_preprocess(path1, config, mean_pose, std_pose)
x_b, x_b_start = load_and_preprocess(path2, config, mean_pose, std_pose)
x_a_batch = x_a.unsqueeze(0).cuda()
x_b_batch = x_b.unsqueeze(0).cuda()
x_ab = ae.cross2d(x_a_batch, x_b_batch, x_a_batch)
x_ab = postprocess(x_ab, mean_pose, std_pose, config.data.unit, start=x_a_start)
np.save(out_path1, x_ab)
motion_data = x_ab
height = 512 # 视频的高度
width = 512 # 视频的宽度
save_path = './an.mp4' # 保存视频的路径
colors = [(255, 0, 0), (0, 255, 0), (0, 0, 255)] # 关节颜色
bg_color = (255, 255, 255) # 背景颜色
fps = 25 # 视频的帧率
# print(motion_data.shape)
# 调用函数生成视频
motion2video(motion_data, height, width, save_path, colors, bg_color=bg_color, transparency=False, fps=fps)
first_frame_image = Image.open('./an-frames/0000.png')
return first_frame_image
# print('hi')
with gr.Blocks() as demo:
gr.Markdown("Upload two `.npy` files to generate motion and visualize the first frame of the output animation.")
with gr.Row():
file1 = gr.File(file_types=[".npy"], label="Upload first .npy file")
file2 = gr.File(file_types=[".npy"], label="Upload second .npy file")
with gr.Row():
generate_btn = gr.Button("Generate Motion")
output_image = gr.Image(label="First Frame of the Generated Animation")
generate_btn.click(
fn=handle_motion_generation,
inputs=[file1, file2],
outputs=output_image
)
if __name__ == "__main__":
# tsp_page.launch(debug = True)
demo.launch() |