Image support

dvclive/data/__init__.py

@@ -1 +1,5 @@
-from .scalar import Scalar  # noqa: F401
+from .image_numpy import ImageNumpy
+from .image_pil import ImagePIL
+from .scalar import Scalar
+
+DATA_TYPES = [ImageNumpy, ImagePIL, Scalar]

dvclive/data/image_numpy.py

@@ -0,0 +1,19 @@
+from .image_pil import ImagePIL
+
+
+class ImageNumpy(ImagePIL):
+    @staticmethod
+    def could_log(val: object) -> bool:
+        if val.__class__.__module__ == "numpy":
+            return True
+        return False
+
+    def dump(self, val, step) -> None:
+        from PIL import Image
+
+        val = Image.fromarray(val)
+        super().dump(val, step)
+
+    @property
+    def summary(self):
+        return {}

dvclive/data/image_pil.py

@@ -0,0 +1,34 @@
+from pathlib import Path
+
+from .base import Data
+
+
+class ImagePIL(Data):
+    subdir = "images"
+    suffixes = [".jpg", ".jpeg", ".gif", ".png"]
+
+    @staticmethod
+    def could_log(val: object) -> bool:
+        if val.__class__.__module__ == "PIL.Image":
+            return True
+        return False
+
+    @property
+    def output_path(self) -> Path:
+        if Path(self.name).suffix not in self.suffixes:
+            raise ValueError(
+                f"Invalid image suffix '{Path(self.name).suffix}'"
+                f" Must be one of {self.suffixes}"
+            )
+        return self.output_folder / self.subdir / self.name
+
+    def dump(self, val, step) -> None:
+        super().dump(val, step)
+        output_path = Path(str(self.output_path).format(step=step))
+        output_path.parent.mkdir(exist_ok=True, parents=True)
+
+        val.save(output_path)
+
+    @property
+    def summary(self):
+        return {}

dvclive/data/scalar.py

@@ -10,6 +10,9 @@
 
 
 class Scalar(Data):
+    subdir = "scalars"
+    suffixes = [".csv", ".tsv"]
+
     @staticmethod
     def could_log(val: object) -> bool:
         if isinstance(val, (int, float)):
@@ -18,7 +21,7 @@ def could_log(val: object) -> bool:
 
     @property
     def output_path(self) -> Path:
-        _path = self.output_folder / self.name
+        _path = self.output_folder / self.subdir / self.name
         _path.parent.mkdir(exist_ok=True, parents=True)
         return _path.with_suffix(".tsv")
 

dvclive/metrics.py

@@ -1,11 +1,12 @@
 import json
 import logging
 import os
+import shutil
 from collections import OrderedDict
 from pathlib import Path
 from typing import Any, Dict, Union
 
-from .data import Scalar
+from .data import DATA_TYPES
 from .dvc import make_checkpoint, make_html
 from .error import InvalidDataTypeError
 
@@ -42,19 +43,29 @@ def __init__(
 
     def _cleanup(self):
 
-        for dvclive_file in Path(self.dir).rglob("*.tsv"):
-            dvclive_file.unlink()
+        for data_type in DATA_TYPES:
+            subdir = Path(self.dir) / data_type.subdir
+            for suffix in data_type.suffixes:
+                for data_file in subdir.rglob(f"*{suffix}"):
+                    data_file.unlink()
 
         if os.path.exists(self.summary_path):
             os.remove(self.summary_path)
 
         if os.path.exists(self.html_path):
-            os.remove(self.html_path)
+            shutil.rmtree(self.html_path, ignore_errors=True)
 
     def _init_paths(self):
         os.makedirs(self.dir, exist_ok=True)
+        for data_type in DATA_TYPES:
+            os.makedirs(
+                os.path.join(self.dir, data_type.subdir), exist_ok=True
+            )
+
         if self._summary:
             self.make_summary()
+        if self._html:
+            os.makedirs(self.html_path, exist_ok=True)
 
     @staticmethod
     def from_env():
@@ -92,11 +103,11 @@ def exists(self):
 
     @property
     def summary_path(self):
-        return self.dir + ".json"
+        return os.path.join(self.dir, "summary.json")
 
     @property
     def html_path(self):
-        return self.dir + "_dvc_plots/index.html"
+        return os.path.join(self.dir, "html")
 
     def get_step(self) -> int:
         return self._step
@@ -115,15 +126,18 @@ def next_step(self):
 
     def log(self, name: str, val: Union[int, float]):
 
+        data = None
         if name in self._data:
             data = self._data[name]
-        elif Scalar.could_log(val):
-            data = Scalar(name, self.dir)
-            self._data[name] = data
         else:
+            for data_type in DATA_TYPES:
+                if data_type.could_log(val):
+                    data = data_type(name, self.dir)
+                    self._data[name] = data
+        if data is None:
             raise InvalidDataTypeError(name, type(val))
-
         data.dump(val, self._step)
+
         if self._summary:
             self.make_summary()
 

setup.py

@@ -43,8 +43,9 @@ def run(self):
 hugginface = ["transformers", "datasets"]
 catalyst = ["catalyst"]
 fastai = ["fastai"]
+image = ["pillow"]
 
-all_libs = mmcv + tf + xgb + lgbm + hugginface + catalyst + fastai
+all_libs = mmcv + tf + xgb + lgbm + hugginface + catalyst + fastai + image
 
 tests_requires = [
     "pylint==2.5.3",

tests/test_catalyst.py

@@ -10,6 +10,7 @@
 
 import dvclive
 from dvclive.catalyst import DvcLiveCallback
+from dvclive.data import Scalar
 
 # pylint: disable=redefined-outer-name, unused-argument
 
@@ -66,8 +67,8 @@ def test_catalyst_callback(tmp_dir, runner, loaders):
 
     assert os.path.exists("dvc_logs")
 
-    train_path = tmp_dir / "dvc_logs/train"
-    valid_path = tmp_dir / "dvc_logs/valid"
+    train_path = tmp_dir / "dvc_logs" / Scalar.subdir / "train"
+    valid_path = tmp_dir / "dvc_logs" / Scalar.subdir / "valid"
 
     assert train_path.is_dir()
     assert valid_path.is_dir()

tests/test_data/test_image.py

@@ -0,0 +1,49 @@
+import os
+
+import numpy as np
+import pytest
+from PIL import Image
+
+# pylint: disable=unused-argument
+import dvclive
+from dvclive.data import ImageNumpy, ImagePIL
+
+
+def test_PIL(tmp_dir):
+    logger = dvclive.init()
+    img = Image.new("RGB", (500, 500), (250, 250, 250))
+    dvclive.log("image.png", img)
+
+    assert (tmp_dir / logger.dir / ImagePIL.subdir / "image.png").exists()
+
+
+def test_invalid_extension(tmp_dir):
+    dvclive.init()
+    img = Image.new("RGB", (500, 500), (250, 250, 250))
+    with pytest.raises(ValueError):
+        dvclive.log("image.foo", img)
+
+
+@pytest.mark.parametrize("shape", [(500, 500), (500, 500, 3), (500, 500, 4)])
+def test_numpy(tmp_dir, shape):
+    logger = dvclive.init()
+    img = np.ones(shape, np.uint8) * 255
+    dvclive.log("image.png", img)
+
+    assert (tmp_dir / logger.dir / ImageNumpy.subdir / "image.png").exists()
+
+
+@pytest.mark.parametrize(
+    "pattern", ["image_{step}.png", str(os.path.join("{step}", "image.png"))]
+)
+def test_step_formatting(tmp_dir, pattern):
+    logger = dvclive.init()
+    img = np.ones((500, 500, 3), np.uint8)
+    for _ in range(3):
+        dvclive.log(pattern, img)
+        dvclive.next_step()
+
+    for step in range(3):
+        assert (
+            tmp_dir / logger.dir / ImagePIL.subdir / pattern.format(step=step)
+        ).exists()

tests/test_fastai.py

@@ -13,6 +13,7 @@
 )
 
 import dvclive
+from dvclive.data import Scalar
 from dvclive.fastai import DvcLiveCallback
 
 # pylint: disable=redefined-outer-name, unused-argument
@@ -49,12 +50,12 @@ def test_fastai_callback(tmp_dir, data_loader):
 
     assert os.path.exists("dvc_logs")
 
-    train_path = tmp_dir / "dvc_logs/train"
-    valid_path = tmp_dir / "dvc_logs/valid"
+    train_path = tmp_dir / "dvc_logs" / Scalar.subdir / "train"
+    valid_path = tmp_dir / "dvc_logs" / Scalar.subdir / "valid"
 
     assert train_path.is_dir()
     assert valid_path.is_dir()
-    assert (tmp_dir / "dvc_logs/accuracy.tsv").exists()
+    assert (tmp_dir / "dvc_logs" / Scalar.subdir / "accuracy.tsv").exists()
 
 
 def test_fastai_model_file(tmp_dir, data_loader):

tests/test_main.py

@@ -10,6 +10,7 @@
 from dvclive import env
 
 # pylint: disable=unused-argument
+from dvclive.data import DATA_TYPES, Scalar
 from dvclive.dvc import SIGNAL_FILE
 from dvclive.error import (
     ConfigMismatchError,
@@ -20,13 +21,16 @@
 
 
 def read_logs(path: str):
-    assert os.path.isdir(path)
+    path = Path(path)
+    assert path.is_dir()
     history = {}
-    for metric_file in Path(path).rglob("*.tsv"):
-        metric_name = str(metric_file).replace(path + os.path.sep, "")
+    for metric_file in (path / Scalar.subdir).rglob("*.tsv"):
+        metric_name = str(metric_file).replace(
+            str(path / Scalar.subdir) + os.path.sep, ""
+        )
         metric_name = metric_name.replace(".tsv", "")
         history[metric_name] = _parse_tsv(metric_file)
-    latest = _parse_json(path + ".json")
+    latest = _parse_json(os.path.join(path, "summary.json"))
     return history, latest
 
 
@@ -57,21 +61,22 @@ def _parse_json(path):
 
 
 @pytest.mark.parametrize("path", ["logs", os.path.join("subdir", "logs")])
-def test_create_logs_dir(tmp_dir, path):
+def test_init_paths(tmp_dir, path):
     dvclive.init(path)
 
     assert (tmp_dir / path).is_dir()
+    for data_type in DATA_TYPES:
+        assert (tmp_dir / path / data_type.subdir).is_dir()
 
 
 @pytest.mark.parametrize("summary", [True, False])
 def test_logging(tmp_dir, summary):
-    dvclive.init("logs", summary=summary)
+    logger = dvclive.init("logs", summary=summary)
 
     dvclive.log("m1", 1)
 
-    assert (tmp_dir / "logs").is_dir()
-    assert (tmp_dir / "logs" / "m1.tsv").is_file()
-    assert (tmp_dir / "logs.json").is_file() == summary
+    assert (tmp_dir / "logs" / Scalar.subdir / "m1.tsv").is_file()
+    assert (tmp_dir / logger.summary_path).is_file() == summary
 
     if summary:
         _, s = read_logs("logs")
@@ -84,11 +89,11 @@ def test_nested_logging(tmp_dir):
     dvclive.log("train/m1", 1)
     dvclive.log("val/val_1/m1", 1)
 
-    assert (tmp_dir / "logs").is_dir()
-    assert (tmp_dir / "logs" / "train").is_dir()
-    assert (tmp_dir / "logs" / "val" / "val_1").is_dir()
-    assert (tmp_dir / "logs" / "train" / "m1.tsv").is_file()
-    assert (tmp_dir / "logs" / "val" / "val_1" / "m1.tsv").is_file()
+    assert (tmp_dir / "logs" / Scalar.subdir / "val" / "val_1").is_dir()
+    assert (tmp_dir / "logs" / Scalar.subdir / "train" / "m1.tsv").is_file()
+    assert (
+        tmp_dir / "logs" / Scalar.subdir / "val" / "val_1" / "m1.tsv"
+    ).is_file()
 
     _, summary = read_logs("logs")
 
@@ -127,22 +132,21 @@ def test_cleanup(tmp_dir, summary, html):
     logger = dvclive.init("logs", summary=summary)
     dvclive.log("m1", 1)
 
-    html_path = tmp_dir / logger.html_path
+    html_path = tmp_dir / logger.html_path / "index.html"
     if html:
-        html_path.parent.mkdir()
         html_path.touch()
 
     (tmp_dir / "logs" / "some_user_file.txt").touch()
 
-    assert (tmp_dir / "logs" / "m1.tsv").is_file()
-    assert (tmp_dir / "logs.json").is_file() == summary
+    assert (tmp_dir / "logs" / Scalar.subdir / "m1.tsv").is_file()
+    assert (tmp_dir / logger.summary_path).is_file() == summary
     assert html_path.is_file() == html
 
     dvclive.init("logs", summary=summary)
 
     assert (tmp_dir / "logs" / "some_user_file.txt").is_file()
-    assert not (tmp_dir / "logs" / "m1.tsv").is_file()
-    assert (tmp_dir / "logs.json").is_file() == summary
+    assert not (tmp_dir / "logs" / Scalar.subdir / "m1.tsv").is_file()
+    assert (tmp_dir / logger.summary_path).is_file() == summary
     assert not (html_path).is_file()