A demonstration of how DVC and MLFlow can be used in the task of data relabeling
"""Dataset."""
import json
import os

from skimage import io
from typing import Tuple, Union, Dict

import tensorflow as tf


def load_class_names() -> Dict[str, str]:
    """Load class names.

    Returns:
        dict: a mapping between class ids and names
    """
    with open("data/class_name_to_id_mapping.txt") as f:
        return json.load(f)


def load_images(dataset_name: str) -> io.ImageCollection:
    """Load images as sklearn-image Image Collection

    Args:
        dataset_name (str): dataset name to load

    Returns:
        io.ImageCollection: image collection
    """
    patterns = [os.path.join("data", dataset_name, class_, "*.jpg") for class_ in
                os.listdir(os.path.join("data", dataset_name))]
    return io.ImageCollection(patterns)


def load_tf_dataset(
        name: str,
        split: bool = True
) -> Union[tf.data.Dataset, Tuple[tf.data.Dataset, tf.data.Dataset]]:
    """Load training dataset and split it into train and validation ones for Keras.

    Note: for the purpose of the exercise paths and sizes are not inferred.

    Returns:
        Union[tf.data.Dataset, Tuple[tf.data.Dataset, tf.data.Dataset]]: a dataset or a pair of train and val datasets
    """
    path = f"data/{name}"
    if split:
        return tf.keras.utils.image_dataset_from_directory(
            path,
            image_size=(128, 128),
            validation_split=0.2,
            subset="both",
            seed=101,
        )
    else:
        return tf.keras.utils.image_dataset_from_directory(
            path,
            image_size=(128, 128),
        )


def relabel_image_in_dataset(dataset_name: str, img_id: str, old_class: str, new_class: str) -> None:
    """Relabel an image by moving it into a different class directory.

    Args:
        dataset_name (str): dataset name
        img_id (str): image to relabel
        old_class (str): old class numeric label
        new_class (str): new class numeric label
    """
    old_path = os.path.join("data", dataset_name, old_class, f"{img_id}.jpg")
    new_path = os.path.join("data", dataset_name, new_class, f"{img_id}.jpg")
    os.rename(old_path, new_path)

bunny_server/dataset.py
d-lowl/bunny-party
"""Module to get DVC specific information.

To date, DVC Python API lacks some things that I'd prefer it to have, such as programmatically getting file versions
(as commit hashes or git tags). Hence we query git directly with https://github.com/gitpython-developers/GitPython.
"""
import datetime
import os
from typing import List

from git import Repo
from pydantic import BaseModel
from typing_extensions import Self


class DVCFileVersion(BaseModel):
    filepath: str
    dot_dvc_filepath: str
    git_revisions: List[str]
    committed_datetime: datetime.datetime

    @classmethod
    def from_filepath(cls, filepath: str) -> Self:
        """Get verbose DVC version for a filepath

        Args:
            filepath (str):

        Returns:
            DVCFileVersion:

        Raises:
             Exception: if the file is not tracked by DVC
        """
        dot_dvc_filepath = filepath + ".dvc"
        if not os.path.exists(dot_dvc_filepath):
            raise Exception(f"{filepath} is not tracked by DVC")

        repo = Repo(".")
        commit_hash = repo.git.rev_list("-1", "HEAD", dot_dvc_filepath)
        commit = repo.commit(commit_hash)
        committed_datetime = commit.committed_datetime
        git_revisions = commit.name_rev.split(" ")

        # Cleanup, see GitPython limitations
        del commit
        del repo

        return cls(
            filepath=filepath,
            dot_dvc_filepath=dot_dvc_filepath,
            git_revisions=git_revisions,
            committed_datetime=committed_datetime
        )



bunny_server/dvc.py
d-lowl/bunny-party
from typing import Dict, List

import numpy as np
import pandas as pd
from skimage import io
from sklearn.decomposition import PCA
from sklearn.metrics import euclidean_distances

from bunny_server.dataset import load_class_names, load_images, relabel_image_in_dataset


def get_vector_database(filepaths: List[str], embedded_images: np.ndarray, class_names: Dict[str, str]) -> pd.DataFrame:
    """Form a vector 'database' using the embeddings and metadata.

    Note: this implementation uses pandas, where's in a real scenario something
    like Chroma would be used directly

    Args:
        filepaths (str): paths to the original image files
        embedded_images (np.ndarray): 2D array of image embeddings
        class_names (Dict[str, str]): a mapping between class ids and names

    Returns:
        pd.DataFrame:
            Columns:
                id (str): image ID
                class_id (str): image's class id
                class_name (str): mapped image's class name
                embedding (np.ndarray): a vector to be used for similarity search
    """

    def get_row_from_filepath(filepath: str) -> pd.Series:
        """Infer class and image filename from its path.

        Returns:
            pd.Series:
                Columns:
                    id (str): image ID
                    class_id (str): image's class id
                    class_name (str): mapped image's class name
        """
        _, dataset_name, class_id, filename = filepath.split("/")
        class_name = class_names[class_id]
        return pd.Series({
            "id": filename.split(".")[0],
            "class_id": class_id,
            "class_name": class_name
        })

    df = pd.DataFrame([get_row_from_filepath(filepath) for filepath in filepaths])
    df["embedding"] = embedded_images.tolist()
    return df


class TerminatorEmbedding:
    """Terminator Embedding for similarity search."""

    _dataset_name: str
    _class_names: Dict[str, str]
    _collection: io.ImageCollection
    _embedding: PCA
    _embedded_images: np.ndarray
    dataframe: pd.DataFrame

    def __init__(self, dataset_name: str) -> None:
        """Initialise Terminator Embedding.

        Note: this embedding class is generated on the fly and stored in-memory due to time limitations
        of the exersice. In the real world, some persistent solution like ChromaDB would have been used.

        Args:
            dataset_name (str): name of the dataset to load
        """
        self._dataset_name = dataset_name
        self._class_names = load_class_names()
        self._collection = load_images(dataset_name)
        self._embedding = PCA(n_components=10)
        no_samples = len(self._collection)
        self._embedded_images = self._embedding.fit_transform(np.reshape(self._collection, (no_samples, -1)))
        self.dataframe = get_vector_database(self._collection.files, self._embedded_images, self._class_names)

    def get_distance(self, id: str) -> pd.DataFrame:
        """Get distance between the image given by ID and each element of the embedded dataset

        Args:
            id (str): image ID

        Returns:
            pd.DataFrame: copy of the dataframe, with `distance: float` column added to it.
        """
        distance = pd.DataFrame(euclidean_distances(
            np.stack(self.dataframe.embedding),
            np.stack(self.dataframe[self.dataframe.id == id].embedding)
        ), columns=["distance"], index=self.dataframe.index)

        return pd.concat([self.dataframe, distance], axis=1)

    def get_n_most_similar(self, img_id: str, n: int, class_id: str) -> pd.DataFrame:
        """Return N images that are most similar to a given image within a class.

        Args:
            img_id (str): image ID to compare with
            n (int): number of returned images
            class_id (str): class ID to search in

        Returns:
            pd.DataFrame: a slice of top N similar images
        """
        df = self.get_distance(img_id)
        return df[(df.id != img_id) & (df.class_id == class_id)].sort_values("distance").head(n)

    def relabel(self, img_id: str, class_id: str) -> None:
        """Relabel an image in-place within the embedding database.

        Args:
            img_id (str): image ID
            class_id (str): class ID
        """
        # Assuming the structure, there's only one entry that satisfies the condition
        old_class_id = self.dataframe[self.dataframe.id == img_id].class_id.iloc[0]
        self.dataframe.loc[self.dataframe.id == img_id, ["class_id"]] = class_id
        self.dataframe.loc[self.dataframe.id == img_id, ["class_name"]] = self._class_names[class_id]
        relabel_image_in_dataset(self._dataset_name, img_id, old_class_id, class_id)

bunny_server/embedding.py
d-lowl/bunny-party
"""Classification model module."""
import mlflow
from typing import Tuple, Dict
from typing_extensions import Self

import tensorflow as tf
from keras import layers, Sequential, losses, Model
from keras.src.callbacks import History
from keras.models import load_model

from bunny_server.dataset import load_tf_dataset
from bunny_server.dvc import DVCFileVersion

MODEL_PATH = "models/classifier.keras"


def get_untrained_model() -> Model:
    """Get a basic sequential model to be fitted.

    Returns:
        Sequential: model to be fitted
    """
    num_classes = 2

    model = Sequential([
        layers.Rescaling(1. / 255, input_shape=(128, 128, 3)),
        layers.Conv2D(16, 3, padding='same', activation='relu'),
        layers.MaxPooling2D(),
        layers.Conv2D(32, 3, padding='same', activation='relu'),
        layers.MaxPooling2D(),
        layers.Conv2D(64, 3, padding='same', activation='relu'),
        layers.MaxPooling2D(),
        layers.Flatten(),
        layers.Dense(128, activation='relu'),
        layers.Dense(num_classes)
    ])

    model.compile(
        optimizer='adam',
        loss=losses.SparseCategoricalCrossentropy(from_logits=True),
        metrics=['accuracy']
    )

    print(model.summary())

    return model


class TerminatorClassifier:
    """Terminator Classifier."""
    def __init__(self, model: Model) -> None:
        """Initialise Terminator classifier.

        Args:
            model (Model): tensorflow model for classification
        """
        self.model = model

    def train(self, train_dataset: tf.data.Dataset, val_dataset: tf.data.Dataset, epochs: int) -> History:
        """Fit a sequential model.

        Args:
            train_dataset (tf.data.Dataset): train set
            val_dataset (tf.data.Dataset): validation set
            epochs (int): number of epochs to fit for

        Returns:
            history of the fitting
        """
        history = self.model.fit(
            train_dataset,
            validation_data=val_dataset,
            epochs=epochs
        )

        return history

    def evaluate(self, test_dataset: tf.data.Dataset) -> Tuple[float, float]:
        """Evaluate model performance.

        Args:
            test_dataset (tf.data.Dataset): test dataset

        Returns:
            Tuple[float, float]: test loss and accuracy
        """
        return self.model.evaluate(test_dataset)

    @classmethod
    def get_untrained(cls) -> Self:
        """Get untrained Terminator Classifier.

        Returns:
            TerminatorClassifier: untrained classifier
        """
        return cls(get_untrained_model())

    @classmethod
    def from_path(cls, path: str) -> Self:
        """Initialise model from a previously trained one.

        Args:
            path (str): path to the saved model

        Return:
            TerminatorClassifier: trained classifier
        """
        return cls(load_model(path))

    def save(self, path: str) -> None:
        """Save model to the filesystem.

        Args:
            path (str): path to save to
        """
        self.model.save(path)


def run_model_training_experiment() -> None:
    """Run end to end experiment and log everything to MLFlow."""
    with mlflow.start_run() as run:
        mlflow.autolog()

        # Log dataset versions
        train_version = DVCFileVersion.from_filepath("data/train")
        mlflow.log_param("train_git_revisions", train_version.git_revisions)
        mlflow.log_param("train_committed_datetime", train_version.committed_datetime)

        test_version = DVCFileVersion.from_filepath("data/test")
        mlflow.log_param("test_git_revisions", test_version.git_revisions)
        mlflow.log_param("test_committed_datetime", test_version.committed_datetime)

        train_dataset, val_dataset = load_tf_dataset("train", split=True)
        test_dataset = load_tf_dataset("test", split=False)
        classifier = TerminatorClassifier.get_untrained()
        history = classifier.train(train_dataset, val_dataset, epochs=10)
        test_loss, test_accuracy = classifier.evaluate(test_dataset)
        mlflow.log_metric("test_loss", test_loss)
        mlflow.log_metric("test_accuracy", test_accuracy)
        classifier.save(MODEL_PATH)


def get_model_evaluation() -> Dict[str, float]:
    """Evaluate previously trained and saved model.

    Returns:
        dict: containing evaluation metrics
    """
    classifier = TerminatorClassifier.from_path(MODEL_PATH)
    test_dataset = load_tf_dataset("test", split=False)
    test_loss, test_accuracy = classifier.evaluate(test_dataset)
    return {
        "test_loss": test_loss,
        "test_accuracy": test_accuracy
    }


bunny_server/model.py
d-lowl/bunny-party
from typing import Dict, Any, List

import uvicorn
from fastapi import FastAPI
from pydantic import BaseModel

from bunny_server.embedding import TerminatorEmbedding
from bunny_server.model import run_model_training_experiment, get_model_evaluation

app = FastAPI()
embedding = TerminatorEmbedding("train")


@app.post("/train_model")
def train_model() -> None:
    """Run the model training experiment."""
    run_model_training_experiment()
    return


@app.get("/evaluate_model")
def evaluate_model() -> Dict[str, float]:
    """Evaluate the current model using the test dataset.

    Returns:
        dict: of evaluation metrics
    """
    return get_model_evaluation()


@app.get("/similar_images")
def get_similar_images(img_id: str, class_id: str, n: int = 3) -> List[Dict[str, Any]]:
    """Get images similar to the given

    Args:
        img_id (str): image ID to search for
        class_id (str): class label ID
        n (int): number of images to return back

    Returns:
        list: of dictionaries describing the similar images
    """
    similar_images = embedding.get_n_most_similar(img_id, n, class_id)
    return similar_images.drop("embedding", axis=1).to_dict(orient="records")


class RelabelRequest(BaseModel):
    """Relabel request data."""
    img_id: str
    class_id: str


@app.post("/relabel")
def relabel(request: RelabelRequest) -> None:
    """Relabeling endpoint."""
    embedding.relabel(img_id=request.img_id, class_id=request.class_id)



if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)

bunny_server/server.py
d-lowl/bunny-party
import hashlib
import json
import os
from typing import List

from PIL import Image
import random
from tqdm import tqdm

SIZE = (128, 128)

TRAIN_SAMPLE_NUMBER = 30
TEST_SAMPLE_NUMBER = 10

ORIGINAL_DATA_LOCATION = os.path.join(os.path.dirname(__file__), "data")
DATA_LOCATION = os.path.join(os.path.dirname(__file__), os.pardir, "data")


DOG = [
    "alvan-nee-brFsZ7qszSY-unsplash.jpg",
    "charlesdeluvio-Mv9hjnEUHR4-unsplash.jpg",
    "marcus-wallis-4zfacTKyZ7w-unsplash.jpg",
    "milli-2l0CWTpcChI-unsplash.jpg",
]

BUNNY = [
    "erik-jan-leusink---SDX4KWIbA-unsplash.jpg",
    "janan-lagerwall-302xfiIGOfE-unsplash.jpg",
    # "melanie-kreutz-IFnknR2Mv5o-unsplash.jpg",
    "satyabratasm-u_kMWN-BWyU-unsplash.jpg",
]

FAKE_BUNNY = "lovely-pet-portrait-isolated.jpg"

CLASSNAMES = {
    "0": "bunny",
    "1": "impostor"
}


def prepare(img: Image) -> Image:
    """Generate a random augmented image by rotation and resizing."""
    angle = random.uniform(0.0, 360.0)
    return img.rotate(angle, expand=True).resize(SIZE)


def get_name(img: Image) -> str:
    """Get md5 name for the new image."""
    return hashlib.md5(img.tobytes()).hexdigest()


def generate_dataset(
        dataset_name: str,
        class_label: str,
        image_names: List[str],
        samples_per_image: int,
) -> None:
    """Generate and write dataset.

    The dataset is written under: data/{dataset_name}/{class_label}

    Args:
        dataset_name (str): Name of the dataset (e.g. train, test).
        class_label (str): Numeric class label, see CLASSNAMES
        image_names (List[str]): List of the original images to generate samples
        samples_per_image (int): number of samples to generate from a given image
    """
    if not os.path.exists(os.path.join(DATA_LOCATION, dataset_name)):
        os.mkdir(os.path.join(DATA_LOCATION, dataset_name))
    if not os.path.exists(os.path.join(DATA_LOCATION, dataset_name, class_label)):
        os.mkdir(os.path.join(DATA_LOCATION, dataset_name, class_label))
    for image_name in tqdm(image_names):
        original = Image.open(os.path.join(ORIGINAL_DATA_LOCATION, "original", image_name))
        for _ in tqdm(range(samples_per_image)):
            prepared = prepare(original)
            img_name = get_name(prepared)
            prepared.save(os.path.join(DATA_LOCATION, dataset_name, class_label, f"{img_name}.jpg"))


if __name__ == "__main__":
    with open(os.path.join(DATA_LOCATION, "class_name_to_id_mapping.txt"), "w") as f:
        json.dump(CLASSNAMES, f, indent=4)

    random.seed(10)

    generate_dataset("train", "0", BUNNY + [FAKE_BUNNY], TRAIN_SAMPLE_NUMBER)
    generate_dataset("train", "1", DOG, TRAIN_SAMPLE_NUMBER)

    generate_dataset("test", "0", BUNNY, TEST_SAMPLE_NUMBER)
    generate_dataset("test", "1", DOG + [FAKE_BUNNY], TEST_SAMPLE_NUMBER)

demo_preparation/prepare_dataset.py
d-lowl/bunny-party
import os.path
from unittest import mock

from bunny_server.dataset import relabel_image_in_dataset, load_tf_dataset, load_class_names, load_images


def test_load_dataset():
    with mock.patch("tensorflow.keras.utils.image_dataset_from_directory") as mock_dataset_loader:
        load_tf_dataset("dataset_name", split=False)
        mock_dataset_loader.assert_called_with("data/dataset_name", image_size=(128, 128))


def test_load_split_dataset():
    with mock.patch("tensorflow.keras.utils.image_dataset_from_directory") as mock_dataset_loader:
        load_tf_dataset("dataset_name", split=True)
        mock_dataset_loader.assert_called_with(
            "data/dataset_name",
            image_size=(128, 128),
            validation_split=0.2,
            subset="both",
            seed=101,
        )


def test_relabel_image_in_dataset():
    with mock.patch("os.rename") as mock_rename:
        relabel_image_in_dataset(
            "test_dataset",
            "image0",
            "1",
            "0"
        )
        mock_rename.assert_called_with(
            "data/test_dataset/1/image0.jpg",
            "data/test_dataset/0/image0.jpg"
        )


def test_load_class_names():
    """Test loading class names."""
    path = "data/class_name_to_id_mapping.txt"  # assumed
    with (mock.patch("builtins.open") as mock_open, mock.patch("json.load") as mock_json_load):
        expected_class_names = {
            "0": "enemy",
            "1": "friend"
        }
        mock_json_load.return_value = expected_class_names
        class_names = load_class_names()
        assert expected_class_names == class_names
        mock_open.assert_called_with(path)


def test_load_images():
    """Test loading image collection."""
    dataset_name = "test_dataset"
    dataset_path = os.path.join("data", dataset_name)
    expected_patterns = [
        "data/test_dataset/0/*.jpg",
        "data/test_dataset/1/*.jpg",
    ]
    with (
        mock.patch("server.dataset.io.ImageCollection") as mock_image_collection,
        mock.patch("os.listdir", return_value=["0", "1"]) as mock_listdir
    ):
        collection = load_images(dataset_name)
        assert collection == mock_image_collection.return_value
        mock_image_collection.assert_called_with(expected_patterns)
        mock_listdir.assert_called_with(dataset_path)


tests/test_dataset.py
d-lowl/bunny-party
import datetime
from unittest import mock
from contextlib import nullcontext as does_not_raise

import pytest

from bunny_server.dvc import DVCFileVersion


@pytest.fixture
def mock_commit() -> mock.MagicMock:
    """Mock a basic commit."""
    with mock.patch("git.Commit") as mock_commit_constructor:
        mock_commit = mock_commit_constructor.return_value
        mock_commit.committed_datetime = datetime.datetime(2010, 11, 15, 10, 10, 10)
        mock_commit.name_rev = "deadbeef"
        yield mock_commit


@pytest.fixture
def mock_commit_with_tag(mock_commit: mock.MagicMock) -> mock.MagicMock:
    """Mock a commit with a tag."""
    mock_commit.name_rev += " tags/v1"
    yield mock_commit


@pytest.fixture
def mock_git_repo(mock_commit: mock.MagicMock) -> mock.MagicMock:
    """Mock a git repo."""
    expected_git_hash = "deadbeef"
    with mock.patch("server.dvc.Repo") as mock_repo_constructor:
        mock_repo = mock_repo_constructor.return_value
        mock_repo.git_dir = "."
        mock_repo.git.rev_list.return_value = expected_git_hash
        mock_repo.commit.return_value = mock_commit
        yield mock_repo


@pytest.fixture
def mock_git_repo_with_tag(mock_git_repo: mock.MagicMock, mock_commit_with_tag: mock.MagicMock) -> mock.MagicMock:
    """Mock a git repo with a tagged commit."""
    mock_git_repo.commit.return_value = mock_commit_with_tag
    yield mock_git_repo


def test_from_filepath_not_tracked():
    """Should throw exception if the file is not tracked."""
    with mock.patch("os.path.exists") as mock_exists:
        mock_exists.return_value = False
        with pytest.raises(Exception):
            version = DVCFileVersion.from_filepath("test/file/path")


def test_from_filepath_commit_hash_only(mock_git_repo: mock.MagicMock):
    """Test the DVC file version is constructed correctly."""
    expected = DVCFileVersion(
        filepath="test/file/path",
        dot_dvc_filepath="test/file/path.dvc",
        git_revisions=["deadbeef"],
        committed_datetime=datetime.datetime(2010, 11, 15, 10, 10, 10)
    )
    with mock.patch("os.path.exists") as mock_exists:
        mock_exists.return_value = True
        with does_not_raise():
            version = DVCFileVersion.from_filepath("test/file/path")
            assert expected == version


def test_from_filepath_commit_with_tag(mock_git_repo_with_tag: mock.MagicMock):
    """Test the DVC file version is constructed correctly."""
    expected = DVCFileVersion(
        filepath="test/file/path",
        dot_dvc_filepath="test/file/path.dvc",
        git_revisions=["deadbeef", "tags/v1"],
        committed_datetime=datetime.datetime(2010, 11, 15, 10, 10, 10)
    )
    with mock.patch("os.path.exists") as mock_exists:
        mock_exists.return_value = True
        with does_not_raise():
            version = DVCFileVersion.from_filepath("test/file/path")
            assert expected == version

tests/test_dvc.py
d-lowl/bunny-party
from typing import Dict, List
from unittest import mock

import numpy as np
import pandas as pd
import pytest

from bunny_server.embedding import get_vector_database, TerminatorEmbedding


@pytest.fixture
def dummy_file_paths() -> List[str]:
    """Dummy file paths."""
    return [
        "data/dataset/0/1.jpg",
        "data/dataset/0/2.jpg",
        "data/dataset/1/3.jpg",
    ]


@pytest.fixture
def class_names() -> Dict[str, str]:
    """Class names mapping for testing."""
    return {
        "0": "bunny",
        "1": "impostor",
    }


@pytest.fixture
def dummy_embedded_images() -> np.ndarray:
    """Dummy embedded images.

    Returns:
        np.ndarray: of 3 images by 5 dimensions
    """
    return np.array([
        [-0.1, -0.2, -0.3, -0.4, -0.5],
        [0.1, 0.2, 0.3, 0.4, 0.5],
        [0.2, 0.2, 0.2, 0.2, 0.2],
    ])


@pytest.fixture
def expected_dataframe() -> pd.DataFrame:
    """Construct a DataFrame expected in embedding tests."""
    return pd.DataFrame([
        {
            "id": "1",
            "class_id": "0",
            "class_name": "bunny",
            "embedding": np.array([-0.1, -0.2, -0.3, -0.4, -0.5]),
        },
        {
            "id": "2",
            "class_id": "0",
            "class_name": "bunny",
            "embedding": np.array([0.1, 0.2, 0.3, 0.4, 0.5]),
        },
        {
            "id": "3",
            "class_id": "1",
            "class_name": "impostor",
            "embedding": np.array([0.2, 0.2, 0.2, 0.2, 0.2]),
        }
    ])


@pytest.fixture
def mock_terminator_embedding(
        dummy_file_paths: List[str],
        dummy_embedded_images: np.ndarray
) -> TerminatorEmbedding:
    """Construct a partially mocked embedding instance."""
    with (
        mock.patch("server.embedding.load_images") as mock_load_images,
        mock.patch("server.embedding.PCA") as mock_pca,
    ):
        mock_collection = mock_load_images.return_value
        mock_collection.__len__.return_value = 3
        mock_collection.files = dummy_file_paths
        mock_embedding = mock_pca.return_value
        mock_embedding.fit_transform.return_value = dummy_embedded_images
        embedding = TerminatorEmbedding("dummy")
        yield embedding


def test_get_vector_database(
        dummy_file_paths: List[str],
        dummy_embedded_images: np.ndarray,
        class_names: Dict[str, str],
        expected_dataframe: pd.DataFrame,
):
    """Test creating a pandas based 'vector database'."""
    got_df = get_vector_database(dummy_file_paths, dummy_embedded_images, class_names)
    pd.testing.assert_frame_equal(expected_dataframe, got_df)


def test_terminator_embedding_constructor(
        mock_terminator_embedding: TerminatorEmbedding,
        expected_dataframe: pd.DataFrame,
):
    """Test TerminatorEmbedding being constructed correctly."""
    # We expect the dataframe within the embedding instance to be equal to the manually constructed one
    pd.testing.assert_frame_equal(expected_dataframe, mock_terminator_embedding.dataframe)


def test_get_distance(
        mock_terminator_embedding: TerminatorEmbedding,
        expected_dataframe: pd.DataFrame
):
    """Test distance calculation."""
    expected_distances = [  # Approximately
        1.161895,
        0.38729833,
        0.0
    ]
    expected_distance_df = expected_dataframe
    expected_dataframe["distance"] = expected_distances
    distance_df = mock_terminator_embedding.get_distance("3")
    pd.testing.assert_frame_equal(expected_distance_df, distance_df)


def test_get_n_most_similar(
        mock_terminator_embedding: TerminatorEmbedding,
        expected_dataframe: pd.DataFrame
):
    """Test getting N most similar images."""
    similar_df = mock_terminator_embedding.get_n_most_similar("3", 3, "0")
    expected_distances = [  # Approximately
        0.38729833,
        1.161895,
    ]
    expected_similar_df = expected_dataframe.iloc[[1, 0]]
    expected_similar_df["distance"] = expected_distances
    pd.testing.assert_frame_equal(expected_similar_df, similar_df)
tests/test_embedding.py
d-lowl/bunny-party
from unittest import mock

import pytest
import tensorflow_datasets as tfds

from bunny_server.model import get_untrained_model, TerminatorClassifier


@pytest.fixture
def mock_dataset() -> mock.MagicMock:
    """Mock TF dataset."""
    with mock.patch("tensorflow.data.Dataset") as mock_dataset_class:
        mock_dataset = mock_dataset_class.return_value
        yield mock_dataset


@pytest.fixture
def mock_tf_model() -> mock.MagicMock:
    """Mock TF model instance."""
    with mock.patch("server.model.Sequential") as model_constructor:
        yield model_constructor.return_value


def test_get_untrained_model():
    """Test that a Sequential model is constructed and compiled"""
    with mock.patch("server.model.Sequential") as mock_sequential:
        returned_model = get_untrained_model()
        mock_sequential.assert_called()
        mock_model = mock_sequential.return_value
        mock_model.compile.assert_called()
        assert mock_model == returned_model


def test_terminator_classifier(mock_tf_model: mock.MagicMock):
    """Test the initialiser."""
    classifier = TerminatorClassifier(mock_tf_model)
    assert classifier.model == mock_tf_model


def test_terminator_classifier_train(
        mock_tf_model: mock.MagicMock,
        mock_dataset: mock.MagicMock
):
    """Test the training procedure."""
    classifier = TerminatorClassifier(mock_tf_model)
    with tfds.testing.mock_data(num_examples=5):
        classifier.train(
            mock_dataset,
            mock_dataset,
            5
        )
        mock_tf_model.fit.assert_called_with(
            mock_dataset,
            validation_data=mock_dataset,
            epochs=5
        )


def test_evaluate(
        mock_tf_model: mock.MagicMock,
        mock_dataset: mock.MagicMock
):
    """Test model evaluation method."""
    expected_metrics = (1.5, 0.5)  # Set mock loss and accuracy
    mock_tf_model.evaluate.return_value = expected_metrics
    classifier = TerminatorClassifier(mock_tf_model)
    got_metrics = classifier.evaluate(mock_dataset)
    assert expected_metrics == got_metrics
    mock_tf_model.evaluate.assert_called_with(mock_dataset)


def test_save(mock_tf_model: mock.MagicMock):
    """Test model saving method"""
    classifier = TerminatorClassifier(mock_tf_model)
    classifier.save("test/model/path.keras")
    mock_tf_model.save.assert_called_with("test/model/path.keras")


def test_from_path():
    """Test initialising Terminator Classifier with the previously saved model"""
    path = "test/model/path.keras"
    with mock.patch("server.model.load_model") as mock_load_model:
        classifier = TerminatorClassifier.from_path(path)
        assert classifier.model == mock_load_model.return_value
        mock_load_model.assert_called_with(path)
tests/test_model.py
d-lowl/bunny-party