layout_data.data.layout 源代码

# -*- encoding: utf-8 -*-
"""
Desc      :   Layout dataset
"""
# File    :   layout.py
# Time    :   2020/04/06 18:02:23
# Author  :   Zweien
# Contact :   278954153@qq.com

import os
from .loadresponse import LoadResponse, mat_loader, LoadResponseH5


class LayoutDataset(LoadResponse):
    """Layout dataset (mutiple files) generated by 'layout-generator'.
    """

    def __init__(
        self, root, train=True, transform=None, target_transform=None,
    ):

        subdir = "train" if train else "test"
        root = os.path.join(root, subdir)
        super().__init__(
            root,
            mat_loader,
            load_name="F",
            resp_name="u",
            extensions="mat",
            transform=transform,
            target_transform=target_transform,
        )


class LayoutDatasetH5(LoadResponseH5):
    """Layout dataset (hdf5 format, single file) generated by 'layout-generator'.
    """

    def __init__(
        self,
        root,
        train=True,
        transform=None,
        target_transform=None,
        train_fn="train.h5",
        test_fn="test.h5",
        train_dir="./",
        test_dir="./"
    ):

        fn = train_fn if train else test_fn
        if train:
            fn = os.path.join(train_dir, train_fn)
        else:
            fn = os.path.join(test_dir, test_fn)
        root = os.path.join(root, fn)
        super().__init__(
            root,
            load_name="F",
            resp_name="u",
            transform=transform,
            target_transform=target_transform,
        )