create wrapper for dataset

data/dataset.py

@@ -1,7 +1,7 @@
 import os
 import sys
 
-# os.environ["CUDA_VISIBLE_DEVICES"] = '6'
+# os.environ["CUDA_VISIBLE_DEVICES"] = '0,1,2,3'
 
 import math
 import time
@@ -14,7 +14,6 @@ import scipy.io as sio
 import plotly.express as ex
 import torchvision
 
-from torch.utils.data import Dataset
 from torchvision.transforms import Resize
 from torchvision.transforms.functional import InterpolationMode
 
@@ -22,6 +21,7 @@ sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
 
 from data.utils import Frame
 from data.cloth3d_io import readPC2, readPC2Frame
+from data.dataset_cape import DatasetCape
 
 
 def read_frames(folder, shuffle=True, is_cape=False, testing=False, frames_file=None):
@@ -131,7 +131,7 @@ def read_frames_file(frames_file, shuffle=False):
     return frames
 
 
-class Dataset(Dataset):
+class DatasetCloth3D(torch.utils.data.Dataset):
     def __init__(self, folder, signature, resolutions=[512], frames_file='', device='cuda', use_cache=False, is_cape=False):
         self.folder = folder
         self.device = device
@@ -229,6 +229,14 @@ class Dataset(Dataset):
         data = { **data, **self.static_cache[frame] }
         return data
 
+def Dataset(folder, signature, resolutions=[512], frames_file='', device='cuda', use_cache=False, is_cape=False):
+    if not is_cape:
+        return DatasetCloth3D(folder, signature, resolutions=resolutions, frames_file=frames_file, device=device, use_cache=use_cache, is_cape=False)
+    else:
+        return DatasetCape(folder, signature, resolutions=resolutions, frames_file=frames_file, device=device, use_cache=use_cache, is_cape=True)
+
+
+
 
 if __name__ == '__main__':
     from torch.utils.data import DataLoader