luni/cvpr_2023
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

save all

Browse Source
This commit is contained in:
luni
2023-08-24 11:36:22 +00:00
parent 59f67a2495
commit afd482522d
36 changed files with 2194 additions and 291 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/calculate.py
+34
View File
@@ -0,0 +1,34 @@
import numpy as np
error_1 = dict()
error_2 = dict()
#end2end with reg
# error_1[0] = {'Jumpsuit': 0.00888, 'Trousers': 0.01193, 'Tshirt': 0.01065, 'Top': 0.00725}
# error_1[1] = {'Jumpsuit': 0.00491, 'Trousers': 0.00704, 'Tshirt': 0.00531, 'Top': 0.00393}
# error_1[2] = {'Jumpsuit': 0.00379, 'Trousers': 0.00547, 'Tshirt': 0.00404, 'Top': 0.00309}
# error_1[3] = {'Jumpsuit': 0.00349, 'Trousers': 0.00502, 'Tshirt': 0.00382, 'Top': 0.00281}
# error_2[0] = {'Dress': 0.01009, 'Skirt': 0.01379}
# error_2[1] = {'Dress': 0.0055, 'Skirt': 0.0077}
# error_2[2] = {'Dress': 0.00418, 'Skirt': 0.00582}
# error_2[3] = {'Dress': 0.00359, 'Skirt': 0.00496}
#end2end without reg
error_1[0] = {'Jumpsuit': 0.00927, 'Trousers': 0.01229, 'Tshirt': 0.01113, 'Top': 0.0076}
error_1[1] = {'Jumpsuit': 0.00503, 'Trousers': 0.00721, 'Tshirt': 0.0054, 'Top': 0.004}
error_1[2] = {'Jumpsuit': 0.00402, 'Trousers': 0.0059, 'Tshirt': 0.00421, 'Top': 0.00319}
error_1[3] = {'Jumpsuit': 0.00347, 'Trousers': 0.00528, 'Tshirt': 0.00365, 'Top': 0.00271}
error_2[0] = {'Dress': 0.01005, 'Skirt': 0.01395}
error_2[1] = {'Dress': 0.0056, 'Skirt': 0.00787}
error_2[2] = {'Dress': 0.00441, 'Skirt': 0.00634}
error_2[3] = {'Dress': 0.00379, 'Skirt': 0.0054}
for i in range(4):
errors_mean = dict()
error = {**error_1[i], **error_2[i]}
for category in error.keys():
errors_mean[category] = round(np.mean(error[category]), 5)
print(f'Level {i}: {errors_mean}')
mean = round(np.mean(list(errors_mean.values())), 5)
print(f'Final level {i}: {mean}')
scripts/calculate_area_diff.py
+47
View File
@@ -0,0 +1,47 @@
import trimesh
import numpy as np
cases = [
# ('06461_237_Trousers', 'temp/sampling_pyramid/06461_237_Trousers', [0,8,12,13]),
# ('01171_213_Dress', 'temp/sampling_pyramid/01171_213_Dress', [0,6,10,15]),
('01852_70_Tshirt_lvl_0', 'temp/sampling_pro_pyramid/01852_70_Tshirt', [15,10,1,5,13]),
('01852_70_Tshirt_lvl_1', 'temp/sampling_pro_pyramid/01852_70_Tshirt', [15,14,3,5,0 ]),
('01852_70_Tshirt_lvl_2', 'temp/sampling_pro_pyramid/01852_70_Tshirt', [15,12,0,8,6]),
('01852_70_Tshirt_lvl_3', 'temp/sampling_pro_pyramid/01852_70_Tshirt', [15,2,8,1,0]),
]
for case in cases:
# Load the original mesh
folder = case[1]
original_mesh = trimesh.load(f'{folder}/{case[0]}_{case[2][0]}.stl')
# Load the variation meshes
variation_meshes = []
for i in case[2][1:]:
variation_meshes.append(trimesh.load(f'{folder}/{case[0]}_{i}.stl'))
# Calculate area of original mesh
# original_area = original_mesh.area
# Calculate area difference for each variation
area_diffs = []
for variation_mesh in variation_meshes:
# calculate per vertex euclid distance
euclid = np.linalg.norm(original_mesh.vertices - variation_mesh.vertices, axis=1) * 1000
euclid_diff = np.mean(euclid)
# area_diff = abs(variation_mesh.area - original_area) / original_area
area_diffs.append(euclid_diff)
# Average area differences
# avg_area_diff = np.mean(area_diffs)
variance = np.var(area_diffs)
# Print results
print(f'Case: {case[0]}')
# print('Original area:', original_area)
print('Euclidean differences:', area_diffs)
print('Euclidean variance:', variance)
# print('Average Euclidean difference:', avg_area_diff)
scripts/calculate_chamfer.py
+168 -39
View File
@@ -4,47 +4,40 @@ sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = '0'
os.environ["CUDA_VISIBLE_DEVICES"] = '6'
import time
import argparse
import wandb
import tqdm
import random
import numpy as np
import plotly.graph_objects as go
import plotly.express as ex
from PIL import Image
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from tqdm import tqdm
from data.dataset import Dataset
from data.utils import Frame
from data.cloth3d_io import readPC2Frame, readPC2
from data.preprocessing import deorient
from utils.checkpoints import CheckpointSaver, load_model
from torch.utils.data import DataLoader
from utils.utils import text_to_html, pickle_load
from utils.postprocess import fix_collision
from losses.losses import collision_loss, normal_loss, kl_loss
from custom_utils.custom_utils import UV2Mesh, Accumulator, plot_3d, get_summary_html, maskFacesBatch, get_summary_str, plot_3d_collision, group_inputs_for_levels
from plotly.subplots import make_subplots
import pytorch3d.loss
from scipy.spatial.distance import cdist
from data.cloth3d_io import readPC2Frame, readPC2, readOBJ
from utils.utils import pickle_load
from scipy.spatial import cKDTree
from smpl.smpl_torch import SMPL
from scipy.spatial.transform import Rotation as R
import trimesh as tm
from losses.losses import collision_loss, normal_loss, kl_loss, collision_error
parser = argparse.ArgumentParser()
parser.add_argument("-s", "--start", help = "start of range", default=0, type=int)
parser.add_argument("-e", "--end", help = "end of range", default=550, type=int)
parser.add_argument("-n", "--name", help = "end of range", default='lhunor')
parser.add_argument('--debug', action='store_true')
args = parser.parse_args()
def chamfer(points1, points2, points2_tree: cKDTree):
tree1 = cKDTree(points1)
dist1, _ = tree1.query(points2, workers=40)
dist2, _ = points2_tree.query(points1, workers=40)
dist1, _ = tree1.query(points2, workers=10)
dist2, _ = points2_tree.query(points1, workers=10)
chamfer_dist = np.mean(dist1) + np.mean(dist2)
return chamfer_dist
@@ -60,34 +53,116 @@ def deorient(verts, pose, inverse=True):
return rotated
def get_collision_depth(m1, m2):
cm = tm.collision.CollisionManager()
cm.add_object('m1', m1)
cm.add_object('m2', m2)
data = cm.in_collision_internal(return_data=True)[1]
# c = np.asarray([d.index('m2') for d in data])
# count = np.unique(c).shape[0]
depth_sum = np.sum([d.depth for d in data])
return depth_sum
def get_area_diff(m1, m2):
return np.abs(m1.area - m2.area)
device = 'cuda'
def get_collision_loss(pred, body, mask):
c_loss, c_ratio = collision_error(torch.from_numpy(pred).unsqueeze(0).to(device), torch.from_numpy(body).unsqueeze(0).to(device), torch.from_numpy(mask).unsqueeze(0).to(device))
return c_loss.cpu().numpy()[0], c_ratio.cpu().numpy()
def get_collision_loss_2(pred, body, mask):
_pred = np.zeros((14475, 3), dtype=np.float32)
_pred[mask] = pred
c_loss, c_ratio = collision_loss(torch.from_numpy(_pred).unsqueeze(0).to(device), torch.from_numpy(body).unsqueeze(0).to(device), torch.from_numpy(mask).unsqueeze(0).to(device), thr=0.0005)
return c_loss.cpu().numpy()[0], c_ratio.cpu().numpy()
def calculate():
errors_hood = []
errors_ours = []
garments = sorted(os.listdir(hood_folder))
collisions_hood = []
areas_hood = []
collision_ratios_hood = []
if calculate_ours:
errors_ours = []
collisions_ours = []
areas_ours = []
collision_ratios_ours = []
garments = sorted(os.listdir(hood_folder))[args.start:args.end]
garments = [g[:-4] for g in garments]
body_faces = np.load('smpl/smpl_faces.npy').astype(np.int32)
for garment in tqdm(garments):
picker = 'skirt' if garment.split('_')[1] in ['Dress', 'Skirt'] else 'nonskirt'
id, garment_type = garment.split('_')
pred_hood = pickle_load(f'{hood_folder}/{garment}.pkl')['verts']
pred_hood_data = pickle_load(f'{hood_folder}/{garment}.pkl')
pred_hood, pred_hood_faces = pred_hood_data['verts'], pred_hood_data['faces']
gt = readPC2(os.path.join(cloth3d_folder, id, garment_type + '.pc16'), True)['V'][2:]
template, template_faces, _, _ = readOBJ(os.path.join(cloth3d_folder, id, garment_type + '.obj'))
mask = np.load(f'/data-local/127_2/cloth3d_processed_test/mesh_fixed/{id}_{garment_type}.npz')['mask']
pose = np.load(f'/data-local/127_2/cloth3d_processed_test/pose/{id}.npy')[2:]
gt = deorient(gt, pose, True)
body = readPC2(f'/data-local/127_2/cloth3d_processed_test/mesh_body_posed_no_orient/{id}.pc16', True)['V'][2:]
fix_data = np.load(f'/data-local/127_2/cloth3d_processed_test/mesh_fixed/{id}_{garment_type}.npz')
faces = fix_data['faces']
m_template = tm.Trimesh(vertices=template, faces=template_faces, process=False, validate=False)
for i in range(pred_hood.shape[0]):
gt_tree = cKDTree(gt[i])
try:
pred_ours = np.load(f'{ours_folder[picker]}/{id}_{i+2}_{garment_type}.npz', allow_pickle=True)['arr_0'].item()['lvl3']
if calculate_ours:
# pred_ours = np.load(f'{ours_folder[picker]}/{id}_{i+2}_{garment_type}.npz', allow_pickle=True)['arr_0'].item()['lvl3']
pred_ours = np.load(f'{ours_folder[picker]}/{id}_{i+2}_{garment_type}.npz', allow_pickle=True)['arr_0']
except:
print('missed one')
continue
# m_body = tm.Trimesh(vertices=body[i], faces=body_faces, process=False, validate=False)
errors_hood.append((f'{id}_{i}_{garment_type}', chamfer(pred_hood[i], gt[i], gt_tree)))
errors_ours.append((f'{id}_{i}_{garment_type}', chamfer(pred_ours[mask,:], gt[i], gt_tree)))
m_hood = tm.Trimesh(vertices=pred_hood[i], faces=pred_hood_faces, process=False, validate=False)
# collision_hood = get_collision_depth(m_body, m_hood)
# collisions_hood.append((f'{id}_{i}_{garment_type}', collision_hood))
area_hood = get_area_diff(m_template, m_hood)
areas_hood.append((f'{id}_{i}_{garment_type}', area_hood))
collision_ratio_hood = get_collision_loss(pred_hood[i], body[i], mask)[1]
collision_ratios_hood.append((f'{id}_{i}_{garment_type}', collision_ratio_hood))
if calculate_ours:
errors_ours.append((f'{id}_{i}_{garment_type}', chamfer(pred_ours[mask], gt[i], gt_tree)))
m_ours = tm.Trimesh(vertices=pred_ours[mask], faces=faces, process=False, validate=False)
# collision_ours = get_collision_depth(m_body, m_ours)
# collisions_ours.append((f'{id}_{i}_{garment_type}', collision_ours))
area_ours = get_area_diff(m_template, m_ours)
areas_ours.append((f'{id}_{i}_{garment_type}', area_ours))
collision_ratio_ours = get_collision_loss(pred_ours[mask], body[i], mask)[1]
# collision_ratio_ours_2 = get_collision_loss_2(pred_ours[mask], body[i], mask)[1]
collision_ratios_ours.append((f'{id}_{i}_{garment_type}', collision_ratio_ours))
# fig = go.Figure(data=[go.Mesh3d(x=pred_ours[mask][:,0], y=pred_ours[mask][:,2], z=pred_ours[mask][:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='blue')])
# fig.add_trace(go.Mesh3d(x=body[i][:,0], y=body[i][:,2], z=body[i][:,1], i=body_faces[:,0], j=body_faces[:,1], k=body_faces[:,2], color='gray'))
# fig.update_layout(scene=dict(xaxis=dict(range=[-1, 1], ), yaxis=dict(range=[-1, 1], ), zaxis=dict(range=[-1, 1], ), ), )
# fig.update_layout(scene_aspectmode='cube', scene_aspectratio=dict(x=1, y=1, z=1))
# fig.show()
# fig = go.Figure(data=[go.Mesh3d(x=gt[i][:,0], y=gt[i][:,2], z=gt[i][:,1], i=pred_hood_faces[:,0], j=pred_hood_faces[:,1], k=pred_hood_faces[:,2], color='blue')])
# fig.add_trace(go.Mesh3d(x=body[i][:,0], y=body[i][:,2], z=body[i][:,1], i=body_faces[:,0], j=body_faces[:,1], k=body_faces[:,2], color='gray'))
# fig.update_layout(scene=dict(xaxis=dict(range=[-1, 1], ), yaxis=dict(range=[-1, 1], ), zaxis=dict(range=[-1, 1], ), ), )
# fig.update_layout(scene_aspectmode='cube', scene_aspectratio=dict(x=1, y=1, z=1))
# fig.show()
# fig = go.Figure(data=[go.Mesh3d(x=pred_hood[i][:,0], y=pred_hood[i][:,2], z=pred_hood[i][:,1], i=pred_hood_faces[:,0], j=pred_hood_faces[:,1], k=pred_hood_faces[:,2], color='blue')])
# fig.add_trace(go.Mesh3d(x=body[i][:,0], y=body[i][:,2], z=body[i][:,1], i=body_faces[:,0], j=body_faces[:,1], k=body_faces[:,2], color='gray'))
# fig.update_layout(scene=dict(xaxis=dict(range=[-1, 1], ), yaxis=dict(range=[-1, 1], ), zaxis=dict(range=[-1, 1], ), ), )
# fig.update_layout(scene_aspectmode='cube', scene_aspectratio=dict(x=1, y=1, z=1))
# fig.show()
# pass
# fig = go.Figure()
# fig.add_trace(go.Scatter3d(x=gt[i][:, 0], y=gt[i][:, 2], z=gt[i][:, 1], mode='markers', marker=dict(size=1)))
@@ -95,6 +170,7 @@ def calculate():
# fig.update_layout(scene_aspectmode='cube')
# fig.update_layout(scene=dict(xaxis=dict(range=[-1, 1]), yaxis=dict(range=[-1, 1]), zaxis=dict(range=[-1, 1])))
# fig.show()
# pass
# fig = go.Figure()
# fig.add_trace(go.Scatter3d(x=gt[i][:, 0], y=gt[i][:, 2], z=gt[i][:, 1], mode='markers', marker=dict(size=1)))
@@ -105,30 +181,83 @@ def calculate():
# save errors to a csv file
with open(f'{results_folder}/hood/{name}_hood.csv', 'w') as f:
for frame, error in errors_hood:
frame_str = ','.join(str(frame).split('_'))
f.write(f'{frame_str},{error}\n')
os.makedirs(f'{results_folder}/hood/{name}_hood', exist_ok=True)
os.makedirs(f'{results_folder}/hood/{name}_hood/chamfer', exist_ok=True)
os.makedirs(f'{results_folder}/hood/{name}_hood/collision', exist_ok=True)
os.makedirs(f'{results_folder}/hood/{name}_hood/area', exist_ok=True)
os.makedirs(f'{results_folder}/hood/{name}_hood/collision_ratio', exist_ok=True)
with open(f'{results_folder}/hood/{name}_ours.csv', 'w') as f:
for frame, error in errors_ours:
frame_str = ','.join(str(frame).split('_'))
f.write(f'{frame_str},{error}\n')
for frame, error in errors_hood:
np.save(f'{results_folder}/hood/{name}_hood/chamfer/{frame}.npy', error)
for frame, error in collisions_hood:
np.save(f'{results_folder}/hood/{name}_hood/collision/{frame}.npy', error)
for frame, error in areas_hood:
np.save(f'{results_folder}/hood/{name}_hood/area/{frame}.npy', error)
for frame, error in collision_ratios_hood:
np.save(f'{results_folder}/hood/{name}_hood/collision_ratio/{frame}.npy', error)
os.makedirs(f'{results_folder}/hood/{name}_ours', exist_ok=True)
os.makedirs(f'{results_folder}/hood/{name}_ours/chamfer', exist_ok=True)
os.makedirs(f'{results_folder}/hood/{name}_ours/collision', exist_ok=True)
os.makedirs(f'{results_folder}/hood/{name}_ours/area', exist_ok=True)
os.makedirs(f'{results_folder}/hood/{name}_ours/collision_ratio', exist_ok=True)
for frame, error in errors_ours:
np.save(f'{results_folder}/hood/{name}_ours/chamfer/{frame}.npy', error)
for frame, error in collisions_ours:
np.save(f'{results_folder}/hood/{name}_ours/collision/{frame}.npy', error)
for frame, error in areas_ours:
np.save(f'{results_folder}/hood/{name}_ours/area/{frame}.npy', error)
for frame, error in collision_ratios_ours:
np.save(f'{results_folder}/hood/{name}_ours/collision_ratio/{frame}.npy', error)
# with open(f'{results_folder}/hood/{name}/chamfer_hood.csv', 'w') as f:
# for frame, error in errors_hood:
# frame_str = ','.join(str(frame).split('_'))
# f.write(f'{frame_str},{error}\n')
# with open(f'{results_folder}/hood/{name}/collision_hood.csv', 'w') as f:
# for frame, error in collisions_hood:
# frame_str = ','.join(str(frame).split('_'))
# f.write(f'{frame_str},{error}\n')
# with open(f'{results_folder}/hood/{name}/area_hood.csv', 'w') as f:
# for frame, error in areas_hood:
# frame_str = ','.join(str(frame).split('_'))
# f.write(f'{frame_str},{error}\n')
# if calculate_ours:
# with open(f'{results_folder}/hood/{name}_chamfer_ours.csv', 'w') as f:
# for frame, error in errors_ours:
# frame_str = ','.join(str(frame).split('_'))
# f.write(f'{frame_str},{error}\n')
# with open(f'{results_folder}/hood/{name}/collision_ours.csv', 'w') as f:
# for frame, error in collisions_ours:
# frame_str = ','.join(str(frame).split('_'))
# f.write(f'{frame_str},{error}\n')
# with open(f'{results_folder}/hood/{name}/area_ours.csv', 'w') as f:
# for frame, error in areas_ours:
# frame_str = ','.join(str(frame).split('_'))
# f.write(f'{frame_str},{error}\n')
if __name__ == '__main__':
name = 'hood_2'
name = 'hood_3_var'
hood_folder = '/data-local/127_2/results/hood/all_2'
hood_folder = '/data-local/127_2/results/hood/all_3'
ours_folder = dict()
ours_folder['skirt'] = '/data-local/127_2/results/pyr_iccvw_64_128_256_512_skirt_e2e_200000'
ours_folder['nonskirt'] = '/data-local/127_2/results/pyr_iccvw_64_128_256_512_e2e_2_250000'
ours_folder['skirt'] = '/data-local/127_2/results/pyr_iccvw_64_128_256_512_skirt_e2e_var_3_450000'
ours_folder['nonskirt'] = '/data-local/127_2/results/pyr_iccvw_64_128_256_512_e2e_var_2_600000'
cloth3d_folder = '/data-net/dataset/cloth3d/test'
results_folder = '/data-local/127_2/results'
out_folder = f'{results_folder}/errors/{name}'
calculate_ours = True
calculate()
scripts/calculate_error_all.py
+8 -6
View File
@@ -10,6 +10,7 @@ parser.add_argument('--config', type=str)
parser.add_argument('--debug', action='store_true')
parser.add_argument('--iteration', type=str, default='200000')
parser.add_argument('-g', '--gpu', type=str, default='0')
parser.add_argument('-f', '--final', action='store_true')
args = vars(parser.parse_args())
if args['debug']:
@@ -68,7 +69,7 @@ def get_normal_loss(pred, gt, mask):
return n_loss.cpu().numpy()
def get_collision_loss(pred, body, mask):
c_loss, c_ratio = collision_loss(torch.from_numpy(pred).unsqueeze(0).to(device), torch.from_numpy(body).unsqueeze(0).to(device), torch.from_numpy(mask).unsqueeze(0).to(device))
c_loss, c_ratio = collision_loss(torch.from_numpy(pred).unsqueeze(0).to(device), torch.from_numpy(body).unsqueeze(0).to(device), torch.from_numpy(mask).unsqueeze(0).to(device), thr=0.0005)
return c_loss.cpu().numpy()[0], c_ratio.cpu().numpy()
@@ -115,10 +116,15 @@ def calculate():
errors_edge.append((str(frame), get_edge_loss(pred, gt, new_mask)))
else:
for i in range(4):
levels = [0,1,2,3] if not args['final'] else [3]
for i in levels:
_pred = pred.item()[f'lvl{i}']
c_loss, c_ratio = get_collision_loss(_pred, body, new_mask)
errors_collision.append((f'{frame}_lvl{i}', c_loss))
errors_collision_ratio.append((f'{frame}_lvl{i}', c_ratio))
errors_normal.append((f'{frame}_lvl{i}', get_normal_loss(_pred, gt, new_mask)))
errors_edge.append((f'{frame}_lvl{i}', get_edge_loss(_pred, gt, new_mask)))
# faces = fix_data['faces']
# body_faces = np.load('smpl/smpl_faces.npy').astype(np.int32)
@@ -143,10 +149,6 @@ def calculate():
# fig.show()
errors_collision.append((f'{frame}_lvl{i}', c_loss))
errors_collision_ratio.append((f'{frame}_lvl{i}', c_ratio))
errors_normal.append((f'{frame}_lvl{i}', get_normal_loss(_pred, gt, new_mask)))
errors_edge.append((f'{frame}_lvl{i}', get_edge_loss(_pred, gt, new_mask)))
# save errors to a csv file
scripts/cvae.py
+50 -2
View File
@@ -11,7 +11,7 @@ parser.add_argument('--debug', action='store_true')
args = vars(parser.parse_args())
if args['debug']:
from configs.iccvw.abl_all_5_var import config
from configs.iccvw.abl_all_5 import config
else:
config_module = importlib.import_module(args['config'])
config = config_module.config
@@ -45,7 +45,7 @@ from data.dataset import Dataset
from models.cond_encoder import AutoEncoder
from utils.checkpoints import CheckpointSaver, load_model
from utils.utils import InfiniteDataLoader as DataLoader
from utils.utils import text_to_html
from utils.utils import text_to_html, save_stl, create_temp_frames
from utils.postprocess import fix_collision
from losses.losses import collision_loss, normal_loss
from pytorch_msssim import MS_SSIM
@@ -352,7 +352,52 @@ def sample():
break
def generate():
from data.utils import Frame
load_model(model, f'checkpoints/iccvw/{config["name"]}')
frame = '01852_70_Tshirt'
# frame = '03849_27_Dress' # 03849_27_Dress 00054_1_Dress
frame = Frame(frame_str=frame)
create_temp_frames(config['train_dir'], frame)
dataset_train = Dataset(config['train_dir'], config['signature'], device=config['device'], frames_file='frames_temp.txt', use_cache=config['dataset_in_memory'])
train_dataloader = DataLoader(dataset_train, batch_size=4, drop_last=True, shuffle=True, )
batch = next(iter(train_dataloader))
if config['use_skinning']:
mesh_orig, uv_orig = batch['mesh_unposed'], batch['uv_unposed']
else:
mesh_orig, uv_orig = batch['mesh_posed'], batch['uv_posed']
# with record_function("forward"):
model.eval()
with torch.inference_mode():
pred, _, _ = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_normals'], batch['uv_mask']])
mesh_pred = uv2mesh(pred)
if config['use_skinning']:
shape = batch['pose_shape'][:, 72:]
gender = torch.greater_equal(shape[:, 10] , 0.5).to(torch.long)
shape_offsets = torch.stack([torch.tensordot(shape[i,:10][None, :], smpl.shapedirs[gender[i]], dims=[[1], [2]])[0,smpl.to_keep,:] for i in range(shape.shape[0])])
# shape_offsets = torch.stack([torch.matmul(shape[i, :10], smpl.shapedirs[gender[i]][:, smpl.to_keep]) for i in range(shape.shape[0])])
mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,None]
mesh_pred_posed = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)[0]
else:
mesh_pred_posed = mesh_pred
fix_data = np.load(Path(dataset_train.folder) / 'mesh_fixed' / f'{frame.id}_{frame.garment}.npz')
mask = fix_data['mask']
faces = fix_data['faces']
pred = mesh_pred_posed.cpu().numpy()[0, mask, :]
body = batch['mesh_body_posed'].cpu().numpy()[0]
body_faces = smpl.f
save_stl(Path('temp') / 'baseline' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_baseline_mesh.stl', pred, faces)
save_stl(Path('temp') / 'baseline' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_baseline_body.stl', body, body_faces)
pass
def train():
@@ -410,6 +455,9 @@ if __name__ == '__main__':
model = CVAE(config)
generate()
exit()
sample()
exit()
scripts/cvae_pyramid.py
+364 -26
View File
@@ -11,7 +11,8 @@ parser.add_argument('--debug', action='store_true')
args = vars(parser.parse_args())
if args['debug']:
from configs.iccvw.pyr_64_128_256_512 import config
# from configs.iccvw.pyr_64_128_256_512 import config
from configs.iccvw.pyr_64_128_256_512_e2e import config
else:
config_module = importlib.import_module(args['config'])
config = config_module.config
@@ -24,7 +25,7 @@ if args['debug']:
import time
import argparse
import wandb
import tqdm
from tqdm import tqdm
import random
import numpy as np
import plotly.graph_objects as go
@@ -501,6 +502,32 @@ def speed_step(levels_input, conds):
return model.sample([levels_input, conds], training=False)
from scipy.spatial import KDTree
import trimesh
def connect_unconnected_vertices(mesh):
# Find unconnected vertices
unconnected_verts = np.where(np.array(mesh.vertex_degree) == 0)[0]
# Create KD-tree from mesh vertices
tree = KDTree(mesh.vertices)
new_faces = mesh.faces.copy()
# Connect unconnected vertices to their two closest neighbors
for vert in unconnected_verts:
# Find the two closest vertices
_, closest_verts = tree.query(mesh.vertices[vert], k=3)
# Create a new face with the unconnected vertex and its two closest neighbors
new_face = [vert, closest_verts[1], closest_verts[2]]
new_faces = np.vstack((new_faces, new_face))
mesh = trimesh.Trimesh(vertices=mesh.vertices, faces=new_faces, process=False)
return mesh
def sample():
from data.utils import Frame
from utils.utils import save_stl
@@ -515,11 +542,10 @@ def sample():
load_model(model, f'checkpoints/iccvw/{config["name"]}')
count = 0
frames = ['01852_70_Tshirt']
frames = ['06461_237_Trousers'] # 03528_132_Tshirt 01852_70_Tshirt 06461_237_Trousers 01893_173_Jumpsuit
frames = ['01171_213_Dress'] # 01039_118_Dress 01171_213_Dress
frames = [Frame(frame_str=frame) for frame in frames]
for frame in frames:
create_temp_frames(config['train_dir'], frame)
dataset = Dataset(config['train_dir'], config['signature'], device=config['device'], frames_file='frames_temp.txt', use_cache=True, resolutions=[64, 128, 256, 512], is_cape=config['is_cape'])
@@ -532,6 +558,8 @@ def sample():
uv_normals = batch['uv_normals']
batch_size = uv_static.size(0)
body = batch['mesh_body_posed'][0].cpu().numpy()
body_faces = smpl.f
for i in range(batch_size):
@@ -541,8 +569,11 @@ def sample():
mask = fix_data['mask']
faces = fix_data['faces']
os.makedirs(Path('temp') / 'sampling_pro_pyramid' / f'{frame.id}_{frame.frame_nr}_{frame.garment}', exist_ok=True)
gt = batch['mesh_posed'][i].cpu().numpy()[mask]
save_stl(Path('temp') / 'sampling_pyramid' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_gt.stl', gt, faces)
save_stl(Path('temp') / 'sampling_pro_pyramid' / f'{frame.id}_{frame.frame_nr}_{frame.garment}' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_gt.stl', gt, faces)
save_stl(Path('temp') / 'sampling_pro_pyramid' / f'{frame.id}_{frame.frame_nr}_{frame.garment}' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_gt_body.stl', body, body_faces)
uv_static_tiled = uv_static[i][None,:,:,:].repeat(4, 1, 1, 1)
uv_body_posed_tiled = uv_body_posed[i][None,:,:,:].repeat(4, 1, 1, 1)
@@ -551,6 +582,8 @@ def sample():
pose_shape = batch['pose_shape'][i][None,:].repeat(4, 1)
body_posed = batch['mesh_body_posed'][i][None,:,:].repeat(4, 1, 1)
mask_tiled = torch.from_numpy(mask[None,:]).to(body_posed.device).repeat(4, 1)
noise = []
for _ in range(4):
with torch.inference_mode():
@@ -570,9 +603,12 @@ def sample():
for k in range(4):
_mesh_posed_fixed = mesh_posed_fixed[k, mask, :]
m = tm.Trimesh(_mesh_posed_fixed, faces, process=False, validate=False)
m = connect_unconnected_vertices(m)
tm.smoothing.filter_humphrey(m, iterations=2)
verts = m.vertices
save_stl(Path('temp') / 'sampling_pyramid' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_{counter}.stl', verts, faces)
save_stl(Path('temp') / 'sampling_pro_pyramid' / f'{frame.id}_{frame.frame_nr}_{frame.garment}' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_{counter}.stl', verts, faces)
save_stl(Path('temp') / 'sampling_pro_pyramid' / f'{frame.id}_{frame.frame_nr}_{frame.garment}' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_{counter}_body.stl', body, body_faces)
counter += 1
pass
@@ -582,6 +618,265 @@ def sample():
break
def sample_pro():
from data.utils import Frame
from utils.utils import save_stl
from pathlib import Path
import trimesh as tm
def create_temp_frames(folder: str, frame: Frame):
with open(Path(folder) / 'frames_temp.txt', 'w') as f:
f.writelines([f'{frame}\n' for _ in range(4)])
model.eval()
load_model(model, f'checkpoints/iccvw/{config["name"]}')
count = 0
# frames = ['06461_237_Trousers'] # 03528_132_Tshirt 01852_70_Tshirt 06461_237_Trousers 01893_173_Jumpsuit
frames = ['01852_70_Tshirt'] # 01039_118_Dress 01171_213_Dress 01852_70_Tshirt
frames = [Frame(frame_str=frame) for frame in frames]
for frame in frames:
create_temp_frames(config['train_dir'], frame)
dataset = Dataset(config['train_dir'], config['signature'], device=config['device'], frames_file='frames_temp.txt', use_cache=True, resolutions=[64, 128, 256, 512], is_cape=config['is_cape'])
dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=True, )
batch = next(iter(dataloader))
uv_static = batch['uv_static']
uv_body_posed = batch['uv_body_posed']
uv_mask = { 512: batch['uv_mask'], 256: batch['uv_mask_256'], 128: batch['uv_mask_128'], 64: batch['uv_mask_64'] }
uv_normals = batch['uv_normals']
batch_size = uv_static.size(0)
body = batch['mesh_body_posed'][0].cpu().numpy()
body_faces = smpl.f
i=0
frame = Frame(batch['id'][i], batch['frame_nr'][i], batch['garment'][i])
fix_data = np.load(Path(dataset_train.folder) / 'mesh_fixed' / f'{frame.id}_{frame.garment}.npz')
mask = fix_data['mask']
faces = fix_data['faces']
os.makedirs(Path('temp') / 'sampling_pro_pyramid' / f'{frame.id}_{frame.frame_nr}_{frame.garment}', exist_ok=True)
gt = batch['mesh_posed'][i].cpu().numpy()[mask]
save_stl(Path('temp') / 'sampling_pro_pyramid' / f'{frame.id}_{frame.frame_nr}_{frame.garment}' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_gt.stl', gt, faces)
save_stl(Path('temp') / 'sampling_pro_pyramid' / f'{frame.id}_{frame.frame_nr}_{frame.garment}' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_gt_body.stl', body, body_faces)
uv_static_tiled = uv_static[i][None,:,:,:].repeat(4, 1, 1, 1)
uv_body_posed_tiled = uv_body_posed[i][None,:,:,:].repeat(4, 1, 1, 1)
uv_mask_tiled = { 512: uv_mask[512][i][None,:,:].repeat(4, 1, 1), 256: uv_mask[256][i][None,:,:].repeat(4, 1, 1), 128: uv_mask[128][i][None,:,:].repeat(4, 1, 1), 64: uv_mask[64][i][None,:,:].repeat(4, 1, 1)}
uv_normals_tiled = uv_normals[i][None,:,:,:].repeat(4, 1, 1, 1)
pose_shape = batch['pose_shape'][i][None,:].repeat(4, 1)
body_posed = batch['mesh_body_posed'][i][None,:,:].repeat(4, 1, 1)
mask_tiled = torch.from_numpy(mask[None,:]).to(body_posed.device).repeat(4, 1)
out = torch.zeros((uv_static_tiled.shape[0], 3, 64, 32)).to(uv_static_tiled.device)
for l in range(4):
with torch.inference_mode():
for b in range(4):
pred_level = model.sample_level(uv_static_tiled, uv_body_posed_tiled, uv_normals_tiled, uv_mask_tiled, out, l)
pred = pred_level
mesh_pred = uv2mesh(pred)
shape = pose_shape[:, 72:]
gender = torch.greater_equal(shape[:, 10] , 0.5).to(torch.long)
shape_offsets = torch.stack([torch.tensordot(shape[i,:10][None, :], smpl.shapedirs[gender[i]], dims=[[1], [2]])[0,smpl.to_keep,:] for i in range(shape.shape[0])])
mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,None]
mesh_pred_posed = smpl(pose_shape[:, :72], pose_shape[:, 72:], mesh_pred_shaped)[0]
mesh_posed_fixed = fix_collision(mesh_pred_posed, body_posed, mask_tiled).cpu().numpy()
for k in range(4):
_mesh_posed_fixed = mesh_posed_fixed[k, mask, :]
m = tm.Trimesh(_mesh_posed_fixed, faces, process=False, validate=False)
m = connect_unconnected_vertices(m)
tm.smoothing.filter_humphrey(m, iterations=2)
verts = m.vertices
save_stl(Path('temp') / 'sampling_pro_pyramid' / f'{frame.id}_{frame.frame_nr}_{frame.garment}' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_lvl_{l}_{b*4+k}.stl', verts, faces)
save_stl(Path('temp') / 'sampling_pro_pyramid' / f'{frame.id}_{frame.frame_nr}_{frame.garment}' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_lvl_{l}_{b*4+k}_body.stl', body, body_faces)
#replace the contentent of out with its own first element tiled
if b == 3:
out = pred_level[3,:,:,:].repeat(4, 1, 1, 1)
pass
# if count > 100:
# break
def pred_pro():
from data.utils import Frame
from utils.utils import save_stl
from pathlib import Path
import trimesh as tm
from typing import List
from models.cvae_norm_model import CVAENorm
from data.downscale import fill
from torchvision.transforms import Resize
from torchvision.transforms.functional import InterpolationMode
def create_temp_frames(folder: str, frames: List[Frame]):
with open(Path(folder) / 'frames_temp.txt', 'w') as f:
f.writelines([f'{frame}\n' for frame in frames])
model.eval()
load_model(model, f'checkpoints/iccvw/{config["name"]}')
count = 0
from configs.iccvw.pyr_64_128_256_512_skirt_e2e import config as config_skirt
model_skirt: Pyramid = Pyramid(config_skirt)
load_model(model_skirt, f'checkpoints/iccvw/{config_skirt["name"]}')
models = {'non-skirt': model, 'skirt': model_skirt}
from configs.iccvw.cvae_norm_5 import config as config_norm
model_norm = CVAENorm(config_norm)
model_norm.eval()
load_model(model_norm, f'checkpoints/iccvw/{config_norm["name"]}')
uv2mesh_skirt = UV2Mesh(res=256, is_skirt=config_skirt['is_skirt'], is_smpl_orig=config['is_cape'])
uv2meshes = {'non-skirt': uv2mesh, 'skirt': uv2mesh_skirt}
# '01320_30_Dress'
frames = ['01852_70_Tshirt', '04717_276_Trousers', '07532_152_Tshirt', '02792_282_Jumpsuit'] # 01039_118_Dress 01171_213_Dress
frames = [Frame(frame_str=frame) for frame in frames]
create_temp_frames(config['train_dir'], frames)
dataset = Dataset(config['train_dir'], config['signature'], device=config['device'], frames_file='frames_temp.txt', use_cache=True, resolutions=[64, 128, 256, 512], is_cape=config['is_cape'])
dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=False, )
batch_orig = next(iter(dataloader))
batch_size = 4
body_faces = smpl.f
masks = []
faces_orig = []
for frame in frames:
fix_data = np.load(Path(dataset_train.folder) / 'mesh_fixed' / f'{frame.id}_{frame.garment}.npz')
masks.append(fix_data['mask'])
faces_orig.append(fix_data['faces'])
masks = np.stack(masks, axis=0)
# faces = np.stack(faces, axis=0)
resizers = { res: Resize(size=(res, int(res/2)), interpolation=InterpolationMode.BICUBIC).to('cuda') for res in [64, 128, 256, 512] }
mode = 'body'
for i in range(batch_size):
batch = batch_orig.copy()
uv_static = batch['uv_static']
uv_body_posed = batch['uv_body_posed']
uv_mask = { 512: batch['uv_mask'], 256: batch['uv_mask_256'], 128: batch['uv_mask_128'], 64: batch['uv_mask_64'] }
uv_normals = batch['uv_normals']
counter = 0
frame = Frame(batch['id'][i], batch['frame_nr'][i], batch['garment'][i])
# fix_data = np.load(Path(dataset_train.folder) / 'mesh_fixed' / f'{frame.id}_{frame.garment}.npz')
# mask = fix_data['mask']
# faces = fix_data['faces']
type = 'skirt' if frame.garment in ['Skirt', 'Dress'] else 'non-skirt'
os.makedirs(Path('temp') / 'pred_pro' / mode / f'{frame.id}_{frame.frame_nr}_{frame.garment}', exist_ok=True)
# gt = batch['mesh_posed'][i].cpu().numpy()[mask]
# save_stl(Path('temp') / 'pred_pro' / f'{frame.id}_{frame.frame_nr}_{frame.garment}' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_gt.stl', gt, faces)
# save_stl(Path('temp') / 'pred_pro' / f'{frame.id}_{frame.frame_nr}_{frame.garment}' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_gt_body.stl', body, body_faces)
uv_static_tiled = uv_static[i][None,:,:,:].repeat(4, 1, 1, 1)
uv_body_posed_tiled = uv_body_posed[i][None,:,:,:].repeat(4, 1, 1, 1)
uv_mask_tiled = { 512: uv_mask[512][i][None,:,:].repeat(4, 1, 1), 256: uv_mask[256][i][None,:,:].repeat(4, 1, 1), 128: uv_mask[128][i][None,:,:].repeat(4, 1, 1), 64: uv_mask[64][i][None,:,:].repeat(4, 1, 1)}
uv_normals_tiled = uv_normals[i][None,:,:,:].repeat(4, 1, 1, 1)
pose_shape = batch['pose_shape'][i][None,:].repeat(4, 1)
body_posed = batch['mesh_body_posed'][i][None,:,:].repeat(4, 1, 1)
# masks_tiled = torch.from_numpy(masks[i][None,:]).to(body_posed.device).repeat(4, 1)
if mode == 'static':
batch['uv_body_posed'] = uv_body_posed_tiled
batch['mesh_body_posed'] = body_posed
batch['pose_shape'] = pose_shape
masks_tiled = masks
for res in [64, 128, 256]:
batch[f'uv_body_posed_{res}'] = batch[f'uv_body_posed_{res}'][i][None,:,:,:].repeat(4, 1, 1, 1)
# orig_img = uv_body_posed_tiled.clone()
# downscaled = {64: [], 128: [], 256: []}
# for i in range(4):
# filled = fill(orig_img[i].cpu().numpy().transpose(1,2,0), uv_mask[512][i].cpu().numpy())
# filled = torch.from_numpy(filled.transpose(2,0,1)).to(orig_img.device)
# for res in [64, 128, 256]:
# downscaled[res].append(resizers[res](filled) * uv_mask[res][i])
# for res in [64, 128, 256]:
# batch[f'uv_body_posed_{res}'] = torch.stack(downscaled[res], dim=0)
faces = faces_orig
elif mode == 'body':
batch['uv_static'] = uv_static_tiled
for res in [64, 128, 256]:
batch[f'uv_static_{res}'] = batch[f'uv_static_{res}'][i][None,:,:,:].repeat(4, 1, 1, 1)
batch['mesh_mask'] = batch['mesh_mask'][i][None,:].repeat(4, 1)
for key in uv_mask_tiled.keys():
batch[f'uv_mask_{key}'] = uv_mask_tiled[key]
batch['uv_mask'] = uv_mask_tiled[512]
masks_tiled = torch.from_numpy(masks[i][None,:]).to(body_posed.device).repeat(4, 1).cpu().numpy()
faces = [faces_orig[i], faces_orig[i], faces_orig[i], faces_orig[i]]
else:
raise Exception('mode not supported')
body_posed = batch['mesh_body_posed']
pose_shape = batch['pose_shape']
with torch.inference_mode():
pred_norm = model_norm.sample(batch['uv_static'], batch['uv_body_posed'], batch['uv_mask']) * 1.4
batch['uv_normals'] = pred_norm
orig_img = pred_norm.clone()
downscaled = {64: [], 128: [], 256: []}
for i in range(4):
filled = fill(orig_img[i].cpu().numpy().transpose(1,2,0), uv_mask[512][i].cpu().numpy())
filled = torch.from_numpy(filled.transpose(2,0,1)).to(orig_img.device)
for res in [64, 128, 256]:
downscaled[res].append(resizers[res](filled) * uv_mask[res][i])
for res in [64, 128, 256]:
batch[f'uv_normals_{res}'] = torch.stack(downscaled[res], dim=0)
levels_input, conds = group_inputs_for_levels(batch, config)
with torch.inference_mode():
out_levels = models[type]([levels_input, conds])
pred, _, _ = out_levels[-1]
mesh_pred = uv2meshes[type](pred)
shape = pose_shape[:, 72:]
gender = torch.greater_equal(shape[:, 10] , 0.5).to(torch.long)
shape_offsets = torch.stack([torch.tensordot(shape[i,:10][None, :], smpl.shapedirs[gender[i]], dims=[[1], [2]])[0,smpl.to_keep,:] for i in range(shape.shape[0])])
mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,None]
mesh_pred_posed = smpl(pose_shape[:, :72], pose_shape[:, 72:], mesh_pred_shaped)[0]
# mesh_posed_fixed = fix_collision(mesh_pred_posed, body_posed, torch.from_numpy(masks).to(body_posed.device)).cpu().numpy()
mesh_posed_fixed = mesh_pred_posed.cpu().numpy()
for k in range(4):
_mesh_posed_fixed = mesh_posed_fixed[k, masks_tiled[k], :]
m = tm.Trimesh(_mesh_posed_fixed, faces[k], process=False, validate=False)
m = connect_unconnected_vertices(m)
tm.smoothing.filter_humphrey(m, iterations=2)
verts = m.vertices
save_stl(Path('temp') / 'pred_pro' / mode / f'{frame.id}_{frame.frame_nr}_{frame.garment}' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_{counter}.stl', verts, faces[k])
save_stl(Path('temp') / 'pred_pro' / mode / f'{frame.id}_{frame.frame_nr}_{frame.garment}' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_{counter}_body.stl', batch['mesh_body_posed'][k].cpu().numpy(), body_faces)
counter += 1
pass
def train():
global g_balancing_coeff
global step
@@ -620,9 +915,44 @@ def train():
step += 1
def speed_test():
from data.utils import Frame
from utils.utils import save_stl
from pathlib import Path
import trimesh as tm
def create_temp_frames(folder: str, frame: Frame):
with open(Path(folder) / 'frames_temp.txt', 'w') as f:
f.writelines([f'{frame}\n' for _ in range(4)])
model.eval()
load_model(model, f'checkpoints/iccvw/{config["name"]}')
frames = ['01852_70_Tshirt']
frames = [Frame(frame_str=frame) for frame in frames]
create_temp_frames(config['train_dir'], frames[0])
dataset = Dataset(config['train_dir'], config['signature'], device=config['device'], frames_file='frames_temp.txt', use_cache=True, resolutions=[64, 128, 256, 512], is_cape=config['is_cape'])
dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=True, )
batch = next(iter(dataloader))
uv_static = batch['uv_static']
uv_body_posed = batch['uv_body_posed']
uv_mask = { 512: batch['uv_mask'], 256: batch['uv_mask_256'], 128: batch['uv_mask_128'], 64: batch['uv_mask_64'] }
uv_normals = batch['uv_normals']
for i in tqdm(range(1000)):
with torch.inference_mode():
if i ==5:
start = time.time()
model.sample(uv_static, uv_body_posed, uv_normals, uv_mask)
print(f'per batch: {(time.time() - start) / 1000}')
if __name__ == '__main__':
if not args['debug']:
init_wandb()
# if not args['debug']:
# init_wandb()
use_amp = True
smpl = SMPL(orig_smpl=config['is_cape'], path='smpl', device=config['device'])
@@ -640,7 +970,14 @@ if __name__ == '__main__':
model: Pyramid = Pyramid(config)
sample()
# count number of parameters
total_params = sum(p.numel() for p in model.parameters())
# speed_test()
# exit()
# sample_pro()
pred_pro()
exit()
optimizer = optim.Adam(model.parameters(), lr=config['lr'])
@@ -652,23 +989,23 @@ if __name__ == '__main__':
reg_coeff = 1
step = 0
if config['pretrained_model'] != '':
if 'ignore_saved_optimizer' in config.keys() and config['ignore_saved_optimizer']:
step = load_model(model, config['pretrained_model'])
else:
step = load_model(model, config['pretrained_model'], optimizer=optimizer, iteration=config['pretrained_model_iteration'])
# if config['pretrained_model'] != '':
# if 'ignore_saved_optimizer' in config.keys() and config['ignore_saved_optimizer']:
# step = load_model(model, config['pretrained_model'])
# else:
# step = load_model(model, config['pretrained_model'], optimizer=optimizer, iteration=config['pretrained_model_iteration'])
if len(config['pyramid']['resolutions']) == 1:
for param in model.static_enc.parameters(): param.requires_grad = False
for param in model.body_enc.parameters(): param.requires_grad = False
if config['use_normals']:
for param in model.norm_enc.parameters(): param.requires_grad = False
else:
load_model(model, config['pyramid']['pretrained_pyramid'], iteration=config['pyramid']['pretrained_pyramid_iteration'])
for param in model.static_enc.parameters(): param.requires_grad = False
for param in model.body_enc.parameters(): param.requires_grad = False
if config['use_normals']:
for param in model.norm_enc.parameters(): param.requires_grad = False
# if len(config['pyramid']['resolutions']) == 1:
# for param in model.static_enc.parameters(): param.requires_grad = False
# for param in model.body_enc.parameters(): param.requires_grad = False
# if config['use_normals']:
# for param in model.norm_enc.parameters(): param.requires_grad = False
# else:
# load_model(model, config['pyramid']['pretrained_pyramid'], iteration=config['pyramid']['pretrained_pyramid_iteration'])
# for param in model.static_enc.parameters(): param.requires_grad = False
# for param in model.body_enc.parameters(): param.requires_grad = False
# if config['use_normals']:
# for param in model.norm_enc.parameters(): param.requires_grad = False
# load_model(model.levels[0].upscale_enhancer, config['pyramid']['levels'][0]['upscale_enhancer'])
# model.levels[0].upscale_enhancer.trainable = False
@@ -691,4 +1028,5 @@ if __name__ == '__main__':
# model = torch.compile(model)
wandb.watch(model, log='all', log_freq=10000, log_graph=True)
model.eval()
train()
scripts/cvae_pyramid_e2e.py
+4 -5
View File
@@ -547,7 +547,7 @@ def train():
valid_accumulator.update(valid_loss)
if step % 1000 == 0 or (step < 1000 and step % 100 == 0):
to_log = dict({ **train_loss, **valid_loss, **{ 'current_step': step } })
to_log = dict({ **train_loss, **valid_loss, **{ 'current_step': step, 'balancing_coeff': g_balancing_coeff } })
if step % 5000 == 0 or (step < 5000 and step % 1000 == 0):
train_figs = plot_everything(batch, 'train')
valid_figs = plot_everything(data_valid, 'valid')
@@ -561,10 +561,9 @@ def train():
valid_accumulator.reset()
print(step)
# if train_loss['train/euclid_error'] < config['balancing_activation_limit']:
# g_balancing_coeff += 1 / 50000
g_balancing_coeff += 1 / 100000
if step % 5000 == 0 and not args['debug']:
if step % 10000 == 0 and not args['debug']:
ckpt_saver(model, step, optimizer, None)
step += 1
@@ -593,7 +592,7 @@ if __name__ == '__main__':
# optimizer = optim.Adam(model.parameters(), lr=config['lr'])
optimizer = optim.Adam(model.parameters(), lr=config['lr'])
_MSSSIM = MS_SSIM(data_range=2, size_average=True, channel=3)
ckpt_saver = CheckpointSaver('checkpoints/iccvw/' + config['name'], top_n=50)
ckpt_saver = CheckpointSaver('checkpoints/iccvw/' + config['name'], top_n=20)
scaler = torch.cuda.amp.GradScaler(enabled=use_amp)
resizer = Resize(size=(512, 256), interpolation=InterpolationMode.BILINEAR).to(config['device'])
scripts/generate_results.py
+1 -1
View File
@@ -14,7 +14,7 @@ parser.add_argument('-l', '--level', type=int, default=-1)
args = vars(parser.parse_args())
if args['debug']:
from configs.iccvw.pyr_64_128_cape_f_2 import config_1
from configs.iccvw.pyr_64_128_256_512_e2e_var import config_1
else:
config_module = importlib.import_module(args['config'])
config_1 = config_module.config
scripts/generate_results_all.py
+3 -1
View File
@@ -10,6 +10,7 @@ parser.add_argument('--config', type=str)
parser.add_argument('--iteration', type=str, default='')
parser.add_argument('--debug', action='store_true')
parser.add_argument('-g', '--gpu', type=str, default='1')
parser.add_argument('-fr', '--frames_file', type=str, default='')
args = vars(parser.parse_args())
if args['debug']:
@@ -117,7 +118,8 @@ if __name__ == '__main__':
# remove all values from config['signature] that start with mesh_
config_1['signature'] = set([x for x in config_1['signature'] if not x.startswith('mesh_')])
dataset_test = Dataset(config_1['test_dir'], config_1['signature'], device=config_1['device'], frames_file=config_1['test_frames'], use_cache=config_1['dataset_in_memory'], resolutions=dataset_resolutions, is_cape=config_1['is_cape'])
frames_file = args['frames_file'] if args['frames_file'] != '' else config_1['test_frames']
dataset_test = Dataset(config_1['test_dir'], config_1['signature'], device=config_1['device'], frames_file=frames_file, use_cache=config_1['dataset_in_memory'], resolutions=dataset_resolutions, is_cape=config_1['is_cape'])
# dataset_test = Dataset(config['test_dir'], config['signature'], device=config['device'], frames_file='frames_temp.txt', use_cache=config['dataset_in_memory'], resolutions=dataset_resolutions, is_cape=config['is_cape'])
test_dataloader = DataLoader(dataset_test, batch_size=4, drop_last=True, shuffle=True, )
scripts/get_error.py
+20 -12
View File
@@ -8,13 +8,16 @@ parser = argparse.ArgumentParser(description='')
parser.add_argument('--name', type=str)
parser.add_argument('--config', type=str)
parser.add_argument('--debug', action='store_true')
parser.add_argument('--iteration', type=str, default='190000')
parser.add_argument('--iteration', type=str, default='')
parser.add_argument('-g', '--gpu', type=str)
parser.add_argument('-n', '--normal', action='store_true')
parser.add_argument('-e', '--edge', action='store_true')
parser.add_argument('-cr', '--collision_ratio', action='store_true')
parser.add_argument('-a', '--all', action='store_true')
args = vars(parser.parse_args())
if args['debug']:
from configs.iccvw.pyr_64_128_256_512_cape_f_e2e import config as config_1
from configs.iccvw.pyr_64_128_256_512_e2e import config as config_1
else:
config_module = importlib.import_module(args['config'])
config_1 = config_module.config
@@ -32,7 +35,7 @@ from data.utils import Frame
def calculate():
def calculate(type=''):
all_levels = args['all']
if all_levels:
errors = {0: dict(), 1: dict(), 2: dict(), 3: dict()}
@@ -41,8 +44,7 @@ def calculate():
errors = dict()
errors_mean = dict()
with open(f'{results_folder}/errors/{name}.csv', 'r') as f:
with open(f'{results_folder}/errors/{name}{type}.csv', 'r') as f:
lines = f.readlines()
for line in lines:
@@ -50,11 +52,11 @@ def calculate():
category = line[2]
if all_levels:
level = int(line[3][-1])
if category not in errors:
if category not in errors[level].keys():
errors[level][category] = []
errors[level][category].append(float(line[4]))
else:
if category not in errors:
if category not in errors.keys():
errors[category] = []
errors[category].append(float(line[3]))
@@ -62,17 +64,18 @@ def calculate():
for i in range(4):
for category in errors[i]:
errors_mean[i][category] = round(np.mean(errors[i][category]), 5)
print(f'Level {i}: {errors_mean[i]}')
print(f'Level {i} error{type}: {errors_mean[i]}')
mean = np.mean(list(errors_mean[i].values()))
print(f'Final error level {i}: {mean}')
print(f'Final error{type} level {i}: {mean}')
pass
else:
for category in errors:
errors_mean[category] = round(np.mean(errors[category]), 5)
print(errors_mean)
mean = np.mean(list(errors_mean.values()))
print(f'Final error: {mean}')
mean = round(np.mean(list(errors_mean.values())), 5)
print(f'Final error{type}: {mean}')
@@ -98,5 +101,10 @@ if __name__ == '__main__':
name = f'{name}_{args["iteration"]}'
out_folder = f'{results_folder}/{name}'
if args['normal']:
calculate('_normal')
if args['edge']:
calculate('_edge')
if args['collision_ratio']:
calculate('_collision_ratio')
calculate()
scripts/get_error_hood.py
+28 -46
View File
@@ -2,30 +2,9 @@ import os
import sys
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
import importlib
import argparse
parser = argparse.ArgumentParser(description='')
parser.add_argument('--name', type=str)
parser.add_argument('--config', type=str)
parser.add_argument('--debug', action='store_true')
parser.add_argument('--iteration', type=str, default='')
parser.add_argument('-g', '--gpu', type=str)
parser.add_argument('-n', '--normal', action='store_true')
parser.add_argument('-e', '--edge', action='store_true')
parser.add_argument('-cr', '--collision_ratio', action='store_true')
parser.add_argument('-a', '--all', action='store_true')
args = vars(parser.parse_args())
if args['debug']:
from configs.iccvw.pyr_64_128_256_512_cape_f_e2e import config as config_1
else:
config_module = importlib.import_module(args['config'])
config_1 = config_module.config
import time
import argparse
import tqdm
from tqdm import tqdm
import random
import numpy as np
import plotly.graph_objects as go
@@ -35,50 +14,53 @@ from data.utils import Frame
def calculate(type=''):
def calculate(type='chamfer'):
errors = dict()
errors_mean = dict()
with open(f'{folder}/{hood_csv}', 'r') as f:
lines_hood = f.readlines()
with open(f'{folder}/{ours_csv}', 'r') as f:
lines_ours = f.readlines()
files_hood = sorted(os.listdir(f'{folder}/{run}_hood/{type}'))
errors_hood = [np.load(f'{folder}/{run}_hood/{type}/{file}') for file in tqdm(files_hood)]
files_ours = sorted(os.listdir(f'{folder}/hood_3_var_ours/{type}'))
errors_ours = [np.load(f'{folder}/hood_3_var_ours/{type}/{file}') for file in tqdm(files_ours)]
for line in lines_hood:
line = line.split(',')
category = line[2]
for error, frame in zip(errors_hood, files_hood):
frame_id, frame_nr, frame_garment = frame.split('_')
category = frame_garment
if category not in errors:
errors[category] = []
errors[category].append(float(line[3]))
errors[category].append(error)
for category in errors:
errors_mean[category] = round(np.mean(errors[category]), 5)
print(errors_mean)
mean = round(np.mean(list(errors_mean.values())), 5)
print(f'Final HOOD error{type}: {mean}')
print(f'Final HOOD error_{type}: {mean}')
for line in lines_ours:
line = line.split(',')
category = line[2]
# errors = dict()
# errors_mean = dict()
if category not in errors:
errors[category] = []
errors[category].append(float(line[3]))
# for error, frame in zip(errors_ours, files_ours):
# frame_id, frame_nr, frame_garment = frame.split('_')
# category = frame_garment
for category in errors:
errors_mean[category] = round(np.mean(errors[category]), 5)
print(errors_mean)
# if category not in errors:
# errors[category] = []
# errors[category].append(error)
mean = round(np.mean(list(errors_mean.values())), 5)
print(f'Final OURS error{type}: {mean}')
# for category in errors:
# errors_mean[category] = round(np.mean(errors[category]), 5)
# print(errors_mean)
# mean = round(np.mean(list(errors_mean.values())), 5)
# print(f'Final OURS error_{type}: {mean}')
if __name__ == '__main__':
folder = '/data-local/127_2/results/hood'
hood_csv = 'hood_2_hood.csv'
ours_csv = 'hood_2_ours.csv'
run = 'hood_3'
calculate()
for type in ['chamfer', 'area', 'collision_ratio']: # ['chamfer', 'collision', 'area', 'collision_ratio']:
calculate(type)
scripts/plot_ablation.py
+43 -28
View File
@@ -25,7 +25,7 @@ import importlib
from data.utils import Frame
from data.dataset import read_frames_file
from utils.utils import save_stl
from utils.utils import save_stl, create_temp_frames
from utils.checkpoints import CheckpointSaver, load_model
from utils.postprocess import fix_collision
from losses.losses import collision_loss, normal_loss, kl_loss
@@ -70,25 +70,21 @@ def step_pyramid(model, batch, config):
return mesh_pred_posed
def create_temp_frames(folder: str, frame: Frame):
with open(Path(folder) / 'frames_temp.txt', 'w') as f:
f.writelines([f'{frame}\n' for _ in range(4)])
if __name__ == '__main__':
results_folder = Path('/data-local/127_2/results')
out_folder = Path('temp/ablation')
out_folder = Path('temp/supplementary')
os.makedirs(out_folder, exist_ok=True)
smpl = SMPL(orig_smpl=config_1['is_cape'], path='smpl', device=config_1['device'])
dataset_resolutions = []
dataset_resolutions = [64, 128, 256, 512]
experiments = {
'final': { 'nonskirt': ('pyr_64_128_256_512', 100000), 'skirt': ('pyr_64_128_256_512_skirt', 200000) },
'e2e': { 'nonskirt': ('pyr_64_128_256_512_e2e', 195000), 'skirt': ('pyr_64_128_256_512_skirt_e2e', 200000) },
'skin': { 'nonskirt': ('pyr_64_128_256_512_skin', 225000), 'skirt': ('pyr_64_128_256_512_skin_skirt', 200000) },
'norm': { 'nonskirt': ('pyr_64_128_256_512_norm', 200000), 'skirt': ('pyr_64_128_256_512_norm_skirt', 200000) },
'reg': { 'nonskirt': ('pyr_64_128_256_512_reg', 200000), 'skirt': ('pyr_64_128_256_512_reg_skirt', 200000) },
# 'inc': { 'nonskirt': ('pyr_64_128_256_512', 100000), 'skirt': ('pyr_64_128_256_512_skirt', 200000) },
'e2e': { 'nonskirt': ('pyr_64_128_256_512_e2e', 250000), 'skirt': ('pyr_64_128_256_512_skirt_e2e', 200000) },
# 'skin': { 'nonskirt': ('pyr_64_128_256_512_skin', 225000), 'skirt': ('pyr_64_128_256_512_skin_skirt', 200000) },
# 'norm': { 'nonskirt': ('pyr_64_128_256_512_norm', 200000), 'skirt': ('pyr_64_128_256_512_norm_skirt', 200000) },
# 'reg': { 'nonskirt': ('pyr_64_128_256_512_reg', 200000), 'skirt': ('pyr_64_128_256_512_reg_skirt', 200000) },
}
configs = dict()
@@ -111,8 +107,23 @@ if __name__ == '__main__':
# frames = ['00054_1_Dress', '00005_1_Jumpsuit']
frames = ['03849_27_Dress', '00054_1_Dress', '01852_70_Tshirt']
frames = [Frame(frame_str=frame) for frame in frames]
# frames = ['03849_27_Dress', '00054_1_Dress', '01852_70_Tshirt']
# frames = ['01852_70_Tshirt'] # , '03849_27_Dress', '00054_1_Dress'
# frames = [Frame(frame_str=frame) for frame in frames]
frames_file = configs['e2e']['nonskirt']['train_frames']
with open(os.path.join(config_1['train_dir'], frames_file), 'r') as f:
frames = f.readlines()
frames_1 = [Frame(frames.split()[0], int(frames.split()[1]), frames.split()[2]) for frames in frames]
frames_file = configs['e2e']['skirt']['train_frames']
with open(os.path.join(config_1['train_dir'], frames_file), 'r') as f:
frames = f.readlines()
frames_2 = [Frame(frames.split()[0], int(frames.split()[1]), frames.split()[2]) for frames in frames]
frames = frames_1 + frames_2
random.shuffle(frames)
body_faces = smpl.f
for frame in frames:
@@ -130,8 +141,8 @@ if __name__ == '__main__':
body = batch['mesh_body_posed'][0].cpu().numpy()
gt = batch['mesh_posed'][0].cpu().numpy()
save_stl(out_folder / f'{frame}_gt.stl', gt[mask, :], faces)
save_stl(out_folder / f'{frame}_gt_body.stl', body, body_faces)
# save_stl(out_folder / f'{frame}_gt.stl', gt[mask, :], faces)
# save_stl(out_folder / f'{frame}_gt_body.stl', body, body_faces)
for experiment in experiments.keys():
@@ -139,20 +150,24 @@ if __name__ == '__main__':
mesh_posed_fixed = fix_collision(pred, batch['mesh_body_posed'], batch['mesh_mask']).cpu().numpy()
mesh_posed_fixed = mesh_posed_fixed[0, mask, :]
m = tm.Trimesh(mesh_posed_fixed, faces, process=False, validate=False)
tm.smoothing.filter_humphrey(m, iterations=2)
verts = m.vertices
try:
m = tm.Trimesh(mesh_posed_fixed, faces, process=False, validate=False)
tm.smoothing.filter_humphrey(m, iterations=2)
verts = m.vertices
except:
verts = mesh_posed_fixed
print('failed to apply smoothing')
save_stl(out_folder / f'{frame}_{experiment}_body.stl', body, body_faces)
save_stl(out_folder / f'{frame}_{experiment}_mesh.stl', verts, faces)
# save_stl(out_folder / f'{frame}_{experiment}_body.stl', body, body_faces)
# save_stl(out_folder / f'{frame}_{experiment}_mesh.stl', verts, faces)
# fig = go.Figure(data=[go.Mesh3d(x=verts[:,0], y=verts[:,2], z=verts[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=1),
# go.Mesh3d(x=body[:,0], y=body[:,2], z=body[:,1], i=body_faces[:,0], j=body_faces[:,1], k=body_faces[:,2], color='gray', opacity=1)
# ])
# fig.update_layout(scene=dict(aspectmode='cube'))
# fig.update_layout(scene=dict(xaxis=dict(range=[-1, 1], ), yaxis=dict(range=[-1, 1], ), zaxis=dict(range=[-1, 1], ), ), )
# fig.update_layout(title_text=f'{frame} {experiment}')
# fig.show()
fig = go.Figure(data=[go.Mesh3d(x=verts[:,0], y=verts[:,2], z=verts[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=1),
go.Mesh3d(x=body[:,0], y=body[:,2], z=body[:,1], i=body_faces[:,0], j=body_faces[:,1], k=body_faces[:,2], color='gray', opacity=1)
])
fig.update_layout(scene=dict(aspectmode='cube'))
fig.update_layout(scene=dict(xaxis=dict(range=[-1, 1], ), yaxis=dict(range=[-1, 1], ), zaxis=dict(range=[-1, 1], ), ), )
fig.update_layout(title_text=f'{frame} {experiment}')
fig.show()
pass
pass
scripts/plot_ablation_pyramid.py
+15 -17
View File
@@ -24,7 +24,7 @@ from pathlib import Path
from data.utils import Frame
from data.dataset import read_frames_file
from utils.utils import save_stl
from utils.utils import save_stl, create_temp_frames
from utils.checkpoints import CheckpointSaver, load_model
from utils.postprocess import fix_collision
from losses.losses import collision_loss, normal_loss, kl_loss
@@ -38,7 +38,6 @@ from torch.utils.data import DataLoader
import trimesh as tm
from zipfile import ZipFile
from configs.iccvw.pyr_64_128_256_512_e2e import config as config_1
from configs.iccvw.pyr_64_128_256_512_skirt_e2e import config as config_2
@@ -76,15 +75,13 @@ def step_pyramid(model, batch, config):
return out
def create_temp_frames(folder: str, frame: Frame):
with open(Path(folder) / 'frames_temp.txt', 'w') as f:
f.writelines([f'{frame}\n' for _ in range(4)])
if __name__ == '__main__':
results_folder = Path('/data-local/127_2/results')
out_folder = Path('temp/pyramid_ablation')
iteration = {'nonskirt': 190000, 'skirt': 200000 }
iteration = {'nonskirt': 250000, 'skirt': 200000 }
smpl = SMPL(orig_smpl=config_1['is_cape'], path='smpl', device=config_1['device'])
dataset_resolutions = []
@@ -99,8 +96,9 @@ if __name__ == '__main__':
model['nonskirt'].eval()
model['skirt'].eval()
frames = ['00054_1_Dress', '01070_13_Jumpsuit', '03849_27_Dress', '06057_152_Dress',
frames = ['01070_13_Jumpsuit', '00054_1_Dress', '03849_27_Dress', '06057_152_Dress',
'03179_274_Dress', '04247_102_Tshirt', '01852_70_Tshirt']
frames = ['00027_1_Dress']
frames_corrupt = ['00300_69_Jumpsuit', '07608_77_Jumpsuit', '06955_117_Dress']
frames = [Frame(frame_str=frame) for frame in frames]
@@ -109,8 +107,8 @@ if __name__ == '__main__':
config = { 'nonskirt': config_1, 'skirt': config_2 }
for frame in tqdm(frames):
create_temp_frames(config_1['train_dir'], frame)
dataset = Dataset(config_1['train_dir'], config_1['signature'], device=config_1['device'], frames_file='frames_temp.txt', use_cache=True, resolutions=dataset_resolutions, is_cape=config_1['is_cape'])
create_temp_frames(config_1['test_dir'], frame)
dataset = Dataset(config_1['test_dir'], config_1['signature'], device=config_1['device'], frames_file='frames_temp.txt', use_cache=True, resolutions=dataset_resolutions, is_cape=config_1['is_cape'])
dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=True, )
picker = 'skirt' if frame.garment in ['Skirt', 'Dress'] else 'nonskirt'
@@ -136,15 +134,15 @@ if __name__ == '__main__':
tm.smoothing.filter_humphrey(m, iterations=2)
verts = m.vertices
save_stl(out_folder / f'{frame}_lvl{i}_body.stl', body, body_faces)
save_stl(out_folder / f'{frame}_lvl{i}_mesh.stl', verts, faces)
# save_stl(out_folder / f'{frame}_lvl{i}_body.stl', body, body_faces)
# save_stl(out_folder / f'{frame}_lvl{i}_mesh.stl', verts, faces)
# fig = go.Figure(data=[go.Mesh3d(x=verts[:,0], y=verts[:,2], z=verts[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=1),
# go.Mesh3d(x=body[:,0], y=body[:,2], z=body[:,1], i=body_faces[:,0], j=body_faces[:,1], k=body_faces[:,2], color='gray', opacity=1)
# ])
# fig.update_layout(scene=dict(aspectmode='cube'))
# fig.update_layout(scene=dict(xaxis=dict(range=[-1, 1], ), yaxis=dict(range=[-1, 1], ), zaxis=dict(range=[-1, 1], ), ), )
# fig.show()
fig = go.Figure(data=[go.Mesh3d(x=verts[:,0], y=verts[:,2], z=verts[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=1),
go.Mesh3d(x=body[:,0], y=body[:,2], z=body[:,1], i=body_faces[:,0], j=body_faces[:,1], k=body_faces[:,2], color='gray', opacity=1)
])
fig.update_layout(scene=dict(aspectmode='cube'))
fig.update_layout(scene=dict(xaxis=dict(range=[-1, 1], ), yaxis=dict(range=[-1, 1], ), zaxis=dict(range=[-1, 1], ), ), )
fig.show()
pass
pass
scripts/plot_cloth3d.py
+148
View File
@@ -0,0 +1,148 @@
import os
import sys
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = '0'
import time
import argparse
import random
import tqdm
import numpy as np
import plotly.graph_objects as go
import plotly.express as ex
from utils.postprocess import fix_collision
import torch
from tqdm import tqdm
from data.utils import Frame
from data.cloth3d_io import readPC2Frame, readPC2, readOBJ
from utils.utils import pickle_load, save_stl
from scipy.spatial import cKDTree
from smpl.smpl_torch import SMPL
from scipy.spatial.transform import Rotation as R
import trimesh as tm
from plotly.subplots import make_subplots
smpl = SMPL(path='smpl', device='cpu')
def deorient(verts, pose, inverse=True):
assert pose.shape[1] == 72
rodr = smpl.rodrigues(torch.tensor(pose, device=smpl.device).reshape(-1, 1, 3)).reshape(verts.shape[0], -1, 3, 3).numpy()
rotated = np.zeros_like(verts)
for frame in range(verts.shape[0]):
r = R.from_matrix(rodr[frame, 0, :, :])
rotated[frame] = r.apply(verts[frame], inverse=inverse)
return rotated
def readBIN(fname):
if '.' not in fname: fname += '.bin'
elif '.' in fname and not fname.endswith('.bin'):
print("File name extension should be '.bin'")
return
with open(fname, 'rb') as f:
F = np.frombuffer(f.read(), dtype=np.uint16).astype(np.int32)
return F.reshape((-1,3))
def calculate():
garments = sorted(os.listdir('/data-local/127_2/cloth3d_processed_test/mesh_mask'))
garments = [g[:-4] for g in garments]
# random.shuffle(garments)
# garments = [('01909_Jumpsuit', 175), ('00960_Dress', 41), ('01573_Dress', 40), ('00079_Tshirt', 190)]
body_faces = np.load('smpl/smpl_faces.npy').astype(np.int32)
for garment in tqdm(garments):
# garment = garment_data[0]
picker = 'skirt' if garment.split('_')[1] in ['Dress', 'Skirt'] else 'nonskirt'
id, garment_type = garment.split('_')
if not os.path.exists(f'temp/cloth3d/{id}/outfit.pc16'):
continue
if not os.path.exists(os.path.join(cloth3d_folder, id, garment_type + '.pc16')):
continue
pred_cloth3d = readPC2(f'temp/cloth3d/{id}/outfit.pc16', True)['V']
pred_cloth3d_faces = readBIN(f'temp/cloth3d/{id}/outfit.bin')
gt = readPC2(os.path.join(cloth3d_folder, id, garment_type + '.pc16'), True)['V']
mask = np.load(f'/data-local/127_2/cloth3d_processed_test/mesh_fixed/{id}_{garment_type}.npz')['mask']
pose = np.load(f'/data-local/127_2/cloth3d_processed_test/pose/{id}.npy')
gt = deorient(gt, pose, True)
# gt = deorient(gt, pose, True)
body = readPC2(f'/data-local/127_2/cloth3d_processed_test/mesh_body_posed_no_orient/{id}.pc16', True)['V']
fix_data = np.load(f'/data-local/127_2/cloth3d_processed_test/mesh_fixed/{id}_{garment_type}.npz')
faces = fix_data['faces']
i = np.random.randint(0, len(pred_cloth3d))
# i = garment_data[1]
try:
if calculate_ours:
pred_ours = np.load(f'{ours_folder[picker]}/{id}_{i+2}_{garment_type}.npz', allow_pickle=True)['arr_0'].item()['lvl3']
except:
print('missed one')
continue
mesh_posed_fixed = fix_collision(torch.from_numpy(pred_ours[None,:,:]).to('cuda'), torch.from_numpy(body[i][None,:,:]).to('cuda'), torch.from_numpy(mask[None,:]).to('cuda')).cpu().numpy()
mesh_posed_fixed = mesh_posed_fixed[0, mask, :]
try:
m = tm.Trimesh(mesh_posed_fixed, faces, process=False, validate=False)
tm.smoothing.filter_humphrey(m, iterations=2)
pred_ours = m.vertices
except:
pred_ours = mesh_posed_fixed
print('failed to apply smoothing')
# fig = go.Figure(data=[go.Mesh3d(x=pred_ours[:,0], y=pred_ours[:,2], z=pred_ours[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='blue')])
# fig.add_trace(go.Mesh3d(x=body[i][:,0], y=body[i][:,2], z=body[i][:,1], i=body_faces[:,0], j=body_faces[:,1], k=body_faces[:,2], color='gray'))
# fig.update_layout(scene=dict(xaxis=dict(range=[-1, 1], ), yaxis=dict(range=[-1, 1], ), zaxis=dict(range=[-1, 1], ), ), )
# fig.update_layout(scene_aspectmode='cube', scene_aspectratio=dict(x=1, y=1, z=1))
# fig.show()
fig = go.Figure(data=[go.Mesh3d(x=pred_cloth3d[i][:,0], y=pred_cloth3d[i][:,1], z=pred_cloth3d[i][:,2], i=pred_cloth3d_faces[:,0], j=pred_cloth3d_faces[:,1], k=pred_cloth3d_faces[:,2], color='blue')])
fig.add_trace(go.Mesh3d(x=body[i][:,0], y=body[i][:,2], z=body[i][:,1], i=body_faces[:,0], j=body_faces[:,1], k=body_faces[:,2], color='gray'))
fig.update_layout(scene=dict(xaxis=dict(range=[-1, 1], ), yaxis=dict(range=[-1, 1], ), zaxis=dict(range=[-1, 1], ), ), )
fig.update_layout(scene_aspectmode='cube', scene_aspectratio=dict(x=1, y=1, z=1))
fig.show()
# fig = go.Figure(data=[go.Mesh3d(x=gt[i][:,0], y=gt[i][:,2], z=gt[i][:,1], i=pred_hood_faces[:,0], j=pred_hood_faces[:,1], k=pred_hood_faces[:,2], color='blue')])
# fig.add_trace(go.Mesh3d(x=body[i][:,0], y=body[i][:,2], z=body[i][:,1], i=body_faces[:,0], j=body_faces[:,1], k=body_faces[:,2], color='gray'))
# fig.update_layout(scene=dict(xaxis=dict(range=[-1, 1], ), yaxis=dict(range=[-1, 1], ), zaxis=dict(range=[-1, 1], ), ), )
# fig.update_layout(scene_aspectmode='cube', scene_aspectratio=dict(x=1, y=1, z=1))
# fig.show()
pass
# save_stl(f'temp/hood/{id}_{i}_{garment_type}_gt.stl', gt[i], pred_hood_faces)
# save_stl(f'temp/hood/{id}_{i}_{garment_type}_gt_body.stl', body[i], body_faces)
# save_stl(f'temp/hood/new/{id}_{i}_{garment_type}_hood.stl', pred_hood[i], pred_hood_faces)
# save_stl(f'temp/hood/new/{id}_{i}_{garment_type}_hood_body.stl', body[i], body_faces)
# save_stl(f'temp/hood/{id}_{i}_{garment_type}_ours.stl', pred_ours, faces)
# save_stl(f'temp/hood/{id}_{i}_{garment_type}_ours_body.stl', body[i], body_faces)
if __name__ == '__main__':
hood_folder = '/data-local/127_2/results/hood/all_2_105k'
ours_folder = dict()
ours_folder['skirt'] = '/data-local/127_2/results/pyr_iccvw_64_128_256_512_skirt_e2e_200000'
ours_folder['nonskirt'] = '/data-local/127_2/results/pyr_iccvw_64_128_256_512_e2e_2_250000'
cloth3d_folder = '/data-net/dataset/cloth3d/test'
results_folder = '/data-local/127_2/results'
calculate_ours = True
calculate()
scripts/plot_hood.py
+131
View File
@@ -0,0 +1,131 @@
import os
import sys
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = '0'
import time
import argparse
import random
import tqdm
import numpy as np
import plotly.graph_objects as go
import plotly.express as ex
from utils.postprocess import fix_collision
import torch
from tqdm import tqdm
from data.utils import Frame
from data.cloth3d_io import readPC2Frame, readPC2, readOBJ
from utils.utils import pickle_load, save_stl
from scipy.spatial import cKDTree
from smpl.smpl_torch import SMPL
from scipy.spatial.transform import Rotation as R
import trimesh as tm
from plotly.subplots import make_subplots
smpl = SMPL(path='smpl', device='cpu')
def deorient(verts, pose, inverse=True):
assert pose.shape[1] == 72
rodr = smpl.rodrigues(torch.tensor(pose, device=smpl.device).reshape(-1, 1, 3)).reshape(verts.shape[0], -1, 3, 3).numpy()
rotated = np.zeros_like(verts)
for frame in range(verts.shape[0]):
r = R.from_matrix(rodr[frame, 0, :, :])
rotated[frame] = r.apply(verts[frame], inverse=inverse)
return rotated
def calculate():
# garments = sorted(os.listdir(hood_folder))
# garments = [g[:-4] for g in garments]
# random.shuffle(garments)
garments = [('01909_Jumpsuit', 175), ('00960_Dress', 41), ('01573_Dress', 40), ('00079_Tshirt', 190)]
body_faces = np.load('smpl/smpl_faces.npy').astype(np.int32)
for garment_data in tqdm(garments):
garment = garment_data[0]
picker = 'skirt' if garment.split('_')[1] in ['Dress', 'Skirt'] else 'nonskirt'
id, garment_type = garment.split('_')
pred_hood_data = pickle_load(f'{hood_folder}/{garment}.pkl')
pred_hood, pred_hood_faces = pred_hood_data['verts'], pred_hood_data['faces']
gt = readPC2(os.path.join(cloth3d_folder, id, garment_type + '.pc16'), True)['V'][2:]
mask = np.load(f'/data-local/127_2/cloth3d_processed_test/mesh_fixed/{id}_{garment_type}.npz')['mask']
pose = np.load(f'/data-local/127_2/cloth3d_processed_test/pose/{id}.npy')[2:]
gt = deorient(gt, pose, True)
body = readPC2(f'/data-local/127_2/cloth3d_processed_test/mesh_body_posed_no_orient/{id}.pc16', True)['V'][2:]
fix_data = np.load(f'/data-local/127_2/cloth3d_processed_test/mesh_fixed/{id}_{garment_type}.npz')
faces = fix_data['faces']
# i = np.random.randint(0, len(pred_hood))
i = garment_data[1]
try:
if calculate_ours:
pred_ours = np.load(f'{ours_folder[picker]}/{id}_{i+2}_{garment_type}.npz', allow_pickle=True)['arr_0'].item()['lvl3']
except:
print('missed one')
continue
mesh_posed_fixed = fix_collision(torch.from_numpy(pred_ours[None,:,:]).to('cuda'), torch.from_numpy(body[i][None,:,:]).to('cuda'), torch.from_numpy(mask[None,:]).to('cuda')).cpu().numpy()
mesh_posed_fixed = mesh_posed_fixed[0, mask, :]
try:
m = tm.Trimesh(mesh_posed_fixed, faces, process=False, validate=False)
tm.smoothing.filter_humphrey(m, iterations=2)
pred_ours = m.vertices
except:
pred_ours = mesh_posed_fixed
print('failed to apply smoothing')
# fig = go.Figure(data=[go.Mesh3d(x=pred_ours[:,0], y=pred_ours[:,2], z=pred_ours[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='blue')])
# fig.add_trace(go.Mesh3d(x=body[i][:,0], y=body[i][:,2], z=body[i][:,1], i=body_faces[:,0], j=body_faces[:,1], k=body_faces[:,2], color='gray'))
# fig.update_layout(scene=dict(xaxis=dict(range=[-1, 1], ), yaxis=dict(range=[-1, 1], ), zaxis=dict(range=[-1, 1], ), ), )
# fig.update_layout(scene_aspectmode='cube', scene_aspectratio=dict(x=1, y=1, z=1))
# fig.show()
fig = go.Figure(data=[go.Mesh3d(x=pred_hood[i][:,0], y=pred_hood[i][:,2], z=pred_hood[i][:,1], i=pred_hood_faces[:,0], j=pred_hood_faces[:,1], k=pred_hood_faces[:,2], color='blue')])
fig.add_trace(go.Mesh3d(x=body[i][:,0], y=body[i][:,2], z=body[i][:,1], i=body_faces[:,0], j=body_faces[:,1], k=body_faces[:,2], color='gray'))
fig.update_layout(scene=dict(xaxis=dict(range=[-1, 1], ), yaxis=dict(range=[-1, 1], ), zaxis=dict(range=[-1, 1], ), ), )
fig.update_layout(scene_aspectmode='cube', scene_aspectratio=dict(x=1, y=1, z=1))
fig.show()
# fig = go.Figure(data=[go.Mesh3d(x=gt[i][:,0], y=gt[i][:,2], z=gt[i][:,1], i=pred_hood_faces[:,0], j=pred_hood_faces[:,1], k=pred_hood_faces[:,2], color='blue')])
# fig.add_trace(go.Mesh3d(x=body[i][:,0], y=body[i][:,2], z=body[i][:,1], i=body_faces[:,0], j=body_faces[:,1], k=body_faces[:,2], color='gray'))
# fig.update_layout(scene=dict(xaxis=dict(range=[-1, 1], ), yaxis=dict(range=[-1, 1], ), zaxis=dict(range=[-1, 1], ), ), )
# fig.update_layout(scene_aspectmode='cube', scene_aspectratio=dict(x=1, y=1, z=1))
# fig.show()
pass
# save_stl(f'temp/hood/2/{id}_{i}_{garment_type}_gt.stl', gt[i], pred_hood_faces)
# save_stl(f'temp/hood/2/{id}_{i}_{garment_type}_gt_body.stl', body[i], body_faces)
# save_stl(f'temp/hood/2/{id}_{i}_{garment_type}_hood.stl', pred_hood[i], pred_hood_faces)
# save_stl(f'temp/hood/ne2w/{id}_{i}_{garment_type}_hood_body.stl', body[i], body_faces)
# save_stl(f'temp/hood/2/{id}_{i}_{garment_type}_ours.stl', pred_ours, faces)
# save_stl(f'temp/hood/2/{id}_{i}_{garment_type}_ours_body.stl', body[i], body_faces)
if __name__ == '__main__':
hood_folder = '/data-local/127_2/results/hood/all_2_105k'
ours_folder = dict()
ours_folder['skirt'] = '/data-local/127_2/results/pyr_iccvw_64_128_256_512_skirt_e2e_200000'
ours_folder['nonskirt'] = '/data-local/127_2/results/pyr_iccvw_64_128_256_512_e2e_2_250000'
cloth3d_folder = '/data-net/dataset/cloth3d/test'
results_folder = '/data-local/127_2/results'
calculate_ours = True
calculate()
scripts/upscaler_train.py
+8 -5
View File
@@ -11,7 +11,7 @@ parser.add_argument('--debug', action='store_true')
args = vars(parser.parse_args())
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = '2'
os.environ["CUDA_VISIBLE_DEVICES"] = '3'
import time
@@ -39,10 +39,10 @@ from utils.utils import InfiniteDataLoader as DataLoader
from pytorch_msssim import MS_SSIM
low_res = 64
is_skirt = False
is_cape = True
is_male = True
low_res = 256
is_skirt = True
is_cape = False
is_male = False
skirt_modifier = '_skirt' if is_skirt else ''
cape_modifier = '_cape' if is_cape else ''
gender_modifier = '' if not is_cape else ('_m' if is_male else '_f')
@@ -244,6 +244,9 @@ def train():
step += 1
if step > 100000:
break
if __name__ == '__main__':
if not args['debug']: