65 lines
2.3 KiB
Python
65 lines
2.3 KiB
Python
from . import common
|
|
from ..data.stanford import config
|
|
from ..data.stanford import read
|
|
from ..data.common import scan
|
|
from typing import Iterable, Optional, Union
|
|
import os
|
|
|
|
|
|
class SingleViewUVScanDataset(common.H5Dataset):
|
|
def __init__(self,
|
|
obj_names : Iterable[str],
|
|
identifiers : Optional[Iterable[str]] = None,
|
|
data_path : Union[str, os.PathLike, None] = None,
|
|
):
|
|
if not obj_names:
|
|
raise ValueError("'obj_names' cannot be empty!")
|
|
if identifiers is None:
|
|
identifiers = read.list_mesh_scan_identifiers()
|
|
if data_path is not None:
|
|
config.DATA_PATH = data_path
|
|
fnames = read.list_mesh_scan_uv_h5_fnames(obj_names, identifiers)
|
|
super().__init__(
|
|
h5_dataclass_cls = scan.SingleViewUVScan,
|
|
fnames = fnames,
|
|
)
|
|
|
|
class AutodecoderSingleViewUVScanDataset(common.AutodecoderDataset):
|
|
def __init__(self,
|
|
obj_names : Iterable[str],
|
|
identifiers : Optional[Iterable[str]] = None,
|
|
data_path : Union[str, os.PathLike, None] = None,
|
|
):
|
|
if identifiers is None:
|
|
identifiers = read.list_mesh_scan_identifiers()
|
|
super().__init__(
|
|
keys = [obj_name for obj_name in obj_names for _ in range(len(identifiers))],
|
|
dataset = SingleViewUVScanDataset(obj_names, identifiers, data_path=data_path),
|
|
)
|
|
|
|
|
|
class SphereScanDataset(common.H5Dataset):
|
|
def __init__(self,
|
|
obj_names : Iterable[str],
|
|
data_path : Union[str, os.PathLike, None] = None,
|
|
):
|
|
if not obj_names:
|
|
raise ValueError("'obj_names' cannot be empty!")
|
|
if data_path is not None:
|
|
config.DATA_PATH = data_path
|
|
fnames = read.list_mesh_sphere_scan_h5_fnames(obj_names)
|
|
super().__init__(
|
|
h5_dataclass_cls = scan.SingleViewUVScan,
|
|
fnames = fnames,
|
|
)
|
|
|
|
class AutodecoderSphereScanDataset(common.AutodecoderDataset):
|
|
def __init__(self,
|
|
obj_names : Iterable[str],
|
|
data_path : Union[str, os.PathLike, None] = None,
|
|
):
|
|
super().__init__(
|
|
keys = obj_names,
|
|
dataset = SphereScanDataset(obj_names, data_path=data_path),
|
|
)
|