Add support for MLFlow logger

.gitignore

@@ -175,3 +175,4 @@ docs/source/_build/
 # logs
 wandb/
 lightning_logs/
+mlruns

notebooks/600_loggers/README.md

@@ -0,0 +1,18 @@
+# Using the MLFlow Logger
+
+## Notebook Contents
+
+This notebook demonstrates how to use the MLFlow logger. The notebook is divided into the following sections:
+
+- Installation
+- Run MLFlow Server
+- Dataset Directory
+- Imports
+- Data Module
+- Model
+- MLFlow Logger
+- Training
+- Testing
+- Demo Track Figure
+- Save Model to MLFLow
+- Load Model from MLFlow

pyproject.toml

@@ -45,6 +45,7 @@ loggers = [
     "gradio>=4",
     "tensorboard",
     "wandb==0.12.17",
+    "mlflow >=1.0.0",
 ]
 notebooks = ["gitpython", "ipykernel", "ipywidgets", "notebook"]
 dev = [

requirements/loggers.txt

@@ -2,5 +2,6 @@ comet-ml>=3.31.7
 gradio>=4
 tensorboard
 wandb==0.12.17
+mlflow >=1.0.0
 GitPython
 ipykernel

src/anomalib/loggers/__init__.py

@@ -15,6 +15,7 @@
 
 try:
     from .comet import AnomalibCometLogger  # noqa: F401
+    from .mlflow import AnomalibMLFlowLogger  # noqa: F401
     from .tensorboard import AnomalibTensorBoardLogger  # noqa: F401
     from .wandb import AnomalibWandbLogger  # noqa: F401
 
@@ -23,6 +24,7 @@
             "AnomalibCometLogger",
             "AnomalibTensorBoardLogger",
             "AnomalibWandbLogger",
+            "AnomalibMLFlowLogger",
         ],
     )
 except ImportError:

src/anomalib/loggers/mlflow.py

@@ -0,0 +1,103 @@
+"""MLFlow logger with add image interface."""
+
+from typing import Literal
+
+import numpy as np
+from lightning.pytorch.loggers.mlflow import MLFlowLogger
+from lightning.pytorch.utilities import rank_zero_only
+from matplotlib.figure import Figure
+
+from anomalib.utils.exceptions.imports import try_import
+
+from .base import ImageLoggerBase
+
+try_import("mlflow")
+
+
+class AnomalibMLFlowLogger(ImageLoggerBase, MLFlowLogger):
+    """Logger for MLFlow.
+
+    Adds interface for ``add_image`` in the logger rather than calling the
+    experiment object.
+
+    .. note::
+        Same as the MLFlowLogger provided by PyTorch Lightning and the doc string is reproduced below.
+
+    Track your parameters, metrics, source code and more using
+    `MLFlow <https://mlflow.org/#core-concepts>`_.
+
+    Install it with pip:
+
+    .. code-block:: bash
+
+        pip install mlflow
+
+    Args:
+        experiment_name: The name of the experiment.
+        run_name: Name of the new run.
+            The `run_name` is internally stored as a ``mlflow.runName`` tag.
+            If the ``mlflow.runName`` tag has already been set in `tags`, the value is overridden by the `run_name`.
+        tracking_uri: Address of local or remote tracking server.
+            If not provided, defaults to `MLFLOW_TRACKING_URI` environment variable if set, otherwise it falls
+            back to `file:<save_dir>`.
+        save_dir: A path to a local directory where the MLflow runs get saved.
+            Defaults to `./mlruns` if `tracking_uri` is not provided.
+            Has no effect if `tracking_uri` is provided.
+        log_model: Log checkpoints created by `ModelCheckpoint` as MLFlow artifacts.
+
+            - if ``log_model == 'all'``, checkpoints are logged during training.
+            - if ``log_model == True``, checkpoints are logged at the end of training, \
+                except when `save_top_k == -1` which also logs every checkpoint during training.
+            - if ``log_model == False`` (default), no checkpoint is logged.
+
+        prefix: A string to put at the beginning of metric keys. Defaults to ``''``.
+        kwargs: Additional arguments like `tags`, `artifact_location` etc. used by
+            `MLFlowExperiment` can be passed as keyword arguments in this logger.
+
+    Example:
+        >>> from anomalib.loggers import AnomalibMLFlowLogger
+        >>> from anomalib.engine import Engine
+        ...
+        >>> mlflow_logger = AnomalibMLFlowLogger()
+        >>> engine = Engine(logger=mlflow_logger)
+
+    See Also:
+        - `MLFlow Documentation <https://mlflow.org/docs/latest/>`_.
+    """
+
+    def __init__(
+        self,
+        experiment_name: str | None = "anomalib_logs",
+        run_name: str | None = None,
+        tracking_uri: str | None = None,
+        save_dir: str | None = "./mlruns",
+        log_model: Literal[True, False, "all"] | None = False,
+        prefix: str | None = "",
+        **kwargs,
+    ) -> None:
+        super().__init__(
+            experiment_name=experiment_name,
+            run_name=run_name,
+            tracking_uri=tracking_uri,
+            save_dir=save_dir,
+            log_model=log_model,
+            prefix=prefix,
+            **kwargs,
+        )
+
+    @rank_zero_only
+    def add_image(self, image: np.ndarray | Figure, name: str | None = None, **kwargs) -> None:
+        """Interface to log images in the mlflow loggers.
+
+        Args:
+            image (np.ndarray | Figure): Image to log.
+            name (str | None): The tag of the image defaults to ``None``.
+            kwargs: Additional keyword arguments that are only used if `image` is of type Figure.
+                These arguments are passed directly to the method that saves the figure.
+                If `image` is a NumPy array, `kwargs` has no effect.
+        """
+        # Need to call different functions of `Experiment` for  Figure vs np.ndarray
+        if isinstance(image, Figure):
+            self.experiment.log_figure(run_id=self.run_id, figure=image, artifact_file=name, **kwargs)
+        else:
+            self.experiment.log_image(run_id=self.run_id, image=image, artifact_file=name)