upscaling pipeline

This commit is contained in:
2023-05-21 13:52:42 +00:00
parent fc27a3cedf
commit 75df8a09cc
5 changed files with 322 additions and 173 deletions
+11 -10
View File
@@ -30,14 +30,15 @@ def read_frames(folder, shuffle=True, is_cape=False, testing=False, frames_file=
if is_cape:
frames = ['.'.join(frame.split('.')[:-1]) for frame in frames] # .split('_')[0]
frames = [Frame(frame.split('.')[0], frame.split('.')[1]) for frame in frames]
else:
frames = [frame.split('.')[0] for frame in frames] # .split('_')[0]
frames = [Frame(frame.split('_')[0], frame.split('_')[1]) for frame in frames]
if testing:
frames = sorted(os.listdir(os.path.join(folder, 'mesh_static') ))
frames = [frame.split('.')[0] + '_4' for frame in frames]
# if testing:
# frames = sorted(os.listdir(os.path.join(folder, 'mesh_static') ))
# frames = [frame.split('.')[0] + '_4' for frame in frames]
frames = [Frame(frame.split('_')[0], frame.split('_')[1]) for frame in frames]
if frames_file is not None:
frames = []
@@ -65,12 +66,12 @@ class StaticCache():
if is_cape:
self.uv_masks = dict()
uv_mask = np.load('custom_data/garment_uv_mask_full.npy').astype(np.int32)
self.uv_masks[512] = uv_mask.astype(np.float32)
uv_mask = torch.tensor(np.load('custom_data/garment_uv_mask_full.npy').astype(np.int32), device=device)
self.uv_masks[512] = uv_mask.to(torch.float32).cpu().numpy()
for res in [32, 64, 128, 256]:
down = Resize(uv_mask[:,:,None], size=(res, int(res/2)), interpolation=InterpolationMode.NEAREST)
down = np.squeeze(down.numpy().astype(np.float32))
self.uv_masks[res] = down
down = Resize(size=(res, int(res/2)), interpolation=InterpolationMode.NEAREST)(uv_mask[None,:,:])
down = torch.squeeze(down).to(torch.float32)
self.uv_masks[res] = down.cpu().numpy()
def read(self, id):
@@ -127,7 +128,7 @@ class Dataset(Dataset):
self.frames = read_frames(folder, frames_file=frames_file, is_cape=is_cape, shuffle=True)
self.body_faces = np.load('smpl/smpl_faces.npy')
self.static_cache = StaticCache(folder, signature, is_cape=is_cape)
self.static_cache = StaticCache(folder, signature, is_cape=is_cape, resolutions=resolutions, device=device)
self.cache = dict() if use_cache else None