Spaces:
Running
on
Zero
Running
on
Zero
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) | |