object-segmentation/tests/test_image.py

"""
Tests for the Image class.
"""

import pytest
import numpy as np
from pathlib import Path
from src.utils.image import Image, ImageLoadError


class TestImage:
    """Test cases for the Image class."""

    def test_load_nonexistent_file(self):
        """Test loading a non-existent file raises ImageLoadError."""
        with pytest.raises(ImageLoadError):
            Image("nonexistent_file.jpg")

    def test_load_unsupported_format(self, tmp_path):
        """Test loading an unsupported format raises ImageLoadError."""
        # Create a dummy file with unsupported extension
        test_file = tmp_path / "test.txt"
        test_file.write_text("not an image")

        with pytest.raises(ImageLoadError):
            Image(test_file)

    def test_supported_extensions(self):
        """Test that supported extensions are correctly defined."""
        expected_extensions = Image.SUPPORTED_EXTENSIONS
        assert Image.SUPPORTED_EXTENSIONS == expected_extensions

    def test_image_properties(self, tmp_path):
        """Test image properties after loading."""
        # Create a simple test image using numpy and cv2
        import cv2

        test_img = np.zeros((100, 200, 3), dtype=np.uint8)
        test_img[:, :] = [255, 0, 0]  # Blue in BGR

        test_file = tmp_path / "test.jpg"
        cv2.imwrite(str(test_file), test_img)

        # Load the image
        img = Image(test_file)

        # Check properties
        assert img.width == 200
        assert img.height == 100
        assert img.channels == 3
        assert img.format == "jpg"
        assert img.shape == (100, 200, 3)
        assert img.size_bytes > 0
        assert img.is_color()
        assert not img.is_grayscale()

    def test_get_rgb(self, tmp_path):
        """Test RGB conversion."""
        import cv2

        # Create BGR image
        test_img = np.zeros((50, 50, 3), dtype=np.uint8)
        test_img[:, :] = [255, 0, 0]  # Blue in BGR

        test_file = tmp_path / "test_rgb.png"
        cv2.imwrite(str(test_file), test_img)

        img = Image(test_file)
        rgb_data = img.get_rgb()

        # RGB should have red channel at 255
        assert rgb_data[0, 0, 0] == 0  # R
        assert rgb_data[0, 0, 1] == 0  # G
        assert rgb_data[0, 0, 2] == 255  # B (was BGR blue)

    def test_get_grayscale(self, tmp_path):
        """Test grayscale conversion."""
        import cv2

        test_img = np.zeros((50, 50, 3), dtype=np.uint8)
        test_img[:, :] = [128, 128, 128]

        test_file = tmp_path / "test_gray.png"
        cv2.imwrite(str(test_file), test_img)

        img = Image(test_file)
        gray_data = img.get_grayscale()

        assert len(gray_data.shape) == 2  # Should be 2D
        assert gray_data.shape == (50, 50)

    def test_copy(self, tmp_path):
        """Test copying image data."""
        import cv2

        test_img = np.zeros((50, 50, 3), dtype=np.uint8)

        test_file = tmp_path / "test_copy.png"
        cv2.imwrite(str(test_file), test_img)

        img = Image(test_file)
        copied = img.copy()

        # Modify copy
        copied[0, 0] = [255, 255, 255]

        # Original should be unchanged
        assert not np.array_equal(img.data[0, 0], copied[0, 0])

    def test_resize(self, tmp_path):
        """Test image resizing."""
        import cv2

        test_img = np.zeros((100, 100, 3), dtype=np.uint8)

        test_file = tmp_path / "test_resize.png"
        cv2.imwrite(str(test_file), test_img)

        img = Image(test_file)
        resized = img.resize(50, 50)

        assert resized.shape == (50, 50, 3)
        # Original should be unchanged
        assert img.width == 100
        assert img.height == 100

    def test_str_repr(self, tmp_path):
        """Test string representation."""
        import cv2

        test_img = np.zeros((100, 200, 3), dtype=np.uint8)

        test_file = tmp_path / "test_str.jpg"
        cv2.imwrite(str(test_file), test_img)

        img = Image(test_file)

        str_repr = str(img)
        assert "test_str.jpg" in str_repr
        assert "100x200x3" in str_repr
        assert "jpg" in str_repr

        repr_str = repr(img)
        assert "Image" in repr_str
        assert "test_str.jpg" in repr_str