feat(app): add get_uncategorized_image_counts method on board_records service

This commit is contained in:
psychedelicious 2024-07-12 13:57:07 +10:00
parent 5a0c99816c
commit c74b9a40af
3 changed files with 37 additions and 1 deletions

View File

@ -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

View File

@ -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.")

View File

@ -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()