File size: 6,981 Bytes
0f8544d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
from typing import Dict, Optional, Tuple, List
from dataclasses import dataclass
import os
import sys
proj_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.append(os.path.join(proj_dir))
import time
import cv2
import gradio as gr
import numpy as np
import torch
import PIL
from PIL import Image
import rembg
from rembg import remove
rembg_session = rembg.new_session()
from segment_anything import sam_model_registry, SamPredictor

import craftsman
from craftsman.systems.base import BaseSystem
from craftsman.utils.config import ExperimentConfig, load_config

parent_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))

def check_input_image(input_image):
    if input_image is None:
        raise gr.Error("No image uploaded!")
    
def load_model(
        ckpt_path: str, 
        config_path: str, 
        device = "cuda"
    ):
    cfg: ExperimentConfig
    cfg = load_config(config_path)
    
    if 'pretrained_model_name_or_path' not in cfg.system.condition_model or cfg.system.condition_model.pretrained_model_name_or_path is None:
        cfg.system.condition_model.config_path = config_path.replace("config.yaml", "clip_config.json")

    system: BaseSystem = craftsman.find(cfg.system_type)(
        cfg.system, 
    )
    
    print(f"Restoring states from the checkpoint path at {ckpt_path} with config {cfg}")
    system.load_state_dict(torch.load(ckpt_path, map_location=torch.device('cpu'))['state_dict'])
    system = system.to(device).eval()

    return system

class RMBG(object):
    def __init__(self, device):
        sam = sam_model_registry["vit_h"](checkpoint=f"{parent_dir}/ckpts/SAM/sam_vit_h_4b8939.pth").to(device)
        self.predictor = SamPredictor(sam)
        
    def rmbg_sam(self, input_image):
        def _sam_segment(predictor, input_image, *bbox_coords):
            bbox = np.array(bbox_coords)
            image = np.asarray(input_image)

            start_time = time.time()
            predictor.set_image(image)

            masks_bbox, scores_bbox, logits_bbox = predictor.predict(box=bbox, multimask_output=True)

            print(f"SAM Time: {time.time() - start_time:.3f}s")
            out_image = np.zeros((image.shape[0], image.shape[1], 4), dtype=np.uint8)
            out_image[:, :, :3] = image
            out_image_bbox = out_image.copy()
            out_image_bbox[:, :, 3] = masks_bbox[-1].astype(np.uint8) * 255
            torch.cuda.empty_cache()
            return Image.fromarray(out_image_bbox, mode='RGBA')

        RES = 1024
        input_image.thumbnail([RES, RES], Image.Resampling.LANCZOS)

        image_rem = input_image.convert('RGBA')
        image_nobg = remove(image_rem, alpha_matting=True)
        arr = np.asarray(image_nobg)[:, :, -1]
        x_nonzero = np.nonzero(arr.sum(axis=0))
        y_nonzero = np.nonzero(arr.sum(axis=1))
        x_min = int(x_nonzero[0].min())
        y_min = int(y_nonzero[0].min())
        x_max = int(x_nonzero[0].max())
        y_max = int(y_nonzero[0].max())
        return _sam_segment(self.predictor, input_image.convert('RGB'), x_min, y_min, x_max, y_max)
        
    def rmbg_rembg(self, input_image):
        def _rembg_remove(
            image: PIL.Image.Image,
            rembg_session = None,
            force: bool = False,
            **rembg_kwargs,
        ) -> PIL.Image.Image:
            do_remove = True
            if image.mode == "RGBA" and image.getextrema()[3][0] < 255:
                # explain why current do not rm bg
                print("alhpa channl not enpty, skip remove background, using alpha channel as mask")
                background = Image.new("RGBA", image.size, (0, 0, 0, 0))
                image = Image.alpha_composite(background, image)
                do_remove = False
            do_remove = do_remove or force
            if do_remove:
                image = rembg.remove(image, session=rembg_session, **rembg_kwargs)
            return image
        return _rembg_remove(input_image, rembg_session, force_remove=True)

    def run(self, rm_type, image, foreground_ratio, background_choice, backgroud_color):
        # image = cv2.resize(np.array(image), (crop_size, crop_size))
        # image = Image.fromarray(image)

        if background_choice == "Alpha as mask":
            background = Image.new("RGBA", image.size, (backgroud_color[0], backgroud_color[1], backgroud_color[2], 0))
            return Image.alpha_composite(background, image)
        elif "Remove" in background_choice:
            if rm_type.upper() == "SAM":
                image = self.rmbg_sam(image)
            elif rm_type.upper() == "REMBG":
                image = self.rmbg_rembg(image)
            else:
                return -1

            image = do_resize_content(image, foreground_ratio)
            image = expand_to_square(image)
            image = add_background(image, backgroud_color)
            return image.convert("RGB")
    
        elif "Original" in background_choice:
            return image
        else:
            return -1
        
def do_resize_content(original_image: Image, scale_rate):
    # resize image content wile retain the original image size
    if scale_rate != 1:
        # Calculate the new size after rescaling
        new_size = tuple(int(dim * scale_rate) for dim in original_image.size)
        # Resize the image while maintaining the aspect ratio
        resized_image = original_image.resize(new_size)
        # Create a new image with the original size and black background
        padded_image = Image.new("RGBA", original_image.size, (0, 0, 0, 0))
        paste_position = ((original_image.width - resized_image.width) // 2, (original_image.height - resized_image.height) // 2)
        padded_image.paste(resized_image, paste_position)
        return padded_image
    else:
        return original_image
    
def expand2square(pil_img, background_color):
    width, height = pil_img.size
    if width == height:
        return pil_img
    elif width > height:
        result = Image.new(pil_img.mode, (width, width), background_color)
        result.paste(pil_img, (0, (width - height) // 2))
        return result
    else:
        result = Image.new(pil_img.mode, (height, height), background_color)
        result.paste(pil_img, ((height - width) // 2, 0))
        return result

def expand_to_square(image, bg_color=(0, 0, 0, 0)):
    # expand image to 1:1
    width, height = image.size
    if width == height:
        return image
    new_size = (max(width, height), max(width, height))
    new_image = Image.new("RGBA", new_size, bg_color)
    paste_position = ((new_size[0] - width) // 2, (new_size[1] - height) // 2)
    new_image.paste(image, paste_position)
    return new_image

def add_background(image, bg_color=(255, 255, 255)):
    # given an RGBA image, alpha channel is used as mask to add background color
    background = Image.new("RGBA", image.size, bg_color)
    return Image.alpha_composite(background, image)