mirror of
https://github.com/storytold/storyteller-ml.git
synced 2026-10-09 00:09:55 +00:00
Justin John
122 lines
3.9 KiB
Python
122 lines
3.9 KiB
Python
from pathlib import Path
|
|
from PIL import Image
|
|
import torch
|
|
import yaml
|
|
import math
|
|
|
|
import torchvision.transforms as T
|
|
from torchvision.io import read_video,write_video
|
|
import os
|
|
import random
|
|
import numpy as np
|
|
from torchvision.io import write_video
|
|
# from kornia.filters import joint_bilateral_blur
|
|
from kornia.geometry.transform import remap
|
|
from kornia.utils.grid import create_meshgrid
|
|
import cv2
|
|
|
|
def save_video_frames(video_path, img_size=(512,512)):
|
|
video, _, _ = read_video(video_path, output_format="TCHW")
|
|
# rotate video -90 degree if video is .mov format. this is a weird bug in torchvision
|
|
if video_path.endswith('.mov'):
|
|
video = T.functional.rotate(video, -90)
|
|
video_name = Path(video_path).stem
|
|
os.makedirs(f'data/{video_name}', exist_ok=True)
|
|
for i in range(len(video)):
|
|
ind = str(i).zfill(5)
|
|
image = T.ToPILImage()(video[i])
|
|
image_resized = image.resize((img_size), resample=Image.Resampling.LANCZOS)
|
|
image_resized.save(f'data/{video_name}/{ind}.png')
|
|
|
|
def video_to_frames(video_path, img_size=(512,512)):
|
|
video, _, _ = read_video(video_path, output_format="TCHW")
|
|
# rotate video -90 degree if video is .mov format. this is a weird bug in torchvision
|
|
if video_path.endswith('.mov'):
|
|
video = T.functional.rotate(video, -90)
|
|
video_name = Path(video_path).stem
|
|
# os.makedirs(f'data/{video_name}', exist_ok=True)
|
|
frames = []
|
|
for i in range(len(video)):
|
|
ind = str(i).zfill(5)
|
|
image = T.ToPILImage()(video[i])
|
|
image_resized = image.resize((img_size), resample=Image.Resampling.LANCZOS)
|
|
# image_resized.save(f'data/{video_name}/{ind}.png')
|
|
frames.append(image_resized)
|
|
return frames
|
|
|
|
def add_dict_to_yaml_file(file_path, key, value):
|
|
data = {}
|
|
|
|
# If the file already exists, load its contents into the data dictionary
|
|
if os.path.exists(file_path):
|
|
with open(file_path, 'r') as file:
|
|
data = yaml.safe_load(file)
|
|
|
|
# Add or update the key-value pair
|
|
data[key] = value
|
|
|
|
# Save the data back to the YAML file
|
|
with open(file_path, 'w') as file:
|
|
yaml.dump(data, file)
|
|
|
|
def isinstance_str(x: object, cls_name: str):
|
|
"""
|
|
Checks whether x has any class *named* cls_name in its ancestry.
|
|
Doesn't require access to the class's implementation.
|
|
|
|
Useful for patching!
|
|
"""
|
|
|
|
for _cls in x.__class__.__mro__:
|
|
if _cls.__name__ == cls_name:
|
|
return True
|
|
|
|
return False
|
|
|
|
|
|
def batch_cosine_sim(x, y):
|
|
if type(x) is list:
|
|
x = torch.cat(x, dim=0)
|
|
if type(y) is list:
|
|
y = torch.cat(y, dim=0)
|
|
x = x / x.norm(dim=-1, keepdim=True)
|
|
y = y / y.norm(dim=-1, keepdim=True)
|
|
similarity = x @ y.T
|
|
return similarity
|
|
|
|
|
|
def load_imgs(data_path, n_frames, device='cuda', pil=False):
|
|
imgs = []
|
|
pils = []
|
|
for i in range(n_frames):
|
|
img_path = os.path.join(data_path, "%05d.jpg" % i)
|
|
if not os.path.exists(img_path):
|
|
img_path = os.path.join(data_path, "%05d.png" % i)
|
|
img_pil = Image.open(img_path)
|
|
pils.append(img_pil)
|
|
img = T.ToTensor()(img_pil).unsqueeze(0)
|
|
imgs.append(img)
|
|
if pil:
|
|
return torch.cat(imgs).to(device), pils
|
|
return torch.cat(imgs).to(device)
|
|
|
|
|
|
def save_video(raw_frames, save_path, fps=10):
|
|
video_codec = "libx264"
|
|
video_options = {
|
|
"crf": "18", # Constant Rate Factor (lower value = higher quality, 18 is a good balance)
|
|
"preset": "slow", # Encoding preset (e.g., ultrafast, superfast, veryfast, faster, fast, medium, slow, slower, veryslow)
|
|
}
|
|
|
|
frames = (raw_frames * 255).to(torch.uint8).cpu().permute(0, 2, 3, 1)
|
|
write_video(save_path, frames, fps=fps, video_codec=video_codec, options=video_options)
|
|
|
|
|
|
def seed_everything(seed):
|
|
torch.manual_seed(seed)
|
|
torch.cuda.manual_seed(seed)
|
|
random.seed(seed)
|
|
np.random.seed(seed)
|
|
|
|
|