feat: support onnx export for domain_specific_object_detection

Link: https://code.alibaba-inc.com/Ali-MaaS/MaaS-lib/codereview/12682688
This commit is contained in:
lee.lcy
2023-05-25 10:44:09 +08:00
committed by xingjun.wxj
parent f2640a5a12
commit a524e01e47
2 changed files with 12 additions and 0 deletions

View File

@@ -13,6 +13,9 @@ from modelscope.metainfo import Models
from modelscope.utils.constant import ModelFile, Tasks
@EXPORTERS.register_module(
Tasks.domain_specific_object_detection,
module_name=Models.tinynas_damoyolo)
@EXPORTERS.register_module(
Tasks.image_object_detection, module_name=Models.tinynas_damoyolo)
class ObjectDetectionDamoyoloExporter(TorchModelExporter):

View File

@@ -27,6 +27,15 @@ class TestExportObjectDetectionDamoyolo(unittest.TestCase):
Exporter.from_model(model).export_onnx(
input_shape=(1, 3, 640, 640), output_dir=self.tmp_dir)
@unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
def test_export_domain_specific_object_detection_damoyolo(self):
model_id = 'damo/cv_tinynas_human-detection_damoyolo'
model = Model.from_pretrained(model_id)
with tempfile.TemporaryDirectory() as tmp_dir:
Exporter.from_model(model).export_onnx(
input_shape=(1, 3, 640, 640), output_dir=tmp_dir)
if __name__ == '__main__':
unittest.main()