change base model

app_diffedit.py

@@ -13,6 +13,10 @@ from diffusers import (
 DEFAULT_SRC_PROMPT = "a woman"
 DEFAULT_EDIT_PROMPT = "a woman, with red lips, 8k, high quality"
 
+BASE_MODEL = "stabilityai/stable-diffusion-xl-base-1.0"
+# BASE_MODEL = "stabilityai/sdxl-turbo"
+# BASE_MODEL = "stabilityai/stable-diffusion-2-1"
+
 basepipeline = StableDiffusionDiffEditPipeline.from_pretrained(
     "stabilityai/stable-diffusion-2-1",
     torch_dtype=torch.float16,

hair_color.py

@@ -0,0 +1,59 @@
+import cv2
+import numpy as np
+import mediapipe as mp
+
+from PIL import Image
+from segment_utils import (
+    segmenter
+)
+
+original_image = Image.open("origin.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))
+
+img = cv2.cvtColor(np.array(original_image), cv2.COLOR_RGB2BGR)
+hair_mask_image = cv2.cvtColor(np.array(mask_image), cv2.COLOR_RGB2BGR)
+
+
+def rgb2scale(bgr):
+    min = 255
+    max = 0
+    for val in bgr:
+        if val < min:
+            min = val
+        if val > max:
+            max = val
+    avg = (min + max) / 2
+    # return [val / max for val in bgr]
+    # return [(val - min) / (max - min) for val in bgr]
+    # return [val / avg for val in bgr]
+    return [val / 127.5 for val in bgr]
+
+bgr = [100, 192, 192]
+scale = [0.3,0.1,1]
+scale = rgb2scale(bgr)
+print(scale)
+
+a = np.where(hair_mask > 0)
+ones = np.ones_like(img)
+ones[a] = scale
+
+
+r = img*ones
+cv2.imwrite("hair_color.jpeg", r)
+
+
+
+green_hair = np.copy(img)
+# boolean indexing and assignment based on mask
+green_hair[(hair_mask_image==255).all(-1)] = [255,0,0]
+
+green_hair_w = cv2.addWeighted(green_hair, 0.3, img, 0.7, 0, green_hair)
+
+cv2.imwrite("green_hair.jpeg", green_hair_w)
+

hair_gray.py

@@ -0,0 +1,60 @@
+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)
+
+