Upload utils.py

apps/utils.py

@@ -0,0 +1,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)