Update app.py

app.py

@@ -14,6 +14,71 @@ pip install -r requirements.txt
 python app.py
 '''
 
+from datasets import load_dataset 
+vid_ds = load_dataset("svjack/video-dataset-genshin-impact-ep-character-organized")
+
+import decord
+import cv2
+import os
+
+def save_video_to_mp4(input_video, output_path):
+    """
+    将 decord.VideoReader 读取的视频或视频文件保存为 MP4 文件。
+
+    参数:
+        input_video (str 或 decord.VideoReader): 视频文件路径或 decord.VideoReader 对象。
+        output_path (str): 输出 MP4 文件的路径。
+    """
+    # 如果输入是路径，则创建 VideoReader 对象
+    if isinstance(input_video, str):
+        vr = decord.VideoReader(input_video)
+    elif isinstance(input_video, decord.VideoReader):
+        vr = input_video
+    else:
+        raise ValueError("输入必须是视频文件路径或 decord.VideoReader 对象")
+
+    # 获取视频的基本信息
+    fps = vr.get_avg_fps()
+    width = vr[0].shape[1]
+    height = vr[0].shape[0]
+
+    # 创建 VideoWriter 对象
+    #fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # 编码器
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # 编码器
+    out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
+
+    # 逐帧读取并写入
+    for frame in vr:
+        frame = frame.asnumpy()  # 将 decord 帧转换为 numpy 数组
+        frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)  # 转换颜色空间
+        out.write(frame)
+
+    # 释放资源
+    out.release()
+    print(f"视频已保存为 {output_path}")
+
+#rd = vid_ds["train"][0]["video"]
+#rd
+#save_video_to_mp4(rd, "output_video.mp4")
+
+out_dir = "examples"
+import shutil 
+from uuid import uuid1
+from tqdm import tqdm
+if os.path.exists(out_dir):
+    shutil.rmtree(out_dir)
+
+vid_l = []
+os.makedirs(out_dir, exist_ok = True)
+for exp in tqdm(vid_ds["train"]):
+    rd = exp["video"]
+    target_path = os.path.join(out_dir, "{}.mp4".format(uuid1()))
+    vid_l.append(target_path)
+    save_video_to_mp4(rd, target_path)
+    if len(vid_l) >= 10:
+        break
+
+
 import subprocess
 import re
 from typing import List, Tuple, Optional
@@ -576,6 +641,24 @@ with gr.Blocks(css=css) as demo:
         ],
         queue = False
     )
+    
+    video_in.change(
+        fn = preprocess_video_in, 
+        inputs = [video_in], 
+        outputs = [
+            first_frame_path, 
+            tracking_points, # update Tracking Points in the gr.State([]) object
+            trackings_input_label, # update Tracking Labels in the gr.State([]) object
+            input_first_frame_image, # hidden component used as ref when clearing points
+            points_map, # Image component where we add new tracking points
+            video_frames_dir, # Array where frames from video_in are deep stored
+            scanned_frames, # Scanned frames by SAM2
+            stored_inference_state, # Sam2 inference state
+            stored_frame_names, # 
+            video_in_drawer, # Accordion to hide uploaded video player
+        ],
+        queue = False
+    )
 
     
     # triggered when we click on image to add new points
@@ -660,4 +743,6 @@ with gr.Blocks(css=css) as demo:
         outputs = [output_propagated, output_video, working_frame, available_frames_to_check, reset_prpgt_brn]
     )
 
+    gr.Examples(vid_l, video_in)
+
 demo.launch(share = True)