|
|
import os |
|
|
import torchvision |
|
|
from PIL import Image, ImageDraw |
|
|
import imageio |
|
|
import cv2 |
|
|
import torch |
|
|
import torch.nn.functional as F |
|
|
import numpy as np |
|
|
import zipfile |
|
|
|
|
|
_gauss_mask_cache = {} |
|
|
|
|
|
|
|
|
def load_gauss_mask(mask_path): |
|
|
if not mask_path: |
|
|
return None |
|
|
abs_path = os.path.abspath(mask_path) |
|
|
mask = _gauss_mask_cache.get(abs_path) |
|
|
if mask is None: |
|
|
mask = torch.load(abs_path, weights_only=False, map_location="cpu") |
|
|
if not torch.is_tensor(mask): |
|
|
mask = torch.tensor(mask) |
|
|
_gauss_mask_cache[abs_path] = mask |
|
|
return mask |
|
|
|
|
|
|
|
|
def apply_alpha_shift(latents, gauss_mask, shift_mean): |
|
|
if gauss_mask is None: |
|
|
return latents |
|
|
mask = gauss_mask |
|
|
if mask.ndim == 3: |
|
|
mask = mask.unsqueeze(0).unsqueeze(0) |
|
|
elif mask.ndim == 4: |
|
|
if mask.shape[0] != 1: |
|
|
mask = mask.unsqueeze(0) |
|
|
if mask.shape[1] != 1: |
|
|
mask = mask.unsqueeze(1) |
|
|
elif mask.ndim != 5: |
|
|
return latents |
|
|
|
|
|
mask = mask.to(device=latents.device, dtype=latents.dtype) |
|
|
target_shape = latents.shape[2:] |
|
|
if mask.shape[-3:] != target_shape: |
|
|
mask = F.interpolate(mask, size=target_shape, mode="trilinear", align_corners=False) |
|
|
shift_mean = torch.as_tensor(shift_mean, dtype=latents.dtype, device=latents.device) |
|
|
return latents + (1.0 - mask) * shift_mean |
|
|
|
|
|
def render_video(tensor_fgr, |
|
|
tensor_pha, |
|
|
nrow=8, |
|
|
normalize=True, |
|
|
value_range=(-1, 1)): |
|
|
def to_tensor(arr_list): |
|
|
tensor_list= [torch.from_numpy(arr).float().div_(127.5).sub_(1) for arr in arr_list] |
|
|
tensor_list = torch.stack(tensor_list, dim = 0).permute(3,0,1,2).unsqueeze(0) |
|
|
return tensor_list |
|
|
|
|
|
if not torch.is_tensor(tensor_fgr): |
|
|
tensor_fgr = to_tensor(tensor_fgr) |
|
|
if not torch.is_tensor(tensor_pha): |
|
|
tensor_pha = to_tensor(tensor_pha) |
|
|
|
|
|
tensor_fgr = tensor_fgr.clamp(min(value_range), max(value_range)) |
|
|
tensor_fgr = torch.stack([ |
|
|
torchvision.utils.make_grid( |
|
|
u, nrow=nrow, normalize=normalize, value_range=value_range) |
|
|
for u in tensor_fgr.unbind(2) |
|
|
], |
|
|
dim=1).permute(1, 2, 3, 0) |
|
|
tensor_fgr = (tensor_fgr * 255).type(torch.uint8).cpu() |
|
|
|
|
|
tensor_pha = tensor_pha.clamp(min(value_range), max(value_range)) |
|
|
tensor_pha = torch.stack([ |
|
|
torchvision.utils.make_grid( |
|
|
u, nrow=nrow, normalize=normalize, value_range=value_range) |
|
|
for u in tensor_pha.unbind(2) |
|
|
], |
|
|
dim=1).permute(1, 2, 3, 0) |
|
|
tensor_pha = (tensor_pha * 255).type(torch.uint8).cpu() |
|
|
|
|
|
frames = [] |
|
|
frames_fgr = [] |
|
|
frames_pha = [] |
|
|
for frame_fgr, frame_pha in zip(tensor_fgr.numpy(), tensor_pha.numpy()): |
|
|
if frame_pha.shape[-1] == 1: |
|
|
frame_pha = frame_pha[:,:,0] |
|
|
else: |
|
|
frame_pha = (0.0 + frame_pha[:,:,0:1] + frame_pha[:,:,1:2] + frame_pha[:,:,2:3]) / 3. |
|
|
frame = np.concatenate([frame_fgr[:,:,::-1], frame_pha.astype(np.uint8)], axis=2) |
|
|
frames.append(frame) |
|
|
frames_fgr.append(frame_fgr) |
|
|
frames_pha.append(frame_pha) |
|
|
|
|
|
def create_checkerboard(size=30, pattern_size=(830, 480), color1=(140, 140, 140), color2=(113, 113, 113)): |
|
|
img = Image.new('RGB', (pattern_size[0], pattern_size[1]), color1) |
|
|
draw = ImageDraw.Draw(img) |
|
|
for i in range(0, pattern_size[0], size): |
|
|
for j in range(0, pattern_size[1], size): |
|
|
if (i + j) // size % 2 == 0: |
|
|
draw.rectangle([i, j, i+size, j+size], fill=color2) |
|
|
return img |
|
|
|
|
|
def blender_background(frame_rgba, checkerboard): |
|
|
alpha_channel = frame_rgba[:, :, 3:] / 255. |
|
|
checkerboard = np.array(checkerboard) |
|
|
checkerboard = cv2.resize(checkerboard, (frame_rgba.shape[1], frame_rgba.shape[0])) |
|
|
|
|
|
frame_rgb = frame_rgba[:, :, :3] * alpha_channel + checkerboard * (1-alpha_channel) |
|
|
return frame_rgb.astype(np.uint8)[:,:,::-1] |
|
|
|
|
|
checkerboard = create_checkerboard() |
|
|
video_checkerboard = [torch.from_numpy(blender_background(f, checkerboard).copy()).float().div_(127.5).sub_(1) for f in frames] |
|
|
video_checkerboard = torch.stack(video_checkerboard ).permute(3, 0, 1, 2) |
|
|
return video_checkerboard, frames |
|
|
|
|
|
def from_BRGA_numpy_to_RGBA_torch(video): |
|
|
video = [torch.from_numpy(f.copy()).float().div_(127.5).sub_(1) for f in video] |
|
|
video = torch.stack(video).permute(3, 0, 1, 2) |
|
|
video[[0, 2], ...] = video[[2, 0], ...] |
|
|
return video |
|
|
|
|
|
def write_zip_file(zip_path, frames): |
|
|
|
|
|
with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf: |
|
|
for idx, img in enumerate(frames): |
|
|
success, buffer = cv2.imencode(".png", img) |
|
|
if not success: |
|
|
print(f"Failed to encode image {idx}, skipping...") |
|
|
continue |
|
|
|
|
|
filename = f"img_{idx:03d}.png" |
|
|
zipf.writestr(filename, buffer.tobytes()) |
|
|
|
