Notes

import numpy as np
import mediapipe as mp
import matplotlib.pyplot as plt
from mediapipe.tasks import python
from mediapipe.tasks.python import vision
import cv2

# Define constants
BG_COLOR = (192, 192, 192) # Gray for background
MASK_COLOR = (255, 255, 255) # White for mask

# Path to the model and image
MODEL_PATH = r'E:Mohithair-coveragehair_segmenter.tflite' # Replace with your model path
IMAGE_PATH = r'E:Mohithair-coverageimages.jpeg' # Replace with your image path

# Create the options for ImageSegmenter
base_options = python.BaseOptions(model_asset_path=MODEL_PATH)
options = vision.ImageSegmenterOptions(
base_options=base_options,
output_category_mask=True
)

# Create the image segmenter
with vision.ImageSegmenter.create_from_options(options) as segmenter:
# Load the image
image = mp.Image.create_from_file(IMAGE_PATH)

# Perform segmentation
segmentation_result = segmenter.segment(image)
category_mask = segmentation_result.category_mask

# Prepare the original and output images
image_data = image.numpy_view()

# Create a condition to apply the mask (detect hair regions)
hair_condition = np.stack((category_mask.numpy_view(),) * 3, axis=-1) > 0.2 # Adjust this threshold as necessary

# Create a blurred version of the entire image
blurred_image = cv2.GaussianBlur(cv2.cvtColor(image_data, cv2.COLOR_BGR2RGB), (55, 55), 0)

# Create the final image where the background is blurred but the hair region stays sharp
final_image = np.where(hair_condition, image_data, blurred_image)

# Calculate hair density heatmap
hair_density = cv2.GaussianBlur(category_mask.numpy_view().astype(np.float32), (55, 55), 0)

# Normalize the hair density values to range [0, 255] for visualization
hair_density_normalized = cv2.normalize(hair_density, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8)

# Apply a colormap to the normalized hair density
heatmap_img = cv2.applyColorMap(hair_density_normalized, cv2.COLORMAP_HSV)

# Superimpose the heatmap on the hair region of the image (only where the hair is detected)
super_imposed_img = np.where(hair_condition, cv2.addWeighted(heatmap_img, 0.5, image_data, 0.5, 0), final_image)

# Plot the original, blurred, and final overlay images
fig, axes = plt.subplots(1, 3, figsize=(18, 6))

# Original image
axes[0].imshow(cv2.cvtColor(image_data, cv2.COLOR_BGR2RGB))
axes[0].set_title("Original Image")
axes[0].axis("off")

# Blurred background image (hair regions are sharp)
axes[1].imshow(final_image)
axes[1].set_title("Blurred Background")
axes[1].axis("off")

# Heatmap overlay image
axes[2].imshow(super_imposed_img)
axes[2].set_title("Heatmap Overlay")
axes[2].axis("off")

# Add legend
plt.figtext(0.5, -0.1, "Heatmap Legend: Blue = Low Density/Scalp, Red = High Density Hair", ha="center", fontsize=12, color="black")

plt.tight_layout()
plt.show()