Adds annotation status filtering to dataset item endpoints (#4976)

Author: A-Artemis

application/backend/app/api/routers/datasets.py

@@ -17,6 +17,7 @@
     SetDatasetItemAnnotations,
 )
 from app.core.models import Pagination
+from app.models import DatasetItemAnnotationStatus
 from app.schemas import ProjectView
 from app.services import DatasetService, ResourceNotFoundError
 from app.services.dataset_service import AnnotationValidationError, InvalidImageError, SubsetAlreadyAssignedError
@@ -103,15 +104,23 @@ def list_dataset_items(
     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,
 ) -> 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)
+    total = dataset_service.count_dataset_items(
+        project=project, start_date=start_date, end_date=end_date, annotation_status=annotation_status
+    )
     dataset_items = dataset_service.list_dataset_items(
-        project=project, limit=limit, offset=offset, start_date=start_date, end_date=end_date
+        project=project,
+        limit=limit,
+        offset=offset,
+        start_date=start_date,
+        end_date=end_date,
+        annotation_status=annotation_status,
     )
     return DatasetItemsWithPagination(
         items=[DatasetItemView.model_validate(dataset_item, from_attributes=True) for dataset_item in dataset_items],

application/backend/app/models/__init__.py

@@ -1,7 +1,13 @@
 # Copyright (C) 2025 Intel Corporation
 # SPDX-License-Identifier: Apache-2.0
 
-from .dataset_item import DatasetItem, DatasetItemAnnotation, DatasetItemFormat, DatasetItemSubset
+from .dataset_item import (
+    DatasetItem,
+    DatasetItemAnnotation,
+    DatasetItemAnnotationStatus,
+    DatasetItemFormat,
+    DatasetItemSubset,
+)
 from .label import Label, LabelReference
 from .shape import FullImage, Point, Polygon, Rectangle, Shape
 from .sink import (
@@ -30,6 +36,7 @@
 __all__ = [
     "DatasetItem",
     "DatasetItemAnnotation",
+    "DatasetItemAnnotationStatus",
     "DatasetItemFormat",
     "DatasetItemSubset",
     "DisconnectedSinkConfig",

application/backend/app/models/dataset_item.py

@@ -27,6 +27,12 @@ class DatasetItemSubset(StrEnum):
     TESTING = "testing"
 
 
+class DatasetItemAnnotationStatus(StrEnum):
+    UNANNOTATED = "unannotated"
+    REVIEWED = "reviewed"
+    TO_REVIEW = "to_review"
+
+
 class DatasetItemAnnotation(BaseModel):
     """
     DatasetItemAnnotation represents an individual annotation within a dataset item.

application/backend/app/repositories/dataset_item_repo.py

@@ -8,7 +8,7 @@
 from sqlalchemy.orm import Session
 
 from app.db.schema import DatasetItemDB, DatasetItemLabelDB
-from app.models import DatasetItemSubset
+from app.models import DatasetItemAnnotationStatus, DatasetItemSubset
 
 
 class UpdateDatasetItemAnnotation(NamedTuple):
@@ -28,8 +28,9 @@ def _base_select(self) -> Select:
         """Create base select statement filtered by project_id."""
         return select(DatasetItemDB).where(DatasetItemDB.project_id == self.project_id)
 
+    @staticmethod
     def _apply_date_filters(
-        self, stmt: Select, start_date: datetime | None = None, end_date: datetime | None = None
+        stmt: Select, start_date: datetime | None = None, end_date: datetime | None = None
     ) -> Select:
         """Apply date range filters to a select statement."""
         if start_date:
@@ -38,22 +39,45 @@ def _apply_date_filters(
             stmt = stmt.where(DatasetItemDB.created_at < end_date)
         return stmt
 
+    @staticmethod
+    def _apply_annotation_status_filter(stmt: Select, annotation_status: str | None = None) -> Select:
+        """Apply annotation status filter to a select statement."""
+        if annotation_status == DatasetItemAnnotationStatus.UNANNOTATED:
+            stmt = stmt.where(DatasetItemDB.annotation_data.is_(None))
+        elif annotation_status == DatasetItemAnnotationStatus.REVIEWED:
+            stmt = stmt.where(DatasetItemDB.annotation_data.is_not(None), DatasetItemDB.user_reviewed.is_(True))
+        elif annotation_status == DatasetItemAnnotationStatus.TO_REVIEW:
+            stmt = stmt.where(DatasetItemDB.annotation_data.is_not(None), DatasetItemDB.user_reviewed.is_(False))
+        return stmt
+
     def save(self, dataset_item_db: DatasetItemDB) -> DatasetItemDB:
         dataset_item_db.updated_at = datetime.now(UTC)
         self.db.add(dataset_item_db)
         self.db.flush()
         return dataset_item_db
 
-    def count(self, start_date: datetime | None = None, end_date: datetime | None = None) -> int:
+    def count(
+        self,
+        start_date: datetime | None = None,
+        end_date: datetime | None = None,
+        annotation_status: str | None = None,
+    ) -> int:
         stmt = select(func.count()).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)
         return self.db.scalar(stmt) or 0
 
     def list_items(
-        self, limit: int, offset: int, start_date: datetime | None = None, end_date: datetime | None = None
+        self,
+        limit: int,
+        offset: int,
+        start_date: datetime | None = None,
+        end_date: datetime | None = None,
+        annotation_status: 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)
         stmt = stmt.order_by(DatasetItemDB.created_at.desc()).offset(offset).limit(limit)
         return list(self.db.scalars(stmt).all())
 

application/backend/app/services/dataset_service.py

@@ -158,10 +158,11 @@ def count_dataset_items(
         project: ProjectView,
         start_date: datetime | None = None,
         end_date: datetime | None = None,
+        annotation_status: str | 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)
+        return repo.count(start_date=start_date, end_date=end_date, annotation_status=annotation_status)
 
     def list_dataset_items(
         self,
@@ -170,12 +171,19 @@ def list_dataset_items(
         offset: int = 0,
         start_date: datetime | None = None,
         end_date: datetime | None = None,
+        annotation_status: str | None = None,
     ) -> list[DatasetItem]:
         """Get information about available dataset items"""
         repo = DatasetItemRepository(project_id=str(project.id), db=self._db_session)
         return [
             DatasetItem.model_validate(db)
-            for db in repo.list_items(limit=limit, offset=offset, start_date=start_date, end_date=end_date)
+            for db in repo.list_items(
+                limit=limit,
+                offset=offset,
+                start_date=start_date,
+                end_date=end_date,
+                annotation_status=annotation_status,
+            )
         ]
 
     def get_dataset_item_by_id(self, project: ProjectView, dataset_item_id: UUID) -> DatasetItem: