mirror of
https://github.com/gaomingqi/Track-Anything.git
synced 2025-12-16 08:27:49 +01:00
61 lines
2.0 KiB
Python
61 lines
2.0 KiB
Python
import torchvision
|
|
import os
|
|
import numpy as np
|
|
import torch
|
|
import cv2
|
|
# from tqdm import tqdm
|
|
from multiprocessing import Pool
|
|
|
|
def ensure_divisible_by_two(image):
|
|
height, width = image.shape[:2]
|
|
if width % 2 != 0:
|
|
width -= 1
|
|
if height % 2 != 0:
|
|
height -= 1
|
|
return cv2.resize(image, (width, height))
|
|
|
|
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 = [ensure_divisible_by_two(image) for image in frames]
|
|
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"))
|
|
sequence_list.remove('list.txt')
|
|
sequence_list.remove('list.txt.bp')
|
|
num_process = 8
|
|
with Pool(num_process) as p:
|
|
video_paths = p.map(process_seq, sequence_list) |