basic test

data/dataset.py

@@ -14,9 +14,9 @@ import plotly.express as ex
 
 from torch.utils.data import Dataset
 
-from render import Render, DepthRenderer
 
 sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
+from data.render import Render, DepthRenderer
 from superSMPL.IO import readPC2Frame, writePC2
 from utils.data_utils import Frame, read_frames_file, mask_faces, get_normals_from_depth
 from cloth3d.read import DataReader
@@ -24,12 +24,13 @@ from cloth3d.util import quads2tris
 
 
 class Dataset(Dataset):
-    def __init__(self, folder, folder_rgb, frames_file='frames_all.txt'):
+    def __init__(self, folder, folder_rgb, frames_file='frames_all.txt', device='cuda'):
         self.folder = folder
         self.folder_rgb = folder_rgb
+        self.device = device
+
         self.transforms = tv.transforms.ToTensor()
         self.frames = read_frames_file(os.path.join(folder, frames_file))
-
         self.render = DepthRenderer()
         self.deodr_render : Render = Render(max_depth=10)
         self.body_faces = np.load('custom_data/smpl_faces.npy')
@@ -58,9 +59,9 @@ class Dataset(Dataset):
         # body_no_trans, _, _, _, _ = self.reader.read_human(frame.id, frame.frame_nr, absolute=False)
         body_faces = quads2tris(body_faces)
 
-        combined = torch.Tensor(np.concatenate((body, mesh_posed[mesh_mask, :]))).to('cuda')
+        combined = torch.Tensor(np.concatenate((body, mesh_posed[mesh_mask, :]))).to(self.device)
         # combined_no_trans = torch.Tensor(np.concatenate((body_no_trans, mesh_posed_no_trans[mesh_mask, :]))).to('cuda')
-        combined_faces = torch.Tensor(np.concatenate((body_faces, faces + body.shape[0]))).to('cuda')
+        combined_faces = torch.Tensor(np.concatenate((body_faces, faces + body.shape[0]))).to(self.device)
 
         start = time.perf_counter()
         # depth_image = self.render.render(mesh_posed[mask, :], faces, info['camLoc'][0])[:,:,0]
@@ -184,7 +185,7 @@ class Dataset(Dataset):
         # ex.imshow(v2i_coverage).show()
 
 
-        rgb_image = tv.io.read_image(os.path.join(self.folder_rgb, frame.id, f'{frame.frame_nr:04d}.png'))
+        rgb_image = tv.io.read_image(os.path.join(self.folder_rgb, frame.id, f'{frame.frame_nr:04d}.png'), mode=tv.io.ImageReadMode.RGB).to(self.device).to(torch.float32)
 
 
         # image = self.transforms(image)