Update app.py

disabled the 10s video restriction, added a button to download the mask pickle

app.py

@@ -1,6 +1,8 @@
 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"]
@@ -38,7 +40,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)
     
@@ -79,20 +81,27 @@ 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的视频限制
         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')
@@ -439,8 +448,15 @@ 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":
@@ -475,6 +491,33 @@ div#component-18, div#component-25, div#component-35, div#component-41{
 }
 """
 
+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([])
@@ -532,6 +575,9 @@ 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():
 
@@ -654,7 +700,7 @@ 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]
+        outputs = [output_propagated, output_video, working_frame, available_frames_to_check, reset_prpgt_brn, mask_final_output, mask_data_file]
     )
 
 demo.launch(show_api=False, show_error=True)