# AUTOGENERATED! DO NOT EDIT! File to edit: resnet_image_classification_mlflow_model.ipynb (unless otherwise specified).

__all__ = [
    "ROOT_PATH",
    "mlflow_save_path",
    "test_data_path",
    "test_samples",
    "test_samples",
    "test_samples_json",
    "images_df",
    "mlflow_predictions",
]

# Cell
import json
import random
from pathlib import Path

import cloudpickle
import mlflow
import numpy as np
import pandas as pd
from pytorch_lightning.utilities.cloud_io import load as pl_load

from roheboam.engine import get_toolbox_lookup
from roheboam.engine.integrations.mlflow.loading import load_mlflow_pyfunc_model
from roheboam.engine.integrations.mlflow.saving import save_mlflow_pyfunc_model
from roheboam.engine.logger import no_engine_logs
from roheboam.engine.pipeline import Pipeline
from roheboam.engine.pipeline.config_loader import load_config_from_paths
from roheboam.engine.pipeline.constants import VARIABLE_KEY
from roheboam.engine.utils.convenience import is_notebook
from roheboam.engine.utils.saving import save_to_temporary_file
from roheboam.engine.utils.testing import all_methods_in_callback_handler_are_called, arrays_are_same, copy_and_flatten_model_weights
from roheboam.engine.vision.tasks.image_classification.data import create_image_classification_sample
from roheboam.engine.vision.tasks.image_classification.deployment import MLFlowPyfuncImageClassificationModel
from roheboam.engine.vision.tasks.image_classification.loaders import create_image_classification_samples_data_for_label_per_folder

# Cell
ROOT_PATH = Path(globals()["_dh"][0]) if is_notebook() else Path(__file__).parent.parent

# Cell
mlflow_save_path = ROOT_PATH / "models" / "model_for_test" / "mlflow"
with no_engine_logs():
    mlflow_model = mlflow.pyfunc.load_model(str(mlflow_save_path))


# Cell
test_data_path = ROOT_PATH / "data" / "hotdog_not_hotdog_test" / "test"
test_samples = create_image_classification_samples_data_for_label_per_folder(test_data_path, label_map={"not_hot_dog": 0, "hot_dog": 1})
test_samples = [create_image_classification_sample(**s) for s in test_samples]
test_samples_json = {
    "columns": ["images", "shapes"],
    "index": [i for i in range(len(test_samples))],
    "data": [[s.image.data.tolist(), s.image.data.shape] for s in test_samples],
}
images_df = pd.read_json(json.dumps(test_samples_json), orient="split")

# Cell
mlflow_predictions = mlflow_model.predict(images_df)
