Update app.py

app.py

@@ -2,6 +2,7 @@ import time
 import shutil
 import gradio as gr
 import os
+import cv2
 import json
 import torch
 import argparse
@@ -36,7 +37,43 @@ def load_and_preprocess(path, config, mean_pose, std_pose):
     motion_proj = torch.from_numpy(motion_proj).float()
 
     return motion_proj, start
-    
+
+def image2video():
+        # 指定包含PNG图像的目录
+    image_directory = './an-frames'
+    # 输出视频的路径
+    output_video_path = './adv.mp4'
+    # 视频的帧率（每秒中的帧数）
+    fps = 24
+
+    # 获取所有PNG图像的文件名并按顺序排序
+    images = sorted([img for img in os.listdir(image_directory) if img.endswith(".png")])
+
+    # 确定视频的分辨率
+    first_image_path = os.path.join(image_directory, images[0])
+    first_frame = cv2.imread(first_image_path)
+    height, width, layers = first_frame.shape
+    size = (width, height)
+
+    # 创建视频写入器
+    out = cv2.VideoWriter(output_video_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, size)
+
+    # 依次读取图像并写入视频
+    for img_name in images:
+        img_path = os.path.join(image_directory, img_name)
+        frame = cv2.imread(img_path)
+        
+        # 确保帧的大小与视频的分辨率匹配
+        if (frame.shape[1], frame.shape[0]) != size:
+            frame = cv2.resize(frame, size)
+        
+        out.write(frame)
+
+    # 关闭视频写入器
+    out.release()
+
+    print(f"Video saved to {output_video_path}")
+
 def handle_motion_generation(npy1,npy2):
     path1 = './data/a.npy'
     path2 = './data/b.npy'
@@ -49,7 +86,7 @@ def handle_motion_generation(npy1,npy2):
 
     config = get_config(config_path)
     ae = get_autoencoder(config)
-    ae.load_state_dict(torch.load(checkpoint_path, map_location=torch.device('cpu')))
+    ae.load_state_dict(torch.load(checkpoint_path))
     # ae.cuda()
     ae.eval()
     mean_pose, std_pose = get_meanpose("test", config.data)
@@ -88,28 +125,28 @@ def handle_motion_generation(npy1,npy2):
     # print(motion_data.shape)
     # 调用函数生成视频
     motion2video(motion_data, height, width, save_path, colors, bg_color=bg_color, transparency=False, fps=fps)
+    image2video()
     first_frame_image = Image.open('./an-frames/0000.png')
-    return first_frame_image
+    return './adv.mp4'
 
+# 创建 Gradio 界面并展示视频而不是图片
 with gr.Blocks() as demo:
-    gr.Markdown("Upload two `.npy` files to generate motion and visualize the first frame of the output animation.")
+    gr.Markdown("Upload two `.npy` files to generate motion and visualize the 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")
+        generate_btn = gr.Button("Generate Animation")
     
-    output_image = gr.Image(label="First Frame of the Generated Animation")
+    output_video = gr.Video(label="Generated Animation",width = 500)
 
     generate_btn.click(
         fn=handle_motion_generation,
         inputs=[file1, file2],
-        outputs=output_image
+        outputs=output_video
     )
 
-
 if __name__ == "__main__":
-    # tsp_page.launch(debug = True)
-    demo.launch()
+    demo.launch()