Keypoint M predefined model (#136)

luxonis · Nov 19, 2024 · 523b775 · 523b775
1 parent 5b28de2
commit 523b775
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/luxonis_train/config/predefined_models/keypoint_detection_model.py b/luxonis_train/config/predefined_models/keypoint_detection_model.py
@@ -30,6 +30,11 @@ def get_variant(variant: VariantLiteral) -> KeypointDetectionVariant:
             backbone_params={"variant": "n"},
             neck_params={"variant": "n", "download_weights": True},
         ),
+        "medium": KeypointDetectionVariant(
+            backbone="EfficientRep",
+            backbone_params={"variant": "s"},
+            neck_params={"variant": "s", "download_weights": True},
+        ),
         "heavy": KeypointDetectionVariant(
             backbone="EfficientRep",
             backbone_params={"variant": "l"},