log image fix (#3286)

* log image fix

Signed-off-by: Jesse Chan <jesse.chan@databricks.com>

* fixed log image tests

Signed-off-by: Jesse Chan <jesse.chan@databricks.com>

* linter

Signed-off-by: Jesse Chan <jesse.chan@databricks.com>

* add simd requirement

* post0?

* versioning yada yada yada

* guh

* import fix?

* update deps

* fix

* fix II

* remove other dependency

* debug statement, remove

* post1?!

* build from source

* whitespace?

* use pillow

* delete a unit test and ignore some types

* s/type/pyright

* formatting

* formatting

* ignore more stuff

* Apply suggestions from code review

* remove rest

* Update setup.py

Co-authored-by: Mihir Patel <mihir.v.patel7@gmail.com>

* try no ignore

* remove intenum

---------

Signed-off-by: Jesse Chan <jesse.chan@databricks.com>
Co-authored-by: Daniel King <43149077+dakinggg@users.noreply.github.com>
Co-authored-by: Milo Cress <milo.cress@databricks.com>
Co-authored-by: Mihir Patel <mihir.v.patel7@gmail.com>

composer/algorithms/utils/augmentation_primitives.py

@@ -30,6 +30,7 @@
 
 import numpy as np
 from PIL import Image, ImageEnhance, ImageOps
+from PIL.Image import Resampling, Transform
 
 AugmentationFn = Callable[[Image.Image, float], Image.Image]
 
@@ -155,7 +156,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)
 
 
 def solarize(pil_img: Image.Image, level: float):
@@ -183,7 +184,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,
+        (1, level, 0, 0, 1, 0),
+        resample=Resampling.BILINEAR,
+    )
 
 
 def shear_y(pil_img: Image.Image, level: float):
@@ -197,7 +203,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,
+        (1, 0, 0, level, 1, 0),
+        resample=Resampling.BILINEAR,
+    )
 
 
 def translate_x(pil_img: Image.Image, level: float):
@@ -211,7 +222,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,
+        (1, 0, level, 0, 1, 0),
+        resample=Resampling.BILINEAR,
+    )
 
 
 def translate_y(pil_img: Image.Image, level: float):
@@ -225,7 +241,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,
+        (1, 0, 0, 0, 1, level),
+        resample=Resampling.BILINEAR,
+    )
 
 
 # The following augmentations overlap with corruptions in the ImageNet-C/CIFAR10-C test

composer/loggers/mlflow_logger.py

@@ -507,8 +507,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,<11',
 ]
 
 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')