Update web-demos/hugging_face/app.py

web-demos/hugging_face/app.py

@@ -365,48 +365,48 @@ def inpaint_video(video_state, *_args):
 
 
 # generate video after vos inference
-def generate_video_from_frames(frames, output_path, fps=30, bitrate=None):
+def generate_video_from_frames(frames, output_path, fps=30):
     """
-    Generates a video from a list of frames.
-    
-    Args:
-        frames (list of numpy arrays): The frames to include in the video.
-        output_path (str): The path to save the generated video.
-        fps (int, optional): The frame rate of the output video. Defaults to 30.
+    Save high-quality video using ffmpeg and CRF.
     """
-    # Приведение fps к обычному float (из np.float64, если нужно)
-    import numpy as np
-    import imageio
-    import torch
-    import torchvision
     import os
-    from fractions import Fraction
-
-    # Convert fps to a clean format
-    if isinstance(fps, np.generic):
-        fps = fps.item()
-    fps = float(fps)
-
-    # Ensure all frames are the same shape
-    assert all(f.shape == frames[0].shape for f in frames), "All frames must have the same shape"
-
-    # Convert to tensor (T, H, W, C)
-    #frames = torch.from_numpy(np.asarray(frames))
-
-    # Ensure output directory exists
-    if not os.path.exists(os.path.dirname(output_path)):
-        os.makedirs(os.path.dirname(output_path))
+    import cv2
+    import numpy as np
+    import subprocess
+    import tempfile
+
+    os.makedirs(os.path.dirname(output_path), exist_ok=True)
+
+    # Убедимся, что все кадры — одинакового размера
+    h, w, _ = frames[0].shape
+    h, w = h // 2 * 2, w // 2 * 2
+    frames = [cv2.resize(f, (w, h)) for f in frames]
+
+    # Временная папка
+    with tempfile.TemporaryDirectory() as tmpdir:
+        frame_dir = os.path.join(tmpdir, "frames")
+        os.makedirs(frame_dir, exist_ok=True)
+
+        # Сохраняем каждый кадр как PNG
+        for idx, frame in enumerate(frames):
+            cv2.imwrite(os.path.join(frame_dir, f"{idx:05d}.png"), frame)
+
+        # Генерируем видео через ffmpeg
+        ffmpeg_cmd = [
+            "ffmpeg",
+            "-y",
+            "-framerate", str(fps),
+            "-i", os.path.join(frame_dir, "%05d.png"),
+            "-c:v", "libx264",
+            "-preset", "slow",             # качество > скорость
+            "-crf", "18",                  # высокое качество (0 — без потерь)
+            "-pix_fmt", "yuv420p",
+            output_path
+        ]
+
+        print("Running ffmpeg:", " ".join(ffmpeg_cmd))
+        subprocess.run(ffmpeg_cmd, check=True)
 
-    # Write the video
-    #torchvision.io.write_video(output_path, frames, fps=fps, video_codec="libx264")
-    #return output_path
-    if bitrate is not None:
-        writer = imageio.get_writer(output_path, fps=fps, codec='libx264', bitrate=bitrate)
-    else:
-        writer = imageio.get_writer(output_path, fps=fps, codec='libx264')
-    for frame in frames:
-        writer.append_data(frame.astype(np.uint8))
-    writer.close()
     return output_path
 
 def restart():