in threestudio/data/images.py [0:0]
def load_images(self):
# load image
assert os.path.exists(
self.cfg.image_path
), f"Could not find image {self.cfg.image_path}!"
rgba = cv2.cvtColor(
cv2.imread(self.cfg.image_path, cv2.IMREAD_UNCHANGED), cv2.COLOR_BGRA2RGBA
)
rgba = (
cv2.resize(
rgba, (self.width, self.height), interpolation=cv2.INTER_AREA
).astype(np.float32)
/ 255.0
)
rgb = rgba[..., :3]
self.rgb: Float[Tensor, "1 H W 3"] = (
torch.from_numpy(rgb).unsqueeze(0).contiguous().to(self.rank)
)
self.mask: Float[Tensor, "1 H W 1"] = (
torch.from_numpy(rgba[..., 3:] > 0.5).unsqueeze(0).to(self.rank)
)
print(
f"[INFO] single image dataset: load image {self.cfg.image_path} {self.rgb.shape}"
)
# load depth
if self.cfg.requires_depth:
depth_path = self.cfg.image_path.replace("_rgba.png", "_depth.png")
assert os.path.exists(depth_path)
depth = cv2.imread(depth_path, cv2.IMREAD_UNCHANGED)
depth = cv2.resize(
depth, (self.width, self.height), interpolation=cv2.INTER_AREA
)
self.depth: Float[Tensor, "1 H W 1"] = (
torch.from_numpy(depth.astype(np.float32) / 255.0)
.unsqueeze(0)
.to(self.rank)
)
print(
f"[INFO] single image dataset: load depth {depth_path} {self.depth.shape}"
)
else:
self.depth = None
# load normal
if self.cfg.requires_normal:
normal_path = self.cfg.image_path.replace("_rgba.png", "_normal.png")
assert os.path.exists(normal_path)
normal = cv2.imread(normal_path, cv2.IMREAD_UNCHANGED)
normal = cv2.resize(
normal, (self.width, self.height), interpolation=cv2.INTER_AREA
)
self.normal: Float[Tensor, "1 H W 3"] = (
torch.from_numpy(normal.astype(np.float32) / 255.0)
.unsqueeze(0)
.to(self.rank)
)
print(
f"[INFO] single image dataset: load normal {normal_path} {self.normal.shape}"
)
else:
self.normal = None