fix: Improve safety of logging images with implicit formats

src/dvclive/live.py

@@ -37,6 +37,7 @@
     clean_and_copy_into,
     env2bool,
     inside_notebook,
+    isinstance_without_import,
     matplotlib_installed,
     open_file_in_browser,
 )
@@ -370,14 +371,25 @@ def log_metric(
         logger.debug(f"Logged {name}: {val}")
 
     def log_image(self, name: str, val):
-        if not Image.could_log(val):
-            raise InvalidDataTypeError(name, type(val))
-
+        # If we're given a path, try loading the image first. This might error out.
         if isinstance(val, (str, Path)):
             from PIL import Image as ImagePIL
 
             val = ImagePIL.open(val)
 
+        # If the provided image name does not have a format on it,
+        # try to infer the format from PIL Image.
+        if (
+            len(name.split(".")) <= 1
+            and isinstance_without_import(val, "PIL.Image", "Image")
+            and f".{str(val.format).lower()}" in Image.suffixes
+        ):
+            name = f"{name}.{str(val.format).lower()}"
+
+        # See if the image format and image name are valid
+        if not Image.could_log(name, val):
+            raise InvalidDataTypeError(name, type(val))
+
         if name in self._images:
             image = self._images[name]
         else:

src/dvclive/plots/image.py

@@ -7,7 +7,7 @@
 
 class Image(Data):
     suffixes = (".jpg", ".jpeg", ".gif", ".png")
-    subfolder = "images"
+    subfolder: str = "images"
 
     @property
     def output_path(self) -> Path:
@@ -16,18 +16,21 @@ def output_path(self) -> Path:
         return _path
 
     @staticmethod
-    def could_log(val: object) -> bool:
+    def could_log(name: str, val: object) -> bool:
         acceptable = {
             ("numpy", "ndarray"),
             ("matplotlib.figure", "Figure"),
             ("PIL.Image", "Image"),
         }
+
+        supported_format = False
         for cls in type(val).mro():
             if any(isinstance_without_import(val, *cls) for cls in acceptable):
-                return True
+                supported_format = True
         if isinstance(val, (PurePath, str)):
-            return True
-        return False
+            supported_format = True
+
+        return supported_format and f".{name.split('.')[-1]}" in Image.suffixes
 
     def dump(self, val, **kwargs) -> None:  # noqa: ARG002
         if isinstance_without_import(val, "numpy", "ndarray"):

tests/plots/test_image.py

@@ -4,6 +4,7 @@
 from PIL import Image
 
 from dvclive import Live
+from dvclive.error import InvalidDataTypeError
 from dvclive.plots import Image as LiveImage
 
 
@@ -24,10 +25,33 @@ def test_pil(tmp_dir):
     assert (tmp_dir / live.plots_dir / LiveImage.subfolder / "image.png").exists()
 
 
+def test_pil_omitting_extension_doesnt_save_without_valid_format(tmp_dir):
+    live = Live()
+    img = Image.new("RGB", (10, 10), (250, 250, 250))
+    with pytest.raises(InvalidDataTypeError, match="has not supported type"):
+        live.log_image("whoops", img)
+
+
+def test_pil_omitting_extension_sets_the_format_if_path_given(tmp_dir):
+    live = Live()
+    img = Image.new("RGB", (10, 10), (250, 250, 250))
+
+    # Save it first, we'll reload it and pass it's path to log_image again
+    live.log_image("saved_with_format.png", img)
+
+    # Now try saving without explicit format and check if the format is set correctly.
+    live.log_image(
+        "whoops",
+        (tmp_dir / live.plots_dir / LiveImage.subfolder / "saved_with_format.png"),
+    )
+
+    assert (tmp_dir / live.plots_dir / LiveImage.subfolder / "whoops.png").exists()
+
+
 def test_invalid_extension(tmp_dir):
     live = Live()
     img = Image.new("RGB", (10, 10), (250, 250, 250))
-    with pytest.raises(ValueError, match="unknown file extension"):
+    with pytest.raises(InvalidDataTypeError, match="has not supported type"):
         live.log_image("image.foo", img)