upscaling pipeline
This commit is contained in:
+11
-10
@@ -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
|
||||
|
||||
|
||||
Reference in New Issue
Block a user