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/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()