mirror of https://github.com/open-mmlab/mmocr.git
106 lines
4.0 KiB
Python
106 lines
4.0 KiB
Python
# Copyright (c) OpenMMLab. All rights reserved.
|
|
from unittest import TestCase
|
|
|
|
import numpy as np
|
|
import torch
|
|
from mmengine.data import InstanceData
|
|
|
|
from mmocr.data import KIEDataSample
|
|
|
|
|
|
class TestTextDetDataSample(TestCase):
|
|
|
|
def _equal(self, a, b):
|
|
if isinstance(a, (torch.Tensor, np.ndarray)):
|
|
return (a == b).all()
|
|
else:
|
|
return a == b
|
|
|
|
def test_init(self):
|
|
meta_info = dict(
|
|
img_size=[256, 256],
|
|
scale_factor=np.array([1.5, 1.5]),
|
|
img_shape=torch.rand(4))
|
|
|
|
kie_data_sample = KIEDataSample(metainfo=meta_info)
|
|
assert 'img_size' in kie_data_sample
|
|
|
|
self.assertListEqual(kie_data_sample.img_size, [256, 256])
|
|
self.assertListEqual(kie_data_sample.get('img_size'), [256, 256])
|
|
|
|
def test_setter(self):
|
|
kie_data_sample = KIEDataSample()
|
|
# test gt_instances
|
|
gt_instances_data = dict(
|
|
bboxes=torch.rand(4, 4),
|
|
labels=torch.rand(4),
|
|
texts=['t1', 't2', 't3', 't4'],
|
|
relations=torch.rand(4, 4),
|
|
edge_labels=torch.randint(0, 4, (4, )))
|
|
gt_instances = InstanceData(**gt_instances_data)
|
|
kie_data_sample.gt_instances = gt_instances
|
|
self.assertIn('gt_instances', kie_data_sample)
|
|
self.assertTrue(
|
|
self._equal(kie_data_sample.gt_instances.bboxes,
|
|
gt_instances_data['bboxes']))
|
|
self.assertTrue(
|
|
self._equal(kie_data_sample.gt_instances.labels,
|
|
gt_instances_data['labels']))
|
|
self.assertTrue(
|
|
self._equal(kie_data_sample.gt_instances.texts,
|
|
gt_instances_data['texts']))
|
|
self.assertTrue(
|
|
self._equal(kie_data_sample.gt_instances.relations,
|
|
gt_instances_data['relations']))
|
|
self.assertTrue(
|
|
self._equal(kie_data_sample.gt_instances.edge_labels,
|
|
gt_instances_data['edge_labels']))
|
|
|
|
# test pred_instances
|
|
pred_instances_data = dict(
|
|
bboxes=torch.rand(4, 4),
|
|
labels=torch.rand(4),
|
|
texts=['t1', 't2', 't3', 't4'],
|
|
relations=torch.rand(4, 4),
|
|
edge_labels=torch.randint(0, 4, (4, )))
|
|
pred_instances = InstanceData(**pred_instances_data)
|
|
kie_data_sample.pred_instances = pred_instances
|
|
assert 'pred_instances' in kie_data_sample
|
|
assert self._equal(kie_data_sample.pred_instances.bboxes,
|
|
pred_instances_data['bboxes'])
|
|
assert self._equal(kie_data_sample.pred_instances.labels,
|
|
pred_instances_data['labels'])
|
|
self.assertTrue(
|
|
self._equal(kie_data_sample.pred_instances.texts,
|
|
pred_instances_data['texts']))
|
|
self.assertTrue(
|
|
self._equal(kie_data_sample.pred_instances.relations,
|
|
pred_instances_data['relations']))
|
|
self.assertTrue(
|
|
self._equal(kie_data_sample.pred_instances.edge_labels,
|
|
pred_instances_data['edge_labels']))
|
|
|
|
# test type error
|
|
with self.assertRaises(AssertionError):
|
|
kie_data_sample.gt_instances = torch.rand(2, 4)
|
|
with self.assertRaises(AssertionError):
|
|
kie_data_sample.pred_instances = torch.rand(2, 4)
|
|
|
|
def test_deleter(self):
|
|
gt_instances_data = dict(
|
|
bboxes=torch.rand(4, 4),
|
|
labels=torch.rand(4),
|
|
)
|
|
|
|
kie_data_sample = KIEDataSample()
|
|
gt_instances = InstanceData(data=gt_instances_data)
|
|
kie_data_sample.gt_instances = gt_instances
|
|
assert 'gt_instances' in kie_data_sample
|
|
del kie_data_sample.gt_instances
|
|
assert 'gt_instances' not in kie_data_sample
|
|
|
|
kie_data_sample.pred_instances = gt_instances
|
|
assert 'pred_instances' in kie_data_sample
|
|
del kie_data_sample.pred_instances
|
|
assert 'pred_instances' not in kie_data_sample
|