|
|
import os
|
|
|
import imageio
|
|
|
import sys
|
|
|
sys.path.append('.')
|
|
|
|
|
|
import numpy as np
|
|
|
import torch
|
|
|
import torch.nn.functional as F
|
|
|
from tqdm import tqdm
|
|
|
from src.visualize.anim import plot_3d_motion_dico, load_anim
|
|
|
import clip
|
|
|
from PIL import Image
|
|
|
import pickle
|
|
|
import src.utils.rotation_conversions as geometry
|
|
|
from textwrap import wrap
|
|
|
import shutil
|
|
|
import subprocess as sp
|
|
|
from copy import deepcopy
|
|
|
|
|
|
GPU_MINIMUM_MEMORY = 5500
|
|
|
|
|
|
def stack_images(real, real_gens, gen):
|
|
|
nleft_cols = len(real_gens) + 1
|
|
|
print("Stacking frames..")
|
|
|
allframes = np.concatenate((real[:, None, ...], *[x[:, None, ...] for x in real_gens], gen), 1)
|
|
|
nframes, nspa, nats, h, w, pix = allframes.shape
|
|
|
blackborder = np.zeros((w//30, h*nats, pix), dtype=allframes.dtype)
|
|
|
frames = []
|
|
|
for frame_idx in tqdm(range(nframes)):
|
|
|
columns = np.vstack(allframes[frame_idx].transpose(1, 2, 3, 4, 0)).transpose(3, 1, 0, 2)
|
|
|
frame = np.concatenate((*columns[0:nleft_cols], blackborder, *columns[nleft_cols:]), 0).transpose(1, 0, 2)
|
|
|
frames.append(frame)
|
|
|
return np.stack(frames)
|
|
|
|
|
|
def stack_gen_and_images(gen, images):
|
|
|
|
|
|
print("Stacking frames..")
|
|
|
allframes = np.concatenate((images, gen), 2)
|
|
|
nframes, nspa, nats, h, w, pix = allframes.shape
|
|
|
blackborder = np.zeros((w//30, h*nats, pix), dtype=allframes.dtype)
|
|
|
frames = []
|
|
|
for frame_idx in tqdm(range(nframes)):
|
|
|
columns = np.vstack(allframes[frame_idx].transpose(1, 2, 3, 4, 0)).transpose(3, 1, 0, 2)
|
|
|
frame = np.concatenate((columns[:]), 0).transpose(1, 0, 2)
|
|
|
frames.append(frame)
|
|
|
return np.stack(frames)
|
|
|
|
|
|
def stack_gen_only(gen):
|
|
|
|
|
|
print("Stacking frames..")
|
|
|
|
|
|
allframes = gen
|
|
|
nframes, nspa, nats, h, w, pix = allframes.shape
|
|
|
blackborder = np.zeros((w//30, h*nats, pix), dtype=allframes.dtype)
|
|
|
frames = []
|
|
|
for frame_idx in tqdm(range(nframes)):
|
|
|
columns = np.vstack(allframes[frame_idx].transpose(1, 2, 3, 4, 0)).transpose(3, 1, 0, 2)
|
|
|
frame = np.concatenate((columns[:]), 0).transpose(1, 0, 2)
|
|
|
frames.append(frame)
|
|
|
return np.stack(frames)
|
|
|
|
|
|
|
|
|
def generate_by_video(visualization, reconstructions, generation,
|
|
|
label_to_action_name, params, nats, nspa, tmp_path, image_pathes=None, mode=None):
|
|
|
|
|
|
|
|
|
fps = params["fps"]
|
|
|
|
|
|
params = params.copy()
|
|
|
|
|
|
if "output_xyz" in visualization or "output_xyz" in generation:
|
|
|
outputkey = "output_xyz"
|
|
|
params["pose_rep"] = "xyz"
|
|
|
else:
|
|
|
outputkey = "poses"
|
|
|
|
|
|
keep = [outputkey, "lengths", "y"]
|
|
|
|
|
|
def _to_np(x):
|
|
|
if type(x).__module__ == np.__name__:
|
|
|
return x
|
|
|
else:
|
|
|
return x.data.cpu().numpy()
|
|
|
|
|
|
visu = {key: _to_np(visualization[key]) for key in keep if key in visualization.keys()}
|
|
|
recons = {mode: {key: _to_np(reconstruction[key]) for key in keep if key in reconstruction.keys()}
|
|
|
for mode, reconstruction in reconstructions.items()}
|
|
|
gener = {key: _to_np(generation[key]) for key in keep if key in generation.keys()}
|
|
|
|
|
|
def get_palette(i, nspa):
|
|
|
if mode == 'edit' and i < 3:
|
|
|
return 'orange'
|
|
|
elif mode == 'interp' and i in [0, nspa-1]:
|
|
|
return 'orange'
|
|
|
return 'blue'
|
|
|
|
|
|
|
|
|
if(len(visu) > 0):
|
|
|
lenmax = max(gener["lengths"].max(),
|
|
|
visu["lengths"].max())
|
|
|
else:
|
|
|
lenmax = gener["lengths"].max()
|
|
|
timesize = lenmax + 5
|
|
|
|
|
|
|
|
|
|
|
|
import multiprocessing
|
|
|
|
|
|
def pool_job_with_desc(pool, iterator, desc, max_, save_path_format, isij):
|
|
|
with tqdm(total=max_, desc=desc.format("Render")) as pbar:
|
|
|
for _ in pool.imap_unordered(plot_3d_motion_dico, iterator):
|
|
|
pbar.update()
|
|
|
if isij:
|
|
|
array = np.stack([[load_anim(save_path_format.format(i, j), timesize)
|
|
|
for j in range(nats)]
|
|
|
for i in tqdm(range(nspa), desc=desc.format("Load"))])
|
|
|
return array.transpose(2, 0, 1, 3, 4, 5)
|
|
|
else:
|
|
|
array = np.stack([load_anim(save_path_format.format(i), timesize)
|
|
|
for i in tqdm(range(nats), desc=desc.format("Load"))])
|
|
|
return array.transpose(1, 0, 2, 3, 4)
|
|
|
|
|
|
with multiprocessing.Pool() as pool:
|
|
|
|
|
|
save_path_format = os.path.join(tmp_path, "gen_{}_{}.gif")
|
|
|
iterator = ((gener[outputkey][i, j],
|
|
|
gener["lengths"][i, j],
|
|
|
save_path_format.format(i, j),
|
|
|
|
|
|
params, {"title": f"{label_to_action_name(gener['y'][i, j])}", "interval": 1000/fps, "palette": get_palette(i, nspa)})
|
|
|
for j in range(nats) for i in range(nspa))
|
|
|
gener["frames"] = pool_job_with_desc(pool, iterator,
|
|
|
"{} the generated samples",
|
|
|
nats*nspa,
|
|
|
save_path_format,
|
|
|
True)
|
|
|
|
|
|
|
|
|
frames_no_title = gener['y'] == ''
|
|
|
gener["frames"][:, frames_no_title] = gener["frames"][:, 0, 0:1, 0:1, 0:1]
|
|
|
|
|
|
|
|
|
if len(visu) > 0:
|
|
|
save_path_format = os.path.join(tmp_path, "real_{}.gif")
|
|
|
iterator = ((visu[outputkey][i],
|
|
|
visu["lengths"][i],
|
|
|
save_path_format.format(i),
|
|
|
params, {"title": f"real: {label_to_action_name(visu['y'][i])}", "interval": 1000/fps})
|
|
|
for i in range(nats))
|
|
|
visu["frames"] = pool_job_with_desc(pool, iterator,
|
|
|
"{} the real samples",
|
|
|
nats,
|
|
|
save_path_format,
|
|
|
False)
|
|
|
for mode, recon in recons.items():
|
|
|
|
|
|
save_path_format = os.path.join(tmp_path, f"reconstructed_{mode}_" + "{}.gif")
|
|
|
iterator = ((recon[outputkey][i],
|
|
|
recon["lengths"][i],
|
|
|
save_path_format.format(i),
|
|
|
params, {"title": f"recons: {label_to_action_name(recon['y'][i])}",
|
|
|
"interval": 1000/fps})
|
|
|
for i in range(nats))
|
|
|
recon["frames"] = pool_job_with_desc(pool, iterator,
|
|
|
"{} the reconstructed samples",
|
|
|
nats,
|
|
|
save_path_format,
|
|
|
False)
|
|
|
if image_pathes is not None:
|
|
|
|
|
|
assert nats == 1
|
|
|
assert nspa == len(image_pathes)
|
|
|
h, w = gener["frames"].shape[3:5]
|
|
|
image_frames = []
|
|
|
for im_path in image_pathes:
|
|
|
im = Image.open(im_path).resize((w, h))
|
|
|
image_frames.append(np.tile(np.expand_dims(np.asarray(im)[..., :3], axis=(0, 1, 2)), (timesize, 1, 1, 1, 1, 1)))
|
|
|
image_frames = np.concatenate(image_frames, axis=1)
|
|
|
assert image_frames.shape == gener["frames"].shape
|
|
|
return stack_gen_and_images(gener["frames"], image_frames)
|
|
|
|
|
|
if len(visu) == 0:
|
|
|
frames = stack_gen_only(gener["frames"])
|
|
|
else:
|
|
|
frames = stack_images(visu["frames"], [recon["frames"] for recon in recons.values()], gener["frames"])
|
|
|
return frames
|
|
|
|
|
|
|
|
|
def generate_by_video_sequences(visualization, label_to_action_name, params, nats, nspa, tmp_path):
|
|
|
|
|
|
|
|
|
fps = params["fps"]
|
|
|
|
|
|
if "output_xyz" in visualization:
|
|
|
outputkey = "output_xyz"
|
|
|
params["pose_rep"] = "xyz"
|
|
|
else:
|
|
|
outputkey = "poses"
|
|
|
|
|
|
keep = [outputkey, "lengths", "y"]
|
|
|
visu = {key: visualization[key].data.cpu().numpy() for key in keep}
|
|
|
lenmax = visu["lengths"].max()
|
|
|
|
|
|
timesize = lenmax + 5
|
|
|
import multiprocessing
|
|
|
|
|
|
def pool_job_with_desc(pool, iterator, desc, max_, save_path_format):
|
|
|
with tqdm(total=max_, desc=desc.format("Render")) as pbar:
|
|
|
for _ in pool.imap_unordered(plot_3d_motion_dico, iterator):
|
|
|
pbar.update()
|
|
|
array = np.stack([[load_anim(save_path_format.format(i, j), timesize)
|
|
|
for j in range(nats)]
|
|
|
for i in tqdm(range(nspa), desc=desc.format("Load"))])
|
|
|
return array.transpose(2, 0, 1, 3, 4, 5)
|
|
|
|
|
|
with multiprocessing.Pool() as pool:
|
|
|
|
|
|
save_path_format = os.path.join(tmp_path, "real_{}_{}.gif")
|
|
|
iterator = ((visu[outputkey][i, j],
|
|
|
visu["lengths"][i, j],
|
|
|
save_path_format.format(i, j),
|
|
|
params, {"title": f"real: {label_to_action_name(visu['y'][i, j])}", "interval": 1000/fps})
|
|
|
for j in range(nats) for i in range(nspa))
|
|
|
visu["frames"] = pool_job_with_desc(pool, iterator,
|
|
|
"{} the real samples",
|
|
|
nats,
|
|
|
save_path_format)
|
|
|
frames = stack_images_sequence(visu["frames"])
|
|
|
return frames
|
|
|
|
|
|
|
|
|
def viz_clip_text(model, text_grid, epoch, params, folder):
|
|
|
""" Generate & viz samples """
|
|
|
|
|
|
|
|
|
model.outputxyz = True
|
|
|
|
|
|
print(f"Visualization of the epoch {epoch}")
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
fact = params["fact_latent"]
|
|
|
figname = params["figname"].format(epoch)
|
|
|
|
|
|
classes = np.array(text_grid, dtype=str)
|
|
|
h, w = classes.shape
|
|
|
|
|
|
texts = classes.reshape([-1])
|
|
|
clip_tokens = clip.tokenize(texts).to(params['device'])
|
|
|
clip_features = model.clip_model.encode_text(clip_tokens).float().unsqueeze(0)
|
|
|
|
|
|
gendurations = torch.ones((h*w, 1), dtype=int) * params['num_frames']
|
|
|
|
|
|
|
|
|
model.eval()
|
|
|
with torch.no_grad():
|
|
|
|
|
|
generation = model.generate(clip_features, gendurations,
|
|
|
is_amass=True,
|
|
|
is_clip_features=True)
|
|
|
generation['y'] = texts
|
|
|
|
|
|
for key, val in generation.items():
|
|
|
if len(generation[key].shape) == 1:
|
|
|
generation[key] = val.reshape(h, w)
|
|
|
else:
|
|
|
generation[key] = val.reshape(h, w, *val.shape[1:])
|
|
|
|
|
|
f_name = params['input_file']
|
|
|
if os.path.isfile(params['input_file']):
|
|
|
f_name = os.path.basename(params['input_file'].replace('.txt', ''))
|
|
|
finalpath = os.path.join(folder, 'clip_text_{}_{}'.format(f_name, 'trans_' if params['vertstrans'] else '') + figname + ".gif")
|
|
|
tmp_path = os.path.join(folder, f"clip_text_subfigures_{figname}")
|
|
|
os.makedirs(tmp_path, exist_ok=True)
|
|
|
|
|
|
|
|
|
|
|
|
print("Generate the videos..")
|
|
|
frames = generate_by_video({}, {}, generation,
|
|
|
lambda x: str(x), params, w, h, tmp_path, mode='text')
|
|
|
|
|
|
|
|
|
print(f"Writing video [{finalpath}]")
|
|
|
imageio.mimsave(finalpath, frames, fps=params["fps"])
|
|
|
|
|
|
|
|
|
def viz_clip_interp(model, datasets, interp_csv, num_stops, epoch, params, folder):
|
|
|
""" Generate & viz samples """
|
|
|
|
|
|
|
|
|
model.outputxyz = True
|
|
|
|
|
|
print(f"Visualization of the epoch {epoch}")
|
|
|
figname = params["figname"].format(epoch)
|
|
|
motion_collection = get_motion_text_mapping(datasets)
|
|
|
|
|
|
|
|
|
all_clip_features = []
|
|
|
all_texts = []
|
|
|
for line in interp_csv:
|
|
|
|
|
|
texts = [line['start'], line['end']]
|
|
|
retrieved_motions = retrieve_motions(datasets, motion_collection, texts, params['device'])
|
|
|
clip_features = encode_motions(model, retrieved_motions, params['device'])
|
|
|
|
|
|
|
|
|
|
|
|
end_factor = np.linspace(0., 1., num=num_stops)
|
|
|
start_factor = 1. - end_factor
|
|
|
interp_features = [(start_factor[i]*clip_features[0]) + (end_factor[i]*clip_features[1]) for i in range(num_stops)]
|
|
|
all_clip_features.append(torch.stack(interp_features))
|
|
|
texts = texts[:1] + [' '] * (num_stops-2) + texts[-1:]
|
|
|
all_texts.append(texts)
|
|
|
|
|
|
all_clip_features = torch.transpose(torch.stack(all_clip_features, axis=0), 0, 1)
|
|
|
all_texts = np.array(all_texts).T
|
|
|
h, w = all_clip_features.shape[:2]
|
|
|
gendurations = torch.ones((h*w, 1), dtype=int) * params['num_frames']
|
|
|
|
|
|
|
|
|
model.eval()
|
|
|
with torch.no_grad():
|
|
|
generation = model.generate(all_clip_features, gendurations,
|
|
|
is_amass=True,
|
|
|
is_clip_features=True)
|
|
|
generation['y'] = all_texts.reshape([-1])
|
|
|
|
|
|
for key, val in generation.items():
|
|
|
if len(generation[key].shape) == 1:
|
|
|
generation[key] = val.reshape(h, w)
|
|
|
else:
|
|
|
generation[key] = val.reshape(h, w, *val.shape[1:])
|
|
|
|
|
|
if os.path.isfile(params['input_file']):
|
|
|
f_name = os.path.basename(params['input_file'].replace('.csv', ''))
|
|
|
finalpath = os.path.join(folder, f'clip_edit_{f_name}_' + figname + ".gif")
|
|
|
tmp_path = os.path.join(folder, f"clip_edit_subfigures_{figname}")
|
|
|
os.makedirs(tmp_path, exist_ok=True)
|
|
|
|
|
|
print("Generate the videos..")
|
|
|
frames = generate_by_video({}, {}, generation,
|
|
|
lambda x: str(x), params, w, h, tmp_path, mode='interp')
|
|
|
|
|
|
print(f"Writing video [{finalpath}]")
|
|
|
imageio.mimsave(finalpath, frames, fps=params["fps"])
|
|
|
|
|
|
|
|
|
def viz_clip_edit(model, datasets, edit_csv, epoch, params, folder):
|
|
|
""" Generate & viz samples """
|
|
|
|
|
|
|
|
|
model.outputxyz = True
|
|
|
|
|
|
print(f"Visualization of the epoch {epoch}")
|
|
|
figname = params["figname"].format(epoch)
|
|
|
motion_collection = get_motion_text_mapping(datasets)
|
|
|
|
|
|
|
|
|
all_clip_features = []
|
|
|
all_texts = []
|
|
|
for line in edit_csv:
|
|
|
|
|
|
texts = [line['base'], line['v_start'], line['v_end']]
|
|
|
if line['motion_source'] == 'data':
|
|
|
retrieved_motions = retrieve_motions(datasets, motion_collection, texts, params['device'])
|
|
|
clip_features = encode_motions(model, retrieved_motions, params['device'])
|
|
|
elif line['motion_source'] == 'text':
|
|
|
clip_tokens = clip.tokenize(texts).to(params['device'])
|
|
|
clip_features = model.clip_model.encode_text(clip_tokens).float()
|
|
|
else:
|
|
|
raise ValueError
|
|
|
|
|
|
|
|
|
result_features = clip_features[0] - clip_features[1] + clip_features[2]
|
|
|
all_clip_features.append(torch.cat([clip_features, result_features.unsqueeze(0)]))
|
|
|
texts.append('Result')
|
|
|
all_texts.append(texts)
|
|
|
|
|
|
all_clip_features = torch.transpose(torch.stack(all_clip_features, axis=0), 0, 1)
|
|
|
all_texts = np.array(all_texts).T
|
|
|
h, w = all_clip_features.shape[:2]
|
|
|
gendurations = torch.ones((h*w, 1), dtype=int) * params['num_frames']
|
|
|
|
|
|
|
|
|
model.eval()
|
|
|
with torch.no_grad():
|
|
|
generation = model.generate(all_clip_features, gendurations,
|
|
|
is_amass=True,
|
|
|
is_clip_features=True)
|
|
|
generation['y'] = all_texts.reshape([-1])
|
|
|
|
|
|
for key, val in generation.items():
|
|
|
if len(generation[key].shape) == 1:
|
|
|
generation[key] = val.reshape(h, w)
|
|
|
else:
|
|
|
generation[key] = val.reshape(h, w, *val.shape[1:])
|
|
|
|
|
|
if os.path.isfile(params['input_file']):
|
|
|
f_name = os.path.basename(params['input_file'].replace('.csv', ''))
|
|
|
finalpath = os.path.join(folder, f'clip_edit_{f_name}_' + figname + ".gif")
|
|
|
tmp_path = os.path.join(folder, f"clip_edit_subfigures_{figname}")
|
|
|
os.makedirs(tmp_path, exist_ok=True)
|
|
|
|
|
|
print("Generate the videos..")
|
|
|
frames = generate_by_video({}, {}, generation,
|
|
|
lambda x: str(x), params, w, h, tmp_path, mode='edit')
|
|
|
|
|
|
print(f"Writing video [{finalpath}]")
|
|
|
imageio.mimsave(finalpath, frames, fps=params["fps"])
|
|
|
|
|
|
|
|
|
def stack_images_sequence(visu):
|
|
|
print("Stacking frames..")
|
|
|
allframes = visu
|
|
|
nframes, nspa, nats, h, w, pix = allframes.shape
|
|
|
frames = []
|
|
|
for frame_idx in tqdm(range(nframes)):
|
|
|
columns = np.vstack(allframes[frame_idx].transpose(1, 2, 3, 4, 0)).transpose(3, 1, 0, 2)
|
|
|
frame = np.concatenate(columns).transpose(1, 0, 2)
|
|
|
frames.append(frame)
|
|
|
return np.stack(frames)
|
|
|
|
|
|
|
|
|
def get_gpu_device():
|
|
|
command = "nvidia-smi --query-gpu=memory.free --format=csv"
|
|
|
memory_free_info = sp.check_output(command.split()).decode('ascii').split('\n')[:-1][1:]
|
|
|
memory_free_values = [int(x.split()[0]) for i, x in enumerate(memory_free_info)]
|
|
|
for gpu_idx, free_mem in enumerate(memory_free_values):
|
|
|
if free_mem > GPU_MINIMUM_MEMORY:
|
|
|
return gpu_idx
|
|
|
Exception('No GPU with required memory')
|
|
|
|
|
|
|
|
|
def get_motion_text_mapping(datasets):
|
|
|
print('Building text-motion mapping...')
|
|
|
split_names = list(datasets.keys())
|
|
|
collection_path = datasets[split_names[0]].datapath.replace('.pt', '_text_labels.txt')
|
|
|
if len(split_names) > 1:
|
|
|
assert split_names[0] in os.path.basename(collection_path)
|
|
|
_base = os.path.basename(collection_path).replace(split_names[0], 'all')
|
|
|
collection_path = os.path.join(os.path.dirname(collection_path), _base)
|
|
|
cache_path = collection_path.replace('.txt', '.npy')
|
|
|
|
|
|
|
|
|
word = 'Loading' if os.path.isfile(cache_path) else 'Saving'
|
|
|
print('{} list of text labels in current dataset to [{}]:'.format(word, collection_path))
|
|
|
print('Look it up next time you want to retrieve new motions using textual labels.')
|
|
|
|
|
|
if os.path.isfile(cache_path):
|
|
|
return np.load(cache_path, allow_pickle=True)[None][0]
|
|
|
|
|
|
motion_collection = {}
|
|
|
for split_name, data in datasets.items():
|
|
|
for i, d in tqdm(enumerate(data)):
|
|
|
motion_collection[d['clip_text']] = motion_collection.get(d['clip_text'], []) + [(split_name, i)]
|
|
|
|
|
|
with open(collection_path, 'w') as fw:
|
|
|
text_labels = sorted(list(motion_collection.keys()))
|
|
|
fw.write('\n'.join(text_labels) + '\n')
|
|
|
np.save(cache_path, motion_collection)
|
|
|
|
|
|
return motion_collection
|
|
|
|
|
|
def retrieve_motions(datasets, motion_collection, texts, device):
|
|
|
retrieved_motions = []
|
|
|
for txt in texts:
|
|
|
_split, _index = motion_collection[txt][0]
|
|
|
retrieved_motions.append(datasets[_split][_index]['inp'].unsqueeze(0).to(device))
|
|
|
return torch.cat(retrieved_motions, axis=0)
|
|
|
|
|
|
def encode_motions(model, motions, device):
|
|
|
return model.encoder({'x': motions,
|
|
|
'y': torch.zeros(motions.shape[0], dtype=int, device=device),
|
|
|
'mask': model.lengths_to_mask(torch.ones(motions.shape[0], dtype=int, device=device) * 60)})["mu"]
|
|
|
|
