Fix keypoint annotation filter logic (#4685)

Author: eugene123tw

lib/src/otx/data/module.py

@@ -98,9 +98,7 @@ def __init__(
         self.save_hyperparameters(ignore=["input_size"])
 
         dataset = DmDataset.import_from(self.data_root, format=self.data_format)
-        if self.task != OTXTaskType.H_LABEL_CLS and not (
-            self.task == OTXTaskType.KEYPOINT_DETECTION and self.data_format == "arrow"
-        ):
+        if self.task != OTXTaskType.H_LABEL_CLS:
             dataset = pre_filtering(
                 dataset,
                 self.data_format,

lib/src/otx/data/utils/pre_filtering.py

@@ -10,7 +10,7 @@
 from functools import partial
 from typing import TYPE_CHECKING
 
-from datumaro.components.annotation import Annotation, Bbox, Ellipse, Polygon
+from datumaro.components.annotation import Annotation, AnnotationType, Bbox, Ellipse, Points, Polygon
 from datumaro.components.dataset import Dataset as DmDataset
 
 from otx.types.task import OTXTaskType
@@ -19,6 +19,14 @@
     from datumaro.components.dataset_base import DatasetItem
 
 
+def get_labels(dataset: DmDataset, task: OTXTaskType) -> list[str]:
+    """Get the labels from the dataset."""
+    # label is funky from arrow dataset
+    if task == OTXTaskType.KEYPOINT_DETECTION:
+        return dataset.categories()[AnnotationType.points][0].labels
+    return dataset.categories()[AnnotationType.label]
+
+
 def pre_filtering(
     dataset: DmDataset,
     data_format: str,
@@ -42,7 +50,16 @@ def pre_filtering(
     used_background_items = set()
     msg = f"There are empty annotation items in train set, Of these, only {unannotated_items_ratio*100}% are used."
     warnings.warn(msg, stacklevel=2)
-    dataset = DmDataset.filter(dataset, partial(is_valid_anno_for_task, task=task), filter_annotations=True)
+
+    labels = get_labels(dataset, task)
+
+    dataset = DmDataset.filter(
+        dataset,
+        partial(is_valid_anno_for_task, task=task, labels=labels),
+        filter_annotations=True,
+    )
+    if task == OTXTaskType.KEYPOINT_DETECTION:
+        return dataset
     dataset = remove_unused_labels(dataset, data_format, ignore_index)
     if unannotated_items_ratio > 0:
         empty_items = [
@@ -61,7 +78,7 @@ def pre_filtering(
     )
 
 
-def is_valid_annot(item: DatasetItem, annotation: Annotation) -> bool:  # noqa: ARG001
+def is_valid_annot(item: DatasetItem, annotation: Annotation, labels: list[str]) -> bool:  # noqa: ARG001
     """Return whether DatasetItem's annotation is valid."""
     if isinstance(annotation, Bbox):
         x1, y1, x2, y2 = annotation.points
@@ -79,28 +96,45 @@ def is_valid_annot(item: DatasetItem, annotation: Annotation) -> bool:  # noqa:
             return True
         msg = "There are invalid polygon, they will be filtered out before training."
         return False
+    if isinstance(annotation, Points):
+        # For keypoint detection, num of (x, y) points should be equal to num of labels
+        if len(annotation.points) == 0:
+            msg = "There are invalid points, they will be filtered out before training."
+            warnings.warn(msg, stacklevel=2)
+            return False
+        return len(annotation.points) // 2 == len(labels)
+
     return True
 
 
-def is_valid_anno_for_task(item: DatasetItem, annotation: Annotation, task: OTXTaskType) -> bool:
+def is_valid_anno_for_task(
+    item: DatasetItem,
+    annotation: Annotation,
+    task: OTXTaskType,
+    labels: list[str],
+) -> bool:
     """Return whether DatasetItem's annotation is valid for a specific task.
 
     Args:
         item (DatasetItem): The item to be checked.
         annotation (Annotation): The annotation to be checked.
         task (OTXTaskType): The task type of the dataset.
+        labels (list[str]): The labels of the dataset.
 
     Returns:
         bool: True if the annotation is valid for the task, False otherwise.
     """
     if task == OTXTaskType.DETECTION:
-        return isinstance(annotation, Bbox) and is_valid_annot(item, annotation)
+        return isinstance(annotation, Bbox) and is_valid_annot(item, annotation, labels)
 
     # Rotated detection is a subset of instance segmentation
     if task in [OTXTaskType.INSTANCE_SEGMENTATION, OTXTaskType.ROTATED_DETECTION]:
-        return isinstance(annotation, (Polygon, Bbox, Ellipse)) and is_valid_annot(item, annotation)
+        return isinstance(annotation, (Polygon, Bbox, Ellipse)) and is_valid_annot(item, annotation, labels)
+
+    if task == OTXTaskType.KEYPOINT_DETECTION:
+        return isinstance(annotation, Points) and is_valid_annot(item, annotation, labels)
 
-    return is_valid_annot(item, annotation)
+    return is_valid_annot(item, annotation, labels)
 
 
 def remove_unused_labels(