Dragreal / utils /vis_trajectory.py
BasicNp's picture
Upload 1672 files
e8aa256 verified
raw
history blame
2.23 kB
import cv2
import numpy as np
from PIL import Image
import json
from tqdm import tqdm
import os
from PIL import Image
# image_path = "./VIPSeg_Video_Generation_Test/imgs/443_hPrsXtTg3lM"
# trajectory = "./VIPSeg_Video_Generation_Test/trajectory_CoTracker/443_hPrsXtTg3lM.json"
save_image = "./VIPSeg_Video_Generation_Test/vis"
# image_path_root = "/mmu-ocr/weijiawu/MovieDiffusion/ShowAnything/data/VIPSeg/VIPSeg_Video_Generation_Test/imgs"
# trajectory_root = "/mmu-ocr/weijiawu/MovieDiffusion/ShowAnything/data/VIPSeg/VIPSeg_Video_Generation_Test/test_traject"
# save_image_root = "/mmu-ocr/weijiawu/MovieDiffusion/ShowAnything/data/VIPSeg/VIPSeg_Video_Generation_Test/trajectory_vis"
image_path_root = "/mmu-ocr/weijiawu/MovieDiffusion/ShowAnything/data/VIPSeg/VIPSeg_Video_Generation_Test/Prediction_Model/DragAnything14frames_OriginalSize"
trajectory_root = "/mmu-ocr/weijiawu/MovieDiffusion/ShowAnything/data/VIPSeg/VIPSeg_Video_Generation_Test/Prediction_Model/trajectory_1024_CoTracker_DragAnything14frames_OriginalSize1"
save_image_root = "/mmu-ocr/weijiawu/MovieDiffusion/ShowAnything/data/VIPSeg/VIPSeg_Video_Generation_Test/trajectory_vis"
def sort_frames(frame_name):
return int(frame_name.split('.')[0])
for video_name in os.listdir(image_path_root):
image_path = os.path.join(image_path_root,video_name)
trajectory = os.path.join(trajectory_root,video_name+".json")
save_image = os.path.join(save_image_root,video_name+".gif")
with open(trajectory, 'r') as json_file:
data = json.load(json_file)
image_files = sorted(os.listdir(image_path), key=sort_frames)
pil_images = []
for idx,images in enumerate(image_files):
image = cv2.imread(os.path.join(image_path,images))
for line in data:
line_data = data[line][:(idx+1)]
print(line_data)
if len(line_data)>=2:
for i in range(len(line_data)-1):
cv2.line(image, line_data[i], line_data[i+1], (0, 255, 0), 3)
cv2.imwrite(os.path.join(save_image,images),image)
pil_images.append(Image.fromarray(image))
pil_images[0].save(save_image, save_all=True, append_images=pil_images[1:], loop=0, duration=110)