add app_gfp

README.md

@@ -10,4 +10,5 @@ pinned: false
 license: mit
 ---
 
+Modified from: https://huggingface.co/spaces/turboedit/turbo_edit
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,10 @@
+import gradio as gr
+
+from app_base import create_demo as create_demo_face
+
+with gr.Blocks(css="style.css") as demo:
+    with gr.Tabs():
+        with gr.Tab(label="Face"):
+            create_demo_face()
+
+demo.launch()

app_gfp.py

@@ -0,0 +1,109 @@
+import os
+import time
+import spaces
+import cv2
+import gradio as gr
+
+from gfpgan.utils import GFPGANer
+
+os.system("pip freeze")
+# download weights
+if not os.path.exists('GFPGANv1.2.pth'):
+    os.system("wget https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.2.pth -P .")
+if not os.path.exists('GFPGANv1.3.pth'):
+    os.system("wget https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth -P .")
+if not os.path.exists('GFPGANv1.4.pth'):
+    os.system("wget https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth -P .")
+if not os.path.exists('RestoreFormer.pth'):
+    os.system("wget https://github.com/TencentARC/GFPGAN/releases/download/v1.3.4/RestoreFormer.pth -P .")
+if not os.path.exists('CodeFormer.pth'):
+    os.system("wget https://github.com/TencentARC/GFPGAN/releases/download/v1.3.4/CodeFormer.pth -P .")
+
+@spaces.GPU(duration=10)
+def enhance(
+    img_path:str,
+    version:str='1.4',
+    scale:int=2,
+):
+    run_task_time = 0
+    time_cost_str = ''
+    run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str)
+    extension = os.path.splitext(os.path.basename(img_path))[1]
+    img = cv2.imread(img_path, cv2.IMREAD_UNCHANGED)
+    if len(img.shape) == 3 and img.shape[2] == 4:
+        img_mode = 'RGBA'
+    elif len(img.shape) == 2:  # for gray inputs
+        img_mode = None
+        img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
+    else:
+        img_mode = None
+
+    h, w = img.shape[0:2]
+    if h < 300:
+        img = cv2.resize(img, (w * 2, h * 2), interpolation=cv2.INTER_LANCZOS4)
+    
+    if version == 'v1.2':
+        face_enhancer = GFPGANer(model_path='GFPGANv1.2.pth', upscale=2, arch='clean', channel_multiplier=2)
+    elif version == 'v1.3':
+        face_enhancer = GFPGANer(model_path='GFPGANv1.3.pth', upscale=2, arch='clean', channel_multiplier=2)
+    elif version == 'v1.4':
+        face_enhancer = GFPGANer(model_path='GFPGANv1.4.pth', upscale=2, arch='clean', channel_multiplier=2)
+    elif version == 'RestoreFormer':
+        face_enhancer = GFPGANer(model_path='RestoreFormer.pth', upscale=2, arch='RestoreFormer', channel_multiplier=2)
+    elif version == 'CodeFormer':
+        face_enhancer = GFPGANer(model_path='CodeFormer.pth', upscale=2, arch='CodeFormer', channel_multiplier=2)
+    elif version == 'RealESR-General-x4v3':
+        face_enhancer = GFPGANer(model_path='realesr-general-x4v3.pth', upscale=2, arch='realesr-general', channel_multiplier=2)
+
+    _, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=True, paste_back=True)
+    if scale != 2:
+        interpolation = cv2.INTER_AREA if scale < 2 else cv2.INTER_LANCZOS4
+        h, w = img.shape[0:2]
+        output = cv2.resize(output, (int(w * scale / 2), int(h * scale / 2)), interpolation=interpolation)
+    
+    if img_mode == 'RGBA':  # RGBA images should be saved in png format
+        extension = 'png'
+    else:
+        extension = 'jpg'
+    save_path = f'output/out.{extension}'
+    cv2.imwrite(save_path, output)
+
+    output = cv2.cvtColor(output, cv2.COLOR_BGR2RGB)
+    run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str)
+    return output, save_path, time_cost_str
+
+
+def get_time_cost(run_task_time, time_cost_str):
+    now_time = int(time.time()*1000)
+    if run_task_time == 0:
+        time_cost_str = 'start'
+    else:
+        if time_cost_str != '': 
+            time_cost_str += f'-->'
+        time_cost_str += f'{now_time - run_task_time}'
+    run_task_time = now_time
+    return run_task_time, time_cost_str
+
+def create_demo() -> gr.Blocks:
+    with gr.Blocks() as demo:
+        with gr.Row():
+            with gr.Column():
+                version = gr.Radio(['v1.2', 'v1.3', 'v1.4', 'RestoreFormer','CodeFormer','RealESR-General-x4v3'], type="value", default='v1.4', label='version')
+                scale = gr.Number(label="Rescaling factor", default=2)
+            with gr.Column():
+                g_btn = gr.Button(label="Enhance")
+        with gr.Row():
+            with gr.Column():
+                input_image = gr.Image(label="Input Image", type="filepath")
+            with gr.Column():
+                restored_image = gr.Image(label="Restored Image", type="numpy", interactive=False)
+                download_path = gr.File(label="Download the output image", interactive=False)
+                restored_cost = gr.Textbox(label="Time cost by step (ms):", visible=True, interactive=False)
+        
+        g_btn.click(
+            fn=enhance,
+            inputs=[input_image, version, scale],
+            outputs=[restored_image, download_path, restored_cost],
+        )
+
+    return demo

croper.py

@@ -0,0 +1,108 @@
+import PIL
+import numpy as np
+
+from PIL import Image
+
+class Croper:
+    def __init__(
+        self,
+        input_image: PIL.Image,
+        target_mask: np.ndarray,
+        mask_size: int = 256,
+        mask_expansion: int = 20,
+    ):
+        self.input_image = input_image
+        self.target_mask = target_mask
+        self.mask_size = mask_size
+        self.mask_expansion = mask_expansion
+    
+    def corp_mask_image(self):
+        target_mask = self.target_mask
+        input_image = self.input_image
+        mask_expansion = self.mask_expansion
+        original_width, original_height = input_image.size
+        mask_indices = np.where(target_mask)
+        start_y = np.min(mask_indices[0])
+        end_y = np.max(mask_indices[0])
+        start_x = np.min(mask_indices[1])
+        end_x = np.max(mask_indices[1])
+        mask_height = end_y - start_y
+        mask_width = end_x - start_x
+        # choose the max side length
+        max_side_length = max(mask_height, mask_width)
+        # expand the mask area
+        height_diff = (max_side_length - mask_height) // 2
+        width_diff = (max_side_length - mask_width) // 2
+        start_y = start_y - mask_expansion - height_diff
+        if start_y < 0:
+            start_y = 0
+        end_y = end_y + mask_expansion + height_diff
+        if end_y > original_height:
+            end_y = original_height
+        start_x = start_x - mask_expansion - width_diff
+        if start_x < 0:
+            start_x = 0
+        end_x = end_x + mask_expansion + width_diff
+        if end_x > original_width:
+            end_x = original_width
+        expanded_height = end_y - start_y
+        expanded_width = end_x - start_x
+        expanded_max_side_length = max(expanded_height, expanded_width)
+        # calculate the crop area
+        crop_mask = target_mask[start_y:end_y, start_x:end_x]
+        crop_mask_start_y = (expanded_max_side_length - expanded_height) // 2
+        crop_mask_end_y = crop_mask_start_y + expanded_height
+        crop_mask_start_x = (expanded_max_side_length - expanded_width) // 2
+        crop_mask_end_x = crop_mask_start_x + expanded_width
+        # create a square mask
+        square_mask = np.zeros((expanded_max_side_length, expanded_max_side_length), dtype=target_mask.dtype)
+        square_mask[crop_mask_start_y:crop_mask_end_y, crop_mask_start_x:crop_mask_end_x] = crop_mask
+        square_mask_image = Image.fromarray((square_mask * 255).astype(np.uint8))
+
+        crop_image = input_image.crop((start_x, start_y, end_x, end_y))
+        square_image = Image.new("RGB", (expanded_max_side_length, expanded_max_side_length))
+        square_image.paste(crop_image, (crop_mask_start_x, crop_mask_start_y))
+
+        self.origin_start_x = start_x
+        self.origin_start_y = start_y
+        self.origin_end_x = end_x
+        self.origin_end_y = end_y
+
+        self.square_start_x = crop_mask_start_x
+        self.square_start_y = crop_mask_start_y
+        self.square_end_x = crop_mask_end_x
+        self.square_end_y = crop_mask_end_y
+
+        self.square_length = expanded_max_side_length
+        self.square_mask_image = square_mask_image
+        self.square_image = square_image
+        self.corp_mask = crop_mask
+
+        mask_size = self.mask_size
+        self.resized_square_mask_image = square_mask_image.resize((mask_size, mask_size))
+        self.resized_square_image = square_image.resize((mask_size, mask_size))
+
+        return self.resized_square_mask_image
+    
+    def restore_result(self, generated_image):
+        square_length = self.square_length
+        generated_image = generated_image.resize((square_length, square_length))
+        square_mask_image = self.square_mask_image
+        cropped_generated_image = generated_image.crop((self.square_start_x, self.square_start_y, self.square_end_x, self.square_end_y))
+        cropped_square_mask_image = square_mask_image.crop((self.square_start_x, self.square_start_y, self.square_end_x, self.square_end_y))
+
+        restored_image = self.input_image.copy()
+        restored_image.paste(cropped_generated_image, (self.origin_start_x, self.origin_start_y), cropped_square_mask_image)
+        
+        return restored_image
+    
+    def restore_result_v2(self, generated_image):
+        square_length = self.square_length
+        generated_image = generated_image.resize((square_length, square_length))
+        cropped_generated_image = generated_image.crop((self.square_start_x, self.square_start_y, self.square_end_x, self.square_end_y))
+
+        restored_image = self.input_image.copy()
+        restored_image.paste(cropped_generated_image, (self.origin_start_x, self.origin_start_y))
+        
+        return restored_image
+        

requirements.txt

@@ -0,0 +1,11 @@
+ml-collections
+gradio
+torch
+diffusers
+transformers
+accelerate
+mediapipe
+spaces
+sentencepiece
+compel
+gfpgan

segment_utils.py

@@ -0,0 +1,88 @@
+import numpy as np
+import mediapipe as mp
+
+from PIL import Image
+from mediapipe.tasks import python
+from mediapipe.tasks.python import vision
+from scipy.ndimage import binary_dilation
+from croper import Croper
+
+segment_model = "checkpoints/selfie_multiclass_256x256.tflite"
+base_options = python.BaseOptions(model_asset_path=segment_model)
+options = vision.ImageSegmenterOptions(base_options=base_options,output_category_mask=True)
+segmenter = vision.ImageSegmenter.create_from_options(options)
+
+def restore_result(croper, category, generated_image):
+    square_length = croper.square_length
+    generated_image = generated_image.resize((square_length, square_length))
+
+    cropped_generated_image = generated_image.crop((croper.square_start_x, croper.square_start_y, croper.square_end_x, croper.square_end_y))
+    cropped_square_mask_image = get_restore_mask_image(croper, category, cropped_generated_image)
+
+    restored_image = croper.input_image.copy()
+    restored_image.paste(cropped_generated_image, (croper.origin_start_x, croper.origin_start_y), cropped_square_mask_image)
+
+    return restored_image
+
+def segment_image(input_image, category, generate_size, mask_expansion, mask_dilation):
+    mask_size = int(generate_size)
+    mask_expansion = int(mask_expansion)
+
+    image = mp.Image(image_format=mp.ImageFormat.SRGB, data=np.asarray(input_image))
+    segmentation_result = segmenter.segment(image)
+    category_mask = segmentation_result.category_mask
+    category_mask_np = category_mask.numpy_view()
+
+    if category == "hair":
+        target_mask = get_hair_mask(category_mask_np, mask_dilation)
+    elif category == "clothes":
+        target_mask = get_clothes_mask(category_mask_np, mask_dilation)
+    elif category == "face":
+        target_mask = get_face_mask(category_mask_np, mask_dilation)
+    else:
+        target_mask = get_face_mask(category_mask_np, mask_dilation)
+    
+    croper = Croper(input_image, target_mask, mask_size, mask_expansion)
+    croper.corp_mask_image()
+    origin_area_image = croper.resized_square_image
+
+    return origin_area_image, croper
+
+def get_face_mask(category_mask_np, dilation=1):
+    face_skin_mask = category_mask_np == 3
+    if dilation > 0:
+        face_skin_mask = binary_dilation(face_skin_mask, iterations=dilation)
+
+    return face_skin_mask
+
+def get_clothes_mask(category_mask_np, dilation=1):
+    body_skin_mask = category_mask_np == 2
+    clothes_mask = category_mask_np == 4
+    combined_mask = np.logical_or(body_skin_mask, clothes_mask)
+    combined_mask = binary_dilation(combined_mask, iterations=4)
+    if dilation > 0:
+        combined_mask = binary_dilation(combined_mask, iterations=dilation)
+    return combined_mask
+
+def get_hair_mask(category_mask_np, dilation=1):
+    hair_mask = category_mask_np == 1
+    if dilation > 0:
+        hair_mask = binary_dilation(hair_mask, iterations=dilation)
+    return hair_mask
+
+def get_restore_mask_image(croper, category, generated_image):
+    image = mp.Image(image_format=mp.ImageFormat.SRGB, data=np.asarray(generated_image))
+    segmentation_result = segmenter.segment(image)
+    category_mask = segmentation_result.category_mask
+    category_mask_np = category_mask.numpy_view()
+
+    if category == "hair":
+        target_mask = get_hair_mask(category_mask_np, 0)
+    elif category == "clothes":
+        target_mask = get_clothes_mask(category_mask_np, 0)
+    elif category == "face":
+        target_mask = get_face_mask(category_mask_np, 0)
+    
+    combined_mask = np.logical_or(target_mask, croper.corp_mask)
+    mask_image = Image.fromarray((combined_mask * 255).astype(np.uint8))
+    return mask_image