[ENH] Local embedders using onnxruntime

orangecontrib/imageanalytics/image_embedder.py

@@ -4,25 +4,109 @@
 from Orange.data import ContinuousVariable, Domain, Table, Variable
 from Orange.misc.utils.embedder_utils import EmbedderCache
 from Orange.util import dummy_callback
+from orangecontrib.imageanalytics.timm_model import InceptionV3, \
+    InceptionV4, ResNet18, ResNet50, ConvNeXt_Small, ConvNeXt_Tiny, ConvNeXt_Atto, \
+    Inception_Next_Atto, Inception_Next_Tiny
 
 from orangecontrib.imageanalytics.local_embedder import LocalEmbedder
 from orangecontrib.imageanalytics.server_embedder import ServerEmbedder
 from orangecontrib.imageanalytics.squeezenet_model import SqueezenetModel
 from orangecontrib.imageanalytics.utils.image_utils import extract_paths
 
 MODELS = {
-    "inception-v3": {
+    "inceptionnext_atto-local": {
+        "name": "InceptionNeXt Atto",
+        "description": "InceptionNeXt image recognition model trained on\nImageNet.",
+        "target_image_size": (224, 224),
+        "layers": ["penultimate"],
+        "order": 0,
+        "is_local": True,
+        "model": Inception_Next_Atto,
+    },
+    "inceptionnext_tiny-local": {
+        "name": "InceptionNeXt Tiny",
+        "description": "InceptionNeXt image recognition model trained on\nImageNet.",
+        "target_image_size": (224, 224),
+        "layers": ["penultimate"],
+        "order": 0,
+        "is_local": True,
+        "model": Inception_Next_Tiny,
+    },
+    "inception-v3-local": {
         "name": "Inception v3",
         "description": "Google's Inception v3 model trained on ImageNet.",
         "target_image_size": (299, 299),
         "layers": ["penultimate"],
         "order": 0,
+        "is_local": True,
+        "model": InceptionV3,
+    },
+    "inception-v3": {
+        "name": "Inception v3",
+        "description": "Google's Inception v3 model trained on ImageNet.",
+        "target_image_size": (299, 299),
+        "layers": ["penultimate"],
+        "order": 1,
         # batch size tell how many images we send in parallel, this number is
         # high for inception since it has many workers, but other embedders
         # send less images since bottleneck are workers, this way we avoid
         # ReadTimeout because of images waiting in a queue at the server
         "batch_size": 500,
     },
+    "inception-v4-local": {
+        "name": "Inception v4",
+        "description": "Google's Inception v4 model trained on ImageNet.",
+        "target_image_size": (299, 299),
+        "layers": ["penultimate"],
+        "order": 1,
+        "is_local": True,
+        "model": InceptionV4,
+    },
+    "resnet18-local": {
+        "name": "ResNet-B (18)",
+        "description": "18-layer image recognition model trained on\nImageNet.",
+        "target_image_size": (224, 224),
+        "layers": ["penultimate"],
+        "order": 2,
+        "is_local": True,
+        "model": ResNet18,
+    },
+    "resnet50-local": {
+        "name": "ResNet-B (50)",
+        "description": "50-layer image recognition model trained on\nImageNet.",
+        "target_image_size": (224, 224),
+        "layers": ["penultimate"],
+        "order": 2,
+        "is_local": True,
+        "model": ResNet50,
+    },
+    "convnext_small-local": {
+        "name": "ConvNeXt Small",
+        "description": "ConvNeXt image recognition model trained on\nImageNet.",
+        "target_image_size": (224, 224),
+        "layers": ["penultimate"],
+        "order": 2,
+        "is_local": True,
+        "model": ConvNeXt_Small,
+    },
+    "convnext_tiny-local": {
+        "name": "ConvNeXt Tiny",
+        "description": "ConvNeXt image recognition model trained on\nImageNet.",
+        "target_image_size": (224, 224),
+        "layers": ["penultimate"],
+        "order": 2,
+        "is_local": True,
+        "model": ConvNeXt_Tiny,
+    },
+    "convnext_atto-local": {
+        "name": "ConvNeXt Atto",
+        "description": "ConvNeXt image recognition model trained on\nImageNet.",
+        "target_image_size": (224, 224),
+        "layers": ["penultimate"],
+        "order": 2,
+        "is_local": True,
+        "model": ConvNeXt_Atto,
+    },
     "painters": {
         "name": "Painters",
         "description": "A model trained to predict painters from artwork\nimages.",
@@ -142,7 +226,8 @@ def _init_embedder(self) -> None:
         Init local or server embedder.
         """
         if self.is_local_embedder():
-            self._embedder = LocalEmbedder(self.model, self._model_settings)
+            bsize = self._model_settings.get("batch_size", LocalEmbedder.DEFAULT_BATCH_SIZE)
+            self._embedder = LocalEmbedder(self.model, self._model_settings, batch_size=bsize)
         else:
             self._embedder = ServerEmbedder(
                 self.model,

orangecontrib/imageanalytics/local_embedder.py

@@ -1,47 +1,106 @@
+from __future__ import annotations
+
+from types import SimpleNamespace
+from concurrent import futures
+from typing import TypeVar, Sequence
+from threading import local
+
+import PIL.Image
+import numpy as np
+from more_itertools import batched
+
 from Orange.misc.utils.embedder_utils import EmbedderCache
 from Orange.util import dummy_callback
 from orangecontrib.imageanalytics.utils.embedder_utils import ImageLoader
 
+T = TypeVar("T")
 
-class LocalEmbedder:
 
+class LocalEmbedder:
+    DEFAULT_BATCH_SIZE = 50
     embedder = None
 
-    def __init__(self, model, model_settings):
+    def __init__(self, model, model_settings, batch_size=DEFAULT_BATCH_SIZE):
         self.embedder = model_settings["model"]()
-
         self._target_image_size = model_settings["target_image_size"]
-
         self._image_loader = ImageLoader()
         self._cache = EmbedderCache(model)
+        self._executor = futures.ThreadPoolExecutor(
+            thread_name_prefix="local-image-embeder-pool"
+        )
+        self.batch_size = batch_size
+        self._tlocal = local()
 
     def embedd_data(self, file_paths, callback=dummy_callback):
         all_embeddings = []
-
-        for i, row in enumerate(file_paths, start=1):
-            all_embeddings.append(self._embed(row))
-            callback(i / len(file_paths))
-
+        i = 0
+        callback(0)
+        with self.embedder:
+            for chunk in batched(file_paths, self.batch_size):
+                embs = self._embed_batch(chunk)
+                all_embeddings.extend(embs)
+                i += len(embs)
+                callback(i / len(file_paths))
         self._cache.persist_cache()
         return all_embeddings
 
-    def _embed(self, file_path):
-        """ Load images and compute cache keys and send requests to
-        an http2 server for valid ones.
-        """
-        image = self._image_loader.load_image_or_none(
-            file_path, self._target_image_size
-        )
-        if image is None:
+    def _load_image(self, path: str) -> PIL.Image.Image | None:
+        # Use a thread local image loader, due to requests-cache/issues/845
+        # (really a cpython sqlite bug)
+        try:
+            self._tlocal.image_loader
+        except AttributeError:
+            self._tlocal.image_loader = ImageLoader()
+        return self._tlocal.image_loader.load_image_or_none(path)
+
+    def _embed_batch(
+            self, paths: Sequence[str | None]
+    ) -> Sequence[np.ndarray | None]:
+        preprocess = self.embedder.preprocess
+        load_image = self._load_image
+        def image(path: str | None) -> np.ndarray | None:
+            if not path:
+                return None
+            img = load_image(path)
+            if img is not None:
+                return preprocess(img)
             return None
-        image = self.embedder.preprocess(image)
 
-        cache_key = self._cache.md5_hash(image)
-        cached_im = self._cache.get_cached_result_or_none(cache_key)
-        if cached_im is not None:
-            return cached_im
+        imgs_f = [self._executor.submit(image, path) for path in paths]
+        done, _ = futures.wait(imgs_f)
+        def Future_result(f: futures.Future[T]) -> T | None:
+            assert f.done()
+            if f.cancelled() or f.exception() is not None:
+                return None
+            return f.result()
+
+        results = [
+            SimpleNamespace(
+                image=img,
+                cache_key=None,
+                embedding=None,
+            )
+            for img in map(Future_result, imgs_f)
+        ]
+
+        def cached_embedding(image: bytes) -> tuple[str, np.ndarray | None]:
+            key = self._cache.md5_hash(image)
+            return key, self._cache.get_cached_result_or_none(key)
 
-        embedded_image = self.embedder.predict(image)
+        # Fill the cached results
+        for r in results:
+            if r.image is not None:
+                r.cache_key, r.embedding = cached_embedding(r.image)
 
-        self._cache.add(cache_key, embedded_image)
-        return embedded_image
+        mask = [r.image is not None and r.embedding is None for r in results]
+        imgs = [r.image for r, masked in zip(results, mask) if masked]
+        if imgs:
+            assert len({img.shape for img in imgs}) == 1
+            imgs = np.stack(imgs, dtype=self.embedder.dtype)
+            embeddings = self.embedder.predict(imgs)
+            embeddings_iter = iter(embeddings)
+            for r, masked in zip(results, mask):
+                if masked:
+                    r.embedding = next(embeddings_iter)
+                    self._cache.add(r.cache_key, r.embedding)
+        return [r.embedding for r in results]

orangecontrib/imageanalytics/onnxrunner.py

@@ -0,0 +1,39 @@
+"""
+A helper module to run onnxruntime inference with multiprocessing in an
+"isolated" environment due to https://www.riverbankcomputing.com/pipermail/pyqt/2025-November/046378.html
+"""
+from __future__ import annotations
+from typing import TYPE_CHECKING
+
+import numpy as np
+
+if TYPE_CHECKING:
+    import onnxruntime as ort
+
+
+class Session:
+    model: ort.InferenceSession |  None
+
+    def __init__(self, model_path: str):
+        self.model_path = model_path
+        self.model = None
+
+    @classmethod
+    def set_global_session(cls, model_path):
+        global session
+        session = Session(model_path)
+
+    @classmethod
+    def global_session_run(cls, *args):
+        global session
+        assert session is not None
+        return session.run(*args)
+
+    def run(self, *args) -> tuple[np.ndarray]:
+        if self.model is None:
+            import onnxruntime as ort
+            self.model = ort.InferenceSession(self.model_path)
+        return self.model.run(*args)
+
+
+session: Session | None = None

orangecontrib/imageanalytics/squeezenet_model.py

@@ -11,6 +11,10 @@ class SqueezenetModel:
     """
     Squeezenet model wrapper.
     """
+    # These are defined for compatibility with LocalEmbederModel
+    dtype = np.float64
+    def __enter__(self): pass
+    def __exit__(self, *args): pass
 
     def __init__(self):
         self.__model = squeezenet(include_softmax=False)
@@ -28,12 +32,11 @@ def preprocess(image: PIL.Image.Image) -> np.ndarray:
         Returns
         -------
         image : np.ndarray
-            An array of size (1, 227, 227, 3).
+            An array of size (227, 227, 3).
         """
+        image = image.resize((227, 227), PIL.Image.LANCZOS)
         mean_pixel = [104.006, 116.669, 122.679]  # imagenet centering
         image = np.array(image, dtype=float)
-        if len(image.shape) < 4:
-            image = image[None, ...]
         swap_img = np.array(image)
         img_out = np.array(swap_img)
         img_out[:, :, 0] = swap_img[:, :, 2]  # from rgb to bgr - caffe mode
@@ -47,11 +50,20 @@ def predict(self, image: np.ndarray) -> np.ndarray:
         Parameters
         ----------
         image : np.ndarray
-            An array of size (1, 227, 227, 3).
+            An array of size (N, 227, 227, 3).
 
         Returns
         -------
         embedding : np.ndarray
-            An array of size (1000,).
+            An array of size (N, 1000,).
         """
-        return self.__model.predict([image])[0][0]
+        res = []
+        for img in image:
+            if img.ndim < 4:
+                img = img[None, :]
+            res.append(self.__model.predict([img])[0][0])
+        return np.stack(res)
+
+    @classmethod
+    def is_cached(cls):
+        return True

orangecontrib/imageanalytics/tests/test_squeezenet_model.py

@@ -14,12 +14,12 @@ def setUp(self):
 
     def test_preprocess(self):
         image = self.embedder.preprocess(self.image)
-        self.assertEqual(image.shape, (1, 227, 227, 3))
+        self.assertEqual(image.shape, (227, 227, 3))
 
     def test_predict(self):
         image = self.embedder.preprocess(self.image)
-        embedding = self.embedder.predict(image)
-        self.assertEqual(embedding.shape, (1000,))
+        embedding = self.embedder.predict(image[None, :])
+        self.assertEqual(embedding.shape, (1, 1000,))
 
 
 if __name__ == "__main__":