mirror of
https://github.com/modelscope/modelscope.git
synced 2026-09-02 03:59:31 +02:00
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:
@@ -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):
|
||||
|
||||
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user