test_data_element.py

# Copyright (c) OpenMMLab. All rights reserved.
import random
from unittest import TestCase

import numpy as np
import pytest
import torch

from mmengine.data import BaseDataElement


class TestBaseDataElement(TestCase):

    def setup_data(self):
        metainfo = dict(
            img_id=random.randint(0, 100),
            img_shape=(random.randint(400, 600), random.randint(400, 600)))
        gt_instances = BaseDataElement(
            bboxes=torch.rand((5, 4)), labels=torch.rand((5, )))
        pred_instances = BaseDataElement(
            bboxes=torch.rand((5, 4)), scores=torch.rand((5, )))
        data = dict(gt_instances=gt_instances, pred_instances=pred_instances)
        return metainfo, data

    def is_equal(self, x, y):
        assert type(x) == type(y)
        if isinstance(
                x, (int, float, str, list, tuple, dict, set, BaseDataElement)):
            return x == y
        elif isinstance(x, (torch.Tensor, np.ndarray)):
            return (x == y).all()

    def check_key_value(self, instances, metainfo=None, data=None):
        # check the existence of keys in metainfo, data, and instances
        if metainfo:
            for k, v in metainfo.items():
                assert k in instances
                assert k in instances.all_keys()
                assert k in instances.metainfo_keys()
                assert k not in instances.keys()
                assert self.is_equal(instances.get(k), v)
                assert self.is_equal(getattr(instances, k), v)
        if data:
            for k, v in data.items():
                assert k in instances
                assert k in instances.keys()
                assert k not in instances.metainfo_keys()
                assert k in instances.all_keys()
                assert self.is_equal(instances.get(k), v)
                assert self.is_equal(getattr(instances, k), v)

    def check_data_device(self, instances, device):
        # assert instances.device == device
        for v in instances.values():
            if isinstance(v, torch.Tensor):
                assert v.device == torch.device(device)
            elif isinstance(v, BaseDataElement):
                self.check_data_device(v, device)

    def check_data_dtype(self, instances, dtype):
        for v in instances.values():
            if isinstance(v, (torch.Tensor, np.ndarray)):
                assert isinstance(v, dtype)
            if isinstance(v, BaseDataElement):
                self.check_data_dtype(v, dtype)

    def check_requires_grad(self, instances):
        for v in instances.values():
            if isinstance(v, torch.Tensor):
                assert v.requires_grad is False
            if isinstance(v, BaseDataElement):
                self.check_requires_grad(v)

    def test_init(self):
        # initialization with no data and metainfo
        metainfo, data = self.setup_data()
        instances = BaseDataElement()
        for k in metainfo:
            assert k not in instances
            assert instances.get(k, None) is None
        for k in data:
            assert k not in instances
            assert instances.get(k, 'abc') == 'abc'

        # initialization with kwargs
        metainfo, data = self.setup_data()
        instances = BaseDataElement(metainfo=metainfo, **data)
        self.check_key_value(instances, metainfo, data)

        # initialization with args
        metainfo, data = self.setup_data()
        instances = BaseDataElement(metainfo=metainfo)
        self.check_key_value(instances, metainfo)
        instances = BaseDataElement(**data)
        self.check_key_value(instances, data=data)

    def test_new(self):
        metainfo, data = self.setup_data()
        instances = BaseDataElement(metainfo=metainfo, **data)

        # test new() with no arguments
        new_instances = instances.new()
        assert type(new_instances) == type(instances)
        # After deepcopy, the address of new data'element will be same as
        # origin, but when change new data' element will not effect the origin
        # element and will have new address
        _, data = self.setup_data()
        new_instances.set_data(data)
        assert not self.is_equal(new_instances.gt_instances,
                                 instances.gt_instances)
        self.check_key_value(new_instances, metainfo, data)

        # test new() with arguments
        metainfo, data = self.setup_data()
        new_instances = instances.new(metainfo=metainfo, **data)
        assert type(new_instances) == type(instances)
        assert id(new_instances.gt_instances) != id(instances.gt_instances)
        _, new_data = self.setup_data()
        new_instances.set_data(new_data)
        assert id(new_instances.gt_instances) != id(data['gt_instances'])
        self.check_key_value(new_instances, metainfo, new_data)

        metainfo, data = self.setup_data()
        new_instances = instances.new(metainfo=metainfo)

    def test_clone(self):
        metainfo, data = self.setup_data()
        instances = BaseDataElement(metainfo=metainfo, **data)
        new_instances = instances.clone()
        assert type(new_instances) == type(instances)

    def test_set_metainfo(self):
        metainfo, _ = self.setup_data()
        instances = BaseDataElement()
        instances.set_metainfo(metainfo)
        self.check_key_value(instances, metainfo=metainfo)

        # test setting existing keys and new keys
        new_metainfo, _ = self.setup_data()
        new_metainfo.update(other=123)
        instances.set_metainfo(new_metainfo)
        self.check_key_value(instances, metainfo=new_metainfo)

        # test have the same key in data
        _, data = self.setup_data()
        instances = BaseDataElement(**data)
        _, data = self.setup_data()
        with self.assertRaises(AttributeError):
            instances.set_metainfo(data)

        with self.assertRaises(AssertionError):
            instances.set_metainfo(123)

    def test_set_data(self):
        metainfo, data = self.setup_data()
        instances = BaseDataElement()

        instances.gt_instances = data['gt_instances']
        instances.pred_instances = data['pred_instances']
        self.check_key_value(instances, data=data)

        # a.xx only set data rather than metainfo
        instances.img_shape = metainfo['img_shape']
        instances.img_id = metainfo['img_id']
        self.check_key_value(instances, data=metainfo)

        metainfo, data = self.setup_data()
        instances = BaseDataElement(metainfo=metainfo, **data)
        with self.assertRaises(AttributeError):
            instances.img_shape = metainfo['img_shape']

        # test set '_metainfo_fields' or '_data_fields'
        with self.assertRaises(AttributeError):
            instances._metainfo_fields = 1
        with self.assertRaises(AttributeError):
            instances._data_fields = 1

        with self.assertRaises(AssertionError):
            instances.set_data(123)

    def test_update(self):
        metainfo, data = self.setup_data()
        instances = BaseDataElement(metainfo=metainfo, **data)
        proposals = BaseDataElement(
            bboxes=torch.rand((5, 4)), scores=torch.rand((5, )))
        new_instances = BaseDataElement(proposals=proposals)
        instances.update(new_instances)
        self.check_key_value(instances, metainfo,
                             data.update(dict(proposals=proposals)))

    def test_delete_modify(self):
        random.seed(10)
        metainfo, data = self.setup_data()
        instances = BaseDataElement(metainfo=metainfo, **data)

        new_metainfo, new_data = self.setup_data()
        # avoid generating same metainfo, data
        while True:
            if new_metainfo['img_id'] == metainfo['img_id'] or new_metainfo[
                    'img_shape'] == metainfo['img_shape']:
                new_metainfo, new_data = self.setup_data()
            else:
                break
        instances.gt_instances = new_data['gt_instances']
        instances.pred_instances = new_data['pred_instances']

        # a.xx only set data rather than metainfo
        instances.set_metainfo(new_metainfo)
        self.check_key_value(instances, new_metainfo, new_data)

        assert not self.is_equal(instances.gt_instances, data['gt_instances'])
        assert not self.is_equal(instances.pred_instances,
                                 data['pred_instances'])
        assert not self.is_equal(instances.img_id, metainfo['img_id'])
        assert not self.is_equal(instances.img_shape, metainfo['img_shape'])

        del instances.gt_instances
        del instances.img_id
        assert not self.is_equal(
            instances.pop('pred_instances', None), data['pred_instances'])
        with self.assertRaises(AttributeError):
            del instances.pred_instances

        assert 'gt_instances' not in instances
        assert 'pred_instances' not in instances
        assert 'img_id' not in instances
        assert instances.pop('gt_instances', None) is None
        # test pop not exist key without default
        with self.assertRaises(KeyError):
            instances.pop('gt_instances')
        assert instances.pop('pred_instances', 'abcdef') == 'abcdef'

        assert instances.pop('img_id', None) is None
        # test pop not exist key without default
        with self.assertRaises(KeyError):
            instances.pop('img_id')
        assert instances.pop('img_shape') == new_metainfo['img_shape']

        # test del '_metainfo_fields' or '_data_fields'
        with self.assertRaises(AttributeError):
            del instances._metainfo_fields
        with self.assertRaises(AttributeError):
            del instances._data_fields

    @pytest.mark.skipif(
        not torch.cuda.is_available(), reason='GPU is required!')
    def test_cuda(self):
        metainfo, data = self.setup_data()
        instances = BaseDataElement(metainfo=metainfo, **data)

        cuda_instances = instances.cuda()
        self.check_data_device(cuda_instances, 'cuda:0')

        # here we further test to convert from cuda to cpu
        cpu_instances = cuda_instances.cpu()
        self.check_data_device(cpu_instances, 'cpu')
        del cuda_instances

        cuda_instances = instances.to('cuda:0')
        self.check_data_device(cuda_instances, 'cuda:0')

    def test_cpu(self):
        metainfo, data = self.setup_data()
        instances = BaseDataElement(metainfo=metainfo, **data)
        self.check_data_device(instances, 'cpu')

        cpu_instances = instances.cpu()
        # assert cpu_instances.device == 'cpu'
        assert cpu_instances.gt_instances.bboxes.device == torch.device('cpu')
        assert cpu_instances.gt_instances.labels.device == torch.device('cpu')

    def test_numpy_tensor(self):
        metainfo, data = self.setup_data()
        instances = BaseDataElement(metainfo=metainfo, **data)

        np_instances = instances.numpy()
        self.check_data_dtype(np_instances, np.ndarray)

        tensor_instances = np_instances.to_tensor()
        self.check_data_dtype(tensor_instances, torch.Tensor)

    def test_detach(self):
        metainfo, data = self.setup_data()
        instances = BaseDataElement(metainfo=metainfo, **data)
        instances.detach()
        self.check_requires_grad(instances)

    def test_repr(self):
        metainfo = dict(img_shape=(800, 1196, 3))
        gt_instances = BaseDataElement(
            metainfo=metainfo, det_labels=torch.LongTensor([0, 1, 2, 3]))
        sample = BaseDataElement(metainfo=metainfo, gt_instances=gt_instances)
        address = hex(id(sample))
        address_gt_instances = hex(id(sample.gt_instances))
        assert repr(sample) == (
            '<BaseDataElement(\n\n'
            '    META INFORMATION\n'
            '    img_shape: (800, 1196, 3)\n\n'
            '    DATA FIELDS\n'
            '    gt_instances: <BaseDataElement(\n        \n'
            '            META INFORMATION\n'
            '            img_shape: (800, 1196, 3)\n        \n'
            '            DATA FIELDS\n'
            '            det_labels: tensor([0, 1, 2, 3])\n'
            f'        ) at {address_gt_instances}>\n'
            f') at {address}>')

    def test_set_fields(self):
        metainfo, data = self.setup_data()
        instances = BaseDataElement(metainfo=metainfo)
        for key, value in data.items():
            instances.set_field(name=key, value=value, dtype=BaseDataElement)
        self.check_key_value(instances, data=data)

        # test type check
        _, data = self.setup_data()
        instances = BaseDataElement()
        for key, value in data.items():
            with self.assertRaises(AssertionError):
                instances.set_field(name=key, value=value, dtype=torch.Tensor)

    def test_inheritance(self):

        class DetDataSample(BaseDataElement):

            @property
            def proposals(self):
                return self._proposals

            @proposals.setter
            def proposals(self, value):
                self.set_field(
                    value=value, name='_proposals', dtype=BaseDataElement)

            @proposals.deleter
            def proposals(self):
                del self._proposals

            @property
            def gt_instances(self):
                return self._gt_instances

            @gt_instances.setter
            def gt_instances(self, value):
                self.set_field(
                    value=value, name='_gt_instances', dtype=BaseDataElement)

            @gt_instances.deleter
            def gt_instances(self):
                del self._gt_instances

            @property
            def pred_instances(self):
                return self._pred_instances

            @pred_instances.setter
            def pred_instances(self, value):
                self.set_field(
                    value=value, name='_pred_instances', dtype=BaseDataElement)

            @pred_instances.deleter
            def pred_instances(self):
                del self._pred_instances

        det_sample = DetDataSample()

        # test set
        proposals = BaseDataElement(bboxes=torch.rand((5, 4)))
        det_sample.proposals = proposals
        assert 'proposals' in det_sample

        # test get
        assert det_sample.proposals == proposals

        # test delete
        del det_sample.proposals
        assert 'proposals' not in det_sample

        # test the data whether meet the requirements
        with self.assertRaises(AssertionError):
            det_sample.proposals = torch.rand((5, 4))

    def test_values(self):
        # test_metainfo_values
        metainfo, data = self.setup_data()
        instances = BaseDataElement(metainfo=metainfo, **data)
        assert len(instances.metainfo_values()) == len(metainfo.values())
        # test_all_values
        assert len(instances.all_values()) == len(metainfo.values()) + len(
            data.values())

        # test_values
        assert len(instances.values()) == len(data.values())

    def test_keys(self):
        # test_metainfo_keys
        metainfo, data = self.setup_data()
        instances = BaseDataElement(metainfo=metainfo, **data)
        assert len(instances.metainfo_keys()) == len(metainfo.keys())

        # test_all_keys
        assert len(
            instances.all_keys()) == len(data.keys()) + len(metainfo.keys())

        # test_keys
        assert len(instances.keys()) == len(data.keys())

    def test_items(self):
        # test_metainfo_items
        metainfo, data = self.setup_data()
        instances = BaseDataElement(metainfo=metainfo, **data)
        assert len(dict(instances.metainfo_items())) == len(
            dict(metainfo.items()))
        # test_all_items
        assert len(dict(instances.all_items())) == len(dict(
            metainfo.items())) + len(dict(data.items()))

        # test_items
        assert len(dict(instances.items())) == len(dict(data.items()))

    def test_to_dict(self):
        metainfo, data = self.setup_data()
        instances = BaseDataElement(metainfo=metainfo, **data)
        dict_instances = instances.to_dict()
        # test convert BaseDataElement to dict
        for k in instances.all_keys():
            # all keys in instances should be in dict_instances
            assert k in dict_instances
        assert isinstance(dict_instances, dict)
        # sub data element should also be converted to dict
        assert isinstance(dict_instances['gt_instances'], dict)
        assert isinstance(dict_instances['pred_instances'], dict)

    def test_metainfo(self):
        # test metainfo property
        metainfo, data = self.setup_data()
        instances = BaseDataElement(metainfo=metainfo, **data)
        self.assertDictEqual(instances.metainfo, metainfo)