File size: 2,169 Bytes
62dbcfb |
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 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 |
from os import path as osp
import msgspec
from traiNNer.data.single_image_dataset import SingleImageDataset
from traiNNer.utils.redux_options import DatasetOptions
def test_singleimagedataset() -> None:
"""Test dataset: SingleImageDataset"""
opt_str = r"""
name: Test
type: SingleImageDataset
dataroot_lq: [
datasets/val/dataset1/lr,
datasets/val/dataset1/lr2
]
io_backend:
type: disk
mean: [0.5, 0.5, 0.5]
std: [0.5, 0.5, 0.5]
"""
opt = msgspec.yaml.decode(opt_str, type=DatasetOptions, strict=True)
dataset = SingleImageDataset(opt)
# ------------------ test scan folder mode -------------------- #
opt.io_backend = {"type": "disk"}
dataset = SingleImageDataset(opt)
assert dataset.io_backend_opt["type"] == "disk" # io backend
assert len(dataset) == 3 # whether to correctly scan folders
# test __getitem__
result = dataset.__getitem__(0)
# check returned keys
expected_keys = ["lq", "lq_path"]
assert set(expected_keys).issubset(set(result.keys()))
assert "lq" in result and "lq_path" in result
# check shape and contents
assert result["lq"].shape == (3, 128, 128)
assert osp.normpath(result["lq_path"]) == osp.normpath(
"datasets/val/dataset1/lr/0007.png"
)
# ------------------ test lmdb backend and with y channel-------------------- #
# TODO
# opt["dataroot_lq"] = "tests/data/lq.lmdb"
# opt["io_backend"] = {"type": "lmdb"}
# opt["color"] = "y"
# opt["mean"] = [0.5]
# opt["std"] = [0.5]
# dataset = SingleImageDataset(opt)
# assert dataset.io_backend_opt["type"] == "lmdb" # io backend
# assert len(dataset) == 2 # whether to read correct meta info
# assert dataset.std == [0.5]
# # test __getitem__
# result = dataset.__getitem__(1)
# # check returned keys
# expected_keys = ["lq", "lq_path"]
# assert set(expected_keys).issubset(set(result.keys()))
# # check shape and contents
# assert "lq" in result and "lq_path" in result
# assert result["lq"].shape == (1, 90, 60)
# assert result["lq_path"] == "comic"
|