Update gradio_app_with_frames.py

gradio_app_with_frames.py

@@ -1,93 +1,158 @@
 import os
-import sys
 import shutil
 import uuid
 import subprocess
 import gradio as gr
-import shutil
+import cv2
+import sys
 from glob import glob
-from huggingface_hub import snapshot_download, hf_hub_download
-from moviepy.editor import VideoFileClip  # Import MoviePy
-
-# Download models
-os.makedirs("pretrained_weights", exist_ok=True)
+from pathlib import Path
 
-# List of subdirectories to create inside "checkpoints"
-subfolders = [
-    "stable-video-diffusion-img2vid-xt"
-]
+# 获取当前Python解释器路径
+#PYTHON_EXECUTABLE = sys.executable
+PYTHON_EXECUTABLE = "python"
 
-# Create each subdirectory
-for subfolder in subfolders:
-    os.makedirs(os.path.join("pretrained_weights", subfolder), exist_ok=True)
+def normalize_path(path: str) -> str:
+    """标准化路径，将Windows路径转换为正斜杠形式"""
+    return str(Path(path).resolve()).replace('\\', '/')
 
-snapshot_download(
-    repo_id="stabilityai/stable-video-diffusion-img2vid",
-    local_dir="./pretrained_weights/stable-video-diffusion-img2vid-xt"
-)
+def check_video_frames(video_path: str) -> int:
+    """检查视频帧数"""
+    video_path = normalize_path(video_path)
+    cap = cv2.VideoCapture(video_path)
+    frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    cap.release()
+    return frame_count
 
-snapshot_download(
-    repo_id="Yhmeng1106/anidoc",
-    local_dir="./pretrained_weights"
-)
-
-hf_hub_download(
-    repo_id="facebook/cotracker",
-    filename="cotracker2.pth",
-    local_dir="./pretrained_weights"
-)
+def preprocess_video(video_path: str) -> str:
+    """预处理视频到14帧"""
+    try:
+        video_path = normalize_path(video_path)
+        unique_id = str(uuid.uuid4())
+        temp_dir = "outputs"
+        output_dir = os.path.join(temp_dir, f"processed_{unique_id}")
+        output_dir = normalize_path(output_dir)
+        os.makedirs(output_dir, exist_ok=True)
+        
+        print(f"Processing video: {video_path}")
+        print(f"Output directory: {output_dir}")
+        
+        # 调用process_video_to_14frames.py处理视频
+        result = subprocess.run(
+            [
+                PYTHON_EXECUTABLE, "process_video_to_14frames.py",
+                "--input", video_path,
+                "--output", output_dir
+            ],
+            check=True,
+            capture_output=True,
+            text=True
+        )
+        
+        if result.stdout:
+            print(f"Preprocessing stdout: {result.stdout}")
+        if result.stderr:
+            print(f"Preprocessing stderr: {result.stderr}")
+        
+        # 获取处理后的视频路径
+        processed_videos = glob(os.path.join(output_dir, "*.mp4"))
+        if not processed_videos:
+            raise gr.Error("Failed to process video: No output video found")
+        return normalize_path(processed_videos[0])
+    except subprocess.CalledProcessError as e:
+        print(f"Preprocessing stderr: {e.stderr}")
+        raise gr.Error(f"Failed to preprocess video: {e.stderr}")
+    except Exception as e:
+        raise gr.Error(f"Error during video preprocessing: {str(e)}")
 
 def generate(control_sequence, ref_image):
-    control_image = control_sequence  # "data_test/sample4.mp4"
-    ref_image = ref_image  # "data_test/sample4.png"
-    unique_id = str(uuid.uuid4())
-    output_dir = f"results_{unique_id}"
-    
     try:
-        # Use MoviePy to get the number of frames in the control_sequence video
-        video_clip = VideoFileClip(control_image)
-        num_frames = int(video_clip.fps * video_clip.duration)  # Calculate total frames
-        video_clip.close()  # Close the video clip to free resources
-
-        # Run the inference command
-        subprocess.run(
+        # 验证输入文件是否存在
+        control_sequence = normalize_path(control_sequence)
+        ref_image = normalize_path(ref_image)
+        
+        if not os.path.exists(control_sequence):
+            raise gr.Error(f"Control sequence file not found: {control_sequence}")
+        if not os.path.exists(ref_image):
+            raise gr.Error(f"Reference image file not found: {ref_image}")
+            
+        # 创建输出目录
+        output_dir = "outputs"
+        os.makedirs(output_dir, exist_ok=True)
+        unique_id = str(uuid.uuid4())
+        result_dir = os.path.join(output_dir, f"results_{unique_id}")
+        result_dir = normalize_path(result_dir)
+        os.makedirs(result_dir, exist_ok=True)
+        
+        print(f"Input control sequence: {control_sequence}")
+        print(f"Input reference image: {ref_image}")
+        print(f"Output directory: {result_dir}")
+        
+        # 检查视频帧数
+        frame_count = check_video_frames(control_sequence)
+        if frame_count != 14:
+            print(f"Video has {frame_count} frames, preprocessing to 14 frames...")
+            control_sequence = preprocess_video(control_sequence)
+            print(f"Preprocessed video saved to: {control_sequence}")
+        
+        # 运行推理命令
+        print(f"Running inference...")
+        result = subprocess.run(
             [
-                "python", "scripts_infer/anidoc_inference.py",
+                PYTHON_EXECUTABLE, "scripts_infer/anidoc_inference.py",
                 "--all_sketch",
                 "--matching",
                 "--tracking",
-                "--control_image", f"{control_image}",
-                "--ref_image", f"{ref_image}",
-                "--output_dir", f"{output_dir}",
+                "--control_image", control_sequence,
+                "--ref_image", ref_image,
+                "--output_dir", result_dir,
                 "--max_point", "10",
-                "--num_frames", str(num_frames)  # Pass the calculated num_frames
             ],
-            check=True
+            check=True,
+            capture_output=True,
+            text=True
         )
+        
+        if result.stdout:
+            print(f"Inference stdout: {result.stdout}")
+        if result.stderr:
+            print(f"Inference stderr: {result.stderr}")
 
-        # Search for the mp4 file in a subfolder of output_dir
-        output_video = glob(os.path.join(output_dir, "*.mp4"))
-        print(output_video)
+        # 搜索输出视频
+        output_video = glob(os.path.join(result_dir, "*.mp4"))
+        print(f"Found output videos: {output_video}")
         
         if output_video:
-            output_video_path = output_video[0]  # Get the first match
+            output_video_path = normalize_path(output_video[0])
+            print(f"Returning output video: {output_video_path}")
         else:
-            output_video_path = None
+            raise gr.Error("No output video generated")
+            
+        # 清理临时文件
+        temp_dirs = glob("outputs/processed_*")
+        for temp_dir in temp_dirs:
+            if os.path.isdir(temp_dir):
+                try:
+                    shutil.rmtree(temp_dir)
+                    print(f"Cleaned up temp directory: {temp_dir}")
+                except Exception as e:
+                    print(f"Warning: Failed to clean up temp directory {temp_dir}: {str(e)}")
         
-        print(output_video_path)
         return output_video_path
-    
+        
     except subprocess.CalledProcessError as e:
-        raise gr.Error(f"Error during inference: {str(e)}")
+        print(f"Inference stderr: {e.stderr}")
+        raise gr.Error(f"Error during inference: {e.stderr}")
     except Exception as e:
-        raise gr.Error(f"Error processing video: {str(e)}")
+        raise gr.Error(f"Error: {str(e)}")
 
-css = """
+css="""
 div#col-container{
     margin: 0 auto;
     max-width: 982px;
 }
 """
+
 with gr.Blocks(css=css) as demo:
     with gr.Column(elem_id="col-container"):
         gr.Markdown("# AniDoc: Animation Creation Made Easier")
@@ -103,12 +168,6 @@ with gr.Blocks(css=css) as demo:
             <a href="https://arxiv.org/pdf/2412.14173">
                 <img src='https://img.shields.io/badge/ArXiv-Paper-red'>
             </a>
-            <a href="https://huggingface.co/spaces/fffiloni/AniDoc?duplicate=true">
-                <img src="https://huggingface.co/datasets/huggingface/badges/resolve/main/duplicate-this-space-sm.svg" alt="Duplicate this Space">
-            </a>
-            <a href="https://huggingface.co/fffiloni">
-                <img src="https://huggingface.co/datasets/huggingface/badges/resolve/main/follow-me-on-HF-sm-dark.svg" alt="Follow me on HF">
-            </a>
         </div>
         """)
         with gr.Row():
@@ -120,19 +179,20 @@ with gr.Blocks(css=css) as demo:
                 video_result = gr.Video(label="Result")
 
                 gr.Examples(
-                    examples=[
+                    examples = [
+                        ["data_test/sample5.mp4", "data_test/sample5.png"],
                         ["data_test/sample1.mp4", "data_test/sample1.png"],
                         ["data_test/sample2.mp4", "data_test/sample2.png"],
                         ["data_test/sample3.mp4", "data_test/sample3.png"],
                         ["data_test/sample4.mp4", "data_test/sample4.png"]
                     ],
-                    inputs=[control_sequence, ref_image]
-                )
-
-    submit_btn.click(
-        fn=generate,
-        inputs=[control_sequence, ref_image],
-        outputs=[video_result]
-    )
+                    inputs = [control_sequence, ref_image]
+        )
+        
+        submit_btn.click(
+        fn = generate,
+        inputs = [control_sequence, ref_image],
+        outputs = [video_result]
+        )
 
-demo.queue().launch(show_api=False, show_error=True, share=True)
+demo.queue().launch(inbrowser=True,show_api=False, show_error=True, share = True)