Spaces:
Sleeping
Sleeping
File size: 4,059 Bytes
24c4def |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 |
# Copyright (c) OpenMMLab. All rights reserved.
from unittest import TestCase
import numpy as np
import torch
from mmengine.structures import InstanceData
from mmocr.structures 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
|