File size: 1,520 Bytes
2c76547 | 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 | # Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
import copy
from dynamic_stereo.models.dynamic_stereo_model import DynamicStereoModel
from pytorch3d.implicitron.tools.config import get_default_args
try:
from dynamic_stereo.models.raft_stereo_model import RAFTStereoModel
MODELS = [RAFTStereoModel, DynamicStereoModel]
except:
MODELS = [DynamicStereoModel]
_MODEL_NAME_TO_MODEL = {model_cls.__name__: model_cls for model_cls in MODELS}
_MODEL_CONFIG_NAME_TO_DEFAULT_CONFIG = {}
for model_cls in MODELS:
_MODEL_CONFIG_NAME_TO_DEFAULT_CONFIG[
model_cls.MODEL_CONFIG_NAME
] = get_default_args(model_cls)
MODEL_NAME_NONE = "NONE"
def model_zoo(model_name: str, **kwargs):
if model_name.upper() == MODEL_NAME_NONE:
return None
model_cls = _MODEL_NAME_TO_MODEL.get(model_name)
if model_cls is None:
raise ValueError(f"No such model name: {model_name}")
model_cls_params = {}
if "model_zoo" in getattr(model_cls, "__dataclass_fields__", []):
model_cls_params["model_zoo"] = model_zoo
print(
f"{model_cls.MODEL_CONFIG_NAME} model configs:",
kwargs.get(model_cls.MODEL_CONFIG_NAME),
)
return model_cls(**model_cls_params, **kwargs.get(model_cls.MODEL_CONFIG_NAME, {}))
def get_all_model_default_configs():
return copy.deepcopy(_MODEL_CONFIG_NAME_TO_DEFAULT_CONFIG)
|