fix rtmdet deploy error (open-mmlab#287)

hanrui1sensetime · Nov 17, 2022 · 5ecaf18 · 5ecaf18
1 parent c043181
commit 5ecaf18
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/mmyolo/deploy/object_detection.py b/mmyolo/deploy/object_detection.py
@@ -16,13 +16,22 @@ class MMYOLO(MMCodebase):
 
     task_registry = MMYOLO_TASK
 
+    @classmethod
+    def register_deploy_modules(cls):
+        """register all rewriters for mmcls."""
+        import mmdeploy.codebase.mmdet.models  # noqa: F401
+        import mmdeploy.codebase.mmdet.ops  # noqa: F401
+        import mmdeploy.codebase.mmdet.structures  # noqa: F401
+
     @classmethod
     def register_all_modules(cls):
         from mmdet.utils.setup_env import \
             register_all_modules as register_all_modules_mmdet
 
         from mmyolo.utils.setup_env import \
             register_all_modules as register_all_modules_mmyolo
+
+        cls.register_deploy_modules()
         register_all_modules_mmyolo(True)
         register_all_modules_mmdet(False)