Skip to content

Commit 5303f48

Browse files
feat(app): add method & route to get uncategorized image counts
1 parent b95be09 commit 5303f48

File tree

4 files changed

+49
-2
lines changed

4 files changed

+49
-2
lines changed

invokeai/app/api/routers/boards.py

+12-1
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
from pydantic import BaseModel, Field
66

77
from invokeai.app.api.dependencies import ApiDependencies
8-
from invokeai.app.services.board_records.board_records_common import BoardChanges
8+
from invokeai.app.services.board_records.board_records_common import BoardChanges, UncategorizedImageCounts
99
from invokeai.app.services.boards.boards_common import BoardDTO
1010
from invokeai.app.services.shared.pagination import OffsetPaginatedResults
1111

@@ -146,3 +146,14 @@ async def list_all_board_image_names(
146146
board_id,
147147
)
148148
return image_names
149+
150+
151+
@boards_router.get(
152+
"/uncategorized/counts",
153+
operation_id="get_uncategorized_image_counts",
154+
response_model=UncategorizedImageCounts,
155+
)
156+
async def get_uncategorized_image_counts() -> UncategorizedImageCounts:
157+
"""Gets count of images and assets for uncategorized images (images with no board assocation)"""
158+
159+
return ApiDependencies.invoker.services.board_records.get_uncategorized_image_counts()

invokeai/app/services/board_records/board_records_base.py

+6-1
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
from abc import ABC, abstractmethod
22

3-
from invokeai.app.services.board_records.board_records_common import BoardChanges, BoardRecord
3+
from invokeai.app.services.board_records.board_records_common import BoardChanges, BoardRecord, UncategorizedImageCounts
44
from invokeai.app.services.shared.pagination import OffsetPaginatedResults
55

66

@@ -48,3 +48,8 @@ def get_many(
4848
def get_all(self, include_archived: bool = False) -> list[BoardRecord]:
4949
"""Gets all board records."""
5050
pass
51+
52+
@abstractmethod
53+
def get_uncategorized_image_counts(self) -> UncategorizedImageCounts:
54+
"""Gets count of images and assets for uncategorized images (images with no board assocation)."""
55+
pass

invokeai/app/services/board_records/board_records_common.py

+5
Original file line numberDiff line numberDiff line change
@@ -79,3 +79,8 @@ class BoardRecordDeleteException(Exception):
7979

8080
def __init__(self, message="Board record not deleted"):
8181
super().__init__(message)
82+
83+
84+
class UncategorizedImageCounts(BaseModel):
85+
image_count: int = Field(description="The number of uncategorized images.")
86+
asset_count: int = Field(description="The number of uncategorized assets.")

invokeai/app/services/board_records/board_records_sqlite.py

+26
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
BoardRecordDeleteException,
1010
BoardRecordNotFoundException,
1111
BoardRecordSaveException,
12+
UncategorizedImageCounts,
1213
deserialize_board_record,
1314
)
1415
from invokeai.app.services.shared.pagination import OffsetPaginatedResults
@@ -228,3 +229,28 @@ def get_all(self, include_archived: bool = False) -> list[BoardRecord]:
228229
raise e
229230
finally:
230231
self._lock.release()
232+
233+
def get_uncategorized_image_counts(self) -> UncategorizedImageCounts:
234+
try:
235+
self._lock.acquire()
236+
query = """
237+
-- Get the count of uncategorized images and assets.
238+
SELECT
239+
CASE
240+
WHEN i.image_category = 'general' THEN 'image_count' -- "Images" are images in the 'general' category
241+
ELSE 'asset_count' -- "Assets" are images in any of the other categories ('control', 'mask', 'user', 'other')
242+
END AS category_type,
243+
COUNT(*) AS unassigned_count
244+
FROM images i
245+
LEFT JOIN board_images bi ON i.image_name = bi.image_name
246+
WHERE bi.board_id IS NULL -- Uncategorized images have no board association
247+
AND i.is_intermediate = 0 -- Omit intermediates from the counts
248+
GROUP BY category_type; -- Group by category_type alias, as derived from the image_category column earlier
249+
"""
250+
self._cursor.execute(query)
251+
results = self._cursor.fetchall()
252+
image_count = dict(results)["image_count"]
253+
asset_count = dict(results)["asset_count"]
254+
return UncategorizedImageCounts(image_count=image_count, asset_count=asset_count)
255+
finally:
256+
self._lock.release()

0 commit comments

Comments
 (0)