|
|
|
|
|
|
|
from .base.mast3r_base_stereo_view_dataset import MASt3RBaseStereoViewDataset |
|
|
|
import mast3r.utils.path_to_dust3r |
|
from dust3r.datasets.arkitscenes import ARKitScenes as DUSt3R_ARKitScenes |
|
from dust3r.datasets.blendedmvs import BlendedMVS as DUSt3R_BlendedMVS |
|
from dust3r.datasets.co3d import Co3d as DUSt3R_Co3d |
|
from dust3r.datasets.megadepth import MegaDepth as DUSt3R_MegaDepth |
|
from dust3r.datasets.scannetpp import ScanNetpp as DUSt3R_ScanNetpp |
|
from dust3r.datasets.staticthings3d import StaticThings3D as DUSt3R_StaticThings3D |
|
from dust3r.datasets.waymo import Waymo as DUSt3R_Waymo |
|
from dust3r.datasets.wildrgbd import WildRGBD as DUSt3R_WildRGBD |
|
|
|
|
|
class ARKitScenes(DUSt3R_ARKitScenes, MASt3RBaseStereoViewDataset): |
|
def __init__(self, *args, split, ROOT, **kwargs): |
|
super().__init__(*args, split=split, ROOT=ROOT, **kwargs) |
|
self.is_metric_scale = True |
|
|
|
|
|
class BlendedMVS(DUSt3R_BlendedMVS, MASt3RBaseStereoViewDataset): |
|
def __init__(self, *args, ROOT, split=None, **kwargs): |
|
super().__init__(*args, ROOT=ROOT, split=split, **kwargs) |
|
self.is_metric_scale = False |
|
|
|
|
|
class Co3d(DUSt3R_Co3d, MASt3RBaseStereoViewDataset): |
|
def __init__(self, mask_bg=True, *args, ROOT, **kwargs): |
|
super().__init__(mask_bg, *args, ROOT=ROOT, **kwargs) |
|
self.is_metric_scale = False |
|
|
|
|
|
class MegaDepth(DUSt3R_MegaDepth, MASt3RBaseStereoViewDataset): |
|
def __init__(self, *args, split, ROOT, **kwargs): |
|
super().__init__(*args, split=split, ROOT=ROOT, **kwargs) |
|
self.is_metric_scale = False |
|
|
|
|
|
class ScanNetpp(DUSt3R_ScanNetpp, MASt3RBaseStereoViewDataset): |
|
def __init__(self, *args, ROOT, **kwargs): |
|
super().__init__(*args, ROOT=ROOT, **kwargs) |
|
self.is_metric_scale = True |
|
|
|
|
|
class StaticThings3D(DUSt3R_StaticThings3D, MASt3RBaseStereoViewDataset): |
|
def __init__(self, ROOT, *args, mask_bg='rand', **kwargs): |
|
super().__init__(ROOT, *args, mask_bg=mask_bg, **kwargs) |
|
self.is_metric_scale = False |
|
|
|
|
|
class Waymo(DUSt3R_Waymo, MASt3RBaseStereoViewDataset): |
|
def __init__(self, *args, ROOT, **kwargs): |
|
super().__init__(*args, ROOT=ROOT, **kwargs) |
|
self.is_metric_scale = True |
|
|
|
|
|
class WildRGBD(DUSt3R_WildRGBD, MASt3RBaseStereoViewDataset): |
|
def __init__(self, mask_bg=True, *args, ROOT, **kwargs): |
|
super().__init__(mask_bg, *args, ROOT=ROOT, **kwargs) |
|
self.is_metric_scale = True |
|
|