def create_object_detection_export_test_case(alt_ssd_export=False, **kwargs):
expected_outputs_dir = os.path.join(os.path.dirname(__file__), '..',
'expected_outputs')
ExportTestCase = create_export_test_case(
'object_detection',
**kwargs,
metric_keys=['bbox'],
expected_outputs_dir=expected_outputs_dir)
if alt_ssd_export:
class ExportWithAltSsdTestCase(ExportTestCase):
def test_alt_ssd_export_on_gpu(self):
skip_if_cuda_not_available()
export_dir = os.path.join(self.output_folder, 'gpu_export')
self.do_export(export_dir, on_gpu=True)
export_dir = os.path.join(export_dir, 'alt_ssd_export')
self.do_evaluation(export_dir)
def test_alt_ssd_export_on_cpu(self):
export_dir = os.path.join(self.output_folder, 'cpu_export')
self.do_export(export_dir, on_gpu=True)
export_dir = os.path.join(export_dir, 'alt_ssd_export')
self.do_evaluation(export_dir)
return ExportWithAltSsdTestCase
return ExportTestCase
def create_image_classification_export_test_case(**kwargs):
expected_outputs_dir = os.path.join(os.path.dirname(__file__), '..',
'expected_outputs')
ExportTestCase = create_export_test_case(
'image_classification',
**kwargs,
metric_keys=['accuracy'],
expected_outputs_dir=expected_outputs_dir)
class ClassificationExportTestCase(ExportTestCase):
@unittest.skipUnless(torch.cuda.is_available(), 'No GPU found')
def test_export_on_gpu(self):
export_dir = os.path.join(self.output_folder, 'gpu_export')
self.do_export(export_dir, on_gpu=True)
def test_export_on_cpu(self):
export_dir = os.path.join(self.output_folder, 'cpu_export')
self.do_export(export_dir, on_gpu=False)
def do_export(self, export_dir, on_gpu):
if not os.path.exists(export_dir):
initial_command = 'export CUDA_VISIBLE_DEVICES=;' if not on_gpu else ''
run_through_shell(f'{initial_command}'
f'cd {os.path.dirname(self.template_file)};'
f'pip install -r requirements.txt;'
f'python export.py'
f' --load-weights snapshot.pth'
f' --save-model-to {export_dir}')
return ClassificationExportTestCase
def create_action_recognition_export_test_case(enable_metrics_eval=True,
**kwargs):
expected_outputs_dir = os.path.join(os.path.dirname(__file__), '..',
'expected_outputs')
ExportTestCase = create_export_test_case(
'action_recognition',
**kwargs,
metric_keys=['accuracy'],
expected_outputs_dir=expected_outputs_dir)
if enable_metrics_eval:
return ExportTestCase
class CustomActionRecognitionExportTestCase(ExportTestCase):
def do_evaluation(self, export_dir):
metrics_path = os.path.join(export_dir, "metrics.yaml")
run_through_shell(
f'cd {os.path.dirname(self.template_file)};'
f'python3 eval.py'
f' --test-ann-files {self.ann_file}'
f' --test-data-roots {self.img_root}'
f' --load-weights {os.path.join(export_dir, "model.bin")}'
f' --save-metrics-to {metrics_path}')
self.assertTrue(os.path.exists(metrics_path))
return CustomActionRecognitionExportTestCase
def create_text_spotting_export_test_case(**kwargs):
expected_outputs_dir = os.path.join(os.path.dirname(__file__), '..', 'expected_outputs')
ExportTestCase = create_export_test_case('text_spotting',
**kwargs,
metric_keys=['f1', 'word_spotting'],
expected_outputs_dir=expected_outputs_dir)
return ExportTestCase
def create_instance_segmentation_export_test_case(**kwargs):
expected_outputs_dir = os.path.join(os.path.dirname(__file__), '..', 'expected_outputs')
ExportTestCase = create_export_test_case('instance_segmentation_2',
**kwargs,
metric_keys=['bbox', 'segm'],
expected_outputs_dir=expected_outputs_dir)
class InstanceSegmenationExportTestCase(ExportTestCase):
@classmethod
def setUpClass(cls):
super().setUpClass()
coco_dir = os.path.abspath(f'{os.path.dirname(__file__)}/../../../../data/coco')
download_and_extract_coco_val2017(coco_dir)
return InstanceSegmenationExportTestCase
def create_image_classification_export_test_case(**kwargs):
expected_outputs_dir = os.path.join(os.path.dirname(__file__), '..',
'expected_outputs')
ExportTestCase = create_export_test_case(
'image_classification',
**kwargs,
metric_keys=['accuracy'],
expected_outputs_dir=expected_outputs_dir)
class ClassificationExportTestCase(ExportTestCase):
def test_export_on_gpu(self):
skip_if_cuda_not_available()
export_dir = os.path.join(self.output_folder, 'gpu_export')
self.do_export(export_dir, on_gpu=True)
def test_export_on_cpu(self):
export_dir = os.path.join(self.output_folder, 'cpu_export')
self.do_export(export_dir, on_gpu=False)
def do_export(self, export_dir, on_gpu):
if not os.path.exists(export_dir):
initial_command = 'export CUDA_VISIBLE_DEVICES=;' if not on_gpu else ''
run_through_shell(f'{initial_command}'
f'cd {os.path.dirname(self.template_file)};'
f'pip install -r requirements.txt;'
f'python3 export.py --openvino'
f' --load-weights snapshot.pth'
f' --save-model-to {export_dir}')
self.assertTrue(
len(list(pathlib.Path(export_dir).rglob('*.onnx'))) > 0,
'Export to onnx failed')
self.assertTrue(
len(list(pathlib.Path(export_dir).rglob('*.bin'))) > 0,
'Export to openvino failed')
return ClassificationExportTestCase