arkindex-base-worker 0.4.0b3__py3-none-any.whl → 0.4.0rc1__py3-none-any.whl

{arkindex_base_worker-0.4.0b3.dist-info → arkindex_base_worker-0.4.0rc1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: arkindex-base-worker
-Version: 0.4.0b3
+Version: 0.4.0rc1
 Summary: Base Worker to easily build Arkindex ML workflows
 Author-email: Teklia <contact@teklia.com>
 Maintainer-email: Teklia <contact@teklia.com>
@@ -40,6 +40,7 @@ Classifier: Programming Language :: Python :: 3.11
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
+Requires-Dist: humanize ==4.9.0
 Requires-Dist: peewee ~=3.17
 Requires-Dist: Pillow ==10.4.0
 Requires-Dist: python-gnupg ==0.5.2
@@ -49,7 +50,7 @@ Requires-Dist: zstandard ==0.22.0
 Provides-Extra: docs
 Requires-Dist: black ==24.4.2 ; extra == 'docs'
 Requires-Dist: mkdocs-material ==9.5.31 ; extra == 'docs'
-Requires-Dist: mkdocstrings-python ==1.10.7 ; extra == 'docs'
+Requires-Dist: mkdocstrings-python ==1.10.8 ; extra == 'docs'
 Provides-Extra: tests
 Requires-Dist: pytest ==8.3.2 ; extra == 'tests'
 Requires-Dist: pytest-mock ==3.14.0 ; extra == 'tests'

{arkindex_base_worker-0.4.0b3.dist-info → arkindex_base_worker-0.4.0rc1.dist-info}/RECORD

@@ -1,51 +1,51 @@
 arkindex_worker/__init__.py,sha256=OlgCtTC9MaWeejviY0a3iQpALcRQGMVArFVVYwTF6I8,162
 arkindex_worker/cache.py,sha256=FTlB0coXofn5zTNRTcVIvh709mcw4a1bPGqkwWjKs3w,11248
-arkindex_worker/image.py,sha256=8Y0PYMbTEsFUv8lCNLBu7UaDy6um5YfHCefyXL2jpnE,14347
+arkindex_worker/image.py,sha256=oEgVCrSHiGh3D5-UXfM6PvT17TttSxC0115irpvB3Dw,18581
 arkindex_worker/models.py,sha256=bPQzGZNs5a6z6DEcygsa8T33VOqPlMUbwKzHqlKzwbw,9923
 arkindex_worker/utils.py,sha256=q1EeLdC6ebYIH-C0LOAqw2cNpjCjVoP-Vbr-39mF4w0,9884
-arkindex_worker/worker/__init__.py,sha256=w1VlDzERabXIp625kkHnojyu5ctCM11WLw4ARh1ja3k,19818
-arkindex_worker/worker/base.py,sha256=JStHpwSP3bis9LLvV2C2n6GTWtLUVIDA9JPgPJEt17o,18717
+arkindex_worker/worker/__init__.py,sha256=Xzn20bD4THFcnDfPjZeE-uU41m_whs_3yA0WjZb9uqk,18195
+arkindex_worker/worker/base.py,sha256=wyEJB5_zcy4cTvqSXMhX8DLaWQVgvIKO77-uovcprq4,19539
 arkindex_worker/worker/classification.py,sha256=ECm1cnQPOj_9m-CoO0e182ElSySAUOoyddHrORbShhc,10951
 arkindex_worker/worker/corpus.py,sha256=s9bCxOszJMwRq1WWAmKjWq888mjDfbaJ18Wo7h-rNOw,1827
 arkindex_worker/worker/dataset.py,sha256=UXElhhARca9m7Himp-yxD5dAqWbdxDKWOUJUGgeCZXI,2934
-arkindex_worker/worker/element.py,sha256=yz7q-emuCIY6MI438QXQk1Cgq991QjYoLewNyUVE4ic,36411
+arkindex_worker/worker/element.py,sha256=1qTnz9Y4nbTSxn274-sRmM2stzT5wJrsbshxXHlBoPw,44789
 arkindex_worker/worker/entity.py,sha256=qGjQvOVXfP84rER0Dkui6q-rb9nTWerHVG0Z5voB8pU,15229
 arkindex_worker/worker/image.py,sha256=t_Az6IGnj0EZyvcA4XxfPikOUjn_pztgsyxTkFZhaXU,621
 arkindex_worker/worker/metadata.py,sha256=VRajtd2kaBvar9GercX4knvR6l1WFYjoCdJWU9ccKgk,7291
 arkindex_worker/worker/task.py,sha256=1O9zrWXxe3na3TOcoHX5Pxn1875v7EU08BSsCPnb62g,1519
 arkindex_worker/worker/training.py,sha256=qnBFEk11JOWWPLTbjF-lZ9iFBdTPpQzZAzQ9a03J1j4,10874
-arkindex_worker/worker/transcription.py,sha256=8ho-8zmF9LgP86oS59ZZLv5I7tfnZ1yNO2A3pY_9GQ8,21353
+arkindex_worker/worker/transcription.py,sha256=52RY9kYsiR1sz9FxOigyo12Ker3VDbQ4U42gK9DpR3g,21146
 arkindex_worker/worker/version.py,sha256=JIT7OI3Mo7RPkNrjOB9hfqrsG-FYygz_zi4l8PbkuAo,1960
 hooks/pre_gen_project.py,sha256=xQJERv3vv9VzIqcBHI281eeWLWREXUF4mMw7PvJHHXM,269
 tests/__init__.py,sha256=6aeTMHf4q_dKY4jIZWg1KT70VKaLvVlzCxh-Uu_cWiQ,241
-tests/conftest.py,sha256=-ZQTV4rg7TgW84-5Ioqndqv8byNILfDOpyUt8wecEiI,21967
-tests/test_base_worker.py,sha256=LdFV0LFdNU2IOyEKlX59MB1kuyxHCuhy4Tm7eE_iPiU,24281
+tests/conftest.py,sha256=KNBZ0xMC9xX2pKQXp_4XwVU07JGeTSFeM4rN2RpipfY,21522
+tests/test_base_worker.py,sha256=2EIYcd_3f9O0zB5WiGIQV0Cn9wndLvnEnSfcAE1qWWU,30607
 tests/test_cache.py,sha256=ii0gyr0DrG7ChEs7pmT8hMdSguAOAcCze4bRMiFQxuk,10640
 tests/test_dataset_worker.py,sha256=d9HG36qnO5HXu9vQ0UTBvdTSRR21FVq1FNoXM-vZbPk,22105
 tests/test_element.py,sha256=2G9M15TLxQRmvrWM9Kw2ucnElh4kSv_oF_5FYwwAxTY,13181
-tests/test_image.py,sha256=Fs9vKYgQ7mEFylbzI4YIO_JyOLeAcs-WxUXpzewxCd8,16188
-tests/test_merge.py,sha256=FMdpsm_ncHNmIvOrJ1vcwlyn8o9-SPcpFTcbAsXwK-w,8320
-tests/test_utils.py,sha256=zbJC24NyTc3slz3Ed3gJDswjRChjkR5oHEgDoQMOBiE,2588
+tests/test_image.py,sha256=J3jqB5OhcdCpB6n0UnwivxrMlne8YjFLXhq1gBMANrs,26711
+tests/test_merge.py,sha256=TuOeUS0UCz66DPOQFFhc4NQBxIjZL9f5czi4XnvGrr4,8270
+tests/test_utils.py,sha256=_WJUPnt-pM_TQ0er4yjPZy-u_LePrHq1lxwk_teky7M,2544
 tests/test_elements_worker/__init__.py,sha256=Fh4nkbbyJSMv_VtjQxnWrOqTnxXaaWI8S9WU0VrzCHs,179
 tests/test_elements_worker/test_classifications.py,sha256=fXZ8cSzIWwZ6LHsY7tKsy9-Pp9fKyKUStIXS4ViBcek,27779
 tests/test_elements_worker/test_cli.py,sha256=a23i1pUDbXi23MUtbWwGEcLLrmc_YlrbDgOG3h66wLM,2620
 tests/test_elements_worker/test_corpus.py,sha256=c_LUHvkJIYgk_wXF06VQPNOoWfiZ06XpjOXrJ7MRiBc,4479
 tests/test_elements_worker/test_dataset.py,sha256=lSXqubhg1EEq2Y2goE8Y2RYaqIpM9Iejq6fGNW2BczU,11411
-tests/test_elements_worker/test_elements.py,sha256=v5MUD-a4gcmuaqG5UHu9AlzSEoRA2dudkht7cEVED_s,93227
+tests/test_elements_worker/test_elements.py,sha256=PBVRIQB8yTCCa22A0VJKIsJSa4gvagDVZVtZT8mlZF0,107199
 tests/test_elements_worker/test_entities.py,sha256=oav2dtvWWavQe1l3Drbxw1Ta2ocUJEVxJfDQ_r6-rYQ,36181
 tests/test_elements_worker/test_image.py,sha256=_E3UGdDOwTo1MW5KMS81PrdeSPBPWinWYoQPNy2F9Ro,2077
 tests/test_elements_worker/test_metadata.py,sha256=cm2NNaXxBYmYMkPexSPVTAqb2skDTB4mliwQCLz8Y98,22293
 tests/test_elements_worker/test_task.py,sha256=7Sr3fbjdgWUXJUhJEiC9CwnbhQIQX3rCInmHMIrmA38,5573
 tests/test_elements_worker/test_training.py,sha256=Qxi9EzGr_uKcn2Fh5aE6jNrq1K8QKLiOiSew4upASPs,8721
 tests/test_elements_worker/test_transcriptions.py,sha256=FNY6E26iTKqe7LP9LO72By4oV4g9hBIZYTU9BAc_w7I,77060
-tests/test_elements_worker/test_worker.py,sha256=AwdP8uSXNQ_SJavXxJV2s3_J3OiCafShVjMV1dgt4xo,17162
+tests/test_elements_worker/test_worker.py,sha256=AuFDyqncIusT-rMMY4sEay9MqGvoNuSuZQq-5rHN02U,10803
 worker-demo/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 worker-demo/tests/conftest.py,sha256=XzNMNeg6pmABUAH8jN6eZTlZSFGLYjS3-DTXjiRN6Yc,1002
 worker-demo/tests/test_worker.py,sha256=3DLd4NRK4bfyatG5P_PK4k9P9tJHx9XQq5_ryFEEFVg,304
 worker-demo/worker_demo/__init__.py,sha256=2BPomV8ZMNf3YXJgloatKeHQCE6QOkwmsHGkO6MkQuM,125
 worker-demo/worker_demo/worker.py,sha256=Rt-DjWa5iBP08k58NDZMfeyPuFbtNcbX6nc5jFX7GNo,440
-arkindex_base_worker-0.4.0b3.dist-info/LICENSE,sha256=NVshRi1efwVezMfW7xXYLrdDr2Li1AfwfGOd5WuH1kQ,1063
-arkindex_base_worker-0.4.0b3.dist-info/METADATA,sha256=KpYeTvNM7sruTB38VaQk_TephTtArTv1I6hrMI9iloM,3270
-arkindex_base_worker-0.4.0b3.dist-info/WHEEL,sha256=R0nc6qTxuoLk7ShA2_Y-UWkN8ZdfDBG2B6Eqpz2WXbs,91
-arkindex_base_worker-0.4.0b3.dist-info/top_level.txt,sha256=58NuslgxQC2vT4DiqZEgO4JqJRrYa2yeNI9QvkbfGQU,40
-arkindex_base_worker-0.4.0b3.dist-info/RECORD,,
+arkindex_base_worker-0.4.0rc1.dist-info/LICENSE,sha256=NVshRi1efwVezMfW7xXYLrdDr2Li1AfwfGOd5WuH1kQ,1063
+arkindex_base_worker-0.4.0rc1.dist-info/METADATA,sha256=22DYiI2CtAzJ9d0P21Y2ZlAoBFX_Ks-yRQMoYlMO5KM,3303
+arkindex_base_worker-0.4.0rc1.dist-info/WHEEL,sha256=Mdi9PDNwEZptOjTlUcAth7XJDFtKrHYaQMPulZeBCiQ,91
+arkindex_base_worker-0.4.0rc1.dist-info/top_level.txt,sha256=58NuslgxQC2vT4DiqZEgO4JqJRrYa2yeNI9QvkbfGQU,40
+arkindex_base_worker-0.4.0rc1.dist-info/RECORD,,

{arkindex_base_worker-0.4.0b3.dist-info → arkindex_base_worker-0.4.0rc1.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (72.1.0)
+Generator: setuptools (73.0.1)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

arkindex_worker/image.py

@@ -2,13 +2,18 @@
 Helper methods to download and open IIIF images, and manage polygons.
 """
 
+import functools
+import os
 import re
+import tempfile
 from collections import namedtuple
+from collections.abc import Generator, Iterator
 from io import BytesIO
 from math import ceil
 from pathlib import Path
 from typing import TYPE_CHECKING
 
+import humanize
 import requests
 from PIL import Image
 from shapely.affinity import rotate, scale, translate
@@ -40,8 +45,57 @@ IIIF_URL = re.compile(r"\w+:\/{2}.+\/.+\/.+\/.+\/(?P<size>.+)\/!?\d+\/\w+\.\w+")
 IIIF_FULL = "full"
 # Maximum size available
 IIIF_MAX = "max"
+# Ratio to resize image
+IMAGE_RATIO = [1, 0.9, 0.85, 0.80, 0.75, 0.70, 0.60, 0.50, 0.40, 0.30]
 
 
+def update_pillow_image_size_limit(func):
+    """
+    Update Pillow Image size limit
+    """
+
+    @functools.wraps(func)
+    def wrapper(
+        *args,
+        max_image_pixels: str | int | None = os.getenv("ARKINDEX_MAX_IMAGE_PIXELS"),
+        **kwargs,
+    ):
+        """
+        Wrapper to update Pillow Image size limit and restore it at the end of the function.
+
+        :param *args: Positional arguments passed to the function.
+        :param max_image_pixels: Pillow Image size limit to use.
+        :param **kwargs: Keyword arguments passed to the function.
+        """
+        MAX_IMAGE_PIXELS = Image.MAX_IMAGE_PIXELS
+
+        # Override Pillow Image size limit
+        if max_image_pixels is not None:
+            max_image_pixels = int(max_image_pixels)
+            # Override Pillow limit for detecting decompression bombs, disabled if set to 0
+            if max_image_pixels == 0:
+                logger.warning(
+                    "Pillow Image size limit is completely disabled, make sure you trust the image source."
+                )
+                Image.MAX_IMAGE_PIXELS = None
+            else:
+                Image.MAX_IMAGE_PIXELS = max_image_pixels
+
+        try:
+            results = func(*args, **kwargs)
+        except:
+            # Restore initial Pillow Image size limit
+            Image.MAX_IMAGE_PIXELS = MAX_IMAGE_PIXELS
+            raise
+
+        # Restore initial Pillow Image size limit
+        Image.MAX_IMAGE_PIXELS = MAX_IMAGE_PIXELS
+        return results
+
+    return wrapper
+
+
+@update_pillow_image_size_limit
 def open_image(
     path: str,
     mode: str | None = "RGB",
@@ -149,6 +203,70 @@ def upload_image(image: Image, url: str) -> requests.Response:
     return resp
 
 
+def resized_images(
+    *args,
+    element: "Element",
+    max_pixels: int | None = None,
+    max_bytes: int | None = None,
+    **kwargs,
+) -> Iterator[Generator[tempfile.NamedTemporaryFile, None, None]]:
+    """
+    Build resized images according to the pixel and byte limits.
+
+    :param *args: Positional arguments passed to [arkindex_worker.models.Element.open_image_tempfile][].
+    :param element: Element whose image needs to be resized.
+    :param max_pixels: Maximum pixel size of the resized images.
+    :param max_bytes: Maximum byte size of the resized images.
+    :param **kwargs: Keyword arguments passed to [arkindex_worker.models.Element.open_image_tempfile][].
+    :returns: An iterator of the temporary file of the resized image.
+    """
+    _, _, element_width, element_height = polygon_bounding_box(element.polygon)
+
+    logger.info(f"This element's image sizes are ({element_width} x {element_height}).")
+    if max_pixels and max(element_width, element_height) > max_pixels:
+        logger.warning(
+            f"Maximum image input size supported is ({max_pixels} x {max_pixels})."
+        )
+        logger.warning("The image will be resized.")
+
+    element_pixel, param = (
+        (element_width, "max_width")
+        if element_width > element_height
+        else (element_height, "max_height")
+    )
+
+    for resized_pixel in sorted(
+        set(
+            min(round(ratio * element_pixel), max_pixels or element_pixel)
+            for ratio in IMAGE_RATIO
+        ),
+        reverse=True,
+    ):
+        with element.open_image_tempfile(
+            *args, **{**kwargs, param: resized_pixel}
+        ) as image:
+            pillow_image = Image.open(image)
+            if (
+                pillow_image.width != element_width
+                or pillow_image.height != element_height
+            ):
+                logger.warning(
+                    f"The image was resized to ({pillow_image.width} x {pillow_image.height})."
+                )
+
+            # The image is still too large
+            image_size = Path(image.name).stat().st_size
+            if max_bytes and image_size > max_bytes:
+                logger.warning(f"The image size is {humanize.naturalsize(image_size)}.")
+                logger.warning(
+                    f"Maximum image input size supported is {humanize.naturalsize(max_bytes)}."
+                )
+                logger.warning("The image will be resized.")
+                continue
+
+            yield image
+
+
 def polygon_bounding_box(polygon: list[list[int | float]]) -> BoundingBox:
     """
     Compute the rectangle bounding box of a polygon.

arkindex_worker/worker/__init__.py

@@ -111,7 +111,7 @@ class ElementsWorker(
             help="One or more Arkindex element ID",
         )
 
-    def list_elements(self) -> Iterable[CachedElement] | list[str]:
+    def get_elements(self) -> Iterable[CachedElement] | list[str]:
         """
         List the elements to be processed, either from the CLI arguments or
         the cache database when enabled.
@@ -173,33 +173,10 @@ class ElementsWorker(
         ), "Worker must be configured to access its process activity state"
         return self.process_information.get("activity_state") == "ready"
 
-    def configure(self):
-        """
-        Setup the worker using CLI arguments and environment variables.
-        """
-        # CLI args are stored on the instance so that implementations can access them
-        self.args = self.parser.parse_args()
-
-        if self.is_read_only:
-            super().configure_for_developers()
-        else:
-            super().configure()
-            super().configure_cache()
-
-        # Retrieve the model configuration
-        if self.model_configuration:
-            self.config.update(self.model_configuration)
-            logger.info("Model version configuration retrieved")
-
-        # Retrieve the user configuration
-        if self.user_configuration:
-            self.config.update(self.user_configuration)
-            logger.info("User configuration retrieved")
-
     def run(self):
         """
         Implements an Arkindex worker that goes through each element returned by
-        [list_elements][arkindex_worker.worker.ElementsWorker.list_elements].
+        [get_elements][arkindex_worker.worker.ElementsWorker.get_elements].
         It calls [process_element][arkindex_worker.worker.ElementsWorker.process_element],
         catching exceptions, and handles saving WorkerActivity updates when enabled.
         """
@@ -207,7 +184,7 @@ class ElementsWorker(
 
         # List all elements either from JSON file
         # or direct list of elements on CLI
-        elements = self.list_elements()
+        elements = self.get_elements()
         if not elements:
             logger.warning("No elements to process, stopping.")
             sys.exit(1)
@@ -224,7 +201,7 @@ class ElementsWorker(
             element = None
             try:
                 if self.use_cache:
-                    # Just use the result of list_elements as the element
+                    # Just use the result of get_elements as the element
                     element = item
                 else:
                     # Load element using the Arkindex API
@@ -397,29 +374,6 @@ class DatasetWorker(BaseWorker, DatasetMixin, TaskMixin):
             default=[],
         )
 
-    def configure(self):
-        """
-        Setup the worker using CLI arguments and environment variables.
-        """
-        # CLI args are stored on the instance so that implementations can access them
-        self.args = self.parser.parse_args()
-
-        if self.is_read_only:
-            super().configure_for_developers()
-        else:
-            super().configure()
-            super().configure_cache()
-
-        # Retrieve the model configuration
-        if self.model_configuration:
-            self.config.update(self.model_configuration)
-            logger.info("Model version configuration retrieved")
-
-        # Retrieve the user configuration
-        if self.user_configuration:
-            self.config.update(self.user_configuration)
-            logger.info("User configuration retrieved")
-
     def cleanup_downloaded_artifact(self) -> None:
         """
         Cleanup the downloaded dataset artifact if any

arkindex_worker/worker/base.py

@@ -219,7 +219,7 @@ class BaseWorker:
         # Load all required secrets
         self.secrets = {name: self.load_secret(Path(name)) for name in required_secrets}
 
-    def configure(self):
+    def configure_worker_run(self):
         """
         Setup the necessary configuration needed using CLI args and environment variables.
         This is the method called when running a worker on Arkindex.
@@ -320,6 +320,29 @@ class BaseWorker:
         else:
             logger.debug("Cache is disabled")
 
+    def configure(self):
+        """
+        Setup the worker using CLI arguments and environment variables.
+        """
+        # CLI args are stored on the instance so that implementations can access them
+        self.args = self.parser.parse_args()
+
+        if self.is_read_only:
+            self.configure_for_developers()
+        else:
+            self.configure_worker_run()
+            self.configure_cache()
+
+        # Retrieve the model configuration
+        if self.model_configuration:
+            self.config.update(self.model_configuration)
+            logger.info("Model version configuration retrieved")
+
+        # Retrieve the user configuration
+        if self.user_configuration:
+            self.config.update(self.user_configuration)
+            logger.info("User configuration retrieved")
+
     def load_secret(self, name: Path):
         """
         Load a Ponos secret by name.