mirror of
https://github.com/invoke-ai/InvokeAI
synced 2024-08-30 20:32:17 +00:00
feat(app): add get_uncategorized_image_counts method on board_records service
This commit is contained in:
parent
5a0c99816c
commit
c74b9a40af
@ -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
|
||||
|
@ -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.")
|
||||
|
@ -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()
|
||||
|
Loading…
Reference in New Issue
Block a user