diff --git a/src/machinevisiontoolbox/BagOfWords.py b/src/machinevisiontoolbox/BagOfWords.py index 33231f3e..207a75f7 100644 --- a/src/machinevisiontoolbox/BagOfWords.py +++ b/src/machinevisiontoolbox/BagOfWords.py @@ -2,21 +2,32 @@ Visual Bag of Words model for image classification and retrieval. """ +from __future__ import annotations + from collections import Counter -from typing import Any +from collections.abc import Iterable, Iterator +from typing import Any, Protocol import cv2 import numpy as np -from machinevisiontoolbox.ImagePointFeatures import BaseFeature2D from machinevisiontoolbox import Image +from machinevisiontoolbox.ImagePointFeatures import BaseFeature2D + + +class _ImageQueries(Protocol): + """Finite image queries, including lists and lazy file collections.""" + + def __len__(self) -> int: ... + + def __iter__(self) -> Iterator[Image]: ... # TODO: remove top N% and bottom M% of words by frequency class BagOfWords: def __init__( self, - images, + images: Iterable[Image] | BaseFeature2D | None, k: int = 2_000, nstopwords: int = 0, attempts: int = 1, @@ -372,7 +383,7 @@ def _remove_stopwords(self, verbose: bool = True) -> None: # rearrange the cluster centroids self._centroids = self._centroids[map] - def similarity(self, query) -> np.ndarray: + def similarity(self, query: Image | _ImageQueries | np.ndarray) -> np.ndarray: """ Compute similarity between bag of words and query @@ -450,7 +461,7 @@ def similarity(self, query) -> np.ndarray: sim = sim[0, :] return sim - def retrieve(self, images) -> tuple[int, float]: + def retrieve(self, images: Image | _ImageQueries | np.ndarray) -> tuple[int, float]: S = self.similarity(images).ravel() k = np.argmax(S) @@ -483,7 +494,7 @@ def occurrence(self, word: int) -> int: return np.sum(self.words == word) @staticmethod - def _word_freq_vector(words, maxwords: int) -> np.ndarray: + def _word_freq_vector(words: np.ndarray | list[int], maxwords: int) -> np.ndarray: # create columns of the W unique, unique_counts = np.unique(words, return_counts=True) # [w,f] = count_unique(words) diff --git a/tests/typing/bag_of_words.py b/tests/typing/bag_of_words.py new file mode 100644 index 00000000..f0f2934c --- /dev/null +++ b/tests/typing/bag_of_words.py @@ -0,0 +1,20 @@ +import numpy as np + +from machinevisiontoolbox import BagOfWords, FileCollection, Image +from machinevisiontoolbox.ImagePointFeatures import BaseFeature2D + + +def accepted( + bag: BagOfWords, image: Image, features: BaseFeature2D, files: FileCollection +) -> None: + BagOfWords(features) + BagOfWords([image]) + BagOfWords(iter([image])) + BagOfWords(None) + bag.similarity(image) + bag.similarity([image]) + bag.similarity(files) + bag.similarity(np.zeros((2, 1))) + bag.retrieve(image) + bag.retrieve(files) + BagOfWords._word_freq_vector([0, 1], 2)