tqdm for visualization -- li

README.md

@@ -90,6 +90,8 @@ Mask save path:
 /ssd2/tracking/vots2023/gt_mask/
 ```
 
+#### 
+
 ## Citation
 If you find this work useful for your research or applications please cite using this BibTeX:
 ```bibtex

app_vots.py

@@ -16,6 +16,7 @@ import torchvision
 import torch 
 import torch.nn.functional as F
 from tools.painter import mask_painter
+from tqdm import tqdm
 import psutil
 import time
 try: 
@@ -365,7 +366,7 @@ def vos_tracking_video(video_state, interactive_state, mask_dropdown):
             os.makedirs(mask_save_dir)
         i = 1
         print("save mask")
-        for mask in video_state["masks"]:
+        for mask in tqdm(video_state["masks"]):
             # np.save(os.path.join(mask_save_dir, '{:05d}.npy'.format(i)), mask)
             Image.fromarray(mask).save(os.path.join(mask_save_dir, '{:08d}.png'.format(i)))
             i+=1
@@ -419,11 +420,13 @@ def generate_video_from_frames(frames_path, output_path, fps=30):
     if os.path.exists(output_path):
         return output_path
     frames = []
-    for file in frames_path:
+    print("read frames from sequence")
+    for file in tqdm(frames_path):
         frames.append(read_image_from_userfolder(file))
     frames = torch.from_numpy(np.asarray(frames))
     if not os.path.exists(os.path.dirname(output_path)):
         os.makedirs(os.path.dirname(output_path))
+    print("generate video from frames for preview")
     torchvision.io.write_video(output_path, frames, fps=fps, video_codec="libx264")
     return output_path
 
@@ -447,21 +450,25 @@ def get_mask_from_vot(video_state, output_path, fps=30):
     masks = video_state["masks"]
     frames = video_state["origin_images"]
     # video_painted_images = [] 
+    height, width = np.asarray(Image.open(video_state["origin_images"][0])).shape[:2]
+    new_size = (width//2, height//2)
     painted_images = []
-    for i in range(len(masks)):
+    print("painting mask")
+    for i in tqdm(range(len(masks))):
         num_objs = masks[i].max()
         painted_image = np.asarray(Image.open(frames[i]).convert('RGB'))
         for obj in range(1, num_objs+1):
             if np.max(masks[i]==obj) == 0:
                 continue
             painted_image = mask_painter(painted_image, (masks[i]==obj).astype('uint8'), mask_color=obj+1)
-        painted_images.append(painted_image)
+        painted_images.append(cv2.resize(painted_image, new_size, interpolation=cv2.INTER_AREA))
         # video_painted_images.append(save_image_to_userfolder(video_state, index=i, image=cv2.cvtColor(np.asarray(painted_image),cv2.COLOR_BGR2RGB), type=False))
 
     painted_images = torch.from_numpy(np.asarray(painted_images))
     # resize for accelerating video generation
     # new_size = [painted_images.size(1)//2, painted_images.size(2)//2]
     # painted_images_resized = F.interpolate(painted_images, size=new_size, mode='bilinear')
+    print("saving result videos")
     if not os.path.exists(os.path.dirname(output_path)):
         os.makedirs(os.path.dirname(output_path))
     torchvision.io.write_video(output_path, painted_images, fps=fps, video_codec="libx264")

generate_video_multipreocess.py

@@ -0,0 +1,46 @@
+import torchvision
+import os
+import numpy as np
+import torch
+import cv2
+# from tqdm import tqdm
+from multiprocessing import Pool
+def read_image_from_userfolder(image_path):
+    # if type:
+    image = cv2.cvtColor(cv2.imread(image_path), cv2.COLOR_BGR2RGB)
+    return image
+def generate_video_from_frames(frames_path, 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.
+    """
+    if os.path.exists(output_path):
+        return output_path
+    frames = []
+    # print("read frames from sequence")
+    for file in frames_path:
+        frames.append(cv2.cvtColor(cv2.imread(file), cv2.COLOR_BGR2RGB))
+    frames = torch.from_numpy(np.asarray(frames))
+    if not os.path.exists(os.path.dirname(output_path)):
+        os.makedirs(os.path.dirname(output_path))
+    # print("generate video from frames for preview")
+    torchvision.io.write_video(output_path, frames, fps=fps, video_codec="libx264")
+    return output_path
+
+def process_seq(seq):
+    frames = [os.path.join(votdir, "sequences", seq, "color", i) for i in os.listdir(os.path.join(votdir, "sequences", seq, "color"))]
+    frames.sort()
+    video_path = generate_video_from_frames(frames, output_path=os.path.join(votdir, "frame2video", "{}.mp4".format(seq)))
+    print("{}.mp4 is ok".format(seq))
+    return video_path
+
+
+votdir = "/home/dataset/vots2023/"
+sequence_list = os.listdir(os.path.join(votdir, "sequences"))
+num_process = 8
+with Pool(num_process) as p:
+    video_paths = p.map(process_seq, sequence_list)