Cleanup _init_test_backbone_with_pan

zhiqwang · Oct 14, 2021 · 9a6f91a · 9a6f91a
1 parent 1327d74
commit 9a6f91a
Showing 1 changed file with 9 additions and 16 deletions.
diff --git a/test/test_models.py b/test/test_models.py
@@ -144,22 +144,15 @@ def _init_test_backbone_with_pan(
     ):
         model_size = get_yolov5_size(depth_multiple, width_multiple)
         backbone_name = f"darknet_{model_size}_{version.replace('.', '_')}"
-        if use_tan:
-            model = darknet_tan_backbone(
-                backbone_name,
-                depth_multiple,
-                width_multiple,
-                version=version,
-                use_p6=use_p6,
-            )
-        else:
-            model = darknet_pan_backbone(
-                backbone_name,
-                depth_multiple,
-                width_multiple,
-                version=version,
-                use_p6=use_p6,
-            )
+        backbone_arch = eval(f"darknet_{'tan' if use_tan else 'pan'}_backbone")
+        assert backbone_arch in [darknet_pan_backbone, darknet_tan_backbone]
+        model = backbone_arch(
+            backbone_name,
+            depth_multiple,
+            width_multiple,
+            version=version,
+            use_p6=use_p6,
+        )
         return model
 
     @pytest.mark.parametrize(