diff --git a/invokeai/app/services/board_records/board_records_base.py b/invokeai/app/services/board_records/board_records_base.py index 9d16dacf60..7bfe6ada6f 100644 --- a/invokeai/app/services/board_records/board_records_base.py +++ b/invokeai/app/services/board_records/board_records_base.py @@ -1,6 +1,6 @@ from abc import ABC, abstractmethod -from invokeai.app.services.board_records.board_records_common import BoardChanges, BoardRecord +from invokeai.app.services.board_records.board_records_common import BoardChanges, BoardRecord, UncategorizedImageCounts from invokeai.app.services.shared.pagination import OffsetPaginatedResults @@ -48,3 +48,8 @@ class BoardRecordStorageBase(ABC): def get_all(self, include_archived: bool = False) -> list[BoardRecord]: """Gets all board records.""" pass + + @abstractmethod + def get_uncategorized_image_counts(self) -> UncategorizedImageCounts: + """Gets count of images and assets for uncategorized images (images with no board assocation).""" + pass diff --git a/invokeai/app/services/board_records/board_records_common.py b/invokeai/app/services/board_records/board_records_common.py index 0dda8a8b6b..1c25aab565 100644 --- a/invokeai/app/services/board_records/board_records_common.py +++ b/invokeai/app/services/board_records/board_records_common.py @@ -79,3 +79,8 @@ class BoardRecordDeleteException(Exception): def __init__(self, message="Board record not deleted"): super().__init__(message) + + +class UncategorizedImageCounts(BaseModel): + image_count: int = Field(description="The number of uncategorized images.") + asset_count: int = Field(description="The number of uncategorized assets.") diff --git a/invokeai/app/services/board_records/board_records_sqlite.py b/invokeai/app/services/board_records/board_records_sqlite.py index c64e060b95..c5167824cd 100644 --- a/invokeai/app/services/board_records/board_records_sqlite.py +++ b/invokeai/app/services/board_records/board_records_sqlite.py @@ -9,6 +9,7 @@ from invokeai.app.services.board_records.board_records_common import ( BoardRecordDeleteException, BoardRecordNotFoundException, BoardRecordSaveException, + UncategorizedImageCounts, deserialize_board_record, ) from invokeai.app.services.shared.pagination import OffsetPaginatedResults @@ -228,3 +229,28 @@ class SqliteBoardRecordStorage(BoardRecordStorageBase): raise e finally: self._lock.release() + + def get_uncategorized_image_counts(self) -> UncategorizedImageCounts: + try: + self._lock.acquire() + query = """ + SELECT + CASE + WHEN i.image_category = 'general' THEN 'images' + ELSE 'assets' + END AS category_type, + COUNT(*) AS unassigned_count + FROM images i + LEFT JOIN board_images bi ON i.image_name = bi.image_name + WHERE i.image_category IN ('general', 'control', 'mask', 'user', 'other') + AND bi.board_id IS NULL + AND i.is_intermediate = 0 + GROUP BY category_type; + """ + self._cursor.execute(query) + results = self._cursor.fetchall() + image_count = results[0][1] + asset_count = results[1][1] + return UncategorizedImageCounts(image_count=image_count, asset_count=asset_count) + finally: + self._lock.release()