import os os.environ["CUDA_VISIBLE_DEVICES"] = "6" 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 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/data/dataset' 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: garment = np.load(os.path.join(folder, exp + skirt_mods[id] + '-70000', frame + '.npy')) 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_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 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 = [] # 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 = ['abl_norm', 'abl_skin', 'abl_curv', 'abl_extra_loss', 'abl_abs'] 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) gts = get_gt_image(frames, bodies, faces, cameras, skirt_mods) results_folder_base = '/home/workspace/data/results' output_folder = os.path.join(results_folder_base, 'ablation_plots') os.makedirs(output_folder, exist_ok=True) 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, images)] for frame, row in zip(frames, rows): custom_imshow(row).write_image(os.path.join(output_folder, frame + '_ablation.png')) # final = np.concatenate(rows, axis=0) # custom_imshow(final).show() print()