import os os.environ["CUDA_VISIBLE_DEVICES"] = "1" import sys import time import cv2 import random import tqdm import numpy as np import trimesh as tm import tensorflow as tf import cmapy as cm import plotly.graph_objects as go import plotly.express as ex from plotly.subplots import make_subplots from scipy.spatial.transform import Rotation import tps from dataset import DatasetFactory, StaticDataset from custom_utils.custom_utils import batch_to_dict, UV2Mesh, _maskFaces, load_model, save_stl from custom_utils.smpl import SMPL from convert_to_uv_api import convert_to_uv_api from models.upscaler import UpscalerBasicConv as Upscaler from postprocessing import fix_collision def color_labels(img): colors = np.array([cm.color('jet', 1 - i) for i in np.linspace(0, 1, img.max()+1)]) / 255.0 color_img = np.zeros((img.shape[0], img.shape[1], 3), np.float32) for i in range(img.max()+1): color_img[img == i, :] = colors[i] return color_img def plot_fixed(garment, body, faces, camera): mesh = np.concatenate((body, garment), axis=0) faces = np.concatenate((F, faces + 14475), axis=0).astype(np.int32) colors = np.concatenate((np.ones_like(body, dtype=np.int32)*100, np.ones_like(garment, dtype=np.int32)*230), axis=0).astype(np.float32) / 255.0 fig = go.Figure() fig.add_trace(go.Mesh3d( x=mesh[:,0], y=mesh[:,2], z=mesh[:,1], i = faces[:,0], j = faces[:,1], k = faces[:,2], text="", vertexcolor=colors, lighting={ 'ambient':0.8, 'diffuse':0.5, 'roughness':0.5 }, ),) fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1)) helper = dict(xaxis = dict(range=[-1,1], showbackground=False, showticklabels=False, showaxeslabels=False, showspikes=False, visible=False), yaxis = dict(range=[-1,1], showbackground=False, showticklabels=False, showaxeslabels=False, showspikes=False, visible=False), zaxis = dict(range=[-1,1], showbackground=False, showticklabels=False, showaxeslabels=False, showspikes=False, visible=False),) fig.update_layout(scene = helper) fig.update_layout(scene_camera=camera, width=800, height=1100, autosize=False, margin=dict(l=0, r=0, b=0, t=0, pad=0) ) return fig def custom_imshow(img): # fig = ex.imshow(np.flip(img, axis=0)) fig = ex.imshow(img) fig.update_layout(margin=dict(l=0, r=0, b=0, t=0, pad=0), width=img.shape[1], height=img.shape[0]) fig.update_xaxes(visible=False) fig.update_yaxes(visible=False) return fig def get_faces(): faces_folder = '/home/workspace/data/results/support/faces' ids = os.listdir(faces_folder) ids = [id.split('.')[0] for id in ids] faces = {id : np.load(faces_folder + '/' + id + '.npy') for id in ids } faces_folder = '/home/workspace/data/results/support/faces_skirt' ids = os.listdir(faces_folder) ids = [id.split('.')[0] for id in ids] faces_skirt = {id : np.load(faces_folder + '/' + id + '.npy') for id in ids } return { **faces, **faces_skirt } def get_bodies(frames, skirt_mods): folder_base = '/home/workspace/local-data' out = dict() for frame in tqdm.tqdm(frames): id = frame.split('_')[0] body = np.load(os.path.join(folder_base, 'test' + skirt_mods[id] + '_ready' + ('_2' if skirt_mods[id]!='' else '' ), 'mesh_posed_body_no_orient', frame + '.npz'))['mesh_posed_body_no_orient'] out[frame] = body return out def postprocess(garment, body, id, is_skirt): mask = np.load(os.path.join('/home/workspace/data/dataset', ('test_skirt_ready_2' if is_skirt else 'test_ready'), 'mesh_mask', id + '.npy')) garment_temp = np.zeros((1,14475,3), np.float32) garment_temp[0,mask,:] = garment garment_fixed = fix_collision(garment_temp, body[np.newaxis,:,:], mask[np.newaxis,:])[0].numpy()[mask, :] return garment_fixed def deorient(mesh, frame, is_skirt): smpl = SMPL() pose = np.load(os.path.join('/home/workspace/data/dataset', ('test_skirt_ready_2' if is_skirt else 'test_ready'), 'pose_shape', frame + '.npy'))[:72].astype(np.float32) rodr = smpl.run_rodrigues(pose[np.newaxis, :]).numpy()[0].astype(np.float32) r = Rotation.from_matrix(rodr[0, :, :]) out = r.apply(mesh, inverse=True) return out def get_experiment_plots(folder, frames, experiments, skirt_mods, bodies, faces, cameras): out = [] for frame in tqdm.tqdm(frames): frame_out = [] id = frame.split('_')[0] for exp in experiments: path = os.path.join(folder, exp[0] + skirt_mods[id] + '-' + exp[1], frame + '.npy') if not os.path.exists(path): path = os.path.join('/home/workspace/local-data_128/results/iccv', exp[0] + skirt_mods[id] + '-' + exp[1], frame + '.npy') assert os.path.exists(path), path garment = np.load(path) if exp == 'abl_orient': garment = deorient(garment, frame, skirt_mods[id] == '_skirt') garment = postprocess(garment, bodies[frame], id, skirt_mods[id] == '_skirt') fig = plot_fixed(garment, bodies[frame], faces[id], cameras[frame]) fig.write_image('/tmp/exp_temp_5.png') img = cv2.imread('/tmp/exp_temp_5.png') frame_out.append(img[:, 100:700, :]) out.append(np.concatenate(frame_out, axis=1)) return out def get_experiment_meshes(folder, frames, experiments, skirt_mods, bodies, faces, cameras): out = [] for frame in tqdm.tqdm(frames): frame_out = [] id = frame.split('_')[0] for exp in experiments: garment = np.load(os.path.join(folder, exp[0] + skirt_mods[id] + '-' + exp[1], frame + '.npy')) if exp == 'abl_orient': garment = deorient(garment, frame, skirt_mods[id] == '_skirt') garment = postprocess(garment, bodies[frame], id, skirt_mods[id] == '_skirt') frame_out.append([garment, bodies[frame], faces[id]]) out.append(frame_out) return out def get_gt_image(frames, bodies, faces, cameras, skirt_mods): out = [] for frame in tqdm.tqdm(frames): id = frame.split('_')[0] mask = np.load(os.path.join('/home/workspace/data/dataset', ('test_skirt_ready_2' if skirt_mods[id] == '_skirt' else 'test_ready'), 'mesh_mask', id + '.npy')) gt = np.load(os.path.join('/home/workspace/data/dataset', ('test_skirt_ready_2' if skirt_mods[id] == '_skirt' else 'test_ready'), 'mesh_posed_no_orient', frame + '.npz'))['mesh_posed_no_orient'].astype(np.float32) fig = plot_fixed(gt[mask,:], bodies[frame], faces[id], cameras[frame]) fig.write_image('/tmp/exp_temp_5.png') img = cv2.imread('/tmp/exp_temp_5.png') out.append(img[:, 100:700, :]) return out def get_gt_meshes(frames, bodies, faces, cameras, skirt_mods): out = [] for frame in tqdm.tqdm(frames): id = frame.split('_')[0] mask = np.load(os.path.join('/home/workspace/data/dataset', ('test_skirt_ready_2' if skirt_mods[id] == '_skirt' else 'test_ready'), 'mesh_mask', id + '.npy')) gt = np.load(os.path.join('/home/workspace/data/dataset', ('test_skirt_ready_2' if skirt_mods[id] == '_skirt' else 'test_ready'), 'mesh_posed_no_orient', frame + '.npz'))['mesh_posed_no_orient'].astype(np.float32) out.append([gt[mask,:], bodies[frame], faces[id]]) return out if __name__ == '__main__': # uv2mesh = UV2Mesh() # smpl = SMPL() F = np.load("custom_utils/smpl_faces.npy") garment_type_table = np.load('custom_data/garment_type_table_test.npy', allow_pickle=True).item() faces = get_faces() frames = ['00079_135', '00505_23', '00688_25', '01693_128', '00408_34'] # frames = [ '00075_4', '00475_273', '00761_295', '00793_3', '00993_220', '01134_112', '01278_108', '00718_38', '00951_34', '00962_60','01065_38','01191_83','01240_168'] rand_frames = [frame.split('.')[0] for frame in os.listdir('/home/workspace/data/dataset/test_ready/pose_shape')] random.shuffle(rand_frames) rand_frames = rand_frames[:25] # frames.extend(rand_frames) # rand_frames = [frame.split('.')[0] for frame in os.listdir('/home/workspace/data/dataset/test_skirt_ready_2/pose_shape')] # random.shuffle(rand_frames) # rand_frames = rand_frames[:25] # frames.extend(rand_frames) camera_1 = dict(up=dict(x=0, y=0, z=1), center=dict(x=0, y=0, z=-0.05), eye=dict(x=0.3, y=1.0, z=0.05)) camera_2 = dict(up=dict(x=0, y=0, z=1), center=dict(x=0, y=0, z=-0.05), eye=dict(x=-0.9, y=0.7, z=0.05)) cameras = {frame: camera_1 for frame in frames} # abl_all replace abl_norm maybe # experiments = [['iccv_pyr_32_64_128_256_512', '120000'], ['iccv_pyr_32_64_128_256_512_curv', '80000'], ['iccv_pyr_32_64_128_256_512_skin', '115000'], ['iccv_pyr_32_64_128_256_512_reg', '75000']] # , 'pyra_32_res' # experiments = [['iccv_pyr_32_64_128_256_512', '120000'], ['iccv_pyr_32_64_128_256_curv', '80000'], ['iccv_pyr_32_64_128_256_skin', '145000'], ['iccv_pyr_32_64_128_256_reg', '35000']] # , 'pyra_32_res' # experiments = [['iccv_pyr_32_64_128_256_512', '120000'], ['iccv_pyr_32_64_128_curv', '20000'], ['iccv_pyr_32_64_128_skin_2', '20000'], ['iccv_pyr_32_64_128_reg_2', '80000']] # , 'pyra_32_res' experiments = [['abl_all_3', '130000']] # , 'pyra_32_res' skirt_mods = { frame.split('_')[0]: ('_skirt' if os.path.exists('/home/workspace/data/dataset/test_skirt_ready_2/pose_shape/' + frame + '.npy') else '') for frame in frames } bodies = get_bodies(frames, skirt_mods) results_folder_base = '/home/workspace/data/results/iccv/' output_folder = os.path.join('examples/ablation/img_3') output_folder_stl = os.path.join('examples/ablation/stl_3') os.makedirs(output_folder, exist_ok=True) os.makedirs(output_folder_stl, exist_ok=True) # gts_img = get_gt_image(frames, bodies, faces, cameras, skirt_mods) # images = get_experiment_plots(results_folder_base, frames, experiments, skirt_mods, bodies, faces, cameras) # rows = [np.concatenate((gt, exp), axis=1) for gt, exp in zip(gts_img, images)] # for frame, row in zip(frames, rows): # custom_imshow(row).write_image(os.path.join(output_folder, frame + '_ablation.png')) gts = get_gt_meshes(frames, bodies, faces, cameras, skirt_mods) meshes = get_experiment_meshes(results_folder_base, frames, experiments, skirt_mods, bodies, faces, cameras) for mesh, frame in zip(meshes, frames): for exp, (garment, body, faces) in zip(experiments, mesh): save_stl(f'{output_folder_stl}/{frame}_{exp[0]}_garment.stl', garment, faces) save_stl(f'{output_folder_stl}/{frame}_{exp[0]}_body.stl', body, F) for mesh, frame in zip(gts, frames): garment, body, faces = mesh save_stl(f'{output_folder_stl}/{frame}_gt_garment.stl', garment, faces) save_stl(f'{output_folder_stl}/{frame}_gt_body.stl', body, F) # final = np.concatenate(rows, axis=0) # custom_imshow(final).show() print()