350 lines
13 KiB
Python
350 lines
13 KiB
Python
import os
|
|
import sys
|
|
import glob
|
|
import tqdm
|
|
import math
|
|
import time
|
|
import random
|
|
import shutil
|
|
import collections
|
|
import argparse
|
|
import threading
|
|
import numpy as np
|
|
import scipy as sp
|
|
import trimesh as tm
|
|
import cmapy as cm
|
|
import scipy.io as sio
|
|
import scipy.interpolate as spi
|
|
import scipy.stats as stat
|
|
import cv2
|
|
import plotly.express as ex
|
|
import plotly.graph_objects as go
|
|
from custom_utils.custom_utils import m2uv_array
|
|
|
|
|
|
def dilate(img):
|
|
uv_mask_2 = img != 0
|
|
m = img.min()
|
|
# norm = img.max() - img.min()
|
|
img = (img - m) * uv_mask_2 #* (np.invert(uv_mask_2) * -1 + uv_mask_2)
|
|
size = 1
|
|
shape = cv2.MORPH_CROSS # cv.MORPH_CROSS cv.MORPH_ELLIPSE MORPH_RECT
|
|
element = cv2.getStructuringElement(shape, (2 * size + 1, 2 * size + 1), (size, size))
|
|
result = cv2.dilate(img, element, iterations=1)
|
|
result = (result + m) * uv_mask
|
|
return result
|
|
|
|
def erode(img):
|
|
# m = img.min()
|
|
# norm = img.max() - img.min()
|
|
# img = (img - m)
|
|
size = 1
|
|
shape = cv2.MORPH_CROSS # cv.MORPH_CROSS cv.MORPH_ELLIPSE MORPH_RECT
|
|
element = cv2.getStructuringElement(shape, (2 * size + 1, 2 * size + 1), (size, size))
|
|
result = cv2.erode(img, element, iterations=2)
|
|
# result = result + m
|
|
return result
|
|
|
|
def chunks(lst, n):
|
|
"""Yield successive n-sized chunks from lst."""
|
|
for i in range(0, len(lst), n):
|
|
yield lst[i:i + n]
|
|
|
|
|
|
def convert_to_uv(res, m2uv, uv_map, verts, neighbours=10,
|
|
mask=None,
|
|
method='rbf',
|
|
garment_uv_mask=None,
|
|
uv_map_mask=None,
|
|
uv_counter=None,
|
|
uv_map_discreet_unique=None,
|
|
rbf_dest=None,
|
|
uv_mask=None):
|
|
|
|
verts_uv = np.zeros((res * 2, res, 3))
|
|
uv_values = np.zeros((uv_map.shape[0], 3))
|
|
uv_values[m2uv[:, 0]] = verts
|
|
uv_values[m2uv[:, 1]] += verts
|
|
uv_values = uv_values[uv_map_mask, :] # add garment mask
|
|
uv_values /= uv_counter[:, np.newaxis]
|
|
return_garment_uv_mask = False
|
|
|
|
if mask is not None:
|
|
garment_mask = np.zeros((uv_map.shape[0]), dtype=bool)
|
|
garment_mask[m2uv[:, 0]] = mask
|
|
garment_mask[m2uv[:, 1]] += mask
|
|
garment_mask = garment_mask[uv_map_mask] # add garment mask
|
|
else:
|
|
garment_mask = np.sum(uv_values, axis=-1) != 0
|
|
|
|
# test
|
|
verts_uv[uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 1], :] = uv_values
|
|
# verts_uv_save = verts_uv.copy()
|
|
|
|
# garment_mask = np.sum(uv_values, axis=-1) != 0
|
|
unique_uv_values = uv_values[garment_mask, :]
|
|
uv_map_discreet_unique_unique = uv_map_discreet_unique[garment_mask, :]
|
|
rbf_dest_current = rbf_dest
|
|
|
|
if garment_uv_mask is None and mask is not None:
|
|
garment_uv_mask = np.zeros((res * 2, res))
|
|
garment_uv_mask[uv_map_discreet_unique_unique[:, 0], uv_map_discreet_unique_unique[:, 1]] = 255
|
|
garment_uv_mask = cv2.morphologyEx(garment_uv_mask, cv2.MORPH_CLOSE, cv2.getStructuringElement(cv2.MORPH_RECT, (4 * 2 + 1, 4 * 2 + 1), (4, 4)), iterations=4)
|
|
garment_uv_mask_border = garment_uv_mask.astype(bool)
|
|
garment_uv_mask = cv2.dilate(garment_uv_mask, cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2).astype(bool)
|
|
garment_uv_mask_border = ~garment_uv_mask_border & garment_uv_mask
|
|
garment_uv_mask = garment_uv_mask * uv_mask[:,:,0]
|
|
rbf_dest_current = rbf_dest * garment_uv_mask
|
|
return_garment_uv_mask = True
|
|
|
|
|
|
rbf_dest_current = np.nonzero(rbf_dest_current)
|
|
rbf_dest_current = np.array(rbf_dest_current).T
|
|
|
|
# uv_verts_uniq = np.array(np.nonzero(np.sum(verts_uv, axis=-1))).T
|
|
if method == 'rbf':
|
|
interpolator = spi.RBFInterpolator(y=uv_map_discreet_unique_unique, d=unique_uv_values, neighbors=neighbours, kernel='linear')
|
|
elif method == 'nearest':
|
|
interpolator = spi.NearestNDInterpolator(x=uv_map_discreet_unique_unique, y=unique_uv_values)
|
|
else:
|
|
raise NotImplementedError
|
|
|
|
final = interpolator(rbf_dest_current)
|
|
|
|
result = np.zeros((res * 2, res, 3), np.float32)
|
|
|
|
result[rbf_dest_current[:, 0], rbf_dest_current[:, 1], :] = final
|
|
result[uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 1], :] = uv_values
|
|
if garment_uv_mask is not None:
|
|
result = result * garment_uv_mask[:,:,np.newaxis]
|
|
|
|
|
|
if return_garment_uv_mask:
|
|
return result, garment_uv_mask, garment_uv_mask_border
|
|
else:
|
|
return result
|
|
|
|
|
|
|
|
def convert_to_uv_batch(res, m2uv, uv_map, verts,
|
|
neighbours=15,
|
|
mask=None,
|
|
method='rbf',
|
|
garment_uv_mask=None,
|
|
garment_uv_mask_border=None,
|
|
id=None,
|
|
uv_map_mask=None,
|
|
uv_counter=None,
|
|
uv_map_discreet_unique=None,
|
|
rbf_dest=None,
|
|
uv_mask=None):
|
|
batch = len(verts)
|
|
assert method == 'rbf'
|
|
|
|
|
|
assert mask is not None
|
|
garment_mask = np.zeros((uv_map.shape[0]), dtype=bool)
|
|
garment_mask[m2uv[:, 0]] = mask
|
|
garment_mask[m2uv[:, 1]] += mask
|
|
garment_mask = garment_mask[uv_map_mask] # add garment mask
|
|
|
|
rbf_dest_current = rbf_dest
|
|
|
|
# add back mesh vertices that are not part of the garment
|
|
uv_map_discreet_unique_unique_inv = uv_map_discreet_unique[np.invert(garment_mask), :]
|
|
rbf_dest_fill = np.zeros((res * 2, res), dtype=bool)
|
|
rbf_dest_fill[uv_map_discreet_unique_unique_inv[:, 0], uv_map_discreet_unique_unique_inv[:, 1]] = True
|
|
|
|
assert not (garment_uv_mask is None and mask is not None)
|
|
assert garment_uv_mask is not None
|
|
rbf_dest_current = (rbf_dest + rbf_dest_fill) * garment_uv_mask + garment_uv_mask_border
|
|
|
|
rbf_dest_current = np.nonzero(rbf_dest_current)
|
|
rbf_dest_current = np.array(rbf_dest_current).T
|
|
|
|
|
|
for i in range(batch):
|
|
print('processing frame: ' + str(i))
|
|
|
|
uv_values = np.zeros((uv_map.shape[0], 3))
|
|
uv_values[m2uv[:, 0]] = verts[i]
|
|
uv_values[m2uv[:, 1]] += verts[i]
|
|
uv_values = uv_values[uv_map_mask, :]
|
|
uv_values /= uv_counter[:, np.newaxis]
|
|
|
|
|
|
unique_uv_values = uv_values[garment_mask, :]
|
|
uv_map_discreet_unique_unique = uv_map_discreet_unique[garment_mask, :]
|
|
|
|
|
|
interpolator = spi.RBFInterpolator(y=uv_map_discreet_unique_unique, d=unique_uv_values[:,:], neighbors=neighbours, kernel='linear')
|
|
result = np.zeros((res * 2, res, 3), np.float32)
|
|
result[rbf_dest_current[:, 0], rbf_dest_current[:, 1], :] = interpolator(rbf_dest_current)
|
|
result[uv_map_discreet_unique_unique[:, 0], uv_map_discreet_unique_unique[:, 1], :] = unique_uv_values[:,:]
|
|
|
|
if garment_uv_mask is not None:
|
|
result *= garment_uv_mask[:,:,np.newaxis]
|
|
|
|
return result
|
|
|
|
assert False
|
|
|
|
|
|
def convert_to_uv_batch_label(res, m2uv, uv_map, values, neighbours=15, mask=None, method='rbf', garment_uv_mask=None, garment_uv_mask_border=None):
|
|
batch = values.shape[0]
|
|
values = (values + 1) * mask
|
|
values_uv = np.zeros((batch, res * 2, res))
|
|
missing = np.zeros((res * 2, res))
|
|
uv_values = np.zeros((batch, uv_map.shape[0]))
|
|
uv_values[:, m2uv[:, 0]] = values
|
|
uv_values[:, m2uv[:, 1]] = values
|
|
uv_values = uv_values[:, uv_map_mask]
|
|
return_garment_uv_mask = False
|
|
# test
|
|
# verts_uv_save = verts_uv.copy()
|
|
|
|
# unique_uv_values_mask = np.sum(uv_values[0, :, :], axis=-1) != 0
|
|
# unique_uv_values = uv_values[:, unique_uv_values_mask, :]
|
|
# uv_map_discreet_unique_unique = uv_map_discreet_unique[unique_uv_values_mask, :]
|
|
|
|
if mask is not None:
|
|
garment_mask = np.zeros((uv_map.shape[0]), dtype=bool)
|
|
garment_mask[m2uv[:, 0]] = mask
|
|
garment_mask[m2uv[:, 1]] += mask
|
|
garment_mask = garment_mask[uv_map_mask] # add garment mask
|
|
else:
|
|
raise Exception
|
|
|
|
# uv_values = uv_values * garment_mask
|
|
|
|
# test
|
|
values_uv[:, uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 1]] = uv_values
|
|
missing[uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 1]] = (uv_values[0,:] == 0)
|
|
|
|
unique_uv_values = uv_values[:, garment_mask]
|
|
uv_map_discreet_unique_unique = uv_map_discreet_unique[garment_mask, :]
|
|
# rbf_dest_current = rbf_dest.astype(bool) + garment_uv_mask_border * uv_mask[:,:,0]
|
|
|
|
if garment_uv_mask is None and mask is not None:
|
|
garment_uv_mask = np.zeros((res * 2, res))
|
|
garment_uv_mask[uv_map_discreet_unique_unique[:, 0], uv_map_discreet_unique_unique[:, 1]] = 255
|
|
garment_uv_mask = cv2.morphologyEx(garment_uv_mask, cv2.MORPH_CLOSE, cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=4)
|
|
garment_uv_mask = cv2.dilate(garment_uv_mask, cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2).astype(bool)
|
|
return_garment_uv_mask = True
|
|
|
|
rbf_dest_current = rbf_dest * garment_uv_mask + missing * garment_uv_mask
|
|
rbf_dest_current = np.nonzero(rbf_dest_current)
|
|
rbf_dest_current = np.array(rbf_dest_current).T
|
|
|
|
# uv_verts_uniq = np.array(np.nonzero(np.sum(verts_uv, axis=-1))).T
|
|
final = []
|
|
for i in range(batch):
|
|
print('processing frame: ' + str(i))
|
|
interpolator = spi.NearestNDInterpolator(x=uv_map_discreet_unique_unique, y=unique_uv_values[i,:])
|
|
final.append(interpolator(rbf_dest_current))
|
|
final = np.asarray(final)
|
|
|
|
result = np.zeros((batch, res * 2, res))
|
|
result[:, uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 1]] = uv_values
|
|
result[:, rbf_dest_current[:, 0], rbf_dest_current[:, 1]] = final
|
|
if garment_uv_mask is not None:
|
|
result = result * garment_uv_mask[np.newaxis,:,:]
|
|
|
|
if return_garment_uv_mask:
|
|
return result, garment_uv_mask
|
|
else:
|
|
return result
|
|
|
|
|
|
# https://stackoverflow.com/questions/38187286/find-unique-pairs-in-list-of-pairs
|
|
def unique_pairs(a):
|
|
a.sort(axis=1)
|
|
b = np.ascontiguousarray(a).view(
|
|
np.dtype((np.void, a.dtype.itemsize * a.shape[1]))
|
|
)
|
|
_, inv, ct = np.unique(b, return_inverse=True, return_counts=True)
|
|
return ct[inv] == 1
|
|
|
|
|
|
def unique_uv_map_mask(uv_map_discreet, uv_res):
|
|
mask = np.zeros((uv_map_discreet.shape[0]), dtype=bool)
|
|
temp = np.zeros((uv_res * 2, uv_res))
|
|
for p in uv_map_discreet:
|
|
temp[p[0], p[1]] += 1
|
|
# temp = temp == 1
|
|
for i, p in enumerate(uv_map_discreet):
|
|
if temp[p[0], p[1]] == 1:
|
|
mask[i] = True
|
|
elif temp[p[0], p[1]] == 2:
|
|
mask[i] = True
|
|
temp[p[0], p[1]] = 0
|
|
return mask
|
|
|
|
|
|
def convert_to_uv_api(to_convert, mask, is_skirt=False, uv_res=256):
|
|
print("skirt" if is_skirt == 1 else "not skirt")
|
|
|
|
F = np.load("custom_utils/smpl_faces" + ("_skirt" if is_skirt == 1 else "") + ".npy")
|
|
uv_map = np.load('custom_utils/stretched_uv_map_' + ("skirt" if is_skirt == 1 else "3") + '.npy')
|
|
m2uv = np.load('custom_utils/m2uv' + ("_skirt" if is_skirt == 1 else "") + '.npy', allow_pickle=True).item()
|
|
m2uv = m2uv_array(m2uv)
|
|
uv_mask = np.load('custom_utils/mask_256x512' + ("_skirt" if is_skirt == 1 else "") + '.npy')
|
|
|
|
uv_mask = (cv2.resize(uv_mask * 255.0, (uv_res, uv_res * 2)) / 255).astype(bool)[:, :, np.newaxis]
|
|
uv_mask = cv2.dilate(uv_mask.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2).astype(bool)[:, :, np.newaxis]
|
|
|
|
uv_map_discreet = np.floor(uv_map * uv_res).astype(np.int32)
|
|
uv_map_mask = unique_uv_map_mask(uv_map_discreet, uv_res)
|
|
|
|
uv_counter = np.zeros((uv_map.shape[0]))
|
|
uv_counter[m2uv[:, 0]] = 1
|
|
uv_counter[m2uv[:, 1]] += 1
|
|
uv_counter = uv_counter[uv_map_mask]
|
|
|
|
uv_map_discreet_unique = uv_map_discreet[uv_map_mask, :]
|
|
|
|
rbf_dest = np.ones((uv_res * 2, uv_res))
|
|
rbf_dest[uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 1]] = 0
|
|
rbf_dest[np.invert(uv_mask[:,:,0])] = 0
|
|
|
|
_, garment_uv_mask, garment_uv_mask_border = convert_to_uv(uv_res, m2uv, uv_map, to_convert,
|
|
neighbours=3,
|
|
mask=mask,
|
|
method='rbf',
|
|
garment_uv_mask=None,
|
|
uv_map_mask=uv_map_mask,
|
|
uv_counter=uv_counter,
|
|
uv_map_discreet_unique=uv_map_discreet_unique,
|
|
rbf_dest=rbf_dest,
|
|
uv_mask=uv_mask)
|
|
|
|
uv = convert_to_uv_batch(uv_res, m2uv, uv_map, to_convert[np.newaxis,:,:],
|
|
neighbours=3,
|
|
mask=mask,
|
|
method='rbf',
|
|
garment_uv_mask=garment_uv_mask,
|
|
garment_uv_mask_border=garment_uv_mask_border,
|
|
uv_map_mask=uv_map_mask,
|
|
uv_counter=uv_counter,
|
|
uv_map_discreet_unique=uv_map_discreet_unique,
|
|
rbf_dest=rbf_dest,
|
|
uv_mask=uv_mask)
|
|
|
|
return uv
|
|
|
|
|
|
if __name__ == '__main__':
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("-n", "--name", help = "end of range", default='lhunor')
|
|
args = parser.parse_args()
|
|
|
|
|
|
start = time.perf_counter()
|
|
|
|
uv = convert_to_uv_api(np.random.random((14475, 3)) - 0.5, np.ones((14475), bool), is_skirt=False)
|
|
|
|
print(time.perf_counter() - start)
|
|
|
|
print()
|
|
|