Source code for jactorch.data.layout

#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : layout.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 03/09/2020
#
# This file is part of Jacinle.
# Distributed under terms of the MIT license.

from typing import Union, Optional
from jacinle.utils.enum import JacEnum

__all__ = ['DataLayoutType', 'DataLayoutSpec', 'DataLayout', 'skip', 'concat', 'pad', 'pad2d', 'padimage']




[docs]
class DataLayoutType(JacEnum):
    SKIP = 'skip'
    STACK = 'stack'
    CONCAT = 'concat'
    PAD = 'pad'
    PAD2D = 'pad2d'
    PADIMAGE = 'padimage'






[docs]
class DataLayoutSpec(object):


[docs]
    def __init__(self, type, **kwargs):
        self.__dict__.update(kwargs)
        self.type = DataLayoutType.from_string(type)












[docs]
def stack():
    return DataLayoutSpec('stack')






[docs]
def concat():
    return DataLayoutSpec('concat')






[docs]
def pad(fill=0):
    return DataLayoutSpec('pad', fill=fill)






[docs]
def pad2d(fill=0):
    return DataLayoutSpec('pad2d', fill=fill)






[docs]
def padimage(fill=0):
    return DataLayoutSpec('pad2d', fill=fill)




_type_to_constructor = {
    'skip': skip,
    'stack': stack,
    'concat': concat,
    'pad': pad,
    'pad2d': pad2d,
    'padimage': padimage
}




[docs]
class DataLayout(object):


[docs]
    def __init__(self, layout: Optional[dict] = None):
        self.layout = dict()

        if layout is not None:
            for k, v in layout.items():
                self.decl(k, v)





[docs]
    def decl(self, key: str, spec: Union[str, DataLayoutSpec]):
        if isinstance(spec, str):
            self.layout[key] = _type_to_constructor[spec]()
        else:
            self.layout[key] = spec



    def __contains__(self, key: str):
        return key in self.layout

    def __getitem__(self, key: str):
        return self.layout.get(key, None)