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 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
|
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]:
|
def get_all(self, include_archived: bool = False) -> list[BoardRecord]:
|
||||||
"""Gets all board records."""
|
"""Gets all board records."""
|
||||||
pass
|
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"):
|
def __init__(self, message="Board record not deleted"):
|
||||||
super().__init__(message)
|
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,
|
BoardRecordDeleteException,
|
||||||
BoardRecordNotFoundException,
|
BoardRecordNotFoundException,
|
||||||
BoardRecordSaveException,
|
BoardRecordSaveException,
|
||||||
|
UncategorizedImageCounts,
|
||||||
deserialize_board_record,
|
deserialize_board_record,
|
||||||
)
|
)
|
||||||
from invokeai.app.services.shared.pagination import OffsetPaginatedResults
|
from invokeai.app.services.shared.pagination import OffsetPaginatedResults
|
||||||
@ -228,3 +229,28 @@ class SqliteBoardRecordStorage(BoardRecordStorageBase):
|
|||||||
raise e
|
raise e
|
||||||
finally:
|
finally:
|
||||||
self._lock.release()
|
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