diff --git a/nerfstudio/data/datamanagers/base_datamanager.py b/nerfstudio/data/datamanagers/base_datamanager.py index b2f23feb7..e48e3621a 100644 --- a/nerfstudio/data/datamanagers/base_datamanager.py +++ b/nerfstudio/data/datamanagers/base_datamanager.py @@ -18,11 +18,11 @@ Datamanager. from __future__ import annotations -import functools from abc import abstractmethod from collections import defaultdict from dataclasses import dataclass, field from pathlib import Path +from functools import cached_property from typing import ( Any, Callable, @@ -35,6 +35,9 @@ from typing import ( Type, Union, cast, + ForwardRef, + get_origin, + get_args, ) import torch @@ -66,6 +69,7 @@ from nerfstudio.engine.callbacks import TrainingCallback, TrainingCallbackAttrib from nerfstudio.model_components.ray_generators import RayGenerator from nerfstudio.utils.misc import IterableWrapper from nerfstudio.utils.rich_utils import CONSOLE +from nerfstudio.utils.misc import get_orig_class def variable_res_collate(batch: List[Dict]) -> Dict: @@ -376,7 +380,6 @@ class VanillaDataManager(DataManager, Generic[TDataset]): local_rank: int = 0, **kwargs, ): - self.dataset_type: Type[TDataset] = kwargs.get("_dataset_type", getattr(TDataset, "__default__")) self.config = config self.device = device self.world_size = world_size @@ -409,15 +412,32 @@ class VanillaDataManager(DataManager, Generic[TDataset]): CONSOLE.print("Variable resolution, using variable_res_collate") self.config.collate_fn = variable_res_collate break - super().__init__() - def __class_getitem__(cls, item): - return type( - cls.__name__, - (cls,), - {"__module__": cls.__module__, "__init__": functools.partialmethod(cls.__init__, _dataset_type=item)}, - ) + @cached_property + def dataset_type(self) -> Type[TDataset]: + """Returns the dataset type passed as the generic argument""" + default: Type[TDataset] = cast(TDataset, TDataset.__default__) # type: ignore + orig_class: Type[VanillaDataManager] = get_orig_class(self, default=None) # type: ignore + if type(self) is VanillaDataManager and orig_class is None: + return default + if orig_class is not None and get_origin(orig_class) is VanillaDataManager: + return get_args(orig_class)[0] + + # For inherited classes, we need to find the correct type to instantiate + for base in getattr(self, "__orig_bases__", []): + if get_origin(base) is VanillaDataManager: + for value in get_args(base): + if isinstance(value, ForwardRef): + if value.__forward_evaluated__: + value = value.__forward_value__ + elif value.__forward_module__ is None: + value.__forward_module__ = type(self).__module__ + value = getattr(value, "_evaluate")(None, None, set()) + assert isinstance(value, type) + if issubclass(value, InputDataset): + return cast(Type[TDataset], value) + return default def create_train_dataset(self) -> TDataset: """Sets up the data loaders for training""" diff --git a/nerfstudio/utils/misc.py b/nerfstudio/utils/misc.py index e95bd4188..3e4c4df86 100644 --- a/nerfstudio/utils/misc.py +++ b/nerfstudio/utils/misc.py @@ -17,6 +17,8 @@ Miscellaneous helper code. """ +from inspect import currentframe +import typing import platform from typing import Any, Callable, Dict, List, Optional, TypeVar, Union import warnings @@ -182,3 +184,32 @@ def torch_compile(*args, **kwargs): return lambda x: x else: return torch.compile(*args, **kwargs) + + +def get_orig_class(obj, default=None): + """Returns the __orig_class__ class of `obj` even when it is not initialized in __init__ (Python>=3.8). + + Workaround for https://github.com/python/typing/issues/658. + Inspired by https://github.com/Stewori/pytypes/pull/53. + """ + try: + return object.__getattribute__(obj, "__orig_class__") + except AttributeError: + cls = object.__getattribute__(obj, "__class__") + try: + is_type_generic = isinstance(cls, typing.GenericMeta) # type: ignore + except AttributeError: # Python 3.8 + is_type_generic = issubclass(cls, typing.Generic) + if is_type_generic: + frame = currentframe().f_back.f_back # type: ignore + try: + while frame: + try: + res = frame.f_locals["self"] + if res.__origin__ is cls: + return res + except (KeyError, AttributeError): + frame = frame.f_back + finally: + del frame + return default diff --git a/tests/data/test_datamanager.py b/tests/data/test_datamanager.py index e1a2a24c0..0547b736e 100644 --- a/tests/data/test_datamanager.py +++ b/tests/data/test_datamanager.py @@ -1,3 +1,4 @@ +import pickle from pathlib import Path from typing import Any @@ -52,18 +53,49 @@ def config(): def test_data_manager_type_inference(config): # Mock for a faster test - assert VanillaDataManager(config).dataset_type is InputDataset assert VanillaDataManager[DepthDataset](config).dataset_type is DepthDataset + assert VanillaDataManager(config).dataset_type is InputDataset + + class tmp2(VanillaDataManager[DepthDataset]): + pass + + assert tmp2(config).dataset_type is DepthDataset class tmp(VanillaDataManager): pass assert tmp(config).dataset_type is InputDataset - class tmp2(VanillaDataManager[DepthDataset]): - pass - assert tmp2(config).dataset_type is DepthDataset +class _pickle_enabled_tmp(VanillaDataManager): + pass + + +class _pickle_enabled_tmp2(VanillaDataManager[DepthDataset]): + pass + + +def test_data_manager_type_can_be_pickled(config): + # Mock for a faster test + assert VanillaDataManager[DepthDataset](config).dataset_type is DepthDataset + obj = pickle.loads(pickle.dumps(VanillaDataManager[DepthDataset](config))) + assert obj.dataset_type is DepthDataset + assert isinstance(obj, VanillaDataManager) + + assert VanillaDataManager(config).dataset_type is InputDataset + obj = pickle.loads(pickle.dumps(VanillaDataManager(config))) + assert obj.dataset_type is InputDataset + assert isinstance(obj, VanillaDataManager) + + assert _pickle_enabled_tmp(config).dataset_type is InputDataset + obj = pickle.loads(pickle.dumps(_pickle_enabled_tmp(config))) + assert obj.dataset_type is InputDataset + assert isinstance(obj, _pickle_enabled_tmp) + + assert _pickle_enabled_tmp2(config).dataset_type is DepthDataset + obj = pickle.loads(pickle.dumps(_pickle_enabled_tmp2(config))) + assert obj.dataset_type is DepthDataset + assert isinstance(obj, _pickle_enabled_tmp2) def test_data_manager_type_can_be_serialized(config):