stablevideo/atlas_utils.py

from PIL import Image
from pathlib import Path
import scipy.interpolate
import torch
from torchvision import transforms
from torchvision.transforms.functional import crop
from tqdm import tqdm
import numpy as np
import cv2

from stablevideo.implicit_neural_networks import IMLP


def load_video(folder: str, resize=(432, 768), num_frames=70):
 resy, resx = resize
 folder = Path(folder)
 input_files = sorted(list(folder.glob("*.jpg")) + list(folder.glob("*.png")))[:num_frames]
 video_tensor = torch.zeros((len(input_files), 3, resy, resx))
 
 for i, file in enumerate(input_files):
 video_tensor[i] = transforms.ToTensor()(Image.open(str(file)).resize((resx, resy), Image.LANCZOS))
 
 return video_tensor


def load_neural_atlases_models(config):
 foreground_mapping = IMLP(
 input_dim=3,
 output_dim=2,
 hidden_dim=256,
 use_positional=False,
 num_layers=6,
 skip_layers=[],
 ).to(config["device"])

 background_mapping = IMLP(
 input_dim=3,
 output_dim=2,
 hidden_dim=256,
 use_positional=False,
 num_layers=4,
 skip_layers=[],
 ).to(config["device"])

 foreground_atlas_model = IMLP(
 input_dim=2,
 output_dim=3,
 hidden_dim=256,
 use_positional=True,
 positional_dim=10,
 num_layers=8,
 skip_layers=[4, 7],
 ).to(config["device"])

 background_atlas_model = IMLP(
 input_dim=2,
 output_dim=3,
 hidden_dim=256,
 use_positional=True,
 positional_dim=10,
 num_layers=8,
 skip_layers=[4, 7],
 ).to(config["device"])

 alpha_model = IMLP(
 input_dim=3,
 output_dim=1,
 hidden_dim=256,
 use_positional=True,
 positional_dim=5,
 num_layers=8,
 skip_layers=[],
 ).to(config["device"])

 checkpoint = torch.load(config["checkpoint_path"], map_location=torch.device('cpu'))
 foreground_mapping.load_state_dict(checkpoint["model_F_mapping1_state_dict"])
 background_mapping.load_state_dict(checkpoint["model_F_mapping2_state_dict"])
 foreground_atlas_model.load_state_dict(checkpoint["F_atlas_state_dict"])
 background_atlas_model.load_state_dict(checkpoint["F_atlas_state_dict"])
 alpha_model.load_state_dict(checkpoint["model_F_alpha_state_dict"])

 foreground_mapping = foreground_mapping.eval().requires_grad_(False)
 background_mapping = background_mapping.eval().requires_grad_(False)
 foreground_atlas_model = foreground_atlas_model.eval().requires_grad_(False)
 background_atlas_model = background_atlas_model.eval().requires_grad_(False)
 alpha_model = alpha_model.eval().requires_grad_(False)

 return foreground_mapping, background_mapping, foreground_atlas_model, background_atlas_model, alpha_model


@torch.no_grad()
def get_frames_data(config, foreground_mapping, background_mapping, alpha_model):
 max_size = max(config["resx"], config["resy"])
 normalizing_factor = torch.tensor([max_size / 2, max_size / 2, config["maximum_number_of_frames"] / 2])
 background_uv_values = torch.zeros(
 size=(config["maximum_number_of_frames"], config["resy"], config["resx"], 2), device=config["device"]
 )
 foreground_uv_values = torch.zeros(
 size=(config["maximum_number_of_frames"], config["resy"], config["resx"], 2), device=config["device"]
 )
 alpha = torch.zeros(
 size=(config["maximum_number_of_frames"], config["resy"], config["resx"], 1), device=config["device"]
 )

 for frame in tqdm(range(config["maximum_number_of_frames"]), leave=False):
 indices = get_grid_indices(0, 0, config["resy"], config["resx"], t=torch.tensor(frame))

 normalized_chunk = (indices / normalizing_factor - 1).to(config["device"])

 # get the atlas UV coordinates from the two mapping networks;
 with torch.no_grad():
 current_background_uv_values = background_mapping(normalized_chunk)
 current_foreground_uv_values = foreground_mapping(normalized_chunk)
 current_alpha = alpha_model(normalized_chunk)

 background_uv_values[frame, indices[:, 1], indices[:, 0]] = current_background_uv_values * 0.5 - 0.5
 foreground_uv_values[frame, indices[:, 1], indices[:, 0]] = current_foreground_uv_values * 0.5 + 0.5
 current_alpha = 0.5 * (current_alpha + 1.0)
 current_alpha = 0.99 * current_alpha + 0.001
 alpha[frame, indices[:, 1], indices[:, 0]] = current_alpha
 # config["return_atlas_alpha"] = True
 if config["return_atlas_alpha"]: # this should take a few minutes
 foreground_atlas_alpha = torch.zeros(
 size=(
 config["maximum_number_of_frames"],
 config["grid_atlas_resolution"],
 config["grid_atlas_resolution"],
 1,
 ),
 )
 # foreground_uv_values: 70 x 432 x 768 x 2
 foreground_uv_values_grid = foreground_uv_values * config["grid_atlas_resolution"]
 # indices: 4000000 x 2
 indices = get_grid_indices(0, 0, config["grid_atlas_resolution"], config["grid_atlas_resolution"])
 for frame in tqdm(range(config["maximum_number_of_frames"]), leave=False):
 interpolated = scipy.interpolate.griddata(
 foreground_uv_values_grid[frame].reshape(-1, 2).cpu().numpy(), # 432 x 768 x 2 -> -1 x 2
 alpha[frame]
 .reshape(
 -1,
 )
 .cpu()
 .numpy(),
 indices.reshape(-1, 2).cpu().numpy(),
 method="linear",
 ).reshape(config["grid_atlas_resolution"], config["grid_atlas_resolution"], 1)
 foreground_atlas_alpha[frame] = torch.from_numpy(interpolated)
 foreground_atlas_alpha[foreground_atlas_alpha.isnan()] = 0.0
 foreground_atlas_alpha = (
 torch.median(foreground_atlas_alpha, dim=0, keepdim=True).values.to(config["device"]).permute(0, 3, 2, 1)
 )
 else:
 foreground_atlas_alpha = None
 return background_uv_values, foreground_uv_values, alpha.permute(0, 3, 1, 2), foreground_atlas_alpha


@torch.no_grad()
def reconstruct_video_layer(uv_values, atlas_model):
 t, h, w, _ = uv_values.shape
 reconstruction = torch.zeros(size=(t, h, w, 3), device=uv_values.device)
 for frame in range(t):
 rgb = (atlas_model(uv_values[frame].reshape(-1, 2)) + 1) * 0.5
 reconstruction[frame] = rgb.reshape(h, w, 3)
 return reconstruction.permute(0, 3, 1, 2)


@torch.no_grad()
def create_uv_mask(config, mapping_model, min_u, min_v, max_u, max_v, uv_shift=-0.5, resolution_shift=1):
 max_size = max(config["resx"], config["resy"])
 normalizing_factor = torch.tensor([max_size / 2, max_size / 2, config["maximum_number_of_frames"] / 2])
 resolution = config["grid_atlas_resolution"]
 uv_mask = torch.zeros(size=(resolution, resolution), device=config["device"])

 for frame in tqdm(range(config["maximum_number_of_frames"]), leave=False):
 indices = get_grid_indices(0, 0, config["resy"], config["resx"], t=torch.tensor(frame))
 for chunk in indices.split(50000, dim=0):
 normalized_chunk = (chunk / normalizing_factor - 1).to(config["device"])

 # get the atlas UV coordinates from the two mapping networks;
 with torch.no_grad():
 uv_values = mapping_model(normalized_chunk)
 uv_values = uv_values * 0.5 + uv_shift
 uv_values = ((uv_values + resolution_shift) * resolution).clip(0, resolution - 1)

 uv_mask[uv_values[:, 1].floor().long(), uv_values[:, 0].floor().long()] = 1
 uv_mask[uv_values[:, 1].floor().long(), uv_values[:, 0].ceil().long()] = 1
 uv_mask[uv_values[:, 1].ceil().long(), uv_values[:, 0].floor().long()] = 1
 uv_mask[uv_values[:, 1].ceil().long(), uv_values[:, 0].ceil().long()] = 1

 uv_mask = crop(uv_mask.unsqueeze(0).unsqueeze(0), min_v, min_u, max_v, max_u)
 return uv_mask.detach().cpu() # shape [1, 1, resolution, resolution]


@torch.no_grad()
def get_high_res_atlas(atlas_model, min_v, min_u, max_v, max_u, resolution, device="cuda", layer="background"):
 inds_grid = get_grid_indices(0, 0, resolution, resolution)
 inds_grid_chunks = inds_grid.split(50000, dim=0)
 if layer == "background":
 shift = -1
 else:
 shift = 0

 rendered_atlas = torch.zeros((resolution, resolution, 3)).to(device) # resy, resx, 3
 with torch.no_grad():
 # reconstruct image row by row
 for chunk in inds_grid_chunks:
 normalized_chunk = torch.stack(
 [
 (chunk[:, 0] / resolution) + shift,
 (chunk[:, 1] / resolution) + shift,
 ],
 dim=-1,
 ).to(device)

 rgb_output = atlas_model(normalized_chunk)
 rendered_atlas[chunk[:, 1], chunk[:, 0], :] = rgb_output
 # move colors to RGB color domain (0,1)
 rendered_atlas = 0.5 * (rendered_atlas + 1)
 rendered_atlas = rendered_atlas.permute(2, 0, 1).unsqueeze(0) # shape (1, 3, resy, resx)
 cropped_atlas = crop(
 rendered_atlas,
 min_v,
 min_u,
 max_v,
 max_u,
 )

 return cropped_atlas


def get_grid_indices(x_start, y_start, h_crop, w_crop, t=None):
 crop_indices = torch.meshgrid(torch.arange(w_crop) + x_start, torch.arange(h_crop) + y_start)
 crop_indices = torch.stack(crop_indices, dim=-1)
 crop_indices = crop_indices.reshape(h_crop * w_crop, crop_indices.shape[-1])
 if t is not None:
 crop_indices = torch.cat([crop_indices, t.repeat(h_crop * w_crop, 1)], dim=1)
 return crop_indices


def get_atlas_crops(uv_values, grid_atlas, augmentation=None):
 if len(uv_values.shape) == 3:
 dims = [0, 1]
 elif len(uv_values.shape) == 4:
 dims = [0, 1, 2]
 else:
 raise ValueError("uv_values should be of shape of len 3 or 4")

 min_u, min_v = uv_values.amin(dim=dims).long()
 max_u, max_v = uv_values.amax(dim=dims).ceil().long()
 # min_u, min_v = uv_values.min(dim=0).values
 # max_u, max_v = uv_values.max(dim=0).values

 h_v = max_v - min_v
 w_u = max_u - min_u
 atlas_crop = crop(grid_atlas, min_v, min_u, h_v, w_u)
 if augmentation is not None:
 atlas_crop = augmentation(atlas_crop)
 return atlas_crop, torch.stack([min_u, min_v]), torch.stack([max_u, max_v])


def get_random_crop_params(input_size, output_size):
 w, h = input_size
 th, tw = output_size

 if h + 1 < th or w + 1 < tw:
 raise ValueError(f"Required crop size {(th, tw)} is larger then input image size {(h, w)}")

 if w == tw and h == th:
 return 0, 0, h, w

 i = torch.randint(0, h - th + 1, size=(1,)).item()
 j = torch.randint(0, w - tw + 1, size=(1,)).item()
 return i, j, th, tw


def get_masks_boundaries(alpha_video, border=20, threshold=0.95, min_crop_size=2 ** 7 + 1):
 resy, resx = alpha_video.shape[-2:]
 num_frames = alpha_video.shape[0]
 masks_borders = torch.zeros((num_frames, 4), dtype=torch.int64)
 for i, file in enumerate(range(num_frames)):
 mask_im = alpha_video[i]
 mask_im[mask_im >= threshold] = 1
 mask_im[mask_im < threshold] = 0
 all_ones = mask_im.squeeze().nonzero()
 min_y, min_x = torch.maximum(all_ones.min(dim=0).values - border, torch.tensor([0, 0]))
 max_y, max_x = torch.minimum(all_ones.max(dim=0).values + border, torch.tensor([resy, resx]))
 h = max_y - min_y
 w = max_x - min_x
 if h < min_crop_size:
 pad = min_crop_size - h
 if max_y + pad > resy:
 min_y -= pad
 else:
 max_y += pad
 h = max_y - min_y
 if w < min_crop_size:
 pad = min_crop_size - w
 if max_x + pad > resx:
 min_x -= pad
 else:
 max_x += pad
 w = max_x - min_x
 masks_borders[i] = torch.tensor([min_y, min_x, h, w])
 return masks_borders


def get_atlas_bounding_box(mask_boundaries, grid_atlas, video_uvs):
 min_uv = torch.tensor(grid_atlas.shape[-2:], device=video_uvs.device)
 max_uv = torch.tensor([0, 0], device=video_uvs.device)
 for boundary, frame in zip(mask_boundaries, video_uvs):
 cropped_uvs = crop(frame.permute(2, 0, 1).unsqueeze(0), *list(boundary)) # 1,2,h,w
 min_uv = torch.minimum(cropped_uvs.amin(dim=[0, 2, 3]), min_uv).floor().int()
 max_uv = torch.maximum(cropped_uvs.amax(dim=[0, 2, 3]), max_uv).ceil().int()

 hw = max_uv - min_uv
 crop_data = [*list(min_uv)[::-1], *list(hw)[::-1]]
 return crop(grid_atlas, *crop_data), crop_data


def tensor2im(input_image, imtype=np.uint8):
 if not isinstance(input_image, np.ndarray):
 if isinstance(input_image, torch.Tensor): # get the data from a variable
 image_tensor = input_image.data
 else:
 return input_image
 image_numpy = image_tensor[0].clamp(0.0, 1.0).cpu().float().numpy() # convert it into a numpy array
 image_numpy = np.transpose(image_numpy, (1, 2, 0)) * 255.0 # post-processing: tranpose and scaling
 else: # if it is a numpy array, do nothing
 image_numpy = input_image
 return image_numpy.astype(imtype)