feat(app): add method & route to get uncategorized image counts

This commit is contained in:
psychedelicious
2024-09-24 15:50:34 +10:00
parent aa9594cb42
commit 845039e544
4 changed files with 49 additions and 2 deletions

View File

@@ -5,7 +5,7 @@ from fastapi.routing import APIRouter
from pydantic import BaseModel, Field
from invokeai.app.api.dependencies import ApiDependencies
from invokeai.app.services.board_records.board_records_common import BoardChanges
from invokeai.app.services.board_records.board_records_common import BoardChanges, UncategorizedImageCounts
from invokeai.app.services.boards.boards_common import BoardDTO
from invokeai.app.services.shared.pagination import OffsetPaginatedResults
@@ -146,3 +146,14 @@ async def list_all_board_image_names(
board_id,
)
return image_names
@boards_router.get(
"/uncategorized/counts",
operation_id="get_uncategorized_image_counts",
response_model=UncategorizedImageCounts,
)
async def get_uncategorized_image_counts() -> UncategorizedImageCounts:
"""Gets count of images and assets for uncategorized images (images with no board assocation)"""
return ApiDependencies.invoker.services.board_records.get_uncategorized_image_counts()

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 = """
-- Get the count of uncategorized images and assets.
SELECT
CASE
WHEN i.image_category = 'general' THEN 'image_count' -- "Images" are images in the 'general' category
ELSE 'asset_count' -- "Assets" are images in any of the other categories ('control', 'mask', 'user', 'other')
END AS category_type,
COUNT(*) AS unassigned_count
FROM images i
LEFT JOIN board_images bi ON i.image_name = bi.image_name
WHERE bi.board_id IS NULL -- Uncategorized images have no board association
AND i.is_intermediate = 0 -- Omit intermediates from the counts
GROUP BY category_type; -- Group by category_type alias, as derived from the image_category column earlier
"""
self._cursor.execute(query)
results = self._cursor.fetchall()
image_count = dict(results)["image_count"]
asset_count = dict(results)["asset_count"]
return UncategorizedImageCounts(image_count=image_count, asset_count=asset_count)
finally:
self._lock.release()