iBrokeTheCode/Multimodal_Product_Classification

1# import os2# import pandas as pd3 4# from src.vision_embeddings_tf import get_embeddings_df5 6 7import numpy as np8import pytest9from PIL import Image10from tensorflow.keras.applications import ResNet5011from transformers import TFConvNextV2Model12 13from src.vision_embeddings_tf import FoundationalCVModel, load_and_preprocess_image14 15# Run tests with CPU and not GPU (custom added)16# os.environ["CUDA_VISIBLE_DEVICES"] = "-1"17 18 19####################################################################################################20#################### Test the foundational CV model and image preprocessing ########################21####################################################################################################22@pytest.fixture23def mock_image(tmp_path):24    """25    Fixture to create a mock image for testing.26    """27    img_path = tmp_path / "test_image.jpg"28    img = Image.new("RGB", (300, 300), color="red")29    img.save(img_path)30    return str(img_path)31 32 33def test_load_and_preprocess_image(mock_image):34    """35    Test loading and preprocessing of an image.36    """37    # Test the load_and_preprocess_image function38    img = load_and_preprocess_image(mock_image, target_size=(224, 224))39 40    # Check if the output is a numpy array41    assert isinstance(img, np.ndarray), "Output is not a numpy array"42 43    # Check if the image has the correct shape44    assert img.shape == (224, 224, 3), (45        f"Image shape is {img.shape}, expected (224, 224, 3)"46    )47 48    # Check if the pixel values are in the range [0, 1]49    assert img.min() >= 0 and img.max() <= 1, (50        "Image pixel values are not in the range [0, 1]"51    )52 53 54@pytest.mark.parametrize(55    "backbone, expected_model_class, expected_output_shape",56    [57        ("resnet50", type(ResNet50()), (2048,)),  # Keras ResNet50 with 2048 features58        (59            "convnextv2_tiny",60            TFConvNextV2Model,61            (768,),62        ),  # ConvNeXt V2 Tiny from Hugging Face with 768 features63    ],64)65def test_foundational_cv_model_generic(66    backbone, expected_model_class, expected_output_shape67):68    """69    Generic test for loading a foundational CV model and making predictions.70 71    This test ensures that:72    - The correct backbone model is loaded.73    - The input shape matches the model's requirements (224x224x3).74    - The output embedding shape matches the expected shape for the backbone.75 76    Parameters:77    ----------78    backbone : str79        The name of the model backbone to test.80    expected_model_class : class81        The expected class of the loaded backbone model (e.g., ResNet50 or TFConvNextV2Model).82    expected_output_shape : tuple83        The expected shape of the output embedding vector.84    """85    # Initialize the model with the provided backbone86    model = FoundationalCVModel(backbone=backbone, mode="eval")87 88    # Check if the model is an instance of the expected model class89    assert isinstance(model.base_model, expected_model_class), (90        f"Expected model class {expected_model_class}, got {type(model.model)}"91    )92 93    # Create a batch of random images (2 images of shape 224x224x3)94    batch_images = np.random.rand(2, 224, 224, 3)95 96    # Ensure that the input shape matches the model's input requirements97    assert model.model.input_shape == (None, 224, 224, 3), (98        f"Expected input shape (None, 224, 224, 3), got {model.model.input_shape}"99    )100 101    # Ensure that the output shape matches the expected output shape without using the model.predict method102    output = model.get_output_shape()103 104    assert output == (None, *expected_output_shape), (105        f"Expected output shape (None, {expected_output_shape}), got {output}"106    )107 108 109if __name__ == "__main__":110    pytest.main()111