stablevideo/atlas_data.py

import random
import json

import torch
import torch.nn.functional as F
from torch.utils.data import Dataset
from torchvision import transforms
from torchvision.transforms.functional import crop

from stablevideo.atlas_utils import (
 load_neural_atlases_models,
 get_frames_data,
 get_high_res_atlas,
 get_atlas_crops,
 reconstruct_video_layer,
 create_uv_mask,
 get_masks_boundaries,
 get_random_crop_params,
 get_atlas_bounding_box,
 load_video
)

class AtlasData():
 def __init__(self, video_name):
 with open(f"data/{video_name}/config.json", "r") as f:
 json_dict = json.load(f)
 try:
 maximum_number_of_frames = json_dict["number_of_frames"]
 except:
 maximum_number_of_frames = json_dict["maximum_number_of_frames"]

 config = {
 "device": "cpu",
 "checkpoint_path": f"data/{video_name}/checkpoint.ckpt",
 "resx": json_dict["resx"],
 "resy": json_dict["resy"],
 "maximum_number_of_frames": maximum_number_of_frames,
 "return_atlas_alpha": False,
 "grid_atlas_resolution": 2000,
 "num_scales": 7,
 "masks_border_expansion": 30,
 "mask_alpha_threshold": 0.99, # 0.95
 "align_corners": False
 }
 self.config = config
 self.device = config["device"]

 self.min_size = min(self.config["resx"], self.config["resy"])
 self.max_size = max(self.config["resx"], self.config["resy"])
 data_folder = f"data/{video_name}/{video_name}"
 self.original_video = load_video(
 data_folder,
 resize=(self.config["resy"], self.config["resx"]),
 num_frames=self.config["maximum_number_of_frames"],
 )
 self.original_video = self.original_video.to(self.device) # tensor

 (
 foreground_mapping,
 background_mapping,
 foreground_atlas_model,
 background_atlas_model,
 alpha_model,
 ) = load_neural_atlases_models(config)
 (
 original_background_all_uvs,
 original_foreground_all_uvs,
 self.all_alpha,
 foreground_atlas_alpha,
 ) = get_frames_data(
 config,
 foreground_mapping,
 background_mapping,
 alpha_model,
 )

 self.background_reconstruction = reconstruct_video_layer(original_background_all_uvs, background_atlas_model)
 # using original video for the foreground layer
 self.foreground_reconstruction = self.original_video * self.all_alpha

 (
 self.background_all_uvs,
 self.scaled_background_uvs,
 self.background_min_u,
 self.background_min_v,
 self.background_max_u,
 self.background_max_v,
 ) = self.preprocess_uv_values(
 original_background_all_uvs, config["grid_atlas_resolution"], device=self.device, layer="background"
 )
 (
 self.foreground_all_uvs,
 self.scaled_foreground_uvs,
 self.foreground_min_u,
 self.foreground_min_v,
 self.foreground_max_u,
 self.foreground_max_v,
 ) = self.preprocess_uv_values(
 original_foreground_all_uvs, config["grid_atlas_resolution"], device=self.device, layer="foreground"
 )

 self.background_uv_mask = create_uv_mask(
 config,
 background_mapping,
 self.background_min_u,
 self.background_min_v,
 self.background_max_u,
 self.background_max_v,
 uv_shift=-0.5,
 resolution_shift=1,
 )
 self.foreground_uv_mask = create_uv_mask(
 config,
 foreground_mapping,
 self.foreground_min_u,
 self.foreground_min_v,
 self.foreground_max_u,
 self.foreground_max_v,
 uv_shift=0.5,
 resolution_shift=0,
 )
 self.background_grid_atlas = get_high_res_atlas(
 background_atlas_model,
 self.background_min_v,
 self.background_min_u,
 self.background_max_v,
 self.background_max_u,
 config["grid_atlas_resolution"],
 device=config["device"],
 layer="background",
 )
 self.foreground_grid_atlas = get_high_res_atlas(
 foreground_atlas_model,
 self.foreground_min_v,
 self.foreground_min_u,
 self.foreground_max_v,
 self.foreground_max_u,
 config["grid_atlas_resolution"],
 device=config["device"],
 layer="foreground",
 )
 if config["return_atlas_alpha"]:
 self.foreground_atlas_alpha = foreground_atlas_alpha # used for visualizations
 self.cnn_min_crop_size = 2 ** self.config["num_scales"] + 1

 self.mask_boundaries = get_masks_boundaries(
 alpha_video=self.all_alpha.cpu(),
 border=self.config["masks_border_expansion"],
 threshold=self.config["mask_alpha_threshold"],
 min_crop_size=self.cnn_min_crop_size,
 )
 self.cropped_foreground_atlas, self.foreground_atlas_bbox = get_atlas_bounding_box(
 self.mask_boundaries, self.foreground_grid_atlas, self.foreground_all_uvs
 )

 self.step = -1
 self.edited_atlas_dict, self.edit_dict, self.uv_mask = {}, {}, {}
 
 @staticmethod
 def preprocess_uv_values(layer_uv_values, resolution, device="cuda", layer="background"):
 if layer == "background":
 shift = 1
 else:
 shift = 0
 uv_values = (layer_uv_values + shift) * resolution
 min_u, min_v = uv_values.reshape(-1, 2).min(dim=0).values.long()
 uv_values -= torch.tensor([min_u, min_v], device=device)
 max_u, max_v = uv_values.reshape(-1, 2).max(dim=0).values.ceil().long()

 edge_size = torch.tensor([max_u, max_v], device=device)
 scaled_uv_values = ((uv_values.reshape(-1, 2) / edge_size) * 2 - 1).unsqueeze(1).unsqueeze(0)

 return uv_values, scaled_uv_values, min_u, min_v, max_u, max_v

 def get_random_crop_data(self, crop_size):
 t = random.randint(0, self.config["maximum_number_of_frames"] - 1)
 y_start, x_start, h_crop, w_crop = get_random_crop_params((self.config["resx"], self.config["resy"]), crop_size)
 return y_start, x_start, h_crop, w_crop, t

 def get_global_crops_multi(self, keyframes, res):
 foreground_atlas_crops = []
 background_atlas_crops = []
 foreground_uvs = []
 background_uvs = []
 background_alpha_crops = []
 foreground_alpha_crops = []
 original_background_crops = []
 original_foreground_crops = []
 output_dict = {}

 self.config["crops_min_cover"] = 0.95
 self.config["grid_atlas_resolution"] = res
 
 for cur_frame in keyframes:
 y_start, x_start, frame_h, frame_w = self.mask_boundaries[cur_frame].tolist()
 crop_size = (
 max(
 random.randint(round(self.config["crops_min_cover"] * frame_h), frame_h),
 self.cnn_min_crop_size,
 ),
 max(
 random.randint(round(self.config["crops_min_cover"] * frame_w), frame_w),
 self.cnn_min_crop_size,
 ),
 )
 y_crop, x_crop, h_crop, w_crop = get_random_crop_params((frame_w, frame_h), crop_size)

 foreground_uv = self.foreground_all_uvs[
 cur_frame,
 y_start + y_crop : y_start + y_crop + h_crop,
 x_start + x_crop : x_start + x_crop + w_crop,
 ]
 alpha = self.all_alpha[
 [cur_frame],
 :,
 y_start + y_crop : y_start + y_crop + h_crop,
 x_start + x_crop : x_start + x_crop + w_crop,
 ]

 original_foreground_crop = self.foreground_reconstruction[
 [cur_frame],
 :,
 y_start + y_crop : y_start + y_crop + h_crop,
 x_start + x_crop : x_start + x_crop + w_crop,
 ]

 foreground_alpha_crops.append(alpha)
 foreground_uvs.append(foreground_uv) # not scaled
 original_foreground_crops.append(original_foreground_crop)

 foreground_max_vals = torch.tensor(
 [self.config["grid_atlas_resolution"]] * 2, device=self.device, dtype=torch.long
 )
 foreground_min_vals = torch.tensor([0] * 2, device=self.device, dtype=torch.long)
 for uv_values in foreground_uvs:
 min_uv = uv_values.amin(dim=[0, 1]).long()
 max_uv = uv_values.amax(dim=[0, 1]).ceil().long()
 foreground_min_vals = torch.minimum(foreground_min_vals, min_uv)
 foreground_max_vals = torch.maximum(foreground_max_vals, max_uv)
 
 h_v = foreground_max_vals[1] - foreground_min_vals[1]
 w_u = foreground_max_vals[0] - foreground_min_vals[0]
 foreground_atlas_crop = crop(
 self.foreground_grid_atlas,
 foreground_min_vals[1],
 foreground_min_vals[0],
 h_v,
 w_u,
 )
 foreground_atlas_crops.append(foreground_atlas_crop)
 
 for i, uv_values in enumerate(foreground_uvs):
 foreground_uvs[i] = (
 2 * (uv_values - foreground_min_vals) / (foreground_max_vals - foreground_min_vals) - 1
 ).unsqueeze(0)


 crop_size = (
 random.randint(round(self.config["crops_min_cover"] * self.min_size), self.min_size),
 random.randint(round(self.config["crops_min_cover"] * self.max_size), self.max_size),
 )
 crop_data = self.get_random_crop_data(crop_size)
 y, x, h, w, _ = crop_data
 background_uv = self.background_all_uvs[keyframes, y : y + h, x : x + w]
 original_background_crop = self.background_reconstruction[
 keyframes, :, y : y + h, x : x + w
 ]
 alpha = self.all_alpha[keyframes, :, y : y + h, x : x + w]


 original_background_crops = [el.unsqueeze(0) for el in original_background_crop]
 background_alpha_crops = [el.unsqueeze(0) for el in alpha]

 background_atlas_crop, background_min_vals, background_max_vals = get_atlas_crops(
 background_uv,
 self.background_grid_atlas,
 )
 background_uv = 2 * (background_uv - background_min_vals) / (background_max_vals - background_min_vals) - 1
 background_atlas_crops = [el.unsqueeze(0) for el in background_atlas_crop]
 background_uvs = [el.unsqueeze(0) for el in background_uv]

 output_dict["foreground_alpha"] = foreground_alpha_crops
 output_dict["foreground_uvs"] = foreground_uvs
 output_dict["original_foreground_crops"] = original_foreground_crops
 output_dict["foreground_atlas_crops"] = foreground_atlas_crops

 output_dict["background_alpha"] = background_alpha_crops
 output_dict["background_uvs"] = background_uvs
 output_dict["original_background_crops"] = original_background_crops
 output_dict["background_atlas_crops"] = background_atlas_crops
 return output_dict