File size: 1,339 Bytes
df2accb
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
# 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