fix cpu ram following graphdeco-inria#667

nnmhuy · Sep 24, 2024 · 894c4e0 · 894c4e0
1 parent bf641e5
commit 894c4e0
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/utils/camera_utils.py b/utils/camera_utils.py
@@ -39,7 +39,8 @@ def loadCam(args, id, cam_info, resolution_scale):
         scale = float(global_down) * float(resolution_scale)
         resolution = (int(orig_w / scale), int(orig_h / scale))
 
-    resized_image_rgb = PILtoTorch(cam_info.image, resolution)
+    # resized_image_rgb = PILtoTorch(cam_info.image, resolution)
+    resized_image_rgb = PILtoTorch(cam_info.image, resolution, args.data_device)
 
     gt_image = resized_image_rgb[:3, ...]
     loaded_mask = None

diff --git a/utils/general_utils.py b/utils/general_utils.py
@@ -18,9 +18,11 @@
 def inverse_sigmoid(x):
     return torch.log(x/(1-x))
 
-def PILtoTorch(pil_image, resolution):
+def PILtoTorch(pil_image, resolution, data_device="cuda"):
     resized_image_PIL = pil_image.resize(resolution)
-    resized_image = torch.from_numpy(np.array(resized_image_PIL)) / 255.0
+    # resized_image = torch.from_numpy(np.array(resized_image_PIL)) / 255.0
+    resized_image = torch.from_numpy(np.array(resized_image_PIL)).to(data_device) / 255.0
+
     if len(resized_image.shape) == 3:
         return resized_image.permute(2, 0, 1)
     else: