Files
cvpr_2023/custom_utils/create_uv_map.py
2023-02-06 16:38:06 +00:00

596 lines
20 KiB
Python

import os
import sys
import shutil
import tqdm
import random
import math
import collections
import tps
import numpy as np
import scipy as sp
import trimesh as tm
import cmapy as cm
import scipy.io as sio
import matplotlib.pyplot as plt
import plotly.express as ex
def m2uv_array(m2uv, is_smpl_orig=False):
result = np.zeros((14475 if not is_smpl_orig else 3638, 2), dtype=np.int32)
for k, v in m2uv.items():
if len(v) == 2:
result[k] = list(v)
elif len(v) == 1:
x = v.pop()
result[k] = [x, x]
else:
assert False
return result
def get_body():
smpl = SMPL('FeatureLines/superSMPL')
pose = np.zeros((1, 72), dtype=np.float32)
beta = np.zeros((1, 11), dtype=np.float32)
R = smpl.run_rodrigues(pose)
T, _, _ = smpl.run(R, beta, True)
return np.squeeze(T)
def plot_3d(x, y=None):
fig = plt.figure(figsize=(15, 15))
ax = fig.add_subplot(111, projection='3d')
ax.set_xlim3d(-1, 1)
ax.set_ylim3d(-1, 1)
ax.set_zlim3d(-1, 1)
ax.set_xlabel("X")
ax.set_ylabel("Y")
ax.set_zlabel("Z")
ax.scatter(x[:, 0], x[:, 1], x[:, 2], s=5)
if y is not None:
ax.scatter(y[:, 0], y[:, 1], y[:, 2], s=5)
plt.show()
def plot_uv(x, y=None):
fig = plt.figure(figsize=(15, 15))
ax = fig.add_subplot(111)
ax.set_xlim(0, 1)
ax.set_ylim(0, 1)
ax.scatter(x[:, 0], x[:, 1])
if y is not None:
ax.scatter(y[:, 0], y[:, 1])
plt.show()
"""
Mesh to UV map
Computes correspondences between 3D mesh and UV map
NOTE: 3D mesh vertices can have multiple correspondences with UV vertices
"""
def mesh2UV(F, Ft):
m2uv = {v: set() for f in F for v in f}
for f, ft in zip(F, Ft):
for v, vt in zip(f, ft):
m2uv[v].add(vt)
# m2uv = {k:list(v) for k,v in m2uv.items()}
return m2uv
def filter_m2u(m2u, to_keep):
for key in m2u:
if len(m2u[key]) == 2:
v = list(m2u[key])
if v[0] in to_keep:
m2u[key].remove(v[1])
else:
m2u[key].remove(v[0])
return m2u
def convert_to_numbers(m2u):
for key in m2u:
m2u[key] = m2u[key].pop()
return m2u
def get_masked_index(mask, index):
m = np.zeros((mask.shape[0]), dtype=bool)
m[index] = True
m = m[mask]
new_index = np.argmax(m)
return new_index
def get_closest_point(verts_1, verts_2):
dist = sp.spatial.distance.cdist(verts_1, verts_2)
closest = np.argmin(dist, axis=-1)
return closest
def get_normal_edge(p1, p2):
d = p2 - p1
return np.array([-d[1], d[0]])
def get_normal_vertex(p1, p2, p3):
n1 = get_normal_edge(p1, p2)
n2 = get_normal_edge(p2, p3)
n = (n1 + n2)
n = n / np.linalg.norm(n) / 20
return n
def get_distance(p1, p2, points):
d = np.array([np.linalg.norm(np.cross(p2 - p1, point - p1)) / np.linalg.norm(p2 - p1) for point in points])
return d
def get_normals(verts):
v = verts.copy()
v = np.concatenate([v[-1][np.newaxis, :], v, v[0][np.newaxis, :]])
normals = np.zeros_like(verts)
for i in range(1, v.shape[0] - 1):
normals[i - 1] = get_normal_vertex(v[i - 1], v[i], v[i + 1])
return normals
def get_closest_points_normals(verts1, verts2):
dist_2 = sp.spatial.distance.cdist(verts1, verts2)
normals = get_normals(verts1)
# fig = plt.figure(figsize=(15, 15))
# ax = fig.add_subplot(111)
# ax.set_xlim(0, 1)
# ax.set_ylim(0, 1)
# ax.set_aspect('equal')
# ax.scatter(verts1[:, 0], verts1[:, 1])
# for p1, p2 in zip(verts1, verts1 + normals):
# ax.plot([p1[0], p2[0]], [p1[1], p2[1]])
# plt.show()
closest = []
i = 0
for p1, p2 in zip(verts1, verts1 + normals):
dist = get_distance(p1, p2, verts2)
closest.append(np.argmin(dist + dist_2[i, :]))
i += 1
return np.array(closest)
def back():
# load data
data_body = np.load("Blender/uv_data_1.npz", allow_pickle=True)
data_back = np.load("Blender/uv_data_plane_1_2.npz", allow_pickle=True)
# translate back uv to match body
uv_back = data_back.f.Vt
# uv_back[:, 0] = uv_back[:, 0] + 0.1
# uv_back[:, 1] = uv_back[:, 1] + 0.42
# find front and back indices in uv
uv_verts = np.append(data_body.f.Vt, np.zeros((data_body.f.Vt.shape[0], 1)), axis=-1)
m = tm.Trimesh(uv_verts, data_body.f.Ft, process=False, validate=False)
front, back = tm.graph.connected_components(m.edges, min_len=1, nodes=None, engine=None)
# get mesh to uv mapping for the back vertices
m2u_body = mesh2UV(data_body.f.F, data_body.f.Ft)
m2u_body = filter_m2u(m2u_body, set(back))
m2u_body = convert_to_numbers(m2u_body)
m2u_back = mesh2UV(data_back.f.F, data_back.f.Ft)
m2u_back = convert_to_numbers(m2u_back)
# create mask for filtering vertices belonging to the back
body_back_mask = np.zeros((data_body.f.Vt.shape[0]), dtype=bool)
body_back_mask[back] = True
# define which points should match after the transform on the mesh
matching_points = np.array([[3284, 3], [9557, 13], [9278, 15], [8095, 1], [10597, 10], [10602, 17], [7114, 9], [6989, 16], [6983, 11], [4219, 0], [5628, 14], [5908, 12], [1447, 2], [1086, 5], [21, 8], [33, 4], [1888, 7], [2931, 6]])
# map matching point to uv coordinates
for i in range(matching_points.shape[0]):
matching_points[i, :] = [get_masked_index(body_back_mask, m2u_body[matching_points[i,0]]), m2u_back[matching_points[i, 1]]]
x = data_body.f.Vt[body_back_mask, :]
y = uv_back
fig = plt.figure(figsize=(15, 15))
ax = fig.add_subplot(121)
ax.set_xlim(0, 1)
ax.set_ylim(0, 1)
ax.set_aspect('equal')
ax.scatter(x[:, 0], x[:, 1])
if y is not None:
ax.scatter(y[:, 0], y[:, 1])
for pair in matching_points:
ax.scatter([x[pair[0], 0], y[pair[1], 0]], [x[pair[0], 1], y[pair[1], 1]], s=50)
# get transformation from body to back contour
source = x[matching_points[:, 0], :]
destination = y[matching_points[:, 1], :]
trans = tps.TPS(source, destination)
transformed = trans(data_body.f.Vt[back, :])
ax = fig.add_subplot(122)
ax.set_xlim(0, 1)
ax.set_ylim(0, 1)
ax.set_aspect('equal')
ax.scatter(transformed[:, 0], transformed[:, 1])
if y is not None:
ax.scatter(y[:, 0], y[:, 1])
for pair in matching_points:
ax.scatter([transformed[pair[0], 0], y[pair[1], 0]], [transformed[pair[0], 1], y[pair[1], 1]], s=50)
plt.show()
data_back_high = np.load("Blender/uv_data_plane_1_high_res.npz", allow_pickle=True)
y_high = data_back_high.f.Vt
# y_high[:, 0] = y_high[:, 0] + 0.1
# y_high[:, 1] = y_high[:, 1] + 0.42
closest_points = get_closest_points_normals(y_high, transformed)
# closest_points = get_closest_point(y_high, transformed)
unique_mask = np.zeros_like(closest_points, dtype=bool)
unique_mask[np.unique(closest_points, return_index=True)[1]] = True
matching_points = np.array([closest_points[unique_mask], np.array(range(y_high.shape[0]))[unique_mask]]).T
# matching_points_2 = []
# for i, p in enumerate(matching_points):
# if i % 10 == 0:
# matching_points_2.append(p)
# matching_points = np.array(matching_points_2)
source_2 = transformed[matching_points[:, 0], :]
destination_2 = y_high[matching_points[:, 1], :]
trans_2 = tps.TPS(source_2, destination_2)
transformed_2 = trans_2(transformed)
fig = plt.figure(figsize=(15, 15))
ax = fig.add_subplot(121)
ax.set_xlim(0, 1)
ax.set_ylim(0, 1)
ax.set_aspect('equal')
ax.scatter(source_2[:, 0], source_2[:, 1])
ax.scatter(destination_2[:, 0], destination_2[:, 1])
for i, pair in enumerate(zip(source_2, destination_2)):
if i % 5 == 0:
s, d = pair
ax.scatter([s[0], d[0]], [s[1], d[1]])
ax = fig.add_subplot(122)
ax.set_xlim(0, 1)
ax.set_ylim(0, 1)
ax.set_aspect('equal')
ax.scatter(transformed_2[:, 0], transformed_2[:, 1])
ax.scatter(y_high[:, 0], y_high[:, 1])
plt.show()
return transformed_2
def match_back_to_front(front_uv):
# load data
data_body = np.load("Blender/uv_data_smpl_2.npz", allow_pickle=True)
# data_back = np.load("Blender/uv_data_plane_1_2.npz", allow_pickle=True)
# translate back uv to match body
# uv_back = data_back.f.Vt
# uv_back[:, 0] = uv_back[:, 0] + 0.1
# uv_back[:, 1] = uv_back[:, 1] + 0.42
# find front and back indices in uv
uv_verts = np.append(data_body.f.Vt, np.zeros((data_body.f.Vt.shape[0], 1)), axis=-1)
m = tm.Trimesh(uv_verts, data_body.f.Ft, process=False, validate=False)
front, back = tm.graph.connected_components(m.edges, min_len=1, nodes=None, engine=None)
# get mesh to uv mapping for the back vertices
m2u_body = mesh2UV(data_body.f.F, data_body.f.Ft)
m2u_body = filter_m2u(m2u_body, set(back))
m2u_body = convert_to_numbers(m2u_body)
# m2u_back = mesh2UV(data_back.f.F, data_back.f.Ft)
# m2u_back = convert_to_numbers(m2u_back)
# create mask for filtering vertices belonging to the back
body_back_mask = np.zeros((data_body.f.Vt.shape[0]), dtype=bool)
body_back_mask[back] = True
body_front_mask = np.zeros((data_body.f.Vt.shape[0]), dtype=bool)
body_front_mask[front] = True
# define which points should match after the transform on the mesh
# matching_points = np.array([[3284, 3], [9557, 13], [9278, 15], [8095, 1], [10597, 10], [10602, 17], [7114, 9], [6989, 16], [6983, 11], [4219, 0], [5628, 14], [5908, 12], [1447, 2], [1086, 5], [21, 8], [33, 4], [1888, 7], [2931, 6]])
# # map matching point to uv coordinates
# for i in range(matching_points.shape[0]):
# matching_points[i, :] = [get_masked_index(body_back_mask, m2u_body[matching_points[i,0]]), m2u_back[matching_points[i, 1]]]
m2u_body_all = mesh2UV(data_body.f.F, data_body.f.Ft)
m2u_body_all = m2uv_array(m2u_body_all)
border_pairs = m2u_body_all[:, 0] != m2u_body_all[:, 1]
border_pairs = m2u_body_all[border_pairs, :]
matching_points = border_pairs
matching_points_sorted = []
for p in matching_points:
if p[0] in front:
matching_points_sorted.append([p[0], p[1]])
elif p[0] in back:
matching_points_sorted.append([p[1], p[0]])
else:
assert False
# add neck V shape control point
matching_points_sorted.append([m2u_body_all[10228][0], m2u_body_all[33][0]])
matching_points = np.array(matching_points_sorted)
mm = matching_points.copy()
for i in range(matching_points.shape[0]):
matching_points[i, :] = [get_masked_index(body_back_mask, matching_points[i,1]), get_masked_index(body_front_mask, matching_points[i,0])]
temp = np.zeros_like(data_body.f.Vt)
temp[front, :] = front_uv
x = data_body.f.Vt[body_back_mask, :]
y = temp[body_front_mask, :] # front_uv #
fig = plt.figure(figsize=(15, 15))
ax = fig.add_subplot(121)
ax.set_xlim(0, 1)
ax.set_ylim(0, 1)
ax.set_aspect('equal')
ax.scatter(x[:, 0], x[:, 1])
if y is not None:
ax.scatter(y[:, 0], y[:, 1])
for pair in matching_points:
ax.scatter([x[pair[0], 0], y[pair[1], 0]], [x[pair[0], 1], y[pair[1], 1]], s=50)
# get transformation from body to back contour
source = x[matching_points[:, 0], :]
destination = y[matching_points[:, 1], :]
trans = tps.TPS(source, destination)
transformed = trans(data_body.f.Vt[back, :])
ax = fig.add_subplot(122)
ax.set_xlim(0, 1)
ax.set_ylim(0, 1)
ax.set_aspect('equal')
ax.scatter(transformed[:, 0], transformed[:, 1])
if y is not None:
ax.scatter(y[:, 0], y[:, 1])
for pair in matching_points:
ax.scatter([transformed[pair[0], 0], y[pair[1], 0]], [transformed[pair[0], 1], y[pair[1], 1]], s=50)
plt.show()
return transformed
def generate_mask():
import cv2
res = 1024
front = np.load("Blender/uv_data_plane_0_high_res_2.npz", allow_pickle=True)
back = np.load("Blender/uv_data_plane_1_high_res_2.npz", allow_pickle=True)
points = np.round(np.concatenate((front.f.Vt, back.f.Vt)) * res).astype(np.int32)
img = np.zeros((res, res, 3))
img[points[:, 0], points[:, 1], :] = 255
cv2.imwrite('MeshUtils/mask_1024.png', img)
# do paint now
img = cv2.imread('MeshUtils/mask_1024_2.png')
img = img[:, :, 0] == 255
np.save('MeshUtils/mask_1024.npy', img)
def generate_mask_2():
import cv2
res = 256
front = np.load("Blender/uv_data_plane_stretched_front_high_extra.npz", allow_pickle=True)
back = np.load("Blender/uv_data_plane_stretched_front_high_extra.npz", allow_pickle=True)
points = np.floor(np.concatenate((front.f.Vt, back.f.Vt + [1, 0])) * res).astype(np.int32)
img = np.zeros((res * 2, res, 3))
img[points[:, 0], points[:, 1], :] = 255
cv2.imwrite('MeshUtils/mask_256x512.png', img)
# do paint now
img = cv2.imread('MeshUtils/mask_256x512.png')
img = img[:, :, 0] > 240
np.save('MeshUtils/mask_256x512.npy', img)
pass
def generate_uv_mask_border():
import cv2
mask = np.load('MeshUtils/mask_256x512.npy')
uv_res = 256
mask_border = (cv2.resize(mask * 255.0, (uv_res, uv_res * 2)) / 255).astype(bool)[:, :, np.newaxis]
mask_border = cv2.erode(mask_border.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=1).astype(bool)
mask[mask_border == True] = 0
# plt.imshow(mask)
idx_front = np.transpose(np.vstack(np.nonzero(mask[:256, :])), [1,0])
test = np.zeros_like(mask)
test[idx_front[:,0], idx_front[:,1]] = 1
idx_back = idx_front.copy()
idx_back[:, 0] += 256
test[idx_back[:,0], idx_back[:,1]] = 1
pass
def front():
# load data
data_body = np.load("custom_data/uv_map/uv_data_smpl_2.npz", allow_pickle=True)
data_front = np.load("custom_data/uv_map/uv_data_plane_stretched_front_low.npz", allow_pickle=True)
# translate back uv to match body
uv_front = data_front.f.Vt
# uv_back[:, 0] = uv_back[:, 0] + 0.1
# uv_back[:, 1] = uv_back[:, 1] + 0.42
# find front and back indices in uv
uv_verts = np.append(data_body.f.Vt, np.zeros((data_body.f.Vt.shape[0], 1)), axis=-1)
m = tm.Trimesh(uv_verts, data_body.f.Ft, process=False, validate=False)
front, back = tm.graph.connected_components(m.edges, min_len=1, nodes=None, engine=None)
# get mesh to uv mapping for the back vertices
m2u_body = mesh2UV(data_body.f.F, data_body.f.Ft)
m2u_body = filter_m2u(m2u_body, set(front))
m2u_body = convert_to_numbers(m2u_body)
m2u_front = mesh2UV(data_front.f.F, data_front.f.Ft)
m2u_front = convert_to_numbers(m2u_front)
# create mask for filtering vertices belonging to the back
body_front_mask = np.zeros((data_body.f.Vt.shape[0]), dtype=bool)
body_front_mask[front] = True
# define which points should match after the transform on the mesh
matching_points = np.array([ [9557, 12], [3284, 2], [3053, 5], [1888, 8], [10228, 4], [21, 7], [1209, 6], [1447, 3], [5908, 13], [4617, 15], [4219, 1], [6983, 10], [6989, 17], [7114, 9], [10602, 16], [10597, 11], [7884, 0], [8280, 14] ])
# matching_points = np.array([ [9557, 12], [3284, 2], [3053, 5], [1888, 8], [10228, 4], [21, 7], [1209, 6], [1447, 3], [5908, 13], [4793, 15], [4219, 1], [6983, 10], [6989, 17], [7114, 9], [10602, 16], [10597, 11], [7884, 0], [8450, 14], [13566, 18], [11767, 19], [8668, 20], [5012, 21], [12716, 22], [10917, 23], [8280, 24], [4617, 25], [2760, 26], [911, 27] ])
# map matching point to uv coordinates
for i in range(matching_points.shape[0]):
matching_points[i, :] = [get_masked_index(body_front_mask, m2u_body[matching_points[i,0]]), m2u_front[matching_points[i, 1]]]
x = data_body.f.Vt[body_front_mask, :]
y = uv_front
fig = plt.figure(figsize=(15, 15))
ax = fig.add_subplot(121)
# ax.set_xlim(0, 1)
# ax.set_ylim(0, 1)
ax.set_aspect('equal')
ax.scatter(x[:, 0], x[:, 1])
if y is not None:
ax.scatter(y[:, 0], y[:, 1])
for pair in matching_points:
ax.scatter([x[pair[0], 0], y[pair[1], 0]], [x[pair[0], 1], y[pair[1], 1]], s=50)
# get transformation from body to front contour
source = x[matching_points[:, 0], :]
destination = y[matching_points[:, 1], :]
trans = tps.TPS(source, destination)
transformed = trans(data_body.f.Vt[front, :])
# transformed = (transformed - transformed.min()) / (transformed.max() - transformed.min())
ax = fig.add_subplot(122)
# ax.set_xlim(0, 1)
# ax.set_ylim(0, 1)
ax.set_aspect('equal')
ax.scatter(transformed[:, 0], transformed[:, 1])
if y is not None:
ax.scatter(y[:, 0], y[:, 1])
for pair in matching_points:
ax.scatter([transformed[pair[0], 0], y[pair[1], 0]], [transformed[pair[0], 1], y[pair[1], 1]], s=50)
plt.show()
data_front_high = np.load("custom_data/uv_map/uv_data_plane_stretched_front_high.npz", allow_pickle=True)
y_high = data_front_high.f.Vt
# y_high[:, 0] = y_high[:, 0] + 0.1
# y_high[:, 1] = y_high[:, 1] + 0.42
closest_points = get_closest_points_normals(y_high, transformed)
# closest_points = get_closest_point(y_high, transformed)
unique_mask = np.zeros_like(closest_points, dtype=bool)
unique_mask[np.unique(closest_points, return_index=True)[1]] = True
matching_points = np.array([closest_points[unique_mask], np.array(range(y_high.shape[0]))[unique_mask]]).T
# matching_points_2 = []
# for i, p in enumerate(matching_points):
# if i % 10 == 0:
# matching_points_2.append(p)
# matching_points = np.array(matching_points_2)
source_2 = transformed[matching_points[:, 0], :]
destination_2 = y_high[matching_points[:, 1], :]
trans_2 = tps.TPS(source_2, destination_2)
transformed_2 = trans_2(transformed)
fig = plt.figure(figsize=(15, 15))
ax = fig.add_subplot(121)
ax.set_xlim(0, 1)
ax.set_ylim(0, 1)
ax.set_aspect('equal')
ax.scatter(source_2[:, 0], source_2[:, 1])
ax.scatter(destination_2[:, 0], destination_2[:, 1])
for i, pair in enumerate(zip(source_2, destination_2)):
if i % 5 == 0:
s, d = pair
ax.scatter([s[0], d[0]], [s[1], d[1]])
ax = fig.add_subplot(122)
ax.set_xlim(0, 1)
ax.set_ylim(0, 1)
ax.set_aspect('equal')
ax.scatter(transformed_2[:, 0], transformed_2[:, 1])
ax.scatter(y_high[:, 0], y_high[:, 1])
plt.show()
return transformed_2, mesh2UV(data_body.f.F, data_body.f.Ft), front, back
def smoothing():
data_body = np.load("Blender/uv_data_smpl_2.npz", allow_pickle=True)
uv_map = np.load('MeshUtils\stretched_uv_map_3.npy')
verts = np.concatenate((uv_map, np.ones((15068, 1))), axis=-1)
m = tm.Trimesh(verts, data_body.f.Ft, process=False, validate=False)
m2u_body_all = mesh2UV(data_body.f.F, data_body.f.Ft)
m2u_body_all = m2uv_array(m2u_body_all)
border_pairs = m2u_body_all[:, 0] != m2u_body_all[:, 1]
border_pairs = m2u_body_all[border_pairs, :]
pinned_vertices = border_pairs.reshape(-1)
lap = tm.smoothing.laplacian_calculation(m, equal_weight=False, pinned_vertices=pinned_vertices)
# tm.smoothing.filter_laplacian(m, lamb=0.5, iterations=100, volume_constraint=False, laplacian_operator=lap)
# tm.smoothing.filter_mut_dif_laplacian(m, lamb=0.5, iterations=100, volume_constraint=False, laplacian_operator=lap)
# tm.smoothing.filter_humphrey(m, iterations=100, laplacian_operator=lap)
tm.smoothing.filter_taubin(m, iterations=100, laplacian_operator=lap)
x = uv_map
y = m.vertices[:, :2]
fig = plt.figure(figsize=(15, 15))
ax = fig.add_subplot(211)
ax.set_xlim(0, 2)
ax.set_ylim(0, 1)
# ax.set_aspect('equal')
ax.scatter(x[:, 0], x[:, 1])
ax = fig.add_subplot(212)
ax.set_xlim(0, 2)
ax.set_ylim(0, 1)
# ax.set_aspect('equal')
ax.scatter(y[:, 0], y[:, 1])
plt.show()
pass
if __name__ == '__main__':
# generate_mask_2()
# smoothing()
# generate_uv_mask_border()
front_uv, m2uv, front_mask, back_mask = front()
back_uv = match_back_to_front(front_uv)
# back_uv = back()
final = np.zeros((15068, 2))
final[front_mask] = front_uv
final[back_mask] = back_uv + [1, 0]
np.save('MeshUtils\stretched_uv_map_3.npy', final)
pass