File size: 2,499 Bytes
57746f1 |
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 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 |
#!/usr/bin/env python3
# Copyright (C) 2024-present Naver Corporation. All rights reserved.
# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
#
# --------------------------------------------------------
# training executable for MASt3R
# --------------------------------------------------------
import sys
sys.path.append('.')
sys.path.append('submodules/mast3r')
from mast3r.model import AsymmetricMASt3R
from mast3r.losses import ConfMatchingLoss, MatchingLoss, APLoss, Regr3D, InfoNCE, Regr3D_ScaleShiftInv
from mast3r.datasets import ARKitScenes, BlendedMVS, Co3d, MegaDepth, ScanNetpp, StaticThings3D, Waymo, WildRGBD
import mast3r.utils.path_to_dust3r # noqa
# add mast3r classes to dust3r imports
import dust3r.training
dust3r.training.AsymmetricMASt3R = AsymmetricMASt3R
dust3r.training.Regr3D = Regr3D
dust3r.training.Regr3D_ScaleShiftInv = Regr3D_ScaleShiftInv
dust3r.training.MatchingLoss = MatchingLoss
dust3r.training.ConfMatchingLoss = ConfMatchingLoss
dust3r.training.InfoNCE = InfoNCE
dust3r.training.APLoss = APLoss
import dust3r.datasets
dust3r.datasets.ARKitScenes = ARKitScenes
dust3r.datasets.BlendedMVS = BlendedMVS
dust3r.datasets.Co3d = Co3d
dust3r.datasets.MegaDepth = MegaDepth
dust3r.datasets.ScanNetpp = ScanNetpp
dust3r.datasets.StaticThings3D = StaticThings3D
dust3r.datasets.Waymo = Waymo
dust3r.datasets.WildRGBD = WildRGBD
from src.datasets.scannet import Scannet
from src.datasets.scannetpp import Scannetpp
from src.datasets.megadepth import MegaDepth
dust3r.datasets.Scannet = Scannet
dust3r.datasets.Scannetpp = Scannetpp
dust3r.datasets.MegaDepth = MegaDepth
from src.model import LSM_MASt3R
dust3r.training.LSM_MASt3R = LSM_MASt3R
from src.losses import GaussianLoss
dust3r.training.GaussianLoss = GaussianLoss
from dust3r.training import get_args_parser as dust3r_get_args_parser # noqa
from dust3r.training import train # noqa
import yaml
def get_args_parser():
parser = dust3r_get_args_parser()
parser.prog = 'LSM_MASt3R training'
# Load the configuration
with open("configs/model_config.yaml", "r") as f:
config = yaml.safe_load(f)
# Convert the config dict to a string of keyword arguments
config_str = ", ".join(f"{k}={v}" for k, v in config.items())
# Set the default model string with parameters
parser.set_defaults(model=f"LSM_MASt3R({config_str})")
return parser
if __name__ == '__main__':
args = get_args_parser()
args = args.parse_args()
train(args)
|