main src files

src/detection.py

@@ -0,0 +1,54 @@
+import numpy as np
+import torch
+import kornia as K
+from kornia.core import Tensor
+from kornia.contrib import FaceDetector, FaceDetectorResult, FaceKeypoint
+
+print('Loading Face Detector...')
+face_detection = FaceDetector()
+print('DONE')
+
+def detect_face(input):
+
+    # Preprocessing
+    img = K.image_to_tensor(np.array(input), keepdim=False)
+    img = K.color.bgr_to_rgb(img.float())
+    
+    with torch.no_grad():
+        dets = face_detection(img)
+        
+    return [FaceDetectorResult(o) for o in dets[0]]
+
+def process_face(dets):
+    vis_threshold = 0.8
+    faces = []
+    hairs = []
+    
+    for b in dets:
+        if b.score  < vis_threshold:
+            continue
+    
+        reye_kpt=b.get_keypoint(FaceKeypoint.EYE_RIGHT).int().tolist()
+        leye_kpt=b.get_keypoint(FaceKeypoint.EYE_LEFT).int().tolist()
+        rmou_kpt=b.get_keypoint(FaceKeypoint.MOUTH_RIGHT).int().tolist()
+        lmou_kpt=b.get_keypoint(FaceKeypoint.MOUTH_LEFT).int().tolist()
+        nose_kpt=b.get_keypoint(FaceKeypoint.NOSE).int().tolist()
+    
+        faces.append([nose_kpt,
+                     rmou_kpt,
+                     lmou_kpt,
+                     reye_kpt,
+                     leye_kpt
+                    ])
+    
+        # point above
+        top=((b.top_right + b.top_left)/2).int().tolist()
+        bot=((b.bottom_right + b.bottom_left)/2).int().tolist()
+        face_h = np.abs(top[1]-bot[1])
+        top_margin=[top[0], top[1]-face_h*0.1]
+    
+        hairs.append([
+                          top_margin
+                    ])
+
+    return faces, hairs

src/masking.py

@@ -0,0 +1,89 @@
+import torch
+from kornia.morphology import dilation, closing
+import requests
+from transformers import SamModel, SamProcessor
+
+print('Loading SAM...')
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model = SamModel.from_pretrained("facebook/sam-vit-huge").to(device)
+processor = SamProcessor.from_pretrained("facebook/sam-vit-huge")
+print('DONE')
+
+def build_mask(image, faces, hairs):
+
+    # 1. Segmentation
+    input_points = faces  # 2D location of the face
+    
+    with torch.no_grad():
+        inputs = processor(image, input_points=input_points, return_tensors="pt").to(device)
+        outputs = model(**inputs)
+        
+        masks = processor.image_processor.post_process_masks(
+            outputs.pred_masks.cpu(), inputs["original_sizes"].cpu(), inputs["reshaped_input_sizes"].cpu()
+        )
+        scores = outputs.iou_scores
+    
+    input_points = hairs  # 2D location of the face
+    
+    with torch.no_grad():
+        inputs = processor(image, input_points=input_points, return_tensors="pt").to(device)
+        outputs = model(**inputs)
+        
+        h_masks = processor.image_processor.post_process_masks(
+            outputs.pred_masks.cpu(), inputs["original_sizes"].cpu(), inputs["reshaped_input_sizes"].cpu()
+        )
+        h_scores = outputs.iou_scores
+
+    # 2. Post-processing
+    mask=masks[0][0].all(0) | h_masks[0][0].all(0)
+    
+    # dilation
+    tensor = mask[None,None,:,:]
+    kernel = torch.ones(3, 3)
+    mask = closing(tensor, kernel)[0,0].bool()
+    
+    return mask
+
+def build_mask_multi(image, faces, hairs):
+
+    all_masks = []
+    
+    for face,hair in zip(faces,hairs):
+        # 1. Segmentation
+        input_points = [face]  # 2D location of the face
+        
+        with torch.no_grad():
+            inputs = processor(image, input_points=input_points, return_tensors="pt").to(device)
+            outputs = model(**inputs)
+            
+            masks = processor.image_processor.post_process_masks(
+                outputs.pred_masks.cpu(), inputs["original_sizes"].cpu(), inputs["reshaped_input_sizes"].cpu()
+            )
+            scores = outputs.iou_scores
+        
+        input_points = [hair]  # 2D location of the face
+        
+        with torch.no_grad():
+            inputs = processor(image, input_points=input_points, return_tensors="pt").to(device)
+            outputs = model(**inputs)
+            
+            h_masks = processor.image_processor.post_process_masks(
+                outputs.pred_masks.cpu(), inputs["original_sizes"].cpu(), inputs["reshaped_input_sizes"].cpu()
+            )
+            h_scores = outputs.iou_scores
+    
+        # 2. Post-processing
+        mask=masks[0][0].all(0) | h_masks[0][0].all(0)
+        
+        # dilation
+        mask_T = mask[None,None,:,:]
+        kernel = torch.ones(3, 3)
+        mask = closing(mask_T, kernel)[0,0].bool()
+
+        all_masks.append(mask)
+
+    mask = all_masks[0]
+    for next_mask in all_masks[1:]:
+        mask = mask | next_mask
+    
+    return mask

src/process.py

@@ -0,0 +1,36 @@
+import gradio as gr
+import cv2
+from PIL import Image
+import numpy as np
+import torch
+
+from .detection import *
+from .masking import *
+from .synthesis import *
+
+def forward(image_cam, image_upload, prompt="", n_prompt=None, num_steps=20, seed=0, original_resolution=False):
+
+    if image_cam is None:
+        image = image_upload
+    else:
+        image = image_cam
+
+    if not original_resolution:
+        w,h = image.size
+        ratio = 512/h
+        new_size = int(w*ratio), int(h*ratio)
+        image = image.resize(new_size)
+        
+    # detect face
+    dets = detect_face(image)
+
+    # segment hair and face
+    faces, hairs = process_face(dets)
+
+    # build mask
+    mask = build_mask_multi(image, faces, hairs)
+
+    # synthesise
+    new_image = synthesis(image,mask, prompt, n_prompt, num_steps=num_steps, seed=seed)
+    
+    return new_image

src/synthesis.py

@@ -0,0 +1,53 @@
+from controlnet_aux import OpenposeDetector
+from diffusers import StableDiffusionInpaintPipeline, ControlNetModel, UniPCMultistepScheduler
+from src.ControlNetInpaint.src.pipeline_stable_diffusion_controlnet_inpaint import *
+from kornia.filters import gaussian_blur2d
+
+if not 'controlnet' in globals():
+    print('Loading ControlNet...')
+    controlnet = ControlNetModel.from_pretrained(
+        "fusing/stable-diffusion-v1-5-controlnet-openpose", torch_dtype=torch.float16
+    )
+
+if 'pipe' not in globals():
+    print('Loading SD...')
+    pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained(
+         "runwayml/stable-diffusion-inpainting", controlnet=controlnet, torch_dtype=torch.float16
+     ).to('cuda')
+    print('DONE')
+    pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
+    
+if 'openpose' not in globals():
+    print('Loading OpenPose...')
+    openpose = OpenposeDetector.from_pretrained('lllyasviel/ControlNet')
+    print('DONE')
+
+def synthesis(image, mask, prompt="", n_prompt="", num_steps=20, seed=0, remix=True):
+    
+    # 1. Get pose
+    with torch.no_grad():
+        pose_image = openpose(image)
+        pose_image=pose_image.resize(image.size)
+    
+    # generate image
+    generator = torch.manual_seed(seed)
+    new_image = pipe(
+        prompt,
+        negative_prompt = n_prompt,
+        generator=generator,
+        num_inference_steps=num_steps,
+        image=image,
+        control_image=pose_image,
+        mask_image=(mask==False).float().numpy(),
+    ).images
+    
+    if remix:
+        for idx in range(len(new_image)):
+            mask =  gaussian_blur2d(1.0*mask[None,None,:,:],
+                                    kernel_size=(11, 11),
+                                    sigma=(29, 29)
+                                   ).squeeze().clip(0,1)
+            
+            new_image[idx] = (mask[:,:,None]*np.asarray(image) + (1-mask[:,:,None])*np.asarray(new_image[idx].resize(image.size))).int().numpy()
+    
+    return new_image