376 lines
13 KiB
Python
376 lines
13 KiB
Python
import os
|
|
import subprocess
|
|
import tempfile
|
|
from pathlib import Path
|
|
from typing import Union
|
|
import shutil
|
|
|
|
import cv2
|
|
import imageio
|
|
import numpy as np
|
|
import torch
|
|
import torchvision
|
|
from decord import VideoReader, cpu
|
|
from einops import rearrange, repeat
|
|
from .iimage import IImage
|
|
from PIL import Image, ImageDraw, ImageFont
|
|
from torchvision.utils import save_image
|
|
|
|
channel_first = 0
|
|
channel_last = -1
|
|
|
|
|
|
def video_naming(prompt, extension, batch_idx, idx):
|
|
prompt_identifier = prompt.replace(" ", "_")
|
|
prompt_identifier = prompt_identifier.replace("/", "_")
|
|
if len(prompt_identifier) > 40:
|
|
prompt_identifier = prompt_identifier[:40]
|
|
filename = f"{batch_idx:04d}_{idx:04d}_{prompt_identifier}.{extension}"
|
|
return filename
|
|
|
|
|
|
def video_naming_chunk(prompt, extension, batch_idx, idx, chunk_idx):
|
|
prompt_identifier = prompt.replace(" ", "_")
|
|
prompt_identifier = prompt_identifier.replace("/", "_")
|
|
if len(prompt_identifier) > 40:
|
|
prompt_identifier = prompt_identifier[:40]
|
|
filename = f"{batch_idx}_{idx}_{chunk_idx}_{prompt_identifier}.{extension}"
|
|
return filename
|
|
|
|
|
|
class ResultProcessor():
|
|
|
|
def __init__(self, fps: int, n_frames: int, logger=None) -> None:
|
|
self.fps = fps
|
|
self.logger = logger
|
|
self.n_frames = n_frames
|
|
|
|
def set_logger(self, logger):
|
|
self.logger = logger
|
|
|
|
def _create_video(self, video, prompt, filename: Union[str, Path], append_video: torch.FloatTensor = None, input_flow=None):
|
|
|
|
if video.ndim == 5:
|
|
# can be batches if we provide list of filenames
|
|
assert video.shape[0] == 1
|
|
video = video[0]
|
|
|
|
if video.shape[0] == 3 and video.shape[1] == self.n_frames:
|
|
video = rearrange(video, "C F W H -> F C W H")
|
|
assert video.shape[1] == 3, f"Wrong video format. Got {video.shape}"
|
|
if isinstance(filename, Path):
|
|
filename = filename.as_posix()
|
|
# assert video.max() <= 1 and video.min() >= 0
|
|
assert video.max() <=1.1 and video.min() >= -0.1, f"video has unexpected range: [{video.min()}, {video.max()}]"
|
|
vid_obj = IImage(video, vmin=0, vmax=1)
|
|
|
|
if prompt is not None:
|
|
vid_obj = vid_obj.append_text(prompt, padding=(0, 50, 0, 0))
|
|
|
|
if append_video is not None:
|
|
if append_video.ndim == 5:
|
|
assert append_video.shape[0] == 1
|
|
append_video = append_video[0]
|
|
if append_video.shape[0] < video.shape[0]:
|
|
append_video = torch.concat([append_video,
|
|
repeat(append_video[-1, None], "F C W H -> (rep F) C W H", rep=video.shape[0]-append_video.shape[0])], dim=0)
|
|
if append_video.ndim == 3 and video.ndim == 4:
|
|
append_video = repeat(
|
|
append_video, "C W H -> F C W H", F=video.shape[0])
|
|
append_video = IImage(append_video, vmin=-1, vmax=1)
|
|
if prompt is not None:
|
|
append_video = append_video.append_text(
|
|
"input_frame", padding=(0, 50, 0, 0))
|
|
vid_obj = vid_obj | append_video
|
|
vid_obj = vid_obj.setFps(self.fps)
|
|
vid_obj.save(filename)
|
|
|
|
def _create_prompt_file(self, prompt, filename, video_path: str = None):
|
|
filename = Path(filename)
|
|
filename = filename.parent / (filename.stem+".txt")
|
|
|
|
with open(filename.as_posix(), "w") as file_writer:
|
|
file_writer.write(prompt)
|
|
file_writer.write("\n")
|
|
if video_path is not None:
|
|
file_writer.write(video_path)
|
|
else:
|
|
file_writer.write(" no_source")
|
|
|
|
def log_video(self, video: torch.FloatTensor, prompt: str, video_id: str, log_folder: str, input_flow=None, video_path_input: str = None, extension: str = "gif", prompt_on_vid: bool = True, append_video: torch.FloatTensor = None):
|
|
|
|
with tempfile.TemporaryDirectory() as tmpdirname:
|
|
storage_fol = Path(tmpdirname)
|
|
filename = f"{video_id}.{extension}".replace("/", "_")
|
|
vid_filename = storage_fol / filename
|
|
self._create_video(
|
|
video, prompt if prompt_on_vid else None, vid_filename, append_video, input_flow=input_flow)
|
|
|
|
prompt_file = storage_fol / f"{video_id}.txt"
|
|
self._create_prompt_file(prompt, prompt_file, video_path_input)
|
|
|
|
if self.logger.experiment.__class__.__name__ == "_DummyExperiment":
|
|
run_fol = Path(self.logger.save_dir) / \
|
|
self.logger.experiment_id / self.logger.run_id / "artifacts" / log_folder
|
|
if not run_fol.exists():
|
|
run_fol.mkdir(parents=True, exist_ok=True)
|
|
shutil.copy(prompt_file.as_posix(),
|
|
(run_fol / f"{video_id}.txt").as_posix())
|
|
shutil.copy(vid_filename,
|
|
(run_fol / filename).as_posix())
|
|
else:
|
|
self.logger.experiment.log_artifact(
|
|
self.logger.run_id, prompt_file.as_posix(), log_folder)
|
|
self.logger.experiment.log_artifact(
|
|
self.logger.run_id, vid_filename, log_folder)
|
|
|
|
def save_to_file(self, video: torch.FloatTensor, prompt: str, video_filename: Union[str, Path], input_flow=None, conditional_video_path: str = None, prompt_on_vid: bool = True, conditional_video: torch.FloatTensor = None):
|
|
self._create_video(
|
|
video, prompt if prompt_on_vid else None, video_filename, conditional_video, input_flow=input_flow)
|
|
self._create_prompt_file(
|
|
prompt, video_filename, conditional_video_path)
|
|
|
|
|
|
def add_text_to_image(image_array, text, position, font_size, text_color, font_path=None):
|
|
|
|
# Convert the NumPy array to PIL Image
|
|
image_pil = Image.fromarray(image_array)
|
|
|
|
# Create a drawing object
|
|
draw = ImageDraw.Draw(image_pil)
|
|
|
|
if font_path is not None:
|
|
font = ImageFont.truetype(font_path, font_size)
|
|
else:
|
|
try:
|
|
# Load the font
|
|
font = ImageFont.truetype(
|
|
"/usr/share/fonts/truetype/liberation/LiberationMono-Regular.ttf", font_size)
|
|
except:
|
|
font = ImageFont.load_default()
|
|
|
|
# Draw the text on the image
|
|
draw.text(position, text, font=font, fill=text_color)
|
|
|
|
# Convert the PIL Image back to NumPy array
|
|
modified_image_array = np.array(image_pil)
|
|
|
|
return modified_image_array
|
|
|
|
|
|
def add_text_to_video(video_path, prompt):
|
|
|
|
outputs_with_overlay = []
|
|
with open(video_path, "rb") as f:
|
|
vr = VideoReader(f, ctx=cpu(0))
|
|
|
|
for i in range(len(vr)):
|
|
frame = vr[i]
|
|
frame = add_text_to_image(frame, prompt, position=(
|
|
10, 10), font_size=15, text_color=(255, 0, 0),)
|
|
outputs_with_overlay.append(frame)
|
|
outputs = outputs_with_overlay
|
|
video_path = video_path.replace("mp4", "gif")
|
|
imageio.mimsave(video_path, outputs, duration=100, loop=0)
|
|
|
|
|
|
def save_videos_grid(videos: torch.Tensor, path: str, rescale=False, n_rows=4, fps=30, prompt=None):
|
|
videos = rearrange(videos, "b c t h w -> t b c h w")
|
|
outputs = []
|
|
for x in videos:
|
|
x = torchvision.utils.make_grid(x, nrow=n_rows)
|
|
x = x.transpose(0, 1).transpose(1, 2).squeeze(-1)
|
|
if rescale:
|
|
x = (x + 1.0) / 2.0 # -1,1 -> 0,1
|
|
x = (x * 255).numpy().astype(np.uint8)
|
|
outputs.append(x)
|
|
|
|
os.makedirs(os.path.dirname(path), exist_ok=True)
|
|
|
|
if prompt is not None:
|
|
outputs_with_overlay = []
|
|
for frame in outputs:
|
|
frame_out = add_text_to_image(
|
|
frame, prompt, position=(10, 10), font_size=10, text_color=(255, 0, 0),)
|
|
outputs_with_overlay.append(frame_out)
|
|
outputs = outputs_with_overlay
|
|
imageio.mimsave(path, outputs, duration=round(1/fps*1000), loop=0)
|
|
# iio.imwrite(path, outputs)
|
|
# optimize(path)
|
|
|
|
|
|
def set_channel_pos(data, shape_dict, channel_pos):
|
|
|
|
assert data.ndim == 5 or data.ndim == 4
|
|
batch_dim = data.shape[0]
|
|
frame_dim = shape_dict["frame_dim"]
|
|
channel_dim = shape_dict["channel_dim"]
|
|
width_dim = shape_dict["width_dim"]
|
|
height_dim = shape_dict["height_dim"]
|
|
|
|
assert batch_dim != frame_dim
|
|
assert channel_dim != frame_dim
|
|
assert channel_dim != batch_dim
|
|
|
|
video_shape = list(data.shape)
|
|
batch_pos = video_shape.index(batch_dim)
|
|
|
|
channel_pos = video_shape.index(channel_dim)
|
|
w_pos = video_shape.index(width_dim)
|
|
h_pos = video_shape.index(height_dim)
|
|
if w_pos == h_pos:
|
|
video_shape[w_pos] = -1
|
|
h_pos = video_shape.index(height_dim)
|
|
pattern_order = {}
|
|
pattern_order[batch_pos] = "B"
|
|
pattern_order[channel_pos] = "C"
|
|
|
|
pattern_order[w_pos] = "W"
|
|
pattern_order[h_pos] = "H"
|
|
|
|
if data.ndim == 5:
|
|
frame_pos = video_shape.index(frame_dim)
|
|
pattern_order[frame_pos] = "F"
|
|
if channel_pos == channel_first:
|
|
pattern = " -> B F C W H"
|
|
else:
|
|
pattern = " -> B F W H C"
|
|
else:
|
|
if channel_pos == channel_first:
|
|
pattern = " -> B C W H"
|
|
else:
|
|
pattern = " -> B W H C"
|
|
pattern_input = [pattern_order[idx] for idx in range(data.ndim)]
|
|
pattern_input = " ".join(pattern_input)
|
|
pattern = pattern_input + pattern
|
|
data = rearrange(data, pattern)
|
|
|
|
|
|
def merge_first_two_dimensions(tensor):
|
|
dims = tensor.ndim
|
|
letters = []
|
|
for letter_idx in range(dims-2):
|
|
letters.append(chr(letter_idx+67))
|
|
latters_pattern = " ".join(letters)
|
|
tensor = rearrange(tensor, "A B "+latters_pattern +
|
|
" -> (A B) "+latters_pattern)
|
|
# TODO merging first two dimensions might be easier with reshape so no need to create letters
|
|
# should be 'tensor.view(*tensor.shape[:2], -1)'
|
|
return tensor
|
|
|
|
|
|
def apply_spatial_function_to_video_tensor(video, shape, func):
|
|
# TODO detect batch, frame, channel, width, and height
|
|
|
|
assert video.ndim == 5
|
|
batch_dim = shape["batch_dim"]
|
|
frame_dim = shape["frame_dim"]
|
|
channel_dim = shape["channel_dim"]
|
|
width_dim = shape["width_dim"]
|
|
height_dim = shape["height_dim"]
|
|
|
|
assert batch_dim != frame_dim
|
|
assert channel_dim != frame_dim
|
|
assert channel_dim != batch_dim
|
|
|
|
video_shape = list(video.shape)
|
|
batch_pos = video_shape.index(batch_dim)
|
|
frame_pos = video_shape.index(frame_dim)
|
|
channel_pos = video_shape.index(channel_dim)
|
|
w_pos = video_shape.index(width_dim)
|
|
h_pos = video_shape.index(height_dim)
|
|
if w_pos == h_pos:
|
|
video_shape[w_pos] = -1
|
|
h_pos = video_shape.index(height_dim)
|
|
pattern_order = {}
|
|
pattern_order[batch_pos] = "B"
|
|
pattern_order[channel_pos] = "C"
|
|
pattern_order[frame_pos] = "F"
|
|
pattern_order[w_pos] = "W"
|
|
pattern_order[h_pos] = "H"
|
|
pattern_order = sorted(pattern_order.items(), key=lambda x: x[1])
|
|
pattern_order = [x[0] for x in pattern_order]
|
|
input_pattern = " ".join(pattern_order)
|
|
video = rearrange(video, input_pattern+" -> (B F) C W H")
|
|
|
|
video = func(video)
|
|
video = rearrange(video, "(B F) C W H -> "+input_pattern, F=frame_dim)
|
|
return video
|
|
|
|
|
|
def dump_frames(videos, as_mosaik, storage_fol, save_image_kwargs):
|
|
|
|
# assume videos is in format B F C H W, range [0,1]
|
|
num_frames = videos.shape[1]
|
|
num_videos = videos.shape[0]
|
|
|
|
if videos.shape[2] != 3 and videos.shape[-1] == 3:
|
|
videos = rearrange(videos, "B F W H C -> B F C W H")
|
|
|
|
frame_counter = 0
|
|
if not isinstance(storage_fol, Path):
|
|
storage_fol = Path(storage_fol)
|
|
|
|
for frame_idx in range(num_frames):
|
|
print(f" Creating frame {frame_idx}")
|
|
batch_frame = videos[:, frame_idx, ...]
|
|
|
|
if as_mosaik:
|
|
filename = storage_fol / f"frame_{frame_counter:03d}.png"
|
|
save_image(batch_frame, fp=filename.as_posix(),
|
|
**save_image_kwargs)
|
|
frame_counter += 1
|
|
else:
|
|
for video_idx in range(num_videos):
|
|
frame = batch_frame[video_idx]
|
|
|
|
filename = storage_fol / f"frame_{frame_counter:03d}.png"
|
|
save_image(frame, fp=filename.as_posix(),
|
|
**save_image_kwargs)
|
|
frame_counter += 1
|
|
|
|
|
|
def gif_from_videos(videos):
|
|
|
|
assert videos.dim() == 5
|
|
assert videos.min() >= 0
|
|
assert videos.max() <= 1
|
|
gif_file = Path("tmp.gif").absolute()
|
|
|
|
with tempfile.TemporaryDirectory() as tmpdirname:
|
|
storage_fol = Path(tmpdirname)
|
|
nrows = min(4, videos.shape[0])
|
|
dump_frames(
|
|
videos=videos, storage_fol=storage_fol, as_mosaik=True, save_image_kwargs={"nrow": nrows})
|
|
cmd = f"ffmpeg -y -f image2 -framerate 4 -i {storage_fol / 'frame_%03d.png'} {gif_file.as_posix()}"
|
|
subprocess.check_call(
|
|
cmd, shell=True, stdout=subprocess.DEVNULL, stderr=subprocess.STDOUT)
|
|
return gif_file
|
|
|
|
|
|
|
|
def add_margin(pil_img, top, right, bottom, left, color):
|
|
width, height = pil_img.size
|
|
new_width = width + right + left
|
|
new_height = height + top + bottom
|
|
result = Image.new(pil_img.mode, (new_width, new_height), color)
|
|
result.paste(pil_img, (left, top))
|
|
return result
|
|
|
|
def resize_to_fit(image, size):
|
|
W, H = size
|
|
w, h = image.size
|
|
if H / h > W / w:
|
|
H_ = int(h * W / w)
|
|
W_ = W
|
|
else:
|
|
W_ = int(w * H / h)
|
|
H_ = H
|
|
return image.resize((W_, H_))
|
|
|
|
def pad_to_fit(image, size):
|
|
W, H = size
|
|
w, h = image.size
|
|
pad_h = (H - h) // 2
|
|
pad_w = (W - w) // 2
|
|
return add_margin(image, pad_h, pad_w, pad_h, pad_w, (0, 0, 0)) |