nicolas/abduction-demo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

(feat) add more highlight methods

Browse Source
This commit is contained in:
Nicolas Sebastian Schuler
2025-07-31 12:43:58 +02:00
parent 64385c1864
commit bbb868b357
2 changed files with 181 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

29
main.py
View File

@@ -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()

153
src/img_utils.py
View File

@@ -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