def visualize():
from models.smpl4garment import SMPL4Garment
garment_class = 'shirt'
gender = 'female'
split = None
style_idx = '000'
shape_idx = '008'
smooth_level = 1
smpl = SMPL4Garment(gender=gender)
# gt_ds = MultiStyleShape(garment_class=garment_class, split=split, smooth_level=smooth_level, gender=gender)
# gt_ds = OneStyleShape(garment_class=garment_class, shape_idx=shape_idx, style_idx=style_idx, split=split,
# smooth_level=smooth_level, gender=gender)
gt_ds = MultiStyleShape(garment_class=garment_class, split=None, smooth_level=smooth_level, gender=gender)
print(len(gt_ds))
gt_ds = MultiStyleShape(garment_class=garment_class, split='train', smooth_level=smooth_level, gender=gender)
print(len(gt_ds))
gt_ds = MultiStyleShape(garment_class=garment_class, split='test', smooth_level=smooth_level, gender=gender)
print(len(gt_ds))
for idx in np.random.randint(0, len(gt_ds), 4):
verts, thetas, betas, gammas, item = gt_ds[idx]
# verts, sverts, thetas, betas, gammas, item = gt_ds[idx]
body_m, gar_m = smpl.run(theta=thetas.numpy(), beta=betas.numpy(), garment_class=garment_class,
garment_d=verts.numpy())
def run_tailornet():
gender = 'female'
garment_class = 'pant'
thetas, betas, gammas = get_single_frame_inputs(garment_class, gender)
# # uncomment the line below to run inference on sequence data
# thetas, betas, gammas = get_sequence_inputs(garment_class, gender)
# load model
tn_runner = get_tn_runner(gender=gender, garment_class=garment_class)
smpl = SMPL4Garment(gender=gender)
# make out directory if doesn't exist
if not os.path.isdir(OUT_PATH):
os.mkdir(OUT_PATH)
# run inference
for i, (theta, beta, gamma) in enumerate(zip(thetas, betas, gammas)):
print(i, len(thetas))
# normalize y-rotation to make it front facing
theta_normalized = normalize_y_rotation(theta)
with torch.no_grad():
pred_verts_d = tn_runner.forward(
thetas=torch.from_numpy(theta_normalized[None, :].astype(np.float32)).cuda(),
betas=torch.from_numpy(beta[None, :].astype(np.float32)).cuda(),
gammas=torch.from_numpy(gamma[None, :].astype(np.float32)).cuda(),
)[0].cpu().numpy()
# get garment from predicted displacements
body, pred_gar = smpl.run(beta=beta, theta=theta, garment_class=garment_class, garment_d=pred_verts_d)
pred_gar = remove_interpenetration_fast(pred_gar, body)
# save body and predicted garment
body.write_ply(os.path.join(OUT_PATH, "body_{:04d}.ply".format(i)))
pred_gar.write_ply(os.path.join(OUT_PATH, "pred_gar_{:04d}.ply".format(i)))
def __init__(self, params):
self.params = params
self.gender = params['gender']
self.garment_class = params['garment_class']
self.bs = params['batch_size']
self.vis_freq = params['vis_freq']
self.model_name = params['model_name']
self.note = params['note']
# log and backup
log_name = os.path.join(params['log_name'],
'{}_{}'.format(self.garment_class, self.gender))
if params['shape_style'] != '':
log_name = os.path.join(log_name, params['shape_style'])
self.log_dir = sio.prepare_log_dir(log_name)
sio.save_params(self.log_dir, params, save_name='params')
self.iter_nums = 0 if 'iter_nums' not in params else params['iter_nums']
# smpl for garment
self.smpl = SMPL4Garment(gender=self.gender)
# garment specific things
with open(os.path.join(global_var.DATA_DIR, global_var.GAR_INFO_FILE), 'rb') as f:
class_info = pickle.load(f)
self.body_f_np = self.smpl.smpl_base.f.astype(np.long)
self.garment_f_np = class_info[self.garment_class]['f']
self.garment_f_torch = torch.tensor(self.garment_f_np.astype(np.long)).long().to(device)
self.vert_indices = np.array(
class_info[self.garment_class]['vert_indices'])
# get dataset and dataloader
self.train_dataset, self.train_loader = self.load_dataset('train')
self.test_dataset, self.test_loader = self.load_dataset('test')
print("Train dataset size", len(self.train_dataset))
print("Test dataset size", len(self.test_dataset))
# model and optimizer
self.model = self.build_model()
self.model.to(device)
self.optimizer = torch.optim.Adam(
self.model.parameters(), lr=params['lr'], weight_decay=params['weight_decay'])
# continue training from checkpoint if provided
if params['checkpoint']:
ckpt_path = params['checkpoint']
print('loading ckpt from {}'.format(ckpt_path))
state_dict = torch.load(os.path.join(ckpt_path, 'lin.pth.tar'))
self.model.load_state_dict(state_dict)
state_dict = torch.load(os.path.join(ckpt_path, 'optimizer.pth.tar'))
self.optimizer.load_state_dict(state_dict)
self.best_error = np.inf
self.best_epoch = -1
# logger
self.logger = tensorboardX.SummaryWriter(os.path.join(self.log_dir))
self.csv_logger = self.get_logger()
def visualize():
from models.smpl4garment import SMPL4Garment
garment_class = 'skirt'
gender = 'female'
split = None
style_idx = '003'
shape_idx = '000'
smooth_level = 1
smpl = SMPL4Garment(gender=gender)
# gt_ds = MultiStyleShape(garment_class=garment_class, split=split, smooth_level=smooth_level, gender=gender)
gt_ds = OneStyleShape(garment_class=garment_class,
shape_idx=shape_idx,
style_idx=style_idx,
split=split,
smooth_level=smooth_level,
gender=gender)
# gt_ds = MultiStyleShape(garment_class=garment_class, split=None, smooth_level=smooth_level, gender=gender)
# print(len(gt_ds))
# gt_ds = MultiStyleShape(garment_class=garment_class, split='train', smooth_level=smooth_level, gender=gender)
# print(len(gt_ds))
# gt_ds = MultiStyleShape(garment_class=garment_class, split='test', smooth_level=smooth_level, gender=gender)
# print(len(gt_ds))
# for idx in np.random.randint(0, len(gt_ds), 4):
for idx in [245, 1017, 1274, 1588]:
verts, thetas, betas, gammas, item = gt_ds[idx]
# verts, sverts, thetas, betas, gammas, item = gt_ds[idx]
# verts = torch.zeros_like(verts)
body_m, gar_m = smpl.run(theta=thetas.numpy(),
beta=betas.numpy(),
garment_class=garment_class,
garment_d=verts.numpy())
# body_m, gar_m = smpl.run(garment_class=garment_class,
# garment_d=verts.numpy())
body_m.write_ply("/BS/cpatel/work/body_{}.ply".format(idx))
gar_m.write_ply("/BS/cpatel/work/gar_{}.ply".format(idx))
def temp_files(gar_class='g1'):
import os
import sys
sys.path.append('/BS/garvita/work/code/cloth_static/TailorNet')
from models.smpl4garment import SMPL4Garment
sys.path.append('/BS/garvita/work/code/registration_ver2-master')
sys.path.extend(['/BS/garvita/work/code/RVH', '/BS/RVH/work/frankengeist'])
from core.geometry.geometry import get_submesh
if gar_class == 'g1' or gar_class == 'g2' or gar_class == 'g3' or gar_class == 'g4':
upper = 'shirt'
lower = 'pant'
if gar_class == 'g5' or gar_class == 'g6':
upper = 't-shirt'
lower = 'short-pant'
if gar_class == 'g7':
upper = 't-shirt'
lower = 'short-pant'
upper_old = 'Vest'
root_out = '/BS/garvita4/static00/sizer_final/mesh_utils/{}'.format(
gar_class)
if not os.path.exists(root_out):
os.makedirs(root_out)
smpl = SMPL4Garment('male')
_, upper_m = smpl.run(garment_class=upper)
_, lower_m = smpl.run(garment_class=lower)
upper_m.write_obj('{}/UpperClothes_hres.obj'.format(root_out))
lower_m.write_obj('{}/Pants_hres.obj'.format(root_out))
#
#
# template_file = '/BS/garvita2/work/template_files/allTemplate_withBoundaries_symm.pkl'
# with open(template_file, 'rb') as f:
# template_data = pkl.load(f, encoding="latin1")
#
# #load hres smpl
smpl_hres = Mesh(
filename='/BS/cloth3d/static00/nasa_data/smpl_sdf/meshes/000/000000.obj'
)
upper_id = smpl.class_info[upper]['vert_indices']
lower_id = smpl.class_info[lower]['vert_indices']
body_vert = range(len(smpl_hres.v))
body_vert2 = [i for i in body_vert if i not in upper_id]
body_vert2 = [i for i in body_vert2 if i not in lower_id]
body_vert = np.array(body_vert2)
v, faces_body, _, _ = get_submesh(smpl_hres.v, smpl_hres.f, body_vert)
body_hres = Mesh(v=v, f=faces_body)
body_hres.write_obj('{}/Body_hres.obj'.format(root_out))
#get vertices of body
#create lres for upper clothing
lres_vert_upper = np.array([vid for vid in upper_id if vid < 6890])
lres_vert_lower = np.array([vid for vid in lower_id if vid < 6890])
lres_body_vert = np.array([vid for vid in body_vert if vid < 6890])
#load lres and create the templates
smpl_lres = '/BS/garvita/work/dataset/smpl_mesh.obj'
m1 = Mesh(filename=smpl_lres)
v, faces_body, _, _ = get_submesh(m1.v, m1.f, lres_vert_upper)
upper_lres = Mesh(v=v, f=faces_body)
v, faces_body, _, _ = get_submesh(m1.v, m1.f, lres_vert_lower)
lower_lres = Mesh(v=v, f=faces_body)
v, faces_body, _, _ = get_submesh(m1.v, m1.f, lres_body_vert)
body_lres = Mesh(v=v, f=faces_body)
upper_lres.write_obj('{}/UpperClothes_lres.obj'.format(root_out))
lower_lres.write_obj('{}/Pants_lres.obj'.format(root_out))
body_lres.write_obj('{}/Body_lres.obj'.format(root_out))
np.save('{}/UpperClothes_hres_id.npy'.format(root_out), upper_id)
np.save('{}/Pants_hres_id.npy'.format(root_out), lower_id)
np.save('{}/Body_hres_id.npy'.format(root_out), body_vert)
np.save('{}/UpperClothes_lres_id.npy'.format(root_out), lres_vert_upper)
np.save('{}/Pants_lres_id.npy'.format(root_out), lres_vert_lower)
np.save('{}/Body_lres_id.npy'.format(root_out), lres_body_vert)
def run_tailornet():
gender = 'male'
garment_class = 'short-pant'
garment_class_pairs = {
'shirt': [],
't-shirt': [],
'pant': ['shirt', 't-shirt'],
'short-pant': ['shirt', 't-shirt'],
'skirt': ['shirt', 't-shirt']
}
thetas, betas, gammas = get_single_frame_inputs(garment_class, gender)
# # uncomment the line below to run inference on sequence data
# thetas, betas, gammas = get_sequence_inputs(garment_class, gender)
# load model
tn_runner = get_tn_runner(gender=gender, garment_class=garment_class)
# from trainer.base_trainer import get_best_runner
# tn_runner = get_best_runner("/BS/cpatel/work/data/learn_anim/tn_baseline/{}_{}/".format(garment_class, gender))
smpl = SMPL4Garment(gender=gender)
# make out directory if doesn't exist
if not os.path.isdir(OUT_PATH):
os.mkdir(OUT_PATH)
# run inference
for i, (theta, beta, gamma) in enumerate(zip(thetas, betas, gammas)):
print(i, len(thetas))
# normalize y-rotation to make it front facing
theta_normalized = normalize_y_rotation(theta)
with torch.no_grad():
pred_verts_d = tn_runner.forward(
thetas=torch.from_numpy(theta_normalized[None, :].astype(
np.float32)).cuda(),
betas=torch.from_numpy(beta[None, :].astype(
np.float32)).cuda(),
gammas=torch.from_numpy(gamma[None, :].astype(
np.float32)).cuda(),
)[0].cpu().numpy()
# get garment from predicted displacements
body, pred_gar = smpl.run(beta=beta,
theta=theta,
garment_class=garment_class,
garment_d=pred_verts_d)
pred_gar = remove_interpenetration_fast(pred_gar, body)
hv, hf, mapping = get_hres(pred_gar.v, pred_gar.f)
pred_gar_hres = Mesh(hv, hf)
for garment_class_pair in garment_class_pairs[garment_class]:
gar_pair_hres = Mesh(filename=os.path.join(
OUT_PATH, "bcnet_model_{}_gar_hres_{}_{:04d}.obj".format(
BCNET_MODEL_NUM, garment_class_pair, i)))
# print(len(body.v), body.f.min(), body.f.max())
# print(len(pred_gar_hres.v), pred_gar_hres.f.min(), pred_gar_hres.f.max())
# print(np.vstack((body.v, pred_gar_hres.v)).shape, np.min(pred_gar_hres.f + len(body.v)), np.max(pred_gar_hres.f + len(body.v)))
# gar_pair_hres = remove_interpenetration_fast(gar_pair_hres, pred_gar_hres)
for _ in range(3):
gar_pair_hres = remove_interpenetration_fast(
gar_pair_hres,
Mesh(np.vstack((body.v, pred_gar_hres.v)),
np.vstack((body.f, pred_gar_hres.f + len(body.v)))))
gar_pair_hres.write_obj(
os.path.join(
OUT_PATH,
"bcnet_model_{}_gar_hres_{}_inter_{}_{:04d}.obj".format(
BCNET_MODEL_NUM, garment_class_pair, garment_class,
i)))
# save body and predicted garment
body.write_obj(
os.path.join(OUT_PATH,
"bcnet_model_{}.obj".format(BCNET_MODEL_NUM)))
pred_gar.write_obj(
os.path.join(
OUT_PATH, "bcnet_model_{}_gar_{}_{:04d}.obj".format(
BCNET_MODEL_NUM, garment_class, i)))
pred_gar_hres.write_obj(
os.path.join(
OUT_PATH, "bcnet_model_{}_gar_hres_{}_{:04d}.obj".format(
BCNET_MODEL_NUM, garment_class, i)))
def evaluate_save():
"""Evaluate TailorNet (or any model for that matter) on test set."""
from dataset.static_pose_shape_final import MultiStyleShape
import torch
from torch.utils.data import DataLoader
from utils.eval import AverageMeter
from models import ops
from models.smpl4garment import SMPL4Garment
import os
gender = 'female'
garment_class = 'skirt'
smpl = SMPL4Garment(gender)
vis_freq = 512
log_dir = "/BS/cpatel/work/code_test2/try"
dataset = MultiStyleShape(garment_class=garment_class,
gender=gender,
split='test')
dataloader = DataLoader(dataset,
batch_size=32,
num_workers=0,
shuffle=False,
drop_last=False)
print(len(dataset))
val_dist = AverageMeter()
from models.tailornet_model import get_best_runner as tn_runner
runner = tn_runner(garment_class, gender)
# from trainer.base_trainer import get_best_runner as baseline_runner
# runner = baseline_runner("/BS/cpatel/work/data/learn_anim/{}_{}_weights/tn_orig_baseline/{}_{}".format(garment_class, gender, garment_class, gender))
device = torch.device('cuda:0')
with torch.no_grad():
for i, inputs in enumerate(dataloader):
gt_verts, thetas, betas, gammas, idxs = inputs
thetas, betas, gammas = ops.mask_inputs(thetas, betas, gammas,
garment_class)
gt_verts = gt_verts.to(device)
thetas = thetas.to(device)
betas = betas.to(device)
gammas = gammas.to(device)
pred_verts = runner.forward(thetas=thetas,
betas=betas,
gammas=gammas).view(gt_verts.shape)
for lidx, idx in enumerate(idxs):
if idx % vis_freq != 0:
continue
theta = thetas[lidx].cpu().numpy()
beta = betas[lidx].cpu().numpy()
pred_vert = pred_verts[lidx].cpu().numpy()
gt_vert = gt_verts[lidx].cpu().numpy()
body_m, pred_m = smpl.run(theta=theta,
garment_d=pred_vert,
beta=beta,
garment_class=garment_class)
_, gt_m = smpl.run(theta=theta,
garment_d=gt_vert,
beta=beta,
garment_class=garment_class)
save_dir = log_dir
pred_m.write_ply(
os.path.join(save_dir, "pred_{}.ply".format(idx)))
gt_m.write_ply(os.path.join(save_dir, "gt_{}.ply".format(idx)))
body_m.write_ply(
os.path.join(save_dir, "body_{}.ply".format(idx)))
print(val_dist.avg)
