Move the point counting to the yolo context

Author: kbolashev

.gitignore

@@ -5,4 +5,5 @@ dist/
 
 .DS_Store
 .idea/
+.vscode/
 __pycache__/

dagshub_annotation_converter/converters/yolo.py

@@ -208,20 +208,7 @@ def export_to_fs(
 
     # For pose annotations, if the number of keypoints is not set, guess it from the passed annotations
     if context.annotation_type == "pose" and context.keypoints_in_annotation is None:
-        max_keypoints = 0
-        for anns in grouped_annotations.values():
-            for ann in anns:
-                if isinstance(ann, IRPoseImageAnnotation):
-                    current_keypoints = len(ann.points)
-                    if current_keypoints > max_keypoints:
-                        if max_keypoints > 0:
-                            logger.warning(
-                                f"Found annotation with {current_keypoints} keypoints, "
-                                f"which is more than previously found maximum of {max_keypoints}.\n"
-                                "Your annotations might be inconsistent."
-                            )
-                        max_keypoints = current_keypoints
-        context.keypoints_in_annotation = max_keypoints
+        context.infer_keypoints_from_annotations(annotations)
 
     yaml_file_path = export_path / meta_file
     with open(yaml_file_path, "w") as yaml_f:

dagshub_annotation_converter/formats/yolo/context.py

@@ -1,12 +1,14 @@
+import logging
 from pathlib import Path
-from typing import Dict, Union, Optional, Literal, Callable
+from typing import Dict, List, Union, Optional, Literal, Callable
 
 import yaml
 
 
 from dagshub_annotation_converter.formats.common import ImageType
 from dagshub_annotation_converter.formats.yolo.categories import Categories
 from dagshub_annotation_converter.ir.image import IRImageAnnotationBase
+from dagshub_annotation_converter.ir.image.annotations.pose import IRPoseImageAnnotation
 from dagshub_annotation_converter.util.pydantic_util import ParentModel
 
 YoloConverterFunction = Callable[
@@ -15,6 +17,8 @@
 
 YoloAnnotationTypes = Literal["bbox", "segmentation", "pose"]
 
+logger = logging.getLogger(__name__)
+
 
 class YoloContext(ParentModel):
     annotation_type: YoloAnnotationTypes
@@ -106,3 +110,19 @@ def get_yaml_content(self, path_override: Optional[Path] = None) -> str:
             content["kpt_shape"] = [self.keypoints_in_annotation, self.keypoint_dim]
 
         return yaml.dump(content)
+
+    def infer_keypoints_from_annotations(self, annotations: List[IRImageAnnotationBase]):
+        if self.annotation_type != "pose":
+            raise ValueError("Keypoints in annotation can only be inferred for pose annotations")
+        self.keypoints_in_annotation = 0
+        for ann in annotations:
+            if isinstance(ann, IRPoseImageAnnotation):
+                current_keypoints = len(ann.points)
+                if current_keypoints > self.keypoints_in_annotation:
+                    if self.keypoints_in_annotation > 0:
+                        logger.warning(
+                            f"Found annotation with {current_keypoints} keypoints, "
+                            f"which is more than previously found maximum of {self.keypoints_in_annotation}.\n"
+                            "Your annotations might be inconsistent."
+                        )
+                    self.keypoints_in_annotation = current_keypoints