| 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())
|
|
|
