(feat) add more highlight methods

main.py

@@ -1,7 +1,15 @@
 import cv2
 
 from src.data import load_data
-from src.img_utils import apply_color_overlay, blur_background, desaturate_background
+from src.img_utils import (
+    apply_color_overlay,
+    apply_composite_overlay,
+    apply_gradient_heatmap_overlay,
+    apply_spotlight_heatmap,
+    blur_background,
+    desaturate_background,
+    draw_border,
+)
 
 
 def main():
@@ -19,6 +27,25 @@ def main():
     desaturate_result = desaturate_background(image, mask)
     cv2.imwrite(".tmp-data/highlight_desaturated.jpg", desaturate_result)
 
+    gradient_heatmap_result = apply_gradient_heatmap_overlay(image, mask, colormap=cv2.COLORMAP_JET, alpha=0.6)
+    cv2.imwrite(".tmp-data/highlight_gradient_heatmap.jpg", gradient_heatmap_result)
+
+    spotlight_result = apply_spotlight_heatmap(image, mask, darkness_factor=0.2)
+    cv2.imwrite(".tmp-data/highlight_spotlight_heatmap.jpg", spotlight_result)
+
+    composite_result = apply_composite_overlay(
+        image, mask, colormap=cv2.COLORMAP_JET, foreground_alpha=0.6, background_alpha=0.5
+    )
+    cv2.imwrite(".tmp-data/highlight_composite.jpg", composite_result)
+
+    bordered_image = draw_border(
+        image,
+        mask,
+        color=(50, 255, 50),  # A bright green color
+        thickness=8,
+    )
+    cv2.imwrite(".tmp-data/highlight_border.jpg", bordered_image)
+
 
 if __name__ == "__main__":
     main()

src/img_utils.py

@@ -80,3 +80,156 @@ def blur_background(image, mask, blur_intensity=(35, 35)):
     highlighted_image = cv2.add(foreground, background)
 
     return highlighted_image
+
+
+def apply_gradient_heatmap_overlay(image, mask, colormap=cv2.COLORMAP_JET, alpha=0.6):
+    """
+    Applies a semi-transparent gradient heatmap overlay to the masked region.
+
+    Args:
+        image (np.ndarray): The original BGR image.
+        mask (np.ndarray): The single-channel black and white mask.
+        colormap (int): The OpenCV colormap to use (e.g., cv2.COLORMAP_JET).
+        alpha (float): The transparency of the heatmap (0.0 to 1.0).
+
+    Returns:
+        np.ndarray: The image with the gradient heatmap overlay.
+    """
+    # 1. Create a distance transform from the mask.
+    # This creates a float32 image where each pixel's value is its distance
+    # to the nearest zero-pixel (the edge of the mask).
+    dist_transform = cv2.distanceTransform(mask, cv2.DIST_L2, 5)
+
+    # 2. Normalize the distance transform to the 0-255 range.
+    # This is necessary so we can apply a colormap.
+    normalized_dist = cv2.normalize(dist_transform, None, 0, 255, cv2.NORM_MINMAX)
+    gradient_mask = normalized_dist.astype(np.uint8)
+
+    # 3. Apply the colormap to the gradient mask.
+    heatmap_color = cv2.applyColorMap(gradient_mask, colormap)
+
+    # 4. Isolate the heatmap to the region of interest using the original mask.
+    heatmap_region = cv2.bitwise_and(heatmap_color, heatmap_color, mask=mask)
+
+    # 5. Blend the isolated heatmap with the original image.
+    highlighted_image = cv2.addWeighted(image, 1 - alpha, heatmap_region, alpha, 0)
+
+    return highlighted_image
+
+
+def apply_spotlight_heatmap(image, mask, colormap=cv2.COLORMAP_JET, alpha=0.6, darkness_factor=0.3):
+    """
+    Applies a gradient heatmap overlay and darkens the background.
+
+    Args:
+        image (np.ndarray): The original BGR image.
+        mask (np.ndarray): The single-channel black and white mask.
+        colormap (int): The OpenCV colormap to use.
+        alpha (float): The transparency of the heatmap.
+        darkness_factor (float): How dark the background should be (0.0 to 1.0).
+
+    Returns:
+        np.ndarray: The image with the spotlight heatmap effect.
+    """
+    # 1. Create the gradient heatmap region (same as the previous method)
+    dist_transform = cv2.distanceTransform(mask, cv2.DIST_L2, 5)
+    normalized_dist = cv2.normalize(dist_transform, None, 0, 255, cv2.NORM_MINMAX)
+    gradient_mask = normalized_dist.astype(np.uint8)
+    heatmap_color = cv2.applyColorMap(gradient_mask, colormap)
+    heatmap_region = cv2.bitwise_and(heatmap_color, heatmap_color, mask=mask)
+
+    # 2. Isolate the original object and blend it with the heatmap
+    foreground_original = cv2.bitwise_and(image, image, mask=mask)
+    highlighted_foreground = cv2.addWeighted(foreground_original, 1 - alpha, heatmap_region, alpha, 0)
+
+    # 3. Create the darkened background
+    dark_image = (image * darkness_factor).astype(np.uint8)
+    inverted_mask = cv2.bitwise_not(mask)
+    darkened_background = cv2.bitwise_and(dark_image, dark_image, mask=inverted_mask)
+
+    # 4. Combine the highlighted foreground and darkened background
+    final_image = cv2.add(highlighted_foreground, darkened_background)
+
+    return final_image
+
+
+def apply_composite_overlay(image, mask, colormap=cv2.COLORMAP_JET, foreground_alpha=0.6, background_alpha=0.5):
+    """
+    Creates a composite image with different overlays for foreground and background.
+    - Foreground: Original image + gradient heatmap overlay.
+    - Background: Original image + solid color overlay (coldest map color).
+
+    Args:
+        image (np.ndarray): The original BGR image.
+        mask (np.ndarray): The single-channel black and white mask.
+        colormap (int): The OpenCV colormap to use.
+        foreground_alpha (float): Transparency of the heatmap on the object.
+        background_alpha (float): Transparency of the color on the background.
+
+    Returns:
+        np.ndarray: The final composite image.
+    """
+    # === Part 1: Create the Highlighted Foreground ===
+
+    # Generate the gradient for the object
+    dist_transform = cv2.distanceTransform(mask, cv2.DIST_L2, 5)
+    normalized_dist = cv2.normalize(dist_transform, None, 0, 255, cv2.NORM_MINMAX)
+    gradient_mask = normalized_dist.astype(np.uint8)
+
+    # Isolate the heatmap and original object regions
+    heatmap_region = cv2.applyColorMap(gradient_mask, colormap)
+    foreground_original = cv2.bitwise_and(image, image, mask=mask)
+
+    # Blend the heatmap onto the original object
+    highlighted_foreground = cv2.addWeighted(
+        foreground_original, 1 - foreground_alpha, heatmap_region, foreground_alpha, 0
+    )
+    # Ensure only the foreground is kept after blending
+    highlighted_foreground = cv2.bitwise_and(highlighted_foreground, highlighted_foreground, mask=mask)
+
+    # === Part 2: Create the Colored Background ===
+
+    # Programmatically get the "coldest" color from the colormap
+    zero_pixel = np.zeros((1, 1), dtype=np.uint8)
+    cold_color = cv2.applyColorMap(zero_pixel, colormap)[0][0].tolist()
+
+    # Create a solid color layer and blend it with the full original image
+    color_layer = np.full(image.shape, cold_color, dtype=np.uint8)
+    blended_background_full = cv2.addWeighted(image, 1 - background_alpha, color_layer, background_alpha, 0)
+
+    # Isolate only the background from this blended result
+    inverted_mask = cv2.bitwise_not(mask)
+    final_background = cv2.bitwise_and(blended_background_full, blended_background_full, mask=inverted_mask)
+
+    # === Part 3: Combine Foreground and Background ===
+
+    final_image = cv2.add(highlighted_foreground, final_background)
+
+    return final_image
+
+
+def draw_border(image, mask, color=(0, 255, 0), thickness=3):
+    """
+    Draws a border around the masked region on the original image.
+
+    Args:
+        image (np.ndarray): The original BGR image.
+        mask (np.ndarray): The single-channel black and white mask.
+        color (tuple): The BGR color for the border.
+        thickness (int): The thickness of the border line.
+
+    Returns:
+        np.ndarray: The image with the border drawn on it.
+    """
+    # Create a copy to avoid modifying the original image
+    output_image = image.copy()
+
+    # 1. Find the contours of the shape in the mask
+    # cv2.RETR_EXTERNAL finds only the outer contours of the shape
+    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+
+    # 2. Draw the found contours on the output image
+    # The '-1' argument draws all found contours
+    cv2.drawContours(output_image, contours, -1, color, thickness)
+
+    return output_image