initial commit

deploy_1.py

@@ -0,0 +1,239 @@
+# -*- coding: utf-8 -*-
+"""deploy_1
+
+Automatically generated by Colaboratory.
+
+Original file is located at
+    https://colab.research.google.com/drive/15bRa4lN0gamY1gSoZhpUGDp61rmTJ0Eg
+
+# Installing Modules
+"""
+
+!pip install mediapipe
+!pip install --upgrade diffusers[torch]
+!pip install transformers
+!pip install accelerate
+!pip install git+https://github.com/huggingface/diffusers
+
+
+
+"""# Importing Modules"""
+
+import mediapipe as mp
+from mediapipe.tasks import python
+from mediapipe.tasks.python import vision
+import cv2
+from google.colab.patches import cv2_imshow
+import math
+import numpy as np
+from PIL import Image
+from cv2 import kmeans, TERM_CRITERIA_MAX_ITER, TERM_CRITERIA_EPS, KMEANS_RANDOM_CENTERS
+from numpy import float32
+from matplotlib.pyplot import scatter, show
+import matplotlib.pyplot as plt
+import requests
+from transformers import pipeline
+import torch
+import PIL
+from diffusers import StableDiffusionInpaintPipeline, StableDiffusionControlNetInpaintPipeline, ControlNetModel, DDPMScheduler
+from diffusers.utils import load_image
+import torch
+
+
+
+"""# Stable Diffusion and ControlNet Pipeline"""
+
+# Stable Diffusion Controlnet Pipeline Class
+class StableDiffusionControlnetPipeline:
+    def __init__(self):
+        self.SELFIE_MULTICLASS_SEGMENTER_MODEL_PATH = "/content/selfie_multiclass_256x256.tflite"
+        self.CONTROLNET_PATH = "lllyasviel/control_v11p_sd15_inpaint"
+        self.MODEL_PATH = "Uminosachi/realisticVisionV51_v51VAE-inpainting"
+        self.device = "cuda"
+        self.hair_color_pipeline = pipeline("image-classification", model="enzostvs/hair-color")
+        self.controlnet = ControlNetModel.from_pretrained(
+            self.CONTROLNET_PATH, torch_dtype=torch.float16
+        ).to(self.device)
+        pipe = StableDiffusionInpaintPipeline.from_pretrained(
+            self.MODEL_PATH,
+            controlnet=self.controlnet,
+            safety_checker=None,
+            requires_safety_checker=False,
+            torch_dtype=torch.float16
+        ).to(self.device)
+        pipe.scheduler = DDPMScheduler.from_config(pipe.scheduler.config)
+        self.pipe = pipe
+
+    def get_hair_dominant_color(self, image_path):
+        hair_img = Image.open(image_path).convert('RGB')
+        results = self.hair_color_pipeline.predict(hair_img)
+        first_score, first_hair_color = results[0]["score"], results[0]["label"]
+        second_score, second_hair_color = results[1]["score"], results[1]["label"]
+        if first_hair_color != "completely bald":
+            return first_hair_color
+        else:
+            return second_hair_color
+
+    def make_inpaint_condition(self, image, image_mask):
+        image = np.array(image.convert("RGB")).astype(np.float32) / 255.0
+        image_mask = np.array(image_mask.convert("L")).astype(np.float32) / 255.0
+        assert image.shape[0:1] == image_mask.shape[0:1], "image and image_mask must have the same image size"
+        image[image_mask > 0.5] = -1.0  # set as masked pixel
+        image = np.expand_dims(image, 0).transpose(0, 3, 1, 2)
+        image = torch.from_numpy(image)
+        return image
+
+    def roundUp(self, input, round):
+        return input + round - (input % round)
+
+    def stable_diffusion_controlnet(self, image_path):
+        HAIR_ROOT_MASK_PATH = self.create_hair_root_mask(image_path, self.SELFIE_MULTICLASS_SEGMENTER_MODEL_PATH)
+        HAIR_COLOR = self.get_hair_dominant_color(image_path)
+        PROMPT = f"({HAIR_COLOR} root:1.2), raw photo, high detail"
+        NEGATIVE_PROMPT = "black hair root"
+        init_image = load_image(image_path)
+        mask_image = load_image(HAIR_ROOT_MASK_PATH)
+        height = self.roundUp(init_image.height, 8)
+        width = self.roundUp(init_image.width, 8)
+        generator = torch.Generator(device=self.device).manual_seed(1)
+        control_image = self.make_inpaint_condition(init_image, mask_image)
+        new_image = self.pipe(
+            prompt=PROMPT,
+            image=init_image,
+            mask_image=mask_image,
+            num_inference_steps=40,
+            generator=generator,
+            control_image=control_image,
+            negative_prompt=NEGATIVE_PROMPT,
+            strength=1,
+            height=height,
+            width=width,
+            padding_mask_crop=40,
+            guidance_scale=3.5
+        ).images
+        hair_root_edited_img = new_image[0]
+        hair_root_edited_img.save("new_img_modified.jpg")
+        return hair_root_edited_img
+
+    def view_result(self, init_image, touched_up_image):
+        fig, axes = plt.subplots(1, 2, figsize=(12, 6))
+        axes[0].imshow(init_image)
+        axes[0].set_title('Original Image')
+        axes[0].axis('off')
+        axes[1].imshow(touched_up_image)
+        axes[1].set_title('Hair Root Touched-up')
+        axes[1].axis('off')
+        plt.show()
+
+    def resize_and_show(self, image, INPUT_HEIGHT=512, INPUT_WIDTH=512):
+        h, w = image.shape[:2]
+        if h < w:
+            img = cv2.resize(image, (INPUT_WIDTH, math.floor(h/(w/INPUT_WIDTH))))
+        else:
+            img = cv2.resize(image, (math.floor(w/(h/INPUT_HEIGHT)), INPUT_HEIGHT))
+        cv2_imshow(img)
+
+    def create_hair_root_mask(self, image_path, SELFIE_MULTICLASS_SEGMENTER_MODEL_PATH):
+        BG_COLOR = (0, 0, 0)  # Background RGB Color
+        MASK_COLOR = (255, 255, 255)  # Mask RGB Color
+        HAIR_CLASS_INDEX = 1  # Index of the Hair Class
+        N_CLUSTERS = 3
+        img = cv2.imread(image_path)
+        base_options = python.BaseOptions(model_asset_path=SELFIE_MULTICLASS_SEGMENTER_MODEL_PATH)
+        options = vision.ImageSegmenterOptions(base_options=base_options, output_category_mask=True)
+        with vision.ImageSegmenter.create_from_options(options) as segmenter:
+            image = mp.Image(image_format=mp.ImageFormat.SRGB, data=img)
+            segmentation_result = segmenter.segment(image)
+            category_mask = segmentation_result.category_mask
+            image_data = image.numpy_view()
+            fg_image = np.zeros(image_data.shape, dtype=np.uint8)
+            fg_image[:] = MASK_COLOR
+            bg_image = np.zeros(image_data.shape, dtype=np.uint8)
+            bg_image[:] = BG_COLOR
+            condition = np.stack((category_mask.numpy_view(),) * 3, axis=-1) == HAIR_CLASS_INDEX
+            output_image = np.where(condition, fg_image, bg_image)
+            cv2.imwrite("hair_mask.png", output_image)
+            hair_mask_cropped = cv2.bitwise_and(img, output_image)
+            coords = np.where(output_image != [255, 255, 255])
+            background = np.full(img.shape, 128, dtype=np.uint8)  # gray background color
+            hair_mask_cropped[coords[0], coords[1], coords[2]] = background[coords[0], coords[1], coords[2]]
+            rgb_img_hair_mask_cropped = cv2.cvtColor(hair_mask_cropped, cv2.COLOR_BGR2RGB)
+            pillow_img = Image.fromarray(rgb_img_hair_mask_cropped)
+            pillow_img.save("hair_mask_cropped.jpg")
+            img_data = rgb_img_hair_mask_cropped.reshape(-1, 3)
+            criteria = (TERM_CRITERIA_MAX_ITER + TERM_CRITERIA_EPS, 100, 0.2)
+            compactness, labels, centers = kmeans(data=img_data.astype(float32), K=N_CLUSTERS, bestLabels=None,
+                                                  criteria=criteria, attempts=10, flags=KMEANS_RANDOM_CENTERS)
+            colours = centers[labels].reshape(-1, 3)
+            img_colours = colours.reshape(rgb_img_hair_mask_cropped.shape)
+            number_labels = np.bincount(labels.flatten())
+            minimum_cluster_class = number_labels.argmin()
+            masked_image = np.copy(rgb_img_hair_mask_cropped)
+            masked_image = masked_image.reshape((-1, 3))
+            labels = labels.flatten()
+            masked_image[labels == minimum_cluster_class] = [255, 255, 255]
+            masked_image = masked_image.reshape(rgb_img_hair_mask_cropped.shape)
+            masked_image = np.copy(rgb_img_hair_mask_cropped)
+            masked_image = masked_image.reshape((-1, 3))
+            for i in range(0, len(number_labels)):
+                masked_image[labels == i] = [0, 0, 0]
+            masked_image[labels == minimum_cluster_class] = [255, 255, 255]
+            masked_image = masked_image.reshape(rgb_img_hair_mask_cropped.shape)
+            cv2.imwrite("hair_root_mask.jpg", masked_image)
+            hair_rost_mask_img = cv2.imread('hair_root_mask.jpg')
+            gray = cv2.cvtColor(hair_rost_mask_img, cv2.COLOR_BGR2GRAY)
+            ret, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY+cv2.THRESH_OTSU)
+            contours, hierarchy = cv2.findContours(binary, mode=cv2.RETR_TREE, method=cv2.CHAIN_APPROX_NONE)
+            image_copy = hair_rost_mask_img.copy()
+            image_copy = cv2.drawContours(image_copy, contours, -1, (255, 255, 255), thickness=3, lineType=cv2.LINE_4)
+            cv2.fillPoly(image_copy, pts=contours, color=(255, 255, 255))
+            (h, w) = image_copy.shape[:2]
+            cut_pixel = int((w // 2) * 0.25)
+            chin_point = ((w // 2) - cut_pixel, (h // 2) - cut_pixel)
+            image_copy[chin_point[0]:, :] = [0, 0, 0]
+            cv2.imwrite("hair_root_mask_mdf.png", image_copy)
+            HAIR_ROOT_MASK_PATH = "/content/hair_root_mask_mdf.png"
+            return HAIR_ROOT_MASK_PATH
+
+"""# Installing Gradio"""
+
+!pip install gradio --upgrade
+
+"""## Calling the StableDiffusionControlnetPipeline for Gradio Interface
+
+"""
+
+import numpy as np
+import gradio as gr
+
+# Assuming StableDiffusionControlnetPipeline class is already defined
+
+# Define the function for Gradio
+def process_image(input_img):
+    # Convert Gradio input image to numpy array
+    input_img_np = np.array(input_img)
+
+    # Save the uploaded image to a temporary file
+    temp_image_path = "/tmp/uploaded_image.jpg"
+    input_img.save(temp_image_path)
+
+    # Instantiate your pipeline with the uploaded image
+    SB_ControlNet_pipeline = StableDiffusionControlnetPipeline()
+
+    # Process the image using your pipeline
+    output_img = SB_ControlNet_pipeline.stable_diffusion_controlnet(temp_image_path)
+
+    return output_img
+
+# Create a Gradio interface
+iface = gr.Interface(
+    fn=process_image,
+    inputs=gr.Image(type="pil", label="Upload Image"),
+    outputs="image",
+    title="Hair Root Touch Up using AI!",
+    description="Upload an image to edit hair roots using Stable Diffusion Controlnet:)"
+)
+
+# Launch the Gradio interface
+iface.launch(debug=True)
+