313 lines
9.7 KiB
Python
Executable File
313 lines
9.7 KiB
Python
Executable File
import importlib
|
|
import os
|
|
import os.path as osp
|
|
import shutil
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import av
|
|
import numpy as np
|
|
import torch
|
|
import torchvision
|
|
from einops import rearrange
|
|
from PIL import Image
|
|
import cv2
|
|
|
|
|
|
def save_checkpoint(model, save_dir, prefix, ckpt_num, logger, total_limit=None):
|
|
save_path = osp.join(save_dir, f"{prefix}-{ckpt_num}.pth")
|
|
|
|
if total_limit is not None:
|
|
checkpoints = os.listdir(save_dir)
|
|
checkpoints = [d for d in checkpoints if d.startswith(prefix)]
|
|
checkpoints = sorted(
|
|
checkpoints, key=lambda x: int(x.split("-")[1].split(".")[0])
|
|
)
|
|
|
|
if len(checkpoints) >= total_limit:
|
|
num_to_remove = len(checkpoints) - total_limit + 1
|
|
removing_checkpoints = checkpoints[0:num_to_remove]
|
|
logger.info(
|
|
f"{len(checkpoints)} checkpoints already exist, removing {len(removing_checkpoints)} checkpoints"
|
|
)
|
|
logger.info(f"removing checkpoints: {', '.join(removing_checkpoints)}")
|
|
|
|
for removing_checkpoint in removing_checkpoints:
|
|
removing_checkpoint = os.path.join(save_dir, removing_checkpoint)
|
|
os.remove(removing_checkpoint)
|
|
|
|
state_dict = model.state_dict()
|
|
torch.save(state_dict, save_path)
|
|
|
|
|
|
def create_code_snapshot(root, dst_path, extensions=(".py", ".h", ".cpp", ".cu", ".cc", ".cuh", ".json", ".sh", ".bat", ".yaml"), exclude=()):
|
|
"""Creates tarball with the source code"""
|
|
import tarfile
|
|
from pathlib import Path
|
|
with tarfile.open(str(dst_path), "w:gz") as tar:
|
|
for path in Path(root).rglob("*"):
|
|
if '.git' in path.parts:
|
|
continue
|
|
exclude_flag = False
|
|
if len(exclude) > 0:
|
|
for k in exclude:
|
|
if k in path.parts:
|
|
exclude_flag = True
|
|
if exclude_flag:
|
|
continue
|
|
if path.suffix.lower() in extensions:
|
|
try:
|
|
tar.add(path.as_posix(), arcname=path.relative_to(
|
|
root).as_posix(), recursive=True)
|
|
except:
|
|
print(path)
|
|
assert False, 'Error occur in create_code_snapshot'
|
|
|
|
def seed_everything(seed):
|
|
import random
|
|
|
|
import numpy as np
|
|
|
|
torch.manual_seed(seed)
|
|
torch.cuda.manual_seed_all(seed)
|
|
np.random.seed(seed % (2**32))
|
|
random.seed(seed)
|
|
|
|
|
|
def import_filename(filename):
|
|
spec = importlib.util.spec_from_file_location("mymodule", filename)
|
|
module = importlib.util.module_from_spec(spec)
|
|
sys.modules[spec.name] = module
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
def delete_additional_ckpt(base_path, num_keep):
|
|
dirs = []
|
|
for d in os.listdir(base_path):
|
|
if d.startswith("checkpoint-"):
|
|
dirs.append(d)
|
|
num_tot = len(dirs)
|
|
if num_tot <= num_keep:
|
|
return
|
|
# ensure ckpt is sorted and delete the ealier!
|
|
del_dirs = sorted(dirs, key=lambda x: int(x.split("-")[-1]))[: num_tot - num_keep]
|
|
for d in del_dirs:
|
|
path_to_dir = osp.join(base_path, d)
|
|
if osp.exists(path_to_dir):
|
|
shutil.rmtree(path_to_dir)
|
|
|
|
|
|
def save_videos_from_pil(pil_images, path, fps=8, crf=None):
|
|
import av
|
|
|
|
save_fmt = Path(path).suffix
|
|
os.makedirs(os.path.dirname(path), exist_ok=True)
|
|
width, height = pil_images[0].size
|
|
|
|
if save_fmt == ".mp4":
|
|
if True:
|
|
codec = "libx264"
|
|
container = av.open(path, "w")
|
|
stream = container.add_stream(codec, rate=fps)
|
|
|
|
stream.width = width
|
|
stream.height = height
|
|
if crf is not None:
|
|
stream.options = {'crf': str(crf)}
|
|
|
|
for pil_image in pil_images:
|
|
# pil_image = Image.fromarray(image_arr).convert("RGB")
|
|
av_frame = av.VideoFrame.from_image(pil_image)
|
|
container.mux(stream.encode(av_frame))
|
|
container.mux(stream.encode())
|
|
container.close()
|
|
else:
|
|
|
|
video_writer = cv2.VideoWriter(
|
|
path.replace('.mp4', '_cv.mp4'), cv2.VideoWriter_fourcc(*'mp4v'), fps, (width, height)
|
|
)
|
|
for pil_image in pil_images:
|
|
img_np = cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)
|
|
video_writer.write(img_np)
|
|
video_writer.release()
|
|
|
|
elif save_fmt == ".gif":
|
|
pil_images[0].save(
|
|
fp=path,
|
|
format="GIF",
|
|
append_images=pil_images[1:],
|
|
save_all=True,
|
|
duration=(1 / fps * 1000),
|
|
loop=0,
|
|
)
|
|
else:
|
|
raise ValueError("Unsupported file type. Use .mp4 or .gif.")
|
|
|
|
|
|
def save_videos_grid(videos_, path: str, rescale=False, n_rows=6, fps=8, crf=None):
|
|
if not isinstance(videos_, list): videos_ = [videos_]
|
|
|
|
outputs = []
|
|
vid_len = videos_[0].shape[2]
|
|
for i in range(vid_len):
|
|
output = []
|
|
for videos in videos_:
|
|
videos = rearrange(videos, "b c t h w -> t b c h w")
|
|
height, width = videos.shape[-2:]
|
|
|
|
x = torchvision.utils.make_grid(videos[i], nrow=n_rows) # (c h w)
|
|
x = x.transpose(0, 1).transpose(1, 2).squeeze(-1) # (h w c)
|
|
if rescale:
|
|
x = (x + 1.0) / 2.0 # -1,1 -> 0,1
|
|
x = (x * 255).numpy().astype(np.uint8)
|
|
output.append(x)
|
|
|
|
output = Image.fromarray(np.concatenate(output, axis=0))
|
|
outputs.append(output)
|
|
|
|
os.makedirs(os.path.dirname(path), exist_ok=True)
|
|
save_videos_from_pil(outputs, path, fps, crf)
|
|
|
|
|
|
def save_videos_grid_ori(videos: torch.Tensor, path: str, rescale=False, n_rows=6, fps=8):
|
|
videos = rearrange(videos, "b c t h w -> t b c h w")
|
|
height, width = videos.shape[-2:]
|
|
outputs = []
|
|
|
|
for x in videos:
|
|
x = torchvision.utils.make_grid(x, nrow=n_rows) # (c h w)
|
|
x = x.transpose(0, 1).transpose(1, 2).squeeze(-1) # (h w c)
|
|
if rescale:
|
|
x = (x + 1.0) / 2.0 # -1,1 -> 0,1
|
|
x = (x * 255).numpy().astype(np.uint8)
|
|
x = Image.fromarray(x)
|
|
|
|
outputs.append(x)
|
|
|
|
os.makedirs(os.path.dirname(path), exist_ok=True)
|
|
|
|
save_videos_from_pil(outputs, path, fps)
|
|
|
|
|
|
def read_frames(video_path):
|
|
container = av.open(video_path)
|
|
|
|
video_stream = next(s for s in container.streams if s.type == "video")
|
|
frames = []
|
|
for packet in container.demux(video_stream):
|
|
for frame in packet.decode():
|
|
image = Image.frombytes(
|
|
"RGB",
|
|
(frame.width, frame.height),
|
|
frame.to_rgb().to_ndarray(),
|
|
)
|
|
frames.append(image)
|
|
|
|
return frames
|
|
|
|
|
|
def get_fps(video_path):
|
|
container = av.open(video_path)
|
|
video_stream = next(s for s in container.streams if s.type == "video")
|
|
fps = video_stream.average_rate
|
|
container.close()
|
|
return fps
|
|
|
|
def draw_keypoints(keypoints, height=512, width=512, device="cuda"):
|
|
colors = torch.tensor([
|
|
[255, 0, 0],
|
|
[255, 255, 0],
|
|
[0, 255, 0],
|
|
[0, 255, 255],
|
|
[0, 0, 255],
|
|
[255, 0, 255],
|
|
[255, 0, 85],
|
|
], device=device, dtype=torch.float32)
|
|
|
|
selected = torch.tensor([1, 2, 3, 4, 12, 15, 20], device=device)
|
|
B = keypoints.shape[0]
|
|
|
|
# [B, len(selected), 2]
|
|
pts = keypoints[:, selected] * 0.5 + 0.5
|
|
pts[..., 0] *= width
|
|
pts[..., 1] *= height
|
|
pts = pts.long()
|
|
|
|
canvas = torch.zeros((B, 3, height, width), device=device)
|
|
radius = 4
|
|
|
|
for i, color in enumerate(colors):
|
|
x = pts[:, i, 0]
|
|
y = pts[:, i, 1]
|
|
mask = (
|
|
(x[:, None, None] - torch.arange(width, device=device)) ** 2
|
|
+ (y[:, None, None] - torch.arange(height, device=device)[:, None]) ** 2
|
|
) <= radius**2
|
|
canvas[:, 0] += color[0] / 255.0 * mask
|
|
canvas[:, 1] += color[1] / 255.0 * mask
|
|
canvas[:, 2] += color[2] / 255.0 * mask
|
|
|
|
return canvas.clamp(0, 1)
|
|
|
|
def get_boxes(keypoints, height=512, width=512):
|
|
selected = torch.tensor([1, 2, 3, 4, 12, 15, 20])
|
|
|
|
# [B, len(selected), 2]
|
|
pts = keypoints[:, selected] * 0.5 + 0.5
|
|
pts[..., 0] *= width
|
|
pts[..., 1] *= height
|
|
pts = pts.long()
|
|
|
|
cx = pts[..., 0].float().mean(dim=1) # [B]
|
|
cy = pts[..., 1].float().mean(dim=1) # [B]
|
|
|
|
min_y = pts[..., 1].float().min(dim=1)[0] # [B]
|
|
|
|
side = (cy - min_y) * 2.0
|
|
side = side * 1.7
|
|
|
|
x1 = (cx - side / 2 * 0.95).clamp(0, width - 1).long()
|
|
y1 = (cy - side / 2 * 0.95).clamp(0, height - 1).long()
|
|
x2 = (cx + side / 2 * 1.05).clamp(0, width - 1).long()
|
|
y2 = (cy + side / 2 * 1.05).clamp(0, height - 1).long()
|
|
|
|
boxes = torch.stack([x1, y1, x2, y2], dim=1) # [B, 4]
|
|
|
|
return boxes
|
|
|
|
def crop_face(image_pil, face_mesh, scale=1.1):
|
|
image = np.array(image_pil)
|
|
h, w = image.shape[:2]
|
|
results = face_mesh.process(image)
|
|
face_landmarks = results.multi_face_landmarks[0]
|
|
coords = [(int(l.x * w), int(l.y * h)) for l in face_landmarks.landmark]
|
|
xs, ys = zip(*coords)
|
|
x1, y1 = min(xs), min(ys)
|
|
x2, y2 = max(xs), max(ys)
|
|
face_box = (x1, y1, x2, y2)
|
|
|
|
left, top, right, bot = scale_bb(face_box, scale=scale, size=image.shape[:2])
|
|
|
|
face_patch = image[int(top) : int(bot), int(left) : int(right)]
|
|
|
|
return face_patch
|
|
|
|
def scale_bb(bbox, scale, size):
|
|
left, top, right, bot = bbox
|
|
width = right - left
|
|
height = bot - top
|
|
length = max(width, height) * scale
|
|
center_X = (left + right) * 0.5
|
|
center_Y = (top + bot) * 0.5
|
|
left, top, right, bot = [
|
|
center_X - length / 2,
|
|
center_Y - length / 2,
|
|
center_X + length / 2,
|
|
center_Y + length / 2,
|
|
]
|
|
left = max(0, left)
|
|
top = max(0, top)
|
|
right = min(size[1] - 1, right)
|
|
bot = min(size[0] - 1, bot)
|
|
return np.array([left, top, right, bot]) |