51 lines
1.3 KiB
Python
51 lines
1.3 KiB
Python
# Copyright (c) 2023 Amphion.
|
|
#
|
|
# This source code is licensed under the MIT license found in the
|
|
# LICENSE file in the root directory of this source tree.
|
|
|
|
import json
|
|
import os
|
|
from abc import abstractmethod
|
|
from pathlib import Path
|
|
|
|
import json5
|
|
import torch
|
|
import yaml
|
|
|
|
|
|
# TODO: for training and validating
|
|
class BaseDataset(torch.utils.data.Dataset):
|
|
r"""Base dataset for training and validating."""
|
|
|
|
def __init__(self, args, cfg, is_valid=False):
|
|
pass
|
|
|
|
|
|
class BaseTestDataset(torch.utils.data.Dataset):
|
|
r"""Test dataset for inference."""
|
|
|
|
def __init__(self, args=None, cfg=None, infer_type="from_dataset"):
|
|
assert infer_type in ["from_dataset", "from_file"]
|
|
|
|
self.args = args
|
|
self.cfg = cfg
|
|
self.infer_type = infer_type
|
|
|
|
@abstractmethod
|
|
def __getitem__(self, index):
|
|
pass
|
|
|
|
def __len__(self):
|
|
return len(self.metadata)
|
|
|
|
def get_metadata(self):
|
|
path = Path(self.args.source)
|
|
if path.suffix == ".json" or path.suffix == ".jsonc":
|
|
metadata = json5.load(open(self.args.source, "r"))
|
|
elif path.suffix == ".yaml" or path.suffix == ".yml":
|
|
metadata = yaml.full_load(open(self.args.source, "r"))
|
|
else:
|
|
raise ValueError(f"Unsupported file type: {path.suffix}")
|
|
|
|
return metadata
|