def test_can_save_masks(self):
class TestExtractor(Extractor):
def __iter__(self):
return iter([
DatasetItem(id=1,
subset='train',
image=np.ones((4, 5, 3)),
annotations=[
Mask(image=np.array([
[1, 0, 0, 1],
[0, 1, 1, 0],
[0, 1, 1, 0],
[1, 0, 0, 1],
]),
label=1),
],
attributes={'source_id': ''}),
])
def categories(self):
label_cat = LabelCategories()
for label in range(10):
label_cat.add('label_' + str(label))
return {
AnnotationType.label: label_cat,
}
with TestDir() as test_dir:
self._test_save_and_load(TestExtractor(),
TfDetectionApiConverter(save_masks=True),
test_dir)
def test_can_save_bboxes(self):
class TestExtractor(Extractor):
def __iter__(self):
return iter([
DatasetItem(id=1,
subset='train',
image=np.ones((16, 16, 3)),
annotations=[
Bbox(0, 4, 4, 8, label=2),
Bbox(0, 4, 4, 4, label=3),
Bbox(2, 4, 4, 4),
],
attributes={'source_id': ''}),
])
def categories(self):
label_cat = LabelCategories()
for label in range(10):
label_cat.add('label_' + str(label))
return {
AnnotationType.label: label_cat,
}
with TestDir() as test_dir:
self._test_save_and_load(TestExtractor(),
TfDetectionApiConverter(save_images=True),
test_dir)
def test_can_save_dataset_with_no_subsets(self):
class TestExtractor(Extractor):
def __iter__(self):
return iter([
DatasetItem(id=1,
image=np.ones((16, 16, 3)),
annotations=[
Bbox(2, 1, 4, 4, label=2),
Bbox(4, 2, 8, 4, label=3),
]),
DatasetItem(id=2,
image=np.ones((8, 8, 3)) * 2,
annotations=[
Bbox(4, 4, 4, 4, label=3),
]),
DatasetItem(
id=3,
image=np.ones((8, 4, 3)) * 3,
),
])
def categories(self):
label_cat = LabelCategories()
for label in range(10):
label_cat.add('label_' + str(label))
return {
AnnotationType.label: label_cat,
}
with TestDir() as test_dir:
self._test_save_and_load(TestExtractor(),
TfDetectionApiConverter(save_images=True),
test_dir)
def test_can_save_dataset_with_image_info(self):
class TestExtractor(Extractor):
def __iter__(self):
return iter([
DatasetItem(id=1, image=Image(path='1/q.e',
size=(10, 15))),
])
def categories(self):
return {AnnotationType.label: LabelCategories()}
with TestDir() as test_dir:
self._test_save_and_load(TestExtractor(),
TfDetectionApiConverter(), test_dir)
def generate_dummy_tfrecord(path):
TfDetectionApiConverter()(TestExtractor(), save_dir=path)