import cv2
import numpy as np
import mediapipe as mp

from PIL import Image
from segment_utils import (
    segmenter
)

def convert_to_grayscale(pil_image: Image):
    gray_image = pil_image.convert('L')
    # return Image.merge('RGB', (gray_image, gray_image, gray_image))
    return gray_image

original_image = Image.open("/Users/zhiweili/Documents/20240830-175314.jpeg")
gray_image = convert_to_grayscale(original_image)
# gray_image.save("gray_image.jpeg")


image = mp.Image(image_format=mp.ImageFormat.SRGB, data=np.asarray(original_image))
segmentation_result = segmenter.segment(image)
category_mask = segmentation_result.category_mask
category_mask_np = category_mask.numpy_view()
hair_mask = category_mask_np == 1
mask_image = Image.fromarray((hair_mask * 255).astype(np.uint8))

original_image.paste(gray_image, (0, 0), mask_image)

original_image.save("tmp.jpeg")

# mask_image.show()

img = cv2.cvtColor(np.array(original_image), cv2.COLOR_RGB2BGR)
# extract the hair area to hair image
hair_image = np.copy(img)
hair_image[~hair_mask] = 0

# gray hair image
gray_hair_image = cv2.cvtColor(hair_image, cv2.COLOR_BGR2GRAY)
# gray_hair_image_3d = np.repeat(gray_hair_image[:, :, np.newaxis], 3, axis=2)
gray_hair_image_3d = cv2.merge([gray_hair_image] * 3)
# paste the gray hair image to the original image by mask
img[hair_mask] = gray_hair_image_3d[hair_mask]

cv2.imwrite("gray_hair.jpeg", img)





hairPiLImage = Image.fromarray(gray_hair_image)
# paste the gray hair image to the original image by mask
original_image.paste(hairPiLImage, (0, 0), mask_image)
original_image.save("finalImage.jpeg")



# hair_mask_image = cv2.cvtColor(np.array(mask_image), cv2.COLOR_RGB2BGR)