import argparse import os import tqdm import cv2 import scipy.interpolate as spi import numpy as np import plotly.graph_objects as go import plotly.express as ex import torch from torchvision.transforms import Resize from torchvision.transforms.functional import InterpolationMode os.environ["CUDA_VISIBLE_DEVICES"] = "0" def get_img_idxs(): img_idxs = dict() img_coords = dict() uv_masks = dict() uv_mask = np.load('custom_data/garment_uv_mask_full.npy') coords = np.transpose(np.stack(np.nonzero((1 - uv_mask))), (1,0)) img_idx = np.transpose(np.stack(np.nonzero((uv_mask))), (1,0)) img_idxs[512] = img_idx img_coords[512] = coords uv_masks[512] = uv_mask.astype(np.int32) for res in [32, 64, 128, 256]: down = resizers_nearest[res](torch.from_numpy(uv_mask[None,:,:]).to(device)).cpu().numpy() down = np.squeeze(down.astype(bool)) coords = np.transpose(np.stack(np.nonzero((1 - down))), (1,0)) img_idx = np.transpose(np.stack(np.nonzero((down))), (1,0)) img_idxs[res] = img_idx img_coords[res] = coords uv_masks[res] = down.astype(np.int32) uv_masks_body = dict() uv_mask_body = np.load('custom_utils/mask_256x512.npy') uv_mask_body = cv2.dilate(uv_mask_body.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2) uv_masks_body[512] = uv_mask_body for res in [32, 64, 128, 256]: down = resizers_nearest[res](torch.from_numpy(uv_mask_body[None,:,:]).to(device)).cpu().numpy() down = np.squeeze(down.astype(bool)) uv_masks_body[res] = down return img_idxs, img_coords, uv_masks, uv_masks_body def fill(img, mask): coords = img_coords[img.shape[0]] img_idx = img_idxs[img.shape[0]] # coords = np.transpose(np.stack(np.nonzero((1 - uv_masks[img.shape[0]]))), (1,0)) # img_idx = np.transpose(np.stack(np.nonzero((uv_masks[img.shape[0]]))), (1,0)) img_data = img[img_idx[:,0], img_idx[:,1], :] interpolator = spi.NearestNDInterpolator(x=img_idx, y=img_data) img_2 = interpolator(coords) res = np.zeros((512,256,3), dtype=np.float32) res[coords[:,0], coords[:,1], :] = img_2 res[img_idx[:,0], img_idx[:,1], :] = img_data return res def fill_label(img, mask): coords = img_coords[img.shape[0]] img_idx = img_idxs[img.shape[0]] img_data = img[img_idx[:,0], img_idx[:,1]] interpolator = spi.NearestNDInterpolator(x=img_idx, y=img_data) img_2 = interpolator(coords) res = np.zeros((512,256), dtype=np.float32) res[coords[:,0], coords[:,1]] = img_2 res[img_idx[:,0], img_idx[:,1]] = img_data return res def generate_downscaled_uv_map(folder, id, size): img = np.load(os.path.join(folder, 'uv_mask', id + '.npz')); img = img['uv_mask'].astype(np.float32) down = tf.image.resize(img[:,:,tf.newaxis], size=(size, int(size/2)), method=tf.image.ResizeMethod.NEAREST_NEIGHBOR) down = np.squeeze(down.numpy().astype(bool)) np.savez_compressed(os.path.join(folder, 'uv_mask' + '_' + str(size), id + '.npz'), uv_mask = down) class UVMaskStore(): def __init__(self, folder, size): self.size = size def __getitem__(self, id) -> dict: return uv_masks[self.size] def downscale_uv_maps_posed_body(folder, type, sizes): for size in sizes: os.makedirs(os.path.join(folder, type + '_' + str(size)), exist_ok=True) frames = os.listdir(os.path.join(folder, type))[int(args.start):int(args.end)] for frame in tqdm.tqdm(frames): img = np.load(os.path.join(folder, type, frame)); img = img['uv_posed_body_no_orient'] img = fill(img, uv_masks_body[512]) for size in sizes: mask = uv_masks_body[size] down = resizers[size](torch.from_numpy(img.transpose((2,0,1))).to(device)).cpu().numpy() down = down * mask[None,:,:] np.savez_compressed(os.path.join(folder, type + '_' + str(size), frame), uv_posed_body_no_orient = down) def downscale_uv_normals(folder, type, sizes): for size in sizes: os.makedirs(os.path.join(folder, type + '_' + str(size)), exist_ok=True) frames = os.listdir(os.path.join(folder, type))[int(args.start):int(args.end)] for frame in tqdm.tqdm(frames): img = np.load(os.path.join(folder, type, frame)); img = img['uv_normals'] img = fill(img, uv_masks_body[512]) for size in sizes: mask = uv_masks_body[size] down = resizers[size](torch.from_numpy(img.transpose((2,0,1))).to(device)).cpu().numpy() down = down * mask[None,:,:] np.savez_compressed(os.path.join(folder, type + '_' + str(size), frame), uv_normals = down) def downscale_uv_maps_unposed(folder, type, sizes): for size in sizes: os.makedirs(os.path.join(folder, type + '_' + str(size)), exist_ok=True) frames = os.listdir(os.path.join(folder, type))[int(args.start):int(args.end)] for frame in tqdm.tqdm(frames): img = np.load(os.path.join(folder, type, frame)); img = img['uv_unposed_no_shape'] mask_full = uv_mask_store[512][None] img = fill(img, mask_full) for size in sizes: mask = uv_mask_store[size][None] down = resizers[size](torch.from_numpy(img.transpose((2,0,1))).to(device)).cpu().numpy() down = down * mask[None,:,:] np.savez_compressed(os.path.join(folder, type + '_' + str(size), frame), uv_unposed_no_shape = down) def downscale_uv_maps_static(folder, type, sizes): for size in sizes: os.makedirs(os.path.join(folder, type + '_' + str(size)), exist_ok=True) frames = os.listdir(os.path.join(folder, type))[int(args.start):int(args.end)] for frame in tqdm.tqdm(frames): img = np.load(os.path.join(folder, type, frame)); img = img['uv_static'] mask_full = uv_mask_store[512][None] img = fill(img, mask_full) for size in sizes: down = resizers[size](torch.from_numpy(img.transpose((2,0,1))).to(device)).cpu().numpy() mask = uv_mask_store[size][None] down = down * mask[None,:,:] np.savez_compressed(os.path.join(folder, type + '_' + str(size), frame), uv_static = down) def downscale_uv_maps_mask(folder, type, size): os.makedirs(os.path.join(folder, type + '_' + str(size)), exist_ok=True) frames = os.listdir(os.path.join(folder, type))[int(args.start):int(args.end)] for frame in tqdm.tqdm(frames): img = np.load(os.path.join(folder, type, frame)); img = img['uv_mask'].astype(np.float32) down = tf.image.resize(img[:,:,tf.newaxis], size=(size, int(size/2)), method=tf.image.ResizeMethod.NEAREST_NEIGHBOR) down = np.squeeze(down.numpy().astype(bool)) np.savez_compressed(os.path.join(folder, type + '_' + str(size), frame), uv_mask = down) if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument("-f", "--folder", help = "end of range", default='/data-local/127_2/cape_processed/train/male') parser.add_argument("-s", "--start", help = "start of range", default=0) parser.add_argument("-e", "--end", help = "end of range", default=100000) parser.add_argument("-n", "--name", help = "end of range", default='lhunor') parser.add_argument("-d", "--debug", help = "end of range", default=False, action='store_true') args = parser.parse_args() device = 'cuda' sizes = [32, 64, 128, 256] folder = args.folder resizers = { res: Resize(size=(res, int(res/2)), interpolation=InterpolationMode.BICUBIC).to(device) for res in sizes } resizers_nearest = { res: Resize(size=(res, int(res/2)), interpolation=InterpolationMode.NEAREST_EXACT).to(device) for res in sizes } img_idxs, img_coords, uv_masks, uv_masks_body = get_img_idxs() uv_mask_store = { size: UVMaskStore(folder, size=size) for size in sizes } uv_mask_store[512] = UVMaskStore(folder, size=512) # downscale_uv_maps_mask(folder=folder, type='uv_mask', size=32) # downscale_uv_maps_mask(folder=folder, type='uv_mask', size=64) # downscale_uv_maps_mask(folder=folder, type='uv_mask', size=128) # downscale_uv_maps_mask(folder=folder, type='uv_mask', size=256) # downscale_uv_maps_static(folder=folder, type='uv_static', sizes=sizes) downscale_uv_maps_unposed(folder=folder, type='uv_unposed_no_shape', sizes=sizes) downscale_uv_normals(folder=folder, type='uv_normals', sizes=sizes) downscale_uv_maps_posed_body(folder=folder, type='uv_posed_body_no_orient', sizes=sizes) print()