crfm-helm 0.4.0__py3-none-any.whl → 0.5.10__py3-none-any.whl

helm/benchmark/metrics/vision_language/emd_utils.py

@@ -0,0 +1,347 @@
+from typing import List, Tuple
+from tqdm import tqdm
+
+import numpy as np
+import math
+
+from helm.common.optional_dependencies import handle_module_not_found_error
+
+try:
+    import cv2
+    from PIL import Image
+except ModuleNotFoundError as e:
+    handle_module_not_found_error(e, suggestions=["images"])
+
+
+def to_gray(img: np.ndarray) -> np.ndarray:
+    return np.matmul(img, np.array([[0.299], [0.587], [0.114]]))
+
+
+def get_most_frequent_color(img: np.ndarray) -> Tuple[np.ndarray, float]:
+    """Get the most frequent color in the image and its frequency.
+
+    Args:
+    img (np.array): Input image array of shape (height, width, channels).
+
+    Returns:
+    Tuple[np.array, float]: Most frequent color and its frequency as a percentage of the total number of pixels.
+    """
+    # Assert to ensure input is a 3D numpy array
+    assert len(img.shape) == 3, "Input image must be a 3D numpy array"
+
+    # Reshape image array to 2D (pixel, RGB)
+    pixels = img.reshape(-1, img.shape[2])
+
+    # Find unique rows (colors) and their counts
+    unique_colors, counts = np.unique(pixels, axis=0, return_counts=True)
+
+    # Find the index of the most frequent color
+    most_frequent_color_index = np.argmax(counts)
+
+    # Most frequent color
+    most_frequent_color = unique_colors[most_frequent_color_index]
+
+    # Calculate frequency percentage
+    frequency = counts[most_frequent_color_index] / pixels.shape[0]
+
+    return most_frequent_color, frequency
+
+
+def img_to_sig_patches(
+    img: np.ndarray,
+    rgb_most_frequent_color: np.ndarray,
+    patch_size: Tuple[int, int],
+    weight_most_frequent_color: float = 0.01,
+):
+    """
+    Convert an RGB image to a signature for cv2.EMD, processing the image in patches.
+
+    Args:
+    - img: A 3D numpy array representing an RGB image (height, width, channels).
+    - rgb_most_frequent_color: The most frequent color in the image.
+    - patch_size: Tuple indicating the height and width of the patches.
+    - weight_most_frequent_color: The weight assigned to the most frequent color in the patches.
+
+    Returns:
+    - A numpy array suitable for cv2.EMD, containing color values and coordinates of each patch.
+        The shape is (num_patches, patch_size[0] * patch_size[1] + 3).
+    """
+    assert len(img.shape) == 3, "Input image must be a 3D numpy array"
+
+    # Ensure img is a numpy array of type float32
+    img = np.array(img, dtype=np.float32)
+
+    # Determine padding needs
+    pad_height = (-img.shape[0]) % patch_size[0]
+    pad_width = (-img.shape[1]) % patch_size[1]
+
+    # Adjust padding for RGB channels
+    padding = ((0, pad_height), (0, pad_width), (0, 0))
+    pad_values = (
+        (rgb_most_frequent_color[0], rgb_most_frequent_color[0]),
+        (rgb_most_frequent_color[1], rgb_most_frequent_color[1]),
+        (rgb_most_frequent_color[2], rgb_most_frequent_color[2]),
+    )
+
+    # Find the most frequent color for padding
+    if pad_height > 0 or pad_width > 0:
+        img = np.pad(img, padding, "constant", constant_values=pad_values)
+    img /= 255.0  # Normalize colors to [0, 1]
+
+    # Collapse color dimensions to grayscale
+    img = to_gray(img)
+
+    # Reshape image into patches and flatten the color dimensions within each patch
+    patches = (
+        img.reshape(
+            (img.shape[0] // patch_size[0], patch_size[0], img.shape[1] // patch_size[1], patch_size[1], img.shape[2])
+        )
+        .transpose(0, 2, 1, 3, 4)
+        .reshape(-1, *patch_size, img.shape[2])
+    )
+
+    # Calculate patch positions
+    patch_positions = (
+        np.mgrid[0 : img.shape[0] // patch_size[0], 0 : img.shape[1] // patch_size[1]].transpose(1, 2, 0).reshape(-1, 2)
+    )
+
+    # Normalize positions
+    patch_positions = patch_positions / np.array([img.shape[0] // patch_size[0], img.shape[1] // patch_size[1]])
+
+    # Compute the weight of each patch
+    # The weight of each point is 1 if the color is not the most frequent color, weight_most_frequent_color otherwise
+    flattened_patches = patches.reshape(patches.shape[0], -1)
+    gray_most_frequent_color: float = float(to_gray(rgb_most_frequent_color).squeeze() / 255.0)
+    weight = weight_most_frequent_color + (1 - weight_most_frequent_color) * np.any(
+        flattened_patches != gray_most_frequent_color, axis=1, keepdims=True
+    ).astype(np.float32)
+    weight /= np.sum(weight)
+
+    # Flatten patches and concatenate with their normalized positions and weights
+    sig = np.hstack((weight, flattened_patches, patch_positions))
+
+    return sig.astype(np.float32)
+
+
+def pad(small_image: Image.Image, large_image: Image.Image, axis: int) -> Image.Image:
+    """Pad the axis of the small image to match the size of the large image."""
+    new_dim: List[int] = list(small_image.size)
+    new_dim[axis] = large_image.size[axis]
+    new_dim_tupe: Tuple[int, int] = tuple(new_dim)  # type: ignore
+    new_image: Image.Image = Image.new("RGB", new_dim_tupe, (255, 255, 255))
+    new_image.paste(small_image, (0, 0))
+    return new_image
+
+
+def reshape_sub_sig_batch(
+    sub_sigs: np.ndarray,
+    patch_size: Tuple[int, int],
+    gray_most_frequent_color: float,
+    weight_most_frequent_color: float = 0.01,
+) -> np.ndarray:
+    """
+    Reshape a patch-based signature of an image (Shape: (num_patches, patch_size[0] * patch_size[1] + 3))
+    to a batch of signatures for each patch (Shape: (num_patches, patch_size[0] * patch_size[1], 4)).
+    Basically goes from a signature on the patch level to a batch of signatures on the pixel level.
+
+    Args:
+    - sub_sigs: A numpy array of shape (num_patches, patch_size[0] * patch_size[1] + 1) representing the
+        sub-signatures. (the spatial info should have been stripped).
+    - patch_size: Tuple indicating the height and width of the patches.
+    - gray_most_frequent_color: The most frequent color in the image.
+        This is used to reduce the weight assigned to the most frequent color in the patches.
+    - weight_most_frequent_color: The weight assigned to the most frequent color in the patches.
+
+    Returns:
+    - A numpy array of shape (num_patches, patch_size[0] * patch_size[1], 4) representing the sub-signatures of
+        each patch (pixel-level signatures).
+    """
+    # Ensure sub_sigs has the correct shape
+    num_patches = sub_sigs.shape[0]
+    flat_patch_size = patch_size[0] * patch_size[1]
+    assert sub_sigs.shape[1] == flat_patch_size + 1, f"Expected {flat_patch_size + 1} columns, got {sub_sigs.shape[1]}."
+
+    # Ensure sub_sigs is reshaped to include an extra dimension for concatenation
+    num_channels: int = int(round(sub_sigs.shape[0] * sub_sigs.shape[1] / (num_patches * flat_patch_size)))
+    assert num_channels == 1, "Only grayscale images are supported for now."
+    sub_sigs_reshaped = sub_sigs[:, 1:].reshape(num_patches, flat_patch_size, num_channels)
+
+    # Generate spatial information
+    x = np.arange(patch_size[0]) / patch_size[0]
+    y = np.arange(patch_size[1]) / patch_size[1]
+    x, y = np.meshgrid(x, y)
+    spatial_info = np.stack((x.ravel(), y.ravel()), axis=1)  # Shape: (flat_patch_size, 2)
+
+    # Repeat spatial_info for each patch
+    spatial_info_repeated = np.repeat(
+        spatial_info[np.newaxis, :, :], num_patches, axis=0
+    )  # Shape: (num_patches, flat_patch_size, 2)
+
+    # The weight of each point is 1 if the color is not the most frequent color, weight_most_frequent_color otherwise
+    # The weight of a pixel is the product of the weight of the patch and the weight of the pixel in the patch
+    local_weights = weight_most_frequent_color + (1 - weight_most_frequent_color) * (
+        sub_sigs_reshaped != gray_most_frequent_color
+    ).astype(np.float32)
+    global_weights = sub_sigs[:, 0:1]
+    local_weights *= global_weights.reshape(-1, 1, 1)
+    local_weights /= np.sum(local_weights, axis=1, keepdims=True)
+
+    # Concatenate sub_sigs with weights and spatial information
+    sub_sigs_with_spatial_info = np.concatenate(
+        (local_weights, sub_sigs_reshaped, spatial_info_repeated), axis=2
+    )  # Shape: (num_patches, flat_patch_size, 4)
+
+    return sub_sigs_with_spatial_info
+
+
+def compute_cost_matrix_on_sig(
+    sig1: np.ndarray,
+    sig2: np.ndarray,
+    gray_most_frequent_color: float,
+    patch_size: Tuple[int, int],
+    dim: Tuple[int, int],
+    weight_most_frequent_color: float = 0.01,
+    use_tqdm: bool = True,
+) -> np.ndarray:
+    """
+    Compute the cost matrix for the EMD between two signatures with pre-reshaping optimization.
+
+    Args:
+    - sig1: A numpy array of shape (num_patches, patch_size[0] * patch_size[1] + 2) representing the first signature.
+    - sig2: A numpy array of shape (num_patches, patch_size[0] * patch_size[1] + 2) representing the second signature.
+    - gray_most_frequent_color: The most frequent color in the images, used to filter out patches that are constant
+        equal to the most frequent color.
+    - patch_size: Tuple indicating the height and width of the patches.
+    - use_tqdm: Boolean indicating whether to display a progress bar.
+
+    Returns:
+    - A numpy array of shape (num_patches, num_patches) representing the cost matrix.
+    """
+    assert sig1.shape == sig2.shape
+
+    # Reshape the sub-signatures at the beginning
+    sig1_reshaped = reshape_sub_sig_batch(
+        sig1[:, :-2], patch_size, gray_most_frequent_color, weight_most_frequent_color
+    ).astype(np.float32)
+    sig2_reshaped = reshape_sub_sig_batch(
+        sig2[:, :-2], patch_size, gray_most_frequent_color, weight_most_frequent_color
+    ).astype(np.float32)
+
+    cost_matrix = np.zeros((sig1.shape[0], sig2.shape[0]))
+    multiplier: float = (patch_size[0] * patch_size[1]) ** 0.5 / (dim[0] + dim[1])
+    for i in tqdm(range(sig1.shape[0]), disable=not use_tqdm):
+        for j in range(sig2.shape[0]):
+            pos_sig1 = sig1[i, -2:]
+            pos_sig2 = sig2[j, -2:]
+            sub_sig1 = sig1_reshaped[i]
+            sub_sig2 = sig2_reshaped[j]
+            emd_value, _, _ = cv2.EMD(sub_sig1, sub_sig2, cv2.DIST_L1)
+            cost_matrix[i, j] = emd_value + np.linalg.norm(pos_sig1 - pos_sig2, 1) * multiplier  # Use L1
+    return cost_matrix.astype(np.float32)
+
+
+def compute_emd_recursive(
+    img1_PIL: Image.Image,
+    img2_PIL: Image.Image,
+    threshold_most_frequent_color: float = 0.5,
+    patch_size: Tuple[int, int] = (8, 8),
+    max_num_patches: int = 100,
+    weight_most_frequent_color: float = 0.001,
+    use_tqdm: bool = False,
+):
+    """
+    Compute the Earth Mover's Distance between two images using a recursive approach.
+    Both images are discretized into patches, and the EMD is computed on the patches.
+    This is done by computing a cost matrix C such that C[i, j] is the cost of moving
+    the patch i of img1 to the patch j of img2.
+
+    Moving a patch to another patch has a cost that is not proportional to the number of pixels
+    as this corresponds to moving an entire part of the image to another part.
+
+    Args:
+    - img1_PIL: A PIL Image representing the first image.
+    - img2_PIL: A PIL Image representing the second image (should be the reference if there is one
+        as it is used to determine the most frequent color).
+    - threshold_most_frequent_color: The threshold under which a color is considered as the most frequent color.
+        Constant patches equal to the most frequent color are ignored if the frequency is above this threshold.
+    - patch_size: Tuple indicating the height and width of the patches.
+    - max_num_patches: The maximum number of patches to use for the EMD computation.
+        This is done to avoid having a too long computation time. The images will be resized if necessary.
+    - weight_most_frequent_color: The weight assigned to the most frequent color in the patches.
+        Should be between 0 and 1 (usually low as the most frequentcolor does not carry much information).
+    - use_tqdm: Boolean indicating whether to display a progress bar.
+
+    Returns:
+    - A float representing the Earth Mover's Distance between the images.
+    """
+    assert img1_PIL.size == img2_PIL.size
+    assert patch_size[0] > 0 and patch_size[1] > 0
+    assert 0 < threshold_most_frequent_color <= 1
+    assert max_num_patches > 0
+    assert 0 < weight_most_frequent_color <= 1
+
+    # Convert the images to RGB first. Some images have 4 channels (RGBA)
+    img1_PIL = img1_PIL.convert("RGB")
+    img2_PIL = img2_PIL.convert("RGB")
+
+    # Resize the images so that there are not too many patches
+    # Try to maintain the aspect ratio and resize to a multiple of the patch size
+    num_patches = math.ceil(img1_PIL.size[0] / patch_size[0]) * math.ceil(img1_PIL.size[1] / patch_size[1])
+    if num_patches > max_num_patches:
+        ideal_divider = (num_patches / max_num_patches) ** 0.5
+        closest_round_width = math.ceil((img1_PIL.size[0] / patch_size[1]) / ideal_divider) * patch_size[1]
+        num_patches_width = closest_round_width / patch_size[0]
+        # Chooses a round height such that:
+        # - (round_width / patch_size[1]) * (round_height / patch_size[0]) <= max_num_patches
+        # - the ratio is as unchanged as possible:
+        #   (original_width / round_width) / (original_height / round_height) is close to 1
+        closest_round_height = math.floor(max_num_patches / num_patches_width) * patch_size[0]
+        # Resize the images
+        img1_PIL = img1_PIL.resize((closest_round_width, closest_round_height))
+        img2_PIL = img2_PIL.resize((closest_round_width, closest_round_height))
+
+    # Convert the images to numpy arrays
+    img1_np = np.array(img1_PIL)
+    img2_np = np.array(img2_PIL)
+
+    # Get the patch-signature of the images.
+    # This is of shape (num_patches, patch_size[0] * patch_size[1] + 3)
+    # Each row is a patch, and the columns are:
+    # - index 0: weight of the patch
+    # - index 1 - 1 + patch_size[0] * patch_size[1]: color values of the patch
+    # - index -2, -1: position of the patch
+    (rgb_most_frequent_color, frequency) = get_most_frequent_color(img2_np)
+    gray_most_frequent_color = float(to_gray(rgb_most_frequent_color).squeeze() / 255.0)
+    sig1 = img_to_sig_patches(img1_np, rgb_most_frequent_color, patch_size, weight_most_frequent_color)
+    sig2 = img_to_sig_patches(img2_np, rgb_most_frequent_color, patch_size, weight_most_frequent_color)
+
+    if frequency > threshold_most_frequent_color:
+        # Ignore patches that are constant equal to the most frequent color
+        mask1 = np.any(sig1[:, 1:-2] != gray_most_frequent_color, axis=1)
+        mask2 = np.any(sig2[:, 1:-2] != gray_most_frequent_color, axis=1)
+        mask = np.logical_or(mask1, mask2)
+
+        if np.any(mask):
+            sig1 = sig1[mask]
+            sig2 = sig2[mask]
+
+    # Normalize the weights
+    weight1 = sig1[:, 0]
+    weight2 = sig2[:, 0]
+    weights = np.maximum(weight1, weight2)
+    weights /= np.sum(weights)
+    sig1[:, 0] = weights
+    sig2[:, 0] = weights
+
+    # Compute EMD
+    cost = compute_cost_matrix_on_sig(
+        sig1=sig1,
+        sig2=sig2,
+        gray_most_frequent_color=gray_most_frequent_color,
+        patch_size=patch_size,
+        dim=img1_PIL.size,
+        weight_most_frequent_color=weight_most_frequent_color,
+        use_tqdm=use_tqdm,
+    )
+    emd_value, _, _ = cv2.EMD(sig1, sig2, cv2.DIST_USER, cost)
+    return emd_value