|
|
| import os
|
| from skimage.morphology import remove_small_objects
|
| from skimage.measure import label
|
| import numpy as np
|
| from PIL import Image
|
| import cv2
|
| from torchvision import transforms
|
| import torch
|
| import torch.nn.functional as F
|
| import torchvision.transforms.functional as TF
|
|
|
| def find_bounding_box(gray_image):
|
| _, binary_image = cv2.threshold(gray_image, 1, 255, cv2.THRESH_BINARY)
|
| contours, _ = cv2.findContours(binary_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
| max_contour = max(contours, key=cv2.contourArea)
|
| x, y, w, h = cv2.boundingRect(max_contour)
|
| return x, y, w, h
|
|
|
| def load_image(img_path, bg_color=None, rmbg_net=None, padding_ratio=0.1):
|
| img = cv2.imread(img_path, cv2.IMREAD_UNCHANGED)
|
| if img is None:
|
| return f"invalid image path {img_path}"
|
|
|
| def is_valid_alpha(alpha, min_ratio = 0.01):
|
| bins = 20
|
| if isinstance(alpha, np.ndarray):
|
| hist = cv2.calcHist([alpha], [0], None, [bins], [0, 256])
|
| else:
|
| hist = torch.histc(alpha, bins=bins, min=0, max=1)
|
| min_hist_val = alpha.shape[0] * alpha.shape[1] * min_ratio
|
| return hist[0] >= min_hist_val and hist[-1] >= min_hist_val
|
|
|
| def rmbg(image: torch.Tensor) -> torch.Tensor:
|
| image = TF.normalize(image, [0.5,0.5,0.5], [1.0,1.0,1.0]).unsqueeze(0)
|
| result=rmbg_net(image)
|
| return result[0][0]
|
|
|
| if len(img.shape) == 2:
|
| num_channels = 1
|
| else:
|
| num_channels = img.shape[2]
|
|
|
|
|
| height, width = img.shape[:2]
|
| if height > width:
|
| scale = 2000 / height
|
| else:
|
| scale = 2000 / width
|
| if scale < 1:
|
| new_size = (int(width * scale), int(height * scale))
|
| img = cv2.resize(img, new_size, interpolation=cv2.INTER_AREA)
|
|
|
| if img.dtype != 'uint8':
|
| img = (img * (255. / np.iinfo(img.dtype).max)).astype(np.uint8)
|
|
|
| rgb_image = None
|
| alpha = None
|
|
|
| if num_channels == 1:
|
| rgb_image = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
|
| elif num_channels == 3:
|
| rgb_image = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
| elif num_channels == 4:
|
| rgb_image = cv2.cvtColor(img, cv2.COLOR_BGRA2RGB)
|
|
|
| b, g, r, alpha = cv2.split(img)
|
| if not is_valid_alpha(alpha):
|
| alpha = None
|
| else:
|
| alpha_gpu = torch.from_numpy(alpha).unsqueeze(0).cuda().float() / 255.
|
| else:
|
| return f"invalid image: channels {num_channels}"
|
|
|
| rgb_image_gpu = torch.from_numpy(rgb_image).cuda().float().permute(2, 0, 1) / 255.
|
| if alpha is None:
|
| resize_transform = transforms.Resize((384, 384), antialias=True)
|
| rgb_image_resized = resize_transform(rgb_image_gpu)
|
| normalize_image = rgb_image_resized * 2 - 1
|
|
|
| mean_color = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1).cuda()
|
| resize_transform = transforms.Resize((1024, 1024), antialias=True)
|
| rgb_image_resized = resize_transform(rgb_image_gpu)
|
| max_value = rgb_image_resized.flatten().max()
|
| if max_value < 1e-3:
|
| return "invalid image: pure black image"
|
| normalize_image = rgb_image_resized / max_value - mean_color
|
| normalize_image = normalize_image.unsqueeze(0)
|
| resize_transform = transforms.Resize((rgb_image_gpu.shape[1], rgb_image_gpu.shape[2]), antialias=True)
|
|
|
|
|
| alpha_gpu_rmbg = rmbg(rgb_image_resized)
|
| alpha_gpu_rmbg = alpha_gpu_rmbg.squeeze(0)
|
| alpha_gpu_rmbg = resize_transform(alpha_gpu_rmbg)
|
| ma, mi = alpha_gpu_rmbg.max(), alpha_gpu_rmbg.min()
|
| alpha_gpu_rmbg = (alpha_gpu_rmbg - mi) / (ma - mi)
|
|
|
| alpha_gpu = alpha_gpu_rmbg
|
|
|
| alpha_gpu_tmp = alpha_gpu * 255
|
| alpha = alpha_gpu_tmp.to(torch.uint8).squeeze().cpu().numpy()
|
|
|
| _, alpha = cv2.threshold(alpha, 0, 255, cv2.THRESH_BINARY+cv2.THRESH_OTSU)
|
| labeled_alpha = label(alpha)
|
| cleaned_alpha = remove_small_objects(labeled_alpha, min_size=200)
|
| cleaned_alpha = (cleaned_alpha > 0).astype(np.uint8)
|
| alpha = cleaned_alpha * 255
|
| alpha_gpu = torch.from_numpy(cleaned_alpha).cuda().float().unsqueeze(0)
|
| x, y, w, h = find_bounding_box(alpha)
|
|
|
|
|
| else:
|
| rows, cols = np.where(alpha > 0)
|
| if rows.size > 0 and cols.size > 0:
|
| x_min = np.min(cols)
|
| y_min = np.min(rows)
|
| x_max = np.max(cols)
|
| y_max = np.max(rows)
|
|
|
| width = x_max - x_min + 1
|
| height = y_max - y_min + 1
|
| x, y, w, h = x_min, y_min, width, height
|
|
|
| if np.all(alpha==0):
|
| raise ValueError(f"input image too small")
|
|
|
| bg_gray = bg_color[0]
|
| bg_color = torch.from_numpy(bg_color).float().cuda().repeat(alpha_gpu.shape[1], alpha_gpu.shape[2], 1).permute(2, 0, 1)
|
| rgb_image_gpu = rgb_image_gpu * alpha_gpu + bg_color * (1 - alpha_gpu)
|
| padding_size = [0] * 6
|
| if w > h:
|
| padding_size[0] = int(w * padding_ratio)
|
| padding_size[2] = int(padding_size[0] + (w - h) / 2)
|
| else:
|
| padding_size[2] = int(h * padding_ratio)
|
| padding_size[0] = int(padding_size[2] + (h - w) / 2)
|
| padding_size[1] = padding_size[0]
|
| padding_size[3] = padding_size[2]
|
| padded_tensor = F.pad(rgb_image_gpu[:, y:(y+h), x:(x+w)], pad=tuple(padding_size), mode='constant', value=bg_gray)
|
|
|
| return padded_tensor
|
|
|
| def prepare_image(image_path, bg_color, rmbg_net=None):
|
| if os.path.isfile(image_path):
|
| img_tensor = load_image(image_path, bg_color=bg_color, rmbg_net=rmbg_net)
|
| img_np = img_tensor.permute(1,2,0).cpu().numpy()
|
| img_pil = Image.fromarray((img_np*255).astype(np.uint8))
|
|
|
| return img_pil |
