wehub-resource-sync
367 lines
13 KiB
Python
367 lines
13 KiB
Python
import json
|
|
import os
|
|
import posixpath
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
import mlflow
|
|
from mlflow.utils.file_utils import local_file_uri_to_path
|
|
from mlflow.utils.time import get_current_time_millis
|
|
|
|
|
|
@pytest.mark.parametrize("subdir", [None, ".", "dir", "dir1/dir2", "dir/.."])
|
|
def test_log_image_numpy(subdir):
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
filename = "image.png"
|
|
artifact_file = filename if subdir is None else posixpath.join(subdir, filename)
|
|
|
|
image = np.random.randint(0, 256, size=(100, 100, 3), dtype=np.uint8)
|
|
|
|
with mlflow.start_run():
|
|
mlflow.log_image(image, artifact_file)
|
|
|
|
artifact_path = None if subdir is None else posixpath.normpath(subdir)
|
|
artifact_uri = mlflow.get_artifact_uri(artifact_path)
|
|
run_artifact_dir = local_file_uri_to_path(artifact_uri)
|
|
assert os.listdir(run_artifact_dir) == [filename]
|
|
|
|
logged_path = os.path.join(run_artifact_dir, filename)
|
|
loaded_image = np.asarray(Image.open(logged_path), dtype=np.uint8)
|
|
np.testing.assert_array_equal(loaded_image, image)
|
|
|
|
|
|
@pytest.mark.parametrize("subdir", [None, ".", "dir", "dir1/dir2", "dir/.."])
|
|
def test_log_image_pillow(subdir):
|
|
from PIL import Image, ImageChops
|
|
|
|
filename = "image.png"
|
|
artifact_file = filename if subdir is None else posixpath.join(subdir, filename)
|
|
|
|
image = Image.new("RGB", (100, 100))
|
|
|
|
with mlflow.start_run():
|
|
mlflow.log_image(image, artifact_file)
|
|
|
|
artifact_path = None if subdir is None else posixpath.normpath(subdir)
|
|
artifact_uri = mlflow.get_artifact_uri(artifact_path)
|
|
run_artifact_dir = local_file_uri_to_path(artifact_uri)
|
|
assert os.listdir(run_artifact_dir) == [filename]
|
|
|
|
logged_path = os.path.join(run_artifact_dir, filename)
|
|
loaded_image = Image.open(logged_path)
|
|
# How to check Pillow image equality: https://stackoverflow.com/a/6204954/6943581
|
|
assert ImageChops.difference(loaded_image, image).getbbox() is None
|
|
|
|
|
|
def test_log_image_raises_for_unsupported_objects():
|
|
with mlflow.start_run():
|
|
with pytest.raises(TypeError, match="Unsupported image object type"):
|
|
mlflow.log_image("not_image", "image.png")
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"size",
|
|
[
|
|
(100, 100), # Grayscale (2D)
|
|
(100, 100, 1), # Grayscale (3D)
|
|
(100, 100, 3), # RGB
|
|
(100, 100, 4), # RGBA
|
|
],
|
|
)
|
|
def test_log_image_numpy_shape(size):
|
|
import numpy as np
|
|
|
|
filename = "image.png"
|
|
image = np.random.randint(0, 256, size=size, dtype=np.uint8)
|
|
|
|
with mlflow.start_run():
|
|
mlflow.log_image(image, filename)
|
|
artifact_uri = mlflow.get_artifact_uri()
|
|
run_artifact_dir = local_file_uri_to_path(artifact_uri)
|
|
assert os.listdir(run_artifact_dir) == [filename]
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"dtype",
|
|
[
|
|
# Ref.: https://numpy.org/doc/stable/user/basics.types.html#array-types-and-conversions-between-types
|
|
"int8",
|
|
"int16",
|
|
"int32",
|
|
"int64",
|
|
"uint8",
|
|
"uint16",
|
|
"uint32",
|
|
"uint64",
|
|
"float16",
|
|
"float32",
|
|
"float64",
|
|
"bool",
|
|
],
|
|
)
|
|
def test_log_image_numpy_dtype(dtype):
|
|
import numpy as np
|
|
|
|
filename = "image.png"
|
|
image = np.random.randint(0, 2, size=(100, 100, 3)).astype(np.dtype(dtype))
|
|
|
|
with mlflow.start_run():
|
|
mlflow.log_image(image, filename)
|
|
artifact_uri = mlflow.get_artifact_uri()
|
|
run_artifact_dir = local_file_uri_to_path(artifact_uri)
|
|
assert os.listdir(run_artifact_dir) == [filename]
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"array",
|
|
# 1 pixel images with out-of-range values
|
|
[[[-1]], [[256]], [[-0.1]], [[1.1]]],
|
|
)
|
|
def test_log_image_numpy_emits_warning_for_out_of_range_values(array):
|
|
import numpy as np
|
|
|
|
image = np.array(array).astype(type(array[0][0]))
|
|
if isinstance(array[0][0], int):
|
|
with (
|
|
mlflow.start_run(),
|
|
pytest.raises(ValueError, match="Integer pixel values out of acceptable range"),
|
|
):
|
|
mlflow.log_image(image, "image.png")
|
|
else:
|
|
with (
|
|
mlflow.start_run(),
|
|
pytest.warns(UserWarning, match="Float pixel values out of acceptable range"),
|
|
):
|
|
mlflow.log_image(image, "image.png")
|
|
|
|
|
|
def test_log_image_numpy_raises_exception_for_invalid_array_data_type():
|
|
import numpy as np
|
|
|
|
with mlflow.start_run(), pytest.raises(TypeError, match="Invalid array data type"):
|
|
mlflow.log_image(np.tile("a", (1, 1, 3)), "image.png")
|
|
|
|
|
|
def test_log_image_numpy_raises_exception_for_invalid_array_shape():
|
|
import numpy as np
|
|
|
|
with mlflow.start_run(), pytest.raises(ValueError, match="`image` must be a 2D or 3D array"):
|
|
mlflow.log_image(np.zeros((1,), dtype=np.uint8), "image.png")
|
|
|
|
|
|
def test_log_image_numpy_raises_exception_for_invalid_channel_length():
|
|
import numpy as np
|
|
|
|
with mlflow.start_run(), pytest.raises(ValueError, match="Invalid channel length"):
|
|
mlflow.log_image(np.zeros((1, 1, 5), dtype=np.uint8), "image.png")
|
|
|
|
|
|
def test_log_image_raises_exception_for_unsupported_image_object_type():
|
|
with mlflow.start_run(), pytest.raises(TypeError, match="Unsupported image object type"):
|
|
mlflow.log_image("not_image", "image.png")
|
|
|
|
|
|
def test_log_image_with_steps():
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
image = np.random.randint(0, 256, size=(100, 100, 3), dtype=np.uint8)
|
|
|
|
with mlflow.start_run():
|
|
mlflow.log_image(image, key="dog", step=0, synchronous=True)
|
|
|
|
logged_path = "images/"
|
|
artifact_uri = mlflow.get_artifact_uri(logged_path)
|
|
run_artifact_dir = local_file_uri_to_path(artifact_uri)
|
|
files = os.listdir(run_artifact_dir)
|
|
|
|
# .png file for the image and .webp file for compressed image
|
|
assert len(files) == 2
|
|
for file in files:
|
|
assert file.startswith("dog+step+0")
|
|
logged_path = os.path.join(run_artifact_dir, file)
|
|
if file.endswith(".png"):
|
|
loaded_image = np.asarray(Image.open(logged_path), dtype=np.uint8)
|
|
np.testing.assert_array_equal(loaded_image, image)
|
|
elif file.endswith(".json"):
|
|
with open(logged_path) as f:
|
|
metadata = json.load(f)
|
|
assert metadata["filepath"].startswith("images/dog+step+0")
|
|
assert metadata["key"] == "dog"
|
|
assert metadata["step"] == 0
|
|
assert metadata["timestamp"] <= get_current_time_millis()
|
|
|
|
|
|
@pytest.mark.parametrize("step", [20, 26, 27])
|
|
def test_log_image_with_url_encoding_prone_steps(step):
|
|
"""Regression test: steps like 20, 26, 27 previously created %20, %26, %27 patterns
|
|
in filenames that got URL-decoded, corrupting the artifact path.
|
|
See https://github.com/mlflow/mlflow/issues/21085
|
|
"""
|
|
image = np.random.randint(0, 256, size=(100, 100, 3), dtype=np.uint8)
|
|
|
|
with mlflow.start_run():
|
|
mlflow.log_image(image, key="dog", step=step, synchronous=True)
|
|
|
|
artifact_uri = mlflow.get_artifact_uri("images/")
|
|
run_artifact_dir = local_file_uri_to_path(artifact_uri)
|
|
files = os.listdir(run_artifact_dir)
|
|
|
|
assert len(files) == 2
|
|
for file in files:
|
|
assert file.startswith(f"dog+step+{step}+timestamp+")
|
|
assert "%" not in file
|
|
|
|
|
|
def test_log_image_with_timestamp():
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
image = np.random.randint(0, 256, size=(100, 100, 3), dtype=np.uint8)
|
|
|
|
with mlflow.start_run():
|
|
mlflow.log_image(image, key="dog", timestamp=100, synchronous=True)
|
|
|
|
logged_path = "images/"
|
|
artifact_uri = mlflow.get_artifact_uri(logged_path)
|
|
run_artifact_dir = local_file_uri_to_path(artifact_uri)
|
|
files = os.listdir(run_artifact_dir)
|
|
|
|
# .png file for the image, and .webp file for compressed image
|
|
assert len(files) == 2
|
|
for file in files:
|
|
assert file.startswith("dog+step+0")
|
|
logged_path = os.path.join(run_artifact_dir, file)
|
|
if file.endswith(".png"):
|
|
loaded_image = np.asarray(Image.open(logged_path), dtype=np.uint8)
|
|
np.testing.assert_array_equal(loaded_image, image)
|
|
elif file.endswith(".json"):
|
|
with open(logged_path) as f:
|
|
metadata = json.load(f)
|
|
assert metadata["filepath"].startswith("images/dog+step+0")
|
|
assert metadata["key"] == "dog"
|
|
assert metadata["step"] == 0
|
|
assert metadata["timestamp"] == 100
|
|
|
|
|
|
def test_duplicated_log_image_with_step():
|
|
"""
|
|
MLflow will save both files if there are multiple calls to log_image
|
|
with the same key and step.
|
|
"""
|
|
import numpy as np
|
|
|
|
image1 = np.random.randint(0, 256, size=(100, 100, 3), dtype=np.uint8)
|
|
image2 = np.random.randint(0, 256, size=(100, 100, 3), dtype=np.uint8)
|
|
|
|
with mlflow.start_run():
|
|
mlflow.log_image(image1, key="dog", step=100, synchronous=True)
|
|
mlflow.log_image(image2, key="dog", step=100, synchronous=True)
|
|
|
|
logged_path = "images/"
|
|
artifact_uri = mlflow.get_artifact_uri(logged_path)
|
|
run_artifact_dir = local_file_uri_to_path(artifact_uri)
|
|
files = os.listdir(run_artifact_dir)
|
|
assert len(files) == 2 * 2 # 2 images and 2 files per image
|
|
|
|
|
|
def test_duplicated_log_image_with_timestamp():
|
|
"""
|
|
MLflow will save both files if there are multiple calls to log_image
|
|
with the same key, step, and timestamp.
|
|
"""
|
|
import numpy as np
|
|
|
|
image1 = np.random.randint(0, 256, size=(100, 100, 3), dtype=np.uint8)
|
|
image2 = np.random.randint(0, 256, size=(100, 100, 3), dtype=np.uint8)
|
|
|
|
with mlflow.start_run():
|
|
mlflow.log_image(image1, key="dog", step=100, timestamp=100, synchronous=True)
|
|
mlflow.log_image(image2, key="dog", step=100, timestamp=100, synchronous=True)
|
|
|
|
logged_path = "images/"
|
|
artifact_uri = mlflow.get_artifact_uri(logged_path)
|
|
run_artifact_dir = local_file_uri_to_path(artifact_uri)
|
|
files = os.listdir(run_artifact_dir)
|
|
assert len(files) == 2 * 2
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"args",
|
|
[
|
|
{"key": "image"},
|
|
{"step": 0},
|
|
{"timestamp": 0},
|
|
{"timestamp": 0, "step": 0},
|
|
["image"],
|
|
["image", 0],
|
|
],
|
|
)
|
|
def test_log_image_raises_exception_for_unexpected_arguments_used(args):
|
|
# It will overwrite if the user wants the exact same timestamp for the logged images
|
|
import numpy as np
|
|
|
|
exception = "The `artifact_file` parameter cannot be used in conjunction"
|
|
if isinstance(args, dict):
|
|
with mlflow.start_run(), pytest.raises(TypeError, match=exception):
|
|
mlflow.log_image(np.zeros((1,), dtype=np.uint8), "image.png", **args)
|
|
elif isinstance(args, list):
|
|
with mlflow.start_run(), pytest.raises(TypeError, match=exception):
|
|
mlflow.log_image(np.zeros((1,), dtype=np.uint8), "image.png", *args)
|
|
|
|
|
|
def test_log_image_raises_exception_for_missing_arguments():
|
|
import numpy as np
|
|
|
|
exception = "Invalid arguments: Please specify exactly one of `artifact_file` or `key`"
|
|
with mlflow.start_run(), pytest.raises(TypeError, match=exception):
|
|
mlflow.log_image(np.zeros((1,), dtype=np.uint8))
|
|
|
|
|
|
def test_async_log_image_flush():
|
|
import numpy as np
|
|
|
|
image1 = np.random.randint(0, 256, size=(100, 100, 3), dtype=np.uint8)
|
|
with mlflow.start_run():
|
|
for i in range(100):
|
|
mlflow.log_image(image1, key="dog", step=i, timestamp=i, synchronous=False)
|
|
|
|
mlflow.flush_artifact_async_logging()
|
|
|
|
logged_path = "images/"
|
|
artifact_uri = mlflow.get_artifact_uri(logged_path)
|
|
run_artifact_dir = local_file_uri_to_path(artifact_uri)
|
|
files = os.listdir(run_artifact_dir)
|
|
assert len(files) == 100 * 2
|
|
|
|
|
|
def test_log_image_with_slash_in_key():
|
|
image = np.random.randint(0, 256, size=(100, 100, 3), dtype=np.uint8)
|
|
|
|
with mlflow.start_run():
|
|
mlflow.log_image(image, key="category/name", step=5, synchronous=True)
|
|
|
|
logged_path = "images/"
|
|
artifact_uri = mlflow.get_artifact_uri(logged_path)
|
|
run_artifact_dir = local_file_uri_to_path(artifact_uri)
|
|
files = os.listdir(run_artifact_dir)
|
|
|
|
assert len(files) == 2
|
|
for file in files:
|
|
# '~' must be used instead of '#' as the separator
|
|
assert "category~name" in file
|
|
assert "#" not in file
|
|
|
|
run_id = mlflow.active_run().info.run_id
|
|
|
|
client = mlflow.MlflowClient()
|
|
artifacts = client.list_artifacts(run_id, path="images")
|
|
assert len(artifacts) == 2
|
|
for artifact in artifacts:
|
|
# download_artifacts must not raise MlflowException about '#' in path
|
|
local_path = client.download_artifacts(run_id, artifact.path)
|
|
assert os.path.exists(local_path)
|