Added dataset filtering by labels (#4989)

A-Artemis · leoll2 · web-flow · commit d1591bbb4a09 · 2025-11-06T15:42:48.000+01:00
Co-authored-by: Leonardo Lai &lt;leonardo.lai@intel.com&gt;
diff --git a/application/backend/app/api/routers/datasets.py b/application/backend/app/api/routers/datasets.py
@@ -97,22 +97,23 @@ def add_dataset_item(
         status.HTTP_200_OK: {"description": "List of available dataset items", "model": DatasetItemsWithPagination},
     },
 )
-def list_dataset_items(
+def list_dataset_items(  # noqa: PLR0913
     project: Annotated[ProjectView, Depends(get_project)],
     dataset_service: Annotated[DatasetService, Depends(get_dataset_service)],
     limit: Annotated[int, Query(ge=1, le=MAX_DATASET_ITEMS_NUMBER_RETURNED)] = DEFAULT_DATASET_ITEMS_NUMBER_RETURNED,
     offset: Annotated[int, Query(ge=0)] = 0,
     start_date: Annotated[datetime | None, Query()] = None,
     end_date: Annotated[datetime | None, Query()] = None,
     annotation_status: Annotated[DatasetItemAnnotationStatus | None, Query()] = None,
+    labels: Annotated[list[UUID] | None, Query()] = None,
 ) -> DatasetItemsWithPagination:
     """List the available dataset items and their metadata. This endpoint supports pagination."""
     if start_date is not None and end_date is not None and start_date > end_date:
         raise HTTPException(
             status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, detail="Start date must be before end date."
         )
     total = dataset_service.count_dataset_items(
-        project=project, start_date=start_date, end_date=end_date, annotation_status=annotation_status
+        project=project, start_date=start_date, end_date=end_date, annotation_status=annotation_status, label_ids=labels
     )
     dataset_items = dataset_service.list_dataset_items(
         project=project,
@@ -121,6 +122,7 @@ def list_dataset_items(
         start_date=start_date,
         end_date=end_date,
         annotation_status=annotation_status,
+        label_ids=labels,
     )
     return DatasetItemsWithPagination(
         items=[DatasetItemView.model_validate(dataset_item, from_attributes=True) for dataset_item in dataset_items],
diff --git a/application/backend/app/repositories/dataset_item_repo.py b/application/backend/app/repositories/dataset_item_repo.py
@@ -61,10 +61,18 @@ def count(
         start_date: datetime | None = None,
         end_date: datetime | None = None,
         annotation_status: str | None = None,
+        label_ids: list[str] | None = None,
     ) -> int:
-        stmt = select(func.count()).select_from(DatasetItemDB).where(DatasetItemDB.project_id == self.project_id)
+        # When the query involves a JOIN (e.g. when filtering by labels), count distinct items to avoid duplicates
+        if label_ids:
+            select_fn = func.count(func.distinct(DatasetItemDB.id))
+        else:
+            select_fn = func.count()
+        stmt = select(select_fn).select_from(DatasetItemDB).where(DatasetItemDB.project_id == self.project_id)
         stmt = self._apply_date_filters(stmt, start_date, end_date)
         stmt = self._apply_annotation_status_filter(stmt, annotation_status)
+        if label_ids:
+            stmt = stmt.join(DatasetItemLabelDB).where(DatasetItemLabelDB.label_id.in_(label_ids))
         return self.db.scalar(stmt) or 0
 
     def list_items(
@@ -74,10 +82,13 @@ def list_items(
         start_date: datetime | None = None,
         end_date: datetime | None = None,
         annotation_status: str | None = None,
+        label_ids: list[str] | None = None,
     ) -> list[DatasetItemDB]:
         stmt = self._base_select()
         stmt = self._apply_date_filters(stmt, start_date, end_date)
         stmt = self._apply_annotation_status_filter(stmt, annotation_status)
+        if label_ids:
+            stmt = stmt.join(DatasetItemLabelDB).where(DatasetItemLabelDB.label_id.in_(label_ids)).distinct()
         stmt = stmt.order_by(DatasetItemDB.created_at.desc()).offset(offset).limit(limit)
         return list(self.db.scalars(stmt).all())
 
diff --git a/application/backend/app/services/dataset_service.py b/application/backend/app/services/dataset_service.py
@@ -159,10 +159,14 @@ def count_dataset_items(
         start_date: datetime | None = None,
         end_date: datetime | None = None,
         annotation_status: str | None = None,
+        label_ids: list[UUID] | None = None,
     ) -> int:
         """Get number of available dataset items (within date range if specified)"""
         repo = DatasetItemRepository(project_id=str(project.id), db=self._db_session)
-        return repo.count(start_date=start_date, end_date=end_date, annotation_status=annotation_status)
+        label_ids_str = [str(label_id) for label_id in label_ids] if label_ids else None
+        return repo.count(
+            start_date=start_date, end_date=end_date, annotation_status=annotation_status, label_ids=label_ids_str
+        )
 
     def list_dataset_items(
         self,
@@ -172,9 +176,11 @@ def list_dataset_items(
         start_date: datetime | None = None,
         end_date: datetime | None = None,
         annotation_status: str | None = None,
+        label_ids: list[UUID] | None = None,
     ) -> list[DatasetItem]:
         """Get information about available dataset items"""
         repo = DatasetItemRepository(project_id=str(project.id), db=self._db_session)
+        label_ids_str = [str(label_id) for label_id in label_ids] if label_ids else None
         return [
             DatasetItem.model_validate(db)
             for db in repo.list_items(
@@ -183,6 +189,7 @@ def list_dataset_items(
                 start_date=start_date,
                 end_date=end_date,
                 annotation_status=annotation_status,
+                label_ids=label_ids_str,
             )
         ]
 
diff --git a/application/backend/pyproject.toml b/application/backend/pyproject.toml
@@ -87,6 +87,9 @@ target-version = "py313"
 line-length = 120
 exclude = [".venv*"]
 
+[tool.ruff.format]
+line-ending = "lf"
+
 [tool.ruff.lint]
 select = ["ARG", "E", "F", "I", "N", "UP", "YTT", "ASYNC", "S", "COM", "C4", "FA", "PIE", "PYI", "Q", "RSE", "RET", "SIM",
     "TID", "PL", "RUF", "C90", "D103", "ANN001", "ANN201", "ANN205"]
diff --git a/application/backend/tests/integration/services/test_dataset_service.py b/application/backend/tests/integration/services/test_dataset_service.py
@@ -258,6 +258,82 @@ def _create_annotations(label_id: UUID) -> list[DatasetItemAnnotation]:
     return _create_annotations
 
 
+@pytest.fixture
+def fxt_project_with_labeled_dataset_items(
+    fxt_project_with_pipeline, db_session
+) -> tuple[ProjectView, list[DatasetItemDB]]:
+    """Fixture to create a project with multiple labeled dataset items for testing label filtering."""
+    project, _ = fxt_project_with_pipeline
+
+    # Ensure we have at least 2 labels
+    assert len(project.task.labels) >= 2, "Project must have at least 2 labels for this fixture"
+
+    label_0_id = str(project.task.labels[0].id)
+    label_1_id = str(project.task.labels[1].id)
+
+    configs = [
+        # Item 0: No annotations
+        {"name": "item_no_labels", "format": "jpg", "size": 1024, "width": 1024, "height": 768, "subset": "unassigned"},
+        # Item 1: Has label_0
+        {
+            "name": "item_label_0",
+            "format": "jpg",
+            "size": 1024,
+            "width": 1024,
+            "height": 768,
+            "subset": "unassigned",
+            "annotation_data": [{"labels": [{"id": label_0_id}], "shape": {"type": "full_image"}}],
+        },
+        # Item 2: Has label_1
+        {
+            "name": "item_label_1",
+            "format": "jpg",
+            "size": 1024,
+            "width": 1024,
+            "height": 768,
+            "subset": "unassigned",
+            "annotation_data": [{"labels": [{"id": label_1_id}], "shape": {"type": "full_image"}}],
+        },
+        # Item 3: Has both label_0 and label_1
+        {
+            "name": "item_both_labels",
+            "format": "jpg",
+            "size": 1024,
+            "width": 1024,
+            "height": 768,
+            "subset": "unassigned",
+            "annotation_data": [
+                {
+                    "labels": [{"id": label_0_id}],
+                    "shape": {"type": "rectangle", "x": 0, "y": 0, "width": 10, "height": 10},
+                },
+                {
+                    "labels": [{"id": label_1_id}],
+                    "shape": {"type": "rectangle", "x": 20, "y": 20, "width": 10, "height": 10},
+                },
+            ],
+        },
+    ]
+
+    db_dataset_items = []
+    for config in configs:
+        dataset_item = DatasetItemDB(**config)
+        dataset_item.project_id = str(project.id)
+        dataset_item.created_at = datetime.fromisoformat("2025-02-01T00:00:00Z")
+        db_dataset_items.append(dataset_item)
+    db_session.add_all(db_dataset_items)
+    db_session.flush()
+
+    # Link labels to dataset items
+    db_session.add(DatasetItemLabelDB(dataset_item_id=db_dataset_items[1].id, label_id=label_0_id))
+    db_session.add(DatasetItemLabelDB(dataset_item_id=db_dataset_items[2].id, label_id=label_1_id))
+    db_session.add(DatasetItemLabelDB(dataset_item_id=db_dataset_items[3].id, label_id=label_0_id))
+    db_session.add(DatasetItemLabelDB(dataset_item_id=db_dataset_items[3].id, label_id=label_1_id))
+    db_session.flush()
+
+    return project, db_dataset_items
+
+
 class TestDatasetServiceIntegration:
     """Integration tests for DatasetService."""
 
@@ -891,3 +967,109 @@ def test_annotation_status_filter_verifies_data_correctness(
         for item in to_review_items:
             assert item.annotation_data is not None
             assert item.user_reviewed is False
+
+    def test_list_dataset_items_filter_by_single_label(
+        self,
+        fxt_dataset_service: DatasetService,
+        fxt_project_with_labeled_dataset_items: tuple[ProjectView, list[DatasetItemDB]],
+    ):
+        """Test listing dataset items filtered by a single label."""
+        project, db_dataset_items = fxt_project_with_labeled_dataset_items
+        label_0_id = project.task.labels[0].id
+
+        # Filter by label_0 - should return items 1 and 3 (item_label_0 and item_both_labels)
+        dataset_items = fxt_dataset_service.list_dataset_items(
+            project=project,
+            label_ids=[label_0_id],
+        )
+
+        assert len(dataset_items) == 2
+        item_names = {item.name for item in dataset_items}
+        assert item_names == {"item_label_0", "item_both_labels"}
+
+    def test_list_dataset_items_filter_by_multiple_labels(
+        self,
+        fxt_dataset_service: DatasetService,
+        fxt_project_with_labeled_dataset_items: tuple[ProjectView, list[DatasetItemDB]],
+    ):
+        """Test listing dataset items filtered by multiple labels (OR logic)."""
+        project, db_dataset_items = fxt_project_with_labeled_dataset_items
+        label_0_id = project.task.labels[0].id
+        label_1_id = project.task.labels[1].id
+
+        # Filter by label_0 OR label_1 - should return items 1, 2, and 3
+        dataset_items = fxt_dataset_service.list_dataset_items(
+            project=project,
+            label_ids=[label_0_id, label_1_id],
+        )
+
+        assert len(dataset_items) == 3
+        item_names = {item.name for item in dataset_items}
+        assert item_names == {"item_label_0", "item_label_1", "item_both_labels"}
+
+    def test_list_dataset_items_filter_by_nonexistent_label(
+        self,
+        fxt_dataset_service: DatasetService,
+        fxt_project_with_labeled_dataset_items: tuple[ProjectView, list[DatasetItemDB]],
+    ):
+        """Test listing dataset items filtered by a nonexistent label."""
+        project, db_dataset_items = fxt_project_with_labeled_dataset_items
+        nonexistent_label_id = uuid4()
+
+        # Filter by nonexistent label - should return empty list
+        dataset_items = fxt_dataset_service.list_dataset_items(
+            project=project,
+            label_ids=[nonexistent_label_id],
+        )
+
+        assert len(dataset_items) == 0
+
+    def test_count_dataset_items_filter_by_single_label(
+        self,
+        fxt_dataset_service: DatasetService,
+        fxt_project_with_labeled_dataset_items: tuple[ProjectView, list[DatasetItemDB]],
+    ):
+        """Test counting dataset items filtered by a single label."""
+        project, db_dataset_items = fxt_project_with_labeled_dataset_items
+        label_0_id = project.task.labels[0].id
+
+        # Count items with label_0 - should return 2
+        count = fxt_dataset_service.count_dataset_items(
+            project=project,
+            label_ids=[label_0_id],
+        )
+
+        assert count == 2
+
+    def test_count_dataset_items_filter_by_multiple_labels(
+        self,
+        fxt_dataset_service: DatasetService,
+        fxt_project_with_labeled_dataset_items: tuple[ProjectView, list[DatasetItemDB]],
+    ):
+        """Test counting dataset items filtered by multiple labels (OR logic)."""
+        project, db_dataset_items = fxt_project_with_labeled_dataset_items
+        label_0_id = project.task.labels[0].id
+        label_1_id = project.task.labels[1].id
+
+        # Count items with label_0 OR label_1 - should return 3
+        count = fxt_dataset_service.count_dataset_items(
+            project=project,
+            label_ids=[label_0_id, label_1_id],
+        )
+
+        assert count == 3
+
+    def test_list_dataset_items_no_label_filter(
+        self,
+        fxt_dataset_service: DatasetService,
+        fxt_project_with_labeled_dataset_items: tuple[ProjectView, list[DatasetItemDB]],
+    ):
+        """Test listing dataset items without label filter returns all items."""
+        project, db_dataset_items = fxt_project_with_labeled_dataset_items
+
+        # No filter - should return all 4 items
+        dataset_items = fxt_dataset_service.list_dataset_items(project=project)
+
+        assert len(dataset_items) == 4
+        item_names = {item.name for item in dataset_items}
+        assert item_names == {"item_no_labels", "item_label_0", "item_label_1", "item_both_labels"}
diff --git a/application/backend/tests/unit/routers/test_datasets.py b/application/backend/tests/unit/routers/test_datasets.py
@@ -108,10 +108,20 @@ def test_list_dataset_items(self, fxt_get_project, fxt_dataset_item, fxt_dataset
 
         assert response.status_code == status.HTTP_200_OK
         fxt_dataset_service.count_dataset_items.assert_called_once_with(
-            project=fxt_get_project, start_date=None, end_date=None, annotation_status=None
+            project=fxt_get_project,
+            start_date=None,
+            end_date=None,
+            annotation_status=None,
+            label_ids=None,
         )
         fxt_dataset_service.list_dataset_items.assert_called_once_with(
-            project=fxt_get_project, limit=10, offset=0, start_date=None, end_date=None, annotation_status=None
+            project=fxt_get_project,
+            limit=10,
+            offset=0,
+            start_date=None,
+            end_date=None,
+            annotation_status=None,
+            label_ids=None,
         )
 
     def test_list_dataset_items_filtering_and_pagination(
@@ -130,6 +140,7 @@ def test_list_dataset_items_filtering_and_pagination(
             start_date=datetime(2025, 1, 9, 0, 0, 0, tzinfo=ZoneInfo("UTC")),
             end_date=datetime(2025, 12, 31, 23, 59, 59, tzinfo=ZoneInfo("UTC")),
             annotation_status=None,
+            label_ids=None,
         )
         fxt_dataset_service.list_dataset_items.assert_called_once_with(
             project=fxt_get_project,
@@ -138,6 +149,7 @@ def test_list_dataset_items_filtering_and_pagination(
             start_date=datetime(2025, 1, 9, 0, 0, 0, tzinfo=ZoneInfo("UTC")),
             end_date=datetime(2025, 12, 31, 23, 59, 59, tzinfo=ZoneInfo("UTC")),
             annotation_status=None,
+            label_ids=None,
         )
 
     @pytest.mark.parametrize("limit", [1000, 0, -20])
@@ -174,7 +186,11 @@ def test_list_dataset_items_with_annotation_status(
 
         assert response.status_code == status.HTTP_200_OK
         fxt_dataset_service.count_dataset_items.assert_called_once_with(
-            project=fxt_get_project, start_date=None, end_date=None, annotation_status=annotation_status
+            project=fxt_get_project,
+            start_date=None,
+            end_date=None,
+            annotation_status=annotation_status,
+            label_ids=None,
         )
         fxt_dataset_service.list_dataset_items.assert_called_once_with(
             project=fxt_get_project,
@@ -183,6 +199,7 @@ def test_list_dataset_items_with_annotation_status(
             start_date=None,
             end_date=None,
             annotation_status=annotation_status,
+            label_ids=None,
         )
 
     @pytest.mark.parametrize(
