Update app.py

previous change invalid

app.py

@@ -2,7 +2,6 @@ import subprocess
 import re
 from typing import List, Tuple, Optional
 import pickle
-from datetime import datetime
 
 # Define the command to be executed
 command = ["python", "setup.py", "build_ext", "--inplace"]
@@ -40,7 +39,7 @@ def get_video_properties(video_path):
     if not cap.isOpened():
         print("Error: Could not open video.")
         return None
-                p
+    
     # Get the FPS of the video
     fps = cap.get(cv2.CAP_PROP_FPS)
     
@@ -81,27 +80,26 @@ def preprocess_video_in(video_path):
     if not cap.isOpened():
         print("Error: Could not open video.")
         return None
-    
+
     # Get the frames per second (FPS) of the video
     fps = cap.get(cv2.CAP_PROP_FPS)
-    """
+    
     # Calculate the number of frames to process (10 seconds of video)
-    # 需解除10s的视频限制
     max_frames = int(fps * 10)
-    """
-
+    
     frame_number = 0
     first_frame = None
     
     while True:
         ret, frame = cap.read()
         """
-        解除10s的视频限制
+        不再将视频裁剪至10s
         if not ret or frame_number >= max_frames:
             break
         """
         if not ret:
             break
+
         
         # Format the frame filename as '00000.jpg'
         frame_filename = os.path.join(extracted_frames_output_dir, f'{frame_number:05d}.jpg')
@@ -351,6 +349,8 @@ def propagate_to_all(video_in, checkpoint, stored_inference_state, stored_frame_
     masks_images = []
 
     # run propagation throughout the video and collect the results in a dict
+    # 添加全局变量保存mask数据
+    global video_segments
     video_segments = {}  # video_segments contains the per-frame segmentation results
     for out_frame_idx, out_obj_ids, out_mask_logits in predictor.propagate_in_video(inference_state):
         video_segments[out_frame_idx] = {
@@ -448,15 +448,8 @@ def propagate_to_all(video_in, checkpoint, stored_inference_state, stored_frame_
         video_writer.release()
         print(f"Mask Video saved at {mask_video_filename}")
 
-    # 在函数末尾添加导出mask数据
-    mask_data_file = export_mask_data(video_segments, video_frames_dir, frame_names)
-    
-    if vis_frame_type == "check":
-        return gr.update(value=jpeg_images), gr.update(value=None), gr.update(choices=available_frames_to_check, value=working_frame, visible=True), available_frames_to_check, gr.update(visible=True), None, mask_data_file
-    elif vis_frame_type == "render":
-        # ... 现有代码 ...
-        return gr.update(value=None), gr.update(value=final_vid_output_path), working_frame, available_frames_to_check, gr.update(visible=True), mask_video_filename, mask_data_file
         
+        return gr.update(value=None), gr.update(value=final_vid_output_path), working_frame, available_frames_to_check, gr.update(visible=True), mask_video_filename
 
 def update_ui(vis_frame_type):
     if vis_frame_type == "check":
@@ -485,39 +478,25 @@ def reset_propagation(first_frame_path, predictor, stored_inference_state):
     # print(f"RESET State: {stored_inference_state} ")
     return first_frame_path, gr.State([]), gr.State([]), gr.update(value=None, visible=False), stored_inference_state, None, ["frame_0.jpg"], first_frame_path, "frame_0.jpg", gr.update(visible=False)
 
+def export_masks():
+    # 导出mask数据
+    global video_segments
+    if not video_segments:
+        raise gr.Error("No mask data available. Please run propagation first!")
+    
+    # 保存为pickle文件
+    filename = "mask_data.pkl"
+    with open(filename, 'wb') as f:
+        pickle.dump(video_segments, f)
+    
+    return filename
+
 css="""
 div#component-18, div#component-25, div#component-35, div#component-41{
     align-items: stretch!important;
 }
 """
 
-def export_mask_data(video_segments, video_frames_dir, scanned_frames):
-    if not video_segments:
-        return None
-        
-    # 准备导出数据
-    export_data = {
-        'video_dir': video_frames_dir,
-        'frame_names': scanned_frames,
-        'masks': {}
-    }
-    
-    # 将每一帧的mask数据添加到导出数据中
-    for frame_idx, masks in video_segments.items():
-        export_data['masks'][frame_idx] = {
-            obj_id: mask.tolist() for obj_id, mask in masks.items()
-        }
-    
-    # 生成唯一的文件名
-    timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
-    output_file = f'mask_data_{timestamp}.pkl'
-    
-    # 保存数据
-    with open(output_file, 'wb') as f:
-        pickle.dump(export_data, f)
-        
-    return output_file
-
 with gr.Blocks(css=css) as demo:
     first_frame_path = gr.State()
     tracking_points = gr.State([])
@@ -575,9 +554,6 @@ with gr.Blocks(css=css) as demo:
                     <img src="https://huggingface.co/datasets/huggingface/badges/resolve/main/duplicate-this-space-lg-dark.svg" alt="Duplicate this Space" />
                 </a> to skip queue and avoid OOM errors from heavy public load
                 """)
-
-                # mask数据下载
-                mask_data_file = gr.File(label="Download Mask Data", visible=True)
             
             with gr.Column():
 
@@ -598,6 +574,9 @@ with gr.Blocks(css=css) as demo:
                 output_video = gr.Video(visible=False)
                 mask_final_output = gr.Video(label="Mask Video")
                 # output_result_mask = gr.Image()
+        # 在输出部分添加导出按钮
+        export_btn = gr.Button("Export Mask Data")
+        mask_download = gr.File(label="Download Mask Data", visible=False)
     
     
 
@@ -700,7 +679,16 @@ with gr.Blocks(css=css) as demo:
     ).then(
         fn = propagate_to_all,
         inputs = [video_in, checkpoint, stored_inference_state, stored_frame_names, video_frames_dir, vis_frame_type, available_frames_to_check, working_frame],
-        outputs = [output_propagated, output_video, working_frame, available_frames_to_check, reset_prpgt_brn, mask_final_output, mask_data_file]
+        outputs = [output_propagated, output_video, working_frame, available_frames_to_check, reset_prpgt_brn, mask_final_output]
+    )
+
+    # 添加导出按钮的事件绑定
+    export_btn.click(
+        fn=export_masks,
+        outputs=[mask_download]
+    ).then(
+        fn=lambda: gr.update(visible=True),
+        outputs=[mask_download]
     )
 
 demo.launch(show_api=False, show_error=True)