log image fix

composer/algorithms/utils/augmentation_primitives.py

@@ -26,11 +26,33 @@
             implementations of enhancing color, contrast, brightness, and sharpness.
         * ``augmentation_sets["all"]`` contains all augmentations.
 """
+from enum import IntEnum
 from typing import Callable
 
 import numpy as np
 from PIL import Image, ImageEnhance, ImageOps
 
+try:
+    from PIL.Image import Resampling, Transform  # pyright: ignore
+except ImportError:  # fix version conflict issues
+    # transforms (also defined in Imaging.h)
+    class Transform(IntEnum):
+        AFFINE = 0
+        EXTENT = 1
+        PERSPECTIVE = 2
+        QUAD = 3
+        MESH = 4
+
+    # resampling filters (also defined in Imaging.h)
+    class Resampling(IntEnum):
+        NEAREST = 0
+        BOX = 4
+        BILINEAR = 2
+        HAMMING = 5
+        BICUBIC = 3
+        LANCZOS = 1
+
+
 AugmentationFn = Callable[[Image.Image, float], Image.Image]
 
 __all__ = [
@@ -155,7 +177,7 @@ def rotate(pil_img: Image.Image, level: float):
     degrees = _int_parameter(_sample_level(level), 30)
     if np.random.uniform() > 0.5:
         degrees = -degrees
-    return pil_img.rotate(degrees, resample=Image.BILINEAR)
+    return pil_img.rotate(degrees, resample=Resampling.BILINEAR)  # pyright: ignore
 
 
 def solarize(pil_img: Image.Image, level: float):
@@ -183,7 +205,12 @@ def shear_x(pil_img: Image.Image, level: float):
     level = _float_parameter(_sample_level(level), 0.3)
     if np.random.uniform() > 0.5:
         level = -level
-    return pil_img.transform(pil_img.size, Image.AFFINE, (1, level, 0, 0, 1, 0), resample=Image.BILINEAR)
+    return pil_img.transform(
+        pil_img.size,
+        Transform.AFFINE,  # pyright: ignore
+        (1, level, 0, 0, 1, 0),
+        resample=Resampling.BILINEAR,  # pyright: ignore
+    )  # pyright: ignore
 
 
 def shear_y(pil_img: Image.Image, level: float):
@@ -197,7 +224,12 @@ def shear_y(pil_img: Image.Image, level: float):
     level = _float_parameter(_sample_level(level), 0.3)
     if np.random.uniform() > 0.5:
         level = -level
-    return pil_img.transform(pil_img.size, Image.AFFINE, (1, 0, 0, level, 1, 0), resample=Image.BILINEAR)
+    return pil_img.transform(
+        pil_img.size,
+        Transform.AFFINE,  # pyright: ignore
+        (1, 0, 0, level, 1, 0),
+        resample=Resampling.BILINEAR,  # pyright: ignore
+    )  # pyright: ignore
 
 
 def translate_x(pil_img: Image.Image, level: float):
@@ -211,7 +243,12 @@ def translate_x(pil_img: Image.Image, level: float):
     level = _int_parameter(_sample_level(level), pil_img.size[0] / 3)
     if np.random.random() > 0.5:
         level = -level
-    return pil_img.transform(pil_img.size, Image.AFFINE, (1, 0, level, 0, 1, 0), resample=Image.BILINEAR)
+    return pil_img.transform(
+        pil_img.size,
+        Transform.AFFINE,  # pyright: ignore
+        (1, 0, level, 0, 1, 0),
+        resample=Resampling.BILINEAR,  # pyright: ignore
+    )  # pyright: ignore
 
 
 def translate_y(pil_img: Image.Image, level: float):
@@ -225,7 +262,12 @@ def translate_y(pil_img: Image.Image, level: float):
     level = _int_parameter(_sample_level(level), pil_img.size[1] / 3)
     if np.random.random() > 0.5:
         level = -level
-    return pil_img.transform(pil_img.size, Image.AFFINE, (1, 0, 0, 0, 1, level), resample=Image.BILINEAR)
+    return pil_img.transform(
+        pil_img.size,
+        Transform.AFFINE,  # pyright: ignore
+        (1, 0, 0, 0, 1, level),
+        resample=Resampling.BILINEAR,  # pyright: ignore
+    )  # pyright: ignore
 
 
 # The following augmentations overlap with corruptions in the ImageNet-C/CIFAR10-C test

composer/loggers/mlflow_logger.py

@@ -500,8 +500,9 @@ def log_images(
                 assert isinstance(self._run_id, str)
                 self._mlflow_client.log_image(
                     image=image,
-                    artifact_file=f'{name}_{step}_{im_ind}.png',
+                    key=f'{name}_{step}_{im_ind}',
                     run_id=self._run_id,
+                    step=step,
                 )
 
     def post_close(self):

setup.py

@@ -142,7 +142,7 @@ def package_files(prefix: str, directory: str, extension: str):
     'cryptography==42.0.8',
     'pytest-httpserver>=1.0.4,<1.1',
     'setuptools<=59.5.0',
-    'pillow==9.3.0',  # Matches the Pillow version listed in the Dockerfile
+    'pillow>=10.3.0',
 ]
 
 extra_deps['system_metrics_monitor'] = {

tests/loggers/test_mlflow_logger.py

@@ -650,8 +650,9 @@ def before_forward(self, state: State, logger: Logger):
     experiment_id = run.info.experiment_id
 
     run_file_path = mlflow_uri / Path(experiment_id) / Path(run_id)
-    im_dir = run_file_path / Path('artifacts')
-    assert len(os.listdir(im_dir)) == expected_num_ims
+    im_dir = run_file_path / Path('artifacts') / Path('images')
+    # 2 (compressed & uncompressed) per image, and two log images calls in ImageLogger
+    assert len(os.listdir(im_dir)) == expected_num_ims * 2 * 2
 
 
 @device('cpu')