Repository: Project-MONAI/MONAI
File: monai/handlers/mlflow_handler.py (Python)

Code before:
        This function will log a PandasDataset, generated from the Pandas DataFrame, to MLFlow
        inputs. For more details about PandasDataset, please refer to this link:
        https://mlflow.org/docs/latest/python_api/mlflow.data.html#mlflow.data.pandas_dataset.PandasDataset

        Please note that it may take a while to record the dataset if it has too many samples.

        Args:
            engine: Ignite Engine, it can be a trainer, validator or evaluator.

        """
        dataloader = getattr(engine, "data_loader", None)
        dataset = getattr(dataloader, "dataset", None) if dataloader else None
        if not dataset:
            raise AttributeError("The dataset of the engine is None. Cannot record it by MLFlow.")

        sample_dict: dict[str, list[str]] = {}
        sample_dict["images"] = []
        for sample in dataset:
            if isinstance(sample, dict):
                image_name = sample["image_meta_dict"]["filename_or_obj"] if "image_meta_dict" in sample else None
            elif isinstance(sample, list):
                # When using a transform like `RandCropByPosNegLabel`, a sample will be a list containing image slices.
                image_name = sample[0]["image_meta_dict"]["filename_or_obj"] if "image_meta_dict" in sample[0] else None
            else:
                image_name = None
                warnings.warn(f"Don't support {type(sample)} type samples when recording the dataset with MLFlow.")

            if not isinstance(image_name, str):
                warnings.warn(
                    f"Expected type string, got type {type(image_name)} of the image name."
                    "May log an empty dataset in MLFlow"
                )
            else:
                sample_dict["images"].append(image_name)
        dataset_type = "train" if isinstance(engine, Trainer) else "nontrain"
        self._log_dataset(sample_dict, dataset_type)


Review comment (refactor):
this could be a generic utility function to guess the filename, please also consider metatensor case, for example https://github.com/Project-MONAI/research-contributions/blob/689d595503822a0622a1e905352dd692d50a90d6/auto3dseg/algorithm_templates/swinunetr/scripts/train.py#L541-L544 and put it in monai.utils

Code after:
        like:
            {
                "dataset_name0": dataset0,
                "dataset_name1": dataset1,
                ......
            }
        The keys stand for names of datasets, which will be logged as prefixes of dataset names in MLFlow.
        The values are PyTorch datasets from which sample names are abstracted to build a Pandas DataFrame.
        If the input dataset_dict is None, this function will directly return and do nothing.

        To use this function, every sample in the input datasets must contain keys specified by the `dataset_keys`
        parameter.
        This function will log a PandasDataset to MLFlow inputs, generated from the Pandas DataFrame.
        For more details about PandasDataset, please refer to this link:
        https://mlflow.org/docs/latest/python_api/mlflow.data.html#mlflow.data.pandas_dataset.PandasDataset

        Please note that it may take a while to record the dataset if it has too many samples.

        Args:
            dataset_dict: a dictionary in which the key is the name of the dataset and the value is a PyTorch
                dataset, that needs to be recorded.

        """

        if dataset_dict is None:
            return
        elif len(dataset_dict) == 0:
            warnings.warn("There is no dataset to log!")

        # Log datasets to MLFlow one by one.
        for dataset_type, dataset in dataset_dict.items():
            if dataset is None:
                raise AttributeError(f"The {dataset_type} dataset of is None. Cannot record it by MLFlow.")

            sample_dict: dict[str, list[str]] = {}
            dataset_samples = getattr(dataset, "data", [])
            for sample in tqdm(dataset_samples, f"Recording the {dataset_type} dataset"):
                for key in self.dataset_keys:
                    if key not in sample_dict:
                        sample_dict[key] = []

                    if key in sample:
                        value_to_log = sample[key]
                    else:
                        raise KeyError(f"Unexpect key '{key}' in the sample.")

                    if not isinstance(value_to_log, str):
                        warnings.warn(
                            f"Expected type string, got type {type(value_to_log)} of the {key} name."
                            "May log an empty dataset in MLFlow"
                        )