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 load_model( ckpt_path: str, config_path: str, scheluder_name: str = None, scheluder_dict : dict = None, 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") # cfg.system.denoise_scheduler= { # 'num_train_timesteps': 1000, # 'beta_start': 0.00085, # 'beta_end': 0.012, # 'beta_schedule': 'scaled_linear', # 'steps_offset': 1 # } 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 def rmbg_sam(iamge, foreground_ratio): return iamge def rmbg_rembg(iamge, foreground_ratio): return iamge class RMBG(object): def __init__(self, device): sam_checkpoint = f"{parent_dir}/ckpts/SAM/sam_vit_h_4b8939.pth" model_type = "vit_h" sam = sam_model_registry[model_type](checkpoint=sam_checkpoint).to(device) self.predictor = SamPredictor(sam) def rmbg_sam(self, input_image, crop_size, foreground_ratio, segment=True, rescale=True): RES = 1024 input_image.thumbnail([RES, RES], Image.Resampling.LANCZOS) if segment: 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()) input_image = sam_segment(self.predictor, input_image.convert('RGB'), x_min, y_min, x_max, y_max) # Rescale and recenter if rescale: image_arr = np.array(input_image) in_w, in_h = image_arr.shape[:2] out_res = min(RES, max(in_w, in_h)) ret, mask = cv2.threshold(np.array(input_image.split()[-1]), 0, 255, cv2.THRESH_BINARY) x, y, w, h = cv2.boundingRect(mask) max_size = max(w, h) side_len = int(max_size / foreground_ratio) padded_image = np.zeros((side_len, side_len, 4), dtype=np.uint8) center = side_len // 2 padded_image[center - h // 2 : center - h // 2 + h, center - w // 2 : center - w // 2 + w] = image_arr[y : y + h, x : x + w] rgba = Image.fromarray(padded_image).resize((out_res, out_res), Image.LANCZOS) rgba_arr = np.array(rgba) / 255.0 rgb = rgba_arr[..., :3] * rgba_arr[..., -1:] + (1 - rgba_arr[..., -1:]) input_image = Image.fromarray((rgb * 255).astype(np.uint8)) else: input_image = expand2square(input_image, (127, 127, 127, 0)) return input_image def rmbg_rembg(self, image, crop_size, foreground_ratio, background_choice, backgroud_color): print(background_choice) if background_choice == "Alpha as mask": background = Image.new("RGBA", image.size, (0, 0, 0, 0)) image = Image.alpha_composite(background, image) else: image = remove_background(image, rembg_session, force_remove=True) image = do_resize_content(image, foreground_ratio) image = expand_to_square(image) image = add_background(image, backgroud_color) return image.convert("RGB") def run(self, rm_type, image, crop_size, foreground_ratio, background_choice, backgroud_color): if "Remove" in background_choice: if rm_type.upper() == "SAM": return self.rmbg_sam(image, crop_size, foreground_ratio, background_choice, backgroud_color) elif rm_type.upper() == "REMBG": return self.rmbg_rembg(image, crop_size, foreground_ratio, background_choice, backgroud_color) else: return -1 elif "Original" in background_choice: return image else: return -1 def save_image(tensor): ndarr = tensor.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to("cpu", torch.uint8).numpy() # pdb.set_trace() im = Image.fromarray(ndarr) return ndarr def prepare_data(single_image, crop_size): from apps.third_party.Wonder3D.mvdiffusion.data.single_image_dataset import SingleImageDataset dataset = SingleImageDataset(root_dir='', num_views=6, img_wh=[256, 256], bg_color='white', crop_size=crop_size, single_image=single_image) return dataset[0] 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 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') 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 check_input_image(input_image): if input_image is None: raise gr.Error("No image uploaded!") def remove_background( 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 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 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)