container.py

class Container:
    """
    Help class for manage boxes, labels, etc...
    Not inherit dict due to `default_collate` will change dict's subclass to dict.
    """

    def __init__(self, *args, **kwargs):
        self._data_dict = dict(*args, **kwargs)

    def __setattr__(self, key, value):
        object.__setattr__(self, key, value)

    def __getitem__(self, key):
        return self._data_dict[key]

    def __iter__(self):
        return self._data_dict.__iter__()

    def __setitem__(self, key, value):
        self._data_dict[key] = value

    def _call(self, name, *args, **kwargs):
        keys = list(self._data_dict.keys())
        for key in keys:
            value = self._data_dict[key]
            if hasattr(value, name):
                self._data_dict[key] = getattr(value, name)(*args, **kwargs)
        return self

    def to(self, *args, **kwargs):
        return self._call('to', *args, **kwargs)

    def numpy(self):
        return self._call('numpy')

    def resize(self, size):
        """resize boxes
        Args:
            size: (width, height)
        Returns:
            self
        """
        img_width = getattr(self, 'img_width', -1)
        img_height = getattr(self, 'img_height', -1)
        assert img_width > 0 and img_height > 0
        assert 'boxes' in self._data_dict
        boxes = self._data_dict['boxes']
        new_width, new_height = size
        boxes[:, 0::2] *= (new_width / img_width)
        boxes[:, 1::2] *= (new_height / img_height)
        return self

    def __repr__(self):
        return self._data_dict.__repr__()