diff --git a/main.py b/main.py index 4bb52a7..c204afe 100644 --- a/main.py +++ b/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() diff --git a/src/img_utils.py b/src/img_utils.py index 4e55b8a..8d0dffe 100644 --- a/src/img_utils.py +++ b/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