tohid4n/PartCrafter

1# -*- coding: utf-8 -*-2import os3from skimage.morphology import remove_small_objects4from skimage.measure import label5import numpy as np6from PIL import Image7import cv28from torchvision import transforms9import torch10import torch.nn.functional as F11import torchvision.transforms.functional as TF12 13def find_bounding_box(gray_image):14    _, binary_image = cv2.threshold(gray_image, 1, 255, cv2.THRESH_BINARY)15    contours, _ = cv2.findContours(binary_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)16    max_contour = max(contours, key=cv2.contourArea)17    x, y, w, h = cv2.boundingRect(max_contour)18    return x, y, w, h19 20def load_image(img_path, bg_color=None, rmbg_net=None, padding_ratio=0.1, device='cuda'):21    img = cv2.imread(img_path, cv2.IMREAD_UNCHANGED)22    if img is None:23        return f"invalid image path {img_path}"24 25    def is_valid_alpha(alpha, min_ratio = 0.01):26        bins = 2027        if isinstance(alpha, np.ndarray):28            hist = cv2.calcHist([alpha], [0], None, [bins], [0, 256])29        else:30            hist = torch.histc(alpha, bins=bins, min=0, max=1) 31        min_hist_val = alpha.shape[0] * alpha.shape[1] * min_ratio32        return hist[0] >= min_hist_val and hist[-1] >= min_hist_val33    34    def rmbg(image: torch.Tensor) -> torch.Tensor:35        image = TF.normalize(image, [0.5,0.5,0.5], [1.0,1.0,1.0]).unsqueeze(0)36        result=rmbg_net(image)37        return result[0][0]38 39    if len(img.shape) == 2:40        num_channels = 141    else:42        num_channels = img.shape[2]43 44    # check if too large45    height, width = img.shape[:2]46    if height > width:47        scale = 2000 / height48    else:49        scale = 2000 / width50    if scale < 1:51        new_size = (int(width * scale), int(height * scale))52        img = cv2.resize(img, new_size, interpolation=cv2.INTER_AREA)53 54    if img.dtype != 'uint8':55        img = (img * (255. / np.iinfo(img.dtype).max)).astype(np.uint8)56 57    rgb_image = None58    alpha = None59 60    if num_channels == 1:  61        rgb_image = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)62    elif num_channels == 3:  63        rgb_image = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)64    elif num_channels == 4:  65        rgb_image = cv2.cvtColor(img, cv2.COLOR_BGRA2RGB)66 67        b, g, r, alpha = cv2.split(img)68        if not is_valid_alpha(alpha):69            alpha = None70        else:71            alpha_gpu = torch.from_numpy(alpha).unsqueeze(0).to(device).float() / 255.72    else:73        return f"invalid image: channels {num_channels}"74    75    rgb_image_gpu = torch.from_numpy(rgb_image).to(device).float().permute(2, 0, 1) / 255.76    if alpha is None:77        resize_transform = transforms.Resize((384, 384), antialias=True)78        rgb_image_resized = resize_transform(rgb_image_gpu)79        normalize_image = rgb_image_resized * 2 - 180 81        mean_color = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1).to(device)82        resize_transform = transforms.Resize((1024, 1024), antialias=True)83        rgb_image_resized = resize_transform(rgb_image_gpu)84        max_value = rgb_image_resized.flatten().max()85        if max_value < 1e-3:86            return "invalid image: pure black image"87        normalize_image = rgb_image_resized / max_value - mean_color88        normalize_image = normalize_image.unsqueeze(0)89        resize_transform = transforms.Resize((rgb_image_gpu.shape[1], rgb_image_gpu.shape[2]), antialias=True)90 91        # seg from rmbg92        alpha_gpu_rmbg = rmbg(rgb_image_resized)93        alpha_gpu_rmbg = alpha_gpu_rmbg.squeeze(0)94        alpha_gpu_rmbg = resize_transform(alpha_gpu_rmbg)95        ma, mi = alpha_gpu_rmbg.max(), alpha_gpu_rmbg.min()96        alpha_gpu_rmbg = (alpha_gpu_rmbg - mi) / (ma - mi)97 98        alpha_gpu = alpha_gpu_rmbg99        100        alpha_gpu_tmp = alpha_gpu * 255101        alpha = alpha_gpu_tmp.to(torch.uint8).squeeze().cpu().numpy()102 103        _, alpha = cv2.threshold(alpha, 0, 255, cv2.THRESH_BINARY+cv2.THRESH_OTSU)104        labeled_alpha = label(alpha)105        cleaned_alpha = remove_small_objects(labeled_alpha, min_size=200)106        cleaned_alpha = (cleaned_alpha > 0).astype(np.uint8)107        alpha = cleaned_alpha * 255108        alpha_gpu = torch.from_numpy(cleaned_alpha).to(device).float().unsqueeze(0)109        x, y, w, h = find_bounding_box(alpha)110 111    # If alpha is provided, the bounds of all foreground are used112    else: 113        rows, cols = np.where(alpha > 0)114        if rows.size > 0 and cols.size > 0:115            x_min = np.min(cols)116            y_min = np.min(rows)117            x_max = np.max(cols)118            y_max = np.max(rows)119 120            width = x_max - x_min + 1121            height = y_max - y_min + 1122        x, y, w, h = x_min, y_min, width, height123 124    if np.all(alpha==0):125        raise ValueError(f"input image too small")126    127    bg_gray = bg_color[0]128    bg_color = torch.from_numpy(bg_color).float().to(device).repeat(alpha_gpu.shape[1], alpha_gpu.shape[2], 1).permute(2, 0, 1)129    rgb_image_gpu = rgb_image_gpu * alpha_gpu + bg_color * (1 - alpha_gpu)130    padding_size = [0] * 6131    if w > h:132        padding_size[0] = int(w * padding_ratio)133        padding_size[2] = int(padding_size[0] + (w - h) / 2)134    else:135        padding_size[2] = int(h * padding_ratio)136        padding_size[0] = int(padding_size[2] + (h - w) / 2)137    padding_size[1] = padding_size[0]138    padding_size[3] = padding_size[2]139    padded_tensor = F.pad(rgb_image_gpu[:, y:(y+h), x:(x+w)], pad=tuple(padding_size), mode='constant', value=bg_gray)140 141    return padded_tensor142 143def prepare_image(image_path, bg_color=np.array([1.0, 1.0, 1.0]), rmbg_net=None, padding_ratio=0.1, device='cuda'):144    if os.path.isfile(image_path):145        img_tensor = load_image(image_path, bg_color=bg_color, rmbg_net=rmbg_net, padding_ratio=padding_ratio, device=device)146        img_np = img_tensor.permute(1,2,0).cpu().numpy()147        img_pil = Image.fromarray((img_np*255).astype(np.uint8))148        149        return img_pil150    else:151        raise ValueError(f"Invalid image path: {image_path}")