log image fix (#3286)

* log image fix Signed-off-by: Jesse Chan <[email protected]> * fixed log image tests Signed-off-by: Jesse Chan <[email protected]> * linter Signed-off-by: Jesse Chan <[email protected]> * 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 <[email protected]> * try no ignore * remove intenum --------- Signed-off-by: Jesse Chan <[email protected]> Co-authored-by: Daniel King <[email protected]> Co-authored-by: Milo Cress <[email protected]> Co-authored-by: Mihir Patel <[email protected]>
mosaicml · Jun 17, 2024 · 3859366 · 3859366
1 parent 2cf9262
commit 3859366
Show file tree

Hide file tree

Showing 5 changed files with 32 additions and 27 deletions.
diff --git a/composer/algorithms/utils/augmentation_primitives.py b/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

diff --git a/composer/loggers/mlflow_logger.py b/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):

diff --git a/setup.py b/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'] = {

diff --git a/tests/loggers/test_mlflow_logger.py b/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')

diff --git a/tests/test_docker.py b/tests/test_docker.py