updated tests

luxonis · Nov 5, 2024 · 4511ce9 · 4511ce9
1 parent 12540ca
commit 4511ce9
Showing 1 changed file with 5 additions and 8 deletions.
diff --git a/tests/test_nn_archive/heads.py b/tests/test_nn_archive/heads.py
@@ -7,9 +7,6 @@
     HeadSegmentationMetadata,
     HeadYOLOMetadata,
 )
-from luxonis_ml.nn_archive.config_building_blocks.enums import (
-    ObjectDetectionSubtypeYOLO,
-)
 
 head_metadata = HeadMetadata(
     postprocessor_path="postprocessor.onnx",
@@ -61,7 +58,7 @@
     n_prototypes=None,
     n_keypoints=None,
     is_softmax=None,
-    subtype=ObjectDetectionSubtypeYOLO.YOLOv6,
+    subtype="yolov6",
     postprocessor_path=None,
     yolo_outputs=["feats"],
     mask_outputs=None,
@@ -80,7 +77,7 @@
     n_prototypes=10,
     n_keypoints=None,
     is_softmax=True,
-    subtype=ObjectDetectionSubtypeYOLO.YOLOv6,
+    subtype="yolov6",
     postprocessor_path="postprocessor.onnx",
     yolo_outputs=["feats"],
     mask_outputs=["mask"],
@@ -99,7 +96,7 @@
     n_prototypes=None,
     n_keypoints=21,
     is_softmax=None,
-    subtype=ObjectDetectionSubtypeYOLO.YOLOv6,
+    subtype="yolov6",
     postprocessor_path=None,
     yolo_outputs=["feats"],
     mask_outputs=None,
@@ -118,7 +115,7 @@
     n_prototypes=None,
     n_keypoints=None,
     is_softmax=None,
-    subtype=ObjectDetectionSubtypeYOLO.YOLOv6,
+    subtype="yolov6",
     postprocessor_path="postprocessor.onnx",
     yolo_outputs=["feats"],
     mask_outputs=None,
@@ -137,7 +134,7 @@
     n_prototypes=10,
     n_keypoints=21,
     is_softmax=False,
-    subtype=ObjectDetectionSubtypeYOLO.YOLOv6,
+    subtype="yolov6",
     postprocessor_path="postprocessor.onnx",
     yolo_outputs=["feats"],
     mask_outputs=["mask"],