Spaces:
Runtime error
Runtime error
File size: 2,225 Bytes
2df809d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 |
import os.path as osp
import sys
sys.path.append(osp.join(osp.dirname(__file__), "..", ".."))
import cv2
import numpy as np
from dust3r.datasets.co3d import Co3d_Multi
from dust3r.utils.image import imread_cv2
class WildRGBD_Multi(Co3d_Multi):
def __init__(self, mask_bg="rand", *args, ROOT, **kwargs):
super().__init__(mask_bg, *args, ROOT=ROOT, **kwargs)
self.dataset_label = "WildRGBD"
self.is_metric = True
# load all scenes
self.scenes.pop(("box", "scenes/scene_257"), None)
self.scene_list = list(self.scenes.keys())
cut_off = (
self.num_views if not self.allow_repeat else max(self.num_views // 3, 3)
)
self.cut_off = cut_off
self.all_ref_imgs = [
(key, value)
for key, values in self.scenes.items()
for value in values[: len(values) - cut_off + 1]
]
self.invalidate = {scene: {} for scene in self.scene_list}
self.invalid_scenes = {scene: False for scene in self.scene_list}
def _get_metadatapath(self, obj, instance, view_idx):
return osp.join(self.ROOT, obj, instance, "metadata", f"{view_idx:0>5d}.npz")
def _get_impath(self, obj, instance, view_idx):
return osp.join(self.ROOT, obj, instance, "rgb", f"{view_idx:0>5d}.jpg")
def _get_depthpath(self, obj, instance, view_idx):
return osp.join(self.ROOT, obj, instance, "depth", f"{view_idx:0>5d}.png")
def _get_maskpath(self, obj, instance, view_idx):
return osp.join(self.ROOT, obj, instance, "masks", f"{view_idx:0>5d}.png")
def _read_depthmap(self, depthpath, input_metadata):
# We store depths in the depth scale of 1000.
# That is, when we load depth image and divide by 1000, we could get depth in meters.
depthmap = imread_cv2(depthpath, cv2.IMREAD_UNCHANGED)
depthmap = depthmap.astype(np.float32) / 1000.0
return depthmap
def _get_views(self, idx, resolution, rng, num_views):
views = super()._get_views(idx, resolution, rng, num_views)
for view in views:
assert view["is_metric"]
view["quantile"] = np.array(0.96, dtype=np.float32)
return views
|