File size: 3,515 Bytes
0a9bdfb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 |
import os
import cv2
import argparse
import numpy as np
from tqdm import tqdm
from PIL import Image
from pose.script.tool import save_videos_from_pil
from pose.script.dwpose import draw_pose
def draw_dwpose(video_path, pose_path, out_path, draw_face):
# capture video info
cap = cv2.VideoCapture(video_path)
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
fps = cap.get(cv2.CAP_PROP_FPS)
fps = int(np.around(fps))
# fps = get_fps(video_path)
cap.release()
# render resolution, short edge = 1024
k = float(1024) / min(width, height)
h_render = int(k*height//2 * 2)
w_render = int(k*width//2 * 2)
# save resolution, short edge = 768
k = float(768) / min(width, height)
h_save = int(k*height//2 * 2)
w_save = int(k*width//2 * 2)
poses = np.load(pose_path, allow_pickle=True)
poses = poses.tolist()
frames = []
for pose in tqdm(poses):
detected_map = draw_pose(pose, h_render, w_render, draw_face)
detected_map = cv2.resize(detected_map, (w_save, h_save), interpolation=cv2.INTER_AREA)
# cv2.imshow('', detected_map)
# cv2.waitKey(0)
detected_map = cv2.cvtColor(detected_map, cv2.COLOR_BGR2RGB)
detected_map = Image.fromarray(detected_map)
frames.append(detected_map)
save_videos_from_pil(frames, out_path, fps)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--video_dir", type=str, default="./UBC_fashion/test", help='dance video dir')
parser.add_argument("--pose_dir", type=str, default=None, help='auto makedir')
parser.add_argument("--save_dir", type=str, default=None, help='auto makedir')
parser.add_argument("--draw_face", type=bool, default=False, help='whether draw face or not')
args = parser.parse_args()
# video dir
video_dir = args.video_dir
# pose dir
if args.pose_dir is None:
pose_dir = args.video_dir + "_dwpose_keypoints"
else:
pose_dir = args.pose_dir
# save dir
if args.save_dir is None:
if args.draw_face == True:
save_dir = args.video_dir + "_dwpose"
else:
save_dir = args.video_dir + "_dwpose_without_face"
else:
save_dir = args.save_dir
if not os.path.exists(save_dir):
os.makedirs(save_dir)
# collect all video_folder paths
video_mp4_paths = set()
for root, dirs, files in os.walk(args.video_dir):
for name in files:
if name.endswith(".mp4"):
video_mp4_paths.add(os.path.join(root, name))
video_mp4_paths = list(video_mp4_paths)
# random.shuffle(video_mp4_paths)
video_mp4_paths.sort()
print("Num of videos:", len(video_mp4_paths))
# draw dwpose
for i in range(len(video_mp4_paths)):
video_path = video_mp4_paths[i]
video_name = os.path.relpath(video_path, video_dir)
base_name = os.path.splitext(video_name)[0]
pose_path = os.path.join(pose_dir, base_name + '.npy')
if not os.path.exists(pose_path):
print('no keypoint file:', pose_path)
out_path = os.path.join(save_dir, base_name + '.mp4')
if os.path.exists(out_path):
print('already have rendered pose:', out_path)
continue
draw_dwpose(video_path, pose_path, out_path, args.draw_face)
print(f"Process {i+1}/{len(video_mp4_paths)} video")
print('all done!') |