test/models/image_classification_model_test.py

from unittest.mock import Mock

import numpy as np
import tensorflow as tf

from deep_learning.models.image_classification import ImageClassificationModelBuilder


def test_image_classification_model_builder_outputs_binary_probability():
    """验证图片分类模型会为每张图片输出一个二分类概率。"""
    artifact = ImageClassificationModelBuilder(
        image_size=(32, 32),
        model_filters=(8,)
    ).build_training_artifact()
    model = artifact.model
    images = tf.zeros((2, 32, 32, 3), dtype=tf.float32)

    outputs = model(images)

    assert outputs.shape == (2, 1)
    assert np.all(outputs.numpy() >= 0.0)
    assert np.all(outputs.numpy() <= 1.0)


def test_image_classification_model_builder_compiles_training_model():
    """验证图片分类模型构建器会使用二分类训练配置编译模型。"""
    model = Mock()
    builder = ImageClassificationModelBuilder(
        image_size=(32, 32),
        model_filters=(8,)
    )

    builder.compile_training_model(model)

    model.compile.assert_called_once_with(
        optimizer="adam",
        loss="binary_crossentropy",
        metrics=["accuracy"]
    )