deploy_1.py

# -*- 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)