def test_can_batch_launch_custom_model(self):
class TestExtractor(Extractor):
def __iter__(self):
for i in range(5):
yield DatasetItem(id=i,
subset='train',
image=np.array([i]))
class TestLauncher(Launcher):
def launch(self, inputs):
for i, inp in enumerate(inputs):
yield [Label(attributes={'idx': i, 'data': inp.item()})]
model_name = 'model'
launcher_name = 'custom_launcher'
project = Project()
project.env.launchers.register(launcher_name, TestLauncher)
project.add_model(model_name, {'launcher': launcher_name})
model = project.make_executable_model(model_name)
extractor = TestExtractor()
batch_size = 3
executor = InferenceWrapper(extractor, model, batch_size=batch_size)
for item in executor:
self.assertEqual(1, len(item.annotations))
self.assertEqual(
int(item.id) % batch_size,
item.annotations[0].attributes['idx'])
self.assertEqual(int(item.id),
item.annotations[0].attributes['data'])
def test_can_batch_launch_custom_model(self):
dataset = Dataset.from_iterable([
DatasetItem(id=i, subset='train', image=np.array([i]))
for i in range(5)
],
categories=['label'])
class TestLauncher(Launcher):
def launch(self, inputs):
for i, inp in enumerate(inputs):
yield [Label(0, attributes={'idx': i, 'data': inp.item()})]
model_name = 'model'
launcher_name = 'custom_launcher'
project = Project()
project.env.launchers.register(launcher_name, TestLauncher)
project.add_model(model_name, {'launcher': launcher_name})
model = project.make_executable_model(model_name)
batch_size = 3
executor = ModelTransform(dataset, model, batch_size=batch_size)
for item in executor:
self.assertEqual(1, len(item.annotations))
self.assertEqual(
int(item.id) % batch_size,
item.annotations[0].attributes['idx'])
self.assertEqual(int(item.id),
item.annotations[0].attributes['data'])
def test_can_do_transform_with_custom_model(self):
class TestExtractorSrc(Extractor):
def __iter__(self):
for i in range(2):
yield DatasetItem(id=i,
image=np.ones([2, 2, 3]) * i,
annotations=[Label(i)])
class TestLauncher(Launcher):
def launch(self, inputs):
for inp in inputs:
yield [Label(inp[0, 0, 0])]
class TestConverter(Converter):
def __call__(self, extractor, save_dir):
for item in extractor:
with open(osp.join(save_dir, '%s.txt' % item.id),
'w') as f:
f.write(str(item.annotations[0].label) + '\n')
class TestExtractorDst(Extractor):
def __init__(self, url):
super().__init__()
self.items = [
osp.join(url, p) for p in sorted(os.listdir(url))
]
def __iter__(self):
for path in self.items:
with open(path, 'r') as f:
index = osp.splitext(osp.basename(path))[0]
label = int(f.readline().strip())
yield DatasetItem(id=index, annotations=[Label(label)])
model_name = 'model'
launcher_name = 'custom_launcher'
extractor_name = 'custom_extractor'
project = Project()
project.env.launchers.register(launcher_name, TestLauncher)
project.env.extractors.register(extractor_name, TestExtractorSrc)
project.env.converters.register(extractor_name, TestConverter)
project.add_model(model_name, {'launcher': launcher_name})
project.add_source('source', {'format': extractor_name})
with TestDir() as test_dir:
project.make_dataset().apply_model(model=model_name,
save_dir=test_dir)
result = Project.load(test_dir)
result.env.extractors.register(extractor_name, TestExtractorDst)
it = iter(result.make_dataset())
item1 = next(it)
item2 = next(it)
self.assertEqual(0, item1.annotations[0].label)
self.assertEqual(1, item2.annotations[0].label)
def test_can_dump_added_model(self):
model_name = 'model'
project = Project()
saved = Model({'launcher': 'name'})
project.add_model(model_name, saved)
with TestDir() as test_dir:
project.save(test_dir)
loaded = Project.load(test_dir)
loaded = loaded.get_model(model_name)
self.assertEqual(saved, loaded)
def test_can_batch_launch_custom_model(self):
class TestExtractor(Extractor):
def __init__(self, url, n=0):
super().__init__(length=n)
self.n = n
def __iter__(self):
for i in range(self.n):
yield DatasetItem(id=i, subset='train', image=i)
def subsets(self):
return ['train']
class TestLauncher(Launcher):
def __init__(self, **kwargs):
pass
def launch(self, inputs):
for i, inp in enumerate(inputs):
yield [LabelObject(attributes={'idx': i, 'data': inp})]
model_name = 'model'
launcher_name = 'custom_launcher'
project = Project()
project.env.launchers.register(launcher_name, TestLauncher)
project.add_model(model_name, {'launcher': launcher_name})
model = project.make_executable_model(model_name)
extractor = TestExtractor('', n=5)
batch_size = 3
executor = InferenceWrapper(extractor, model, batch_size=batch_size)
for item in executor:
self.assertEqual(1, len(item.annotations))
self.assertEqual(
int(item.id) % batch_size,
item.annotations[0].attributes['idx'])
self.assertEqual(int(item.id),
item.annotations[0].attributes['data'])
def test_can_do_transform_with_custom_model(self):
class TestExtractorSrc(Extractor):
def __init__(self, url, n=2):
super().__init__(length=n)
self.n = n
def __iter__(self):
for i in range(self.n):
yield DatasetItem(id=i,
subset='train',
image=i,
annotations=[LabelObject(i)])
def subsets(self):
return ['train']
class TestLauncher(Launcher):
def __init__(self, **kwargs):
pass
def launch(self, inputs):
for inp in inputs:
yield [LabelObject(inp)]
class TestConverter(Converter):
def __call__(self, extractor, save_dir):
for item in extractor:
with open(osp.join(save_dir, '%s.txt' % item.id),
'w+') as f:
f.write(str(item.subset) + '\n')
f.write(str(item.annotations[0].label) + '\n')
class TestExtractorDst(Extractor):
def __init__(self, url):
super().__init__()
self.items = [
osp.join(url, p) for p in sorted(os.listdir(url))
]
def __iter__(self):
for path in self.items:
with open(path, 'r') as f:
index = osp.splitext(osp.basename(path))[0]
subset = f.readline()[:-1]
label = int(f.readline()[:-1])
assert (subset == 'train')
yield DatasetItem(id=index,
subset=subset,
annotations=[LabelObject(label)])
def __len__(self):
return len(self.items)
def subsets(self):
return ['train']
model_name = 'model'
launcher_name = 'custom_launcher'
extractor_name = 'custom_extractor'
project = Project()
project.env.launchers.register(launcher_name, TestLauncher)
project.env.extractors.register(extractor_name, TestExtractorSrc)
project.env.converters.register(extractor_name, TestConverter)
project.add_model(model_name, {'launcher': launcher_name})
project.add_source('source', {'format': extractor_name})
with TestDir() as test_dir:
project.make_dataset().transform(model_name, test_dir.path)
result = Project.load(test_dir.path)
result.env.extractors.register(extractor_name, TestExtractorDst)
it = iter(result.make_dataset())
item1 = next(it)
item2 = next(it)
self.assertEqual(0, item1.annotations[0].label)
self.assertEqual(1, item2.annotations[0].label)
