fantasyfish/dreamgaussian

1import os2import glob3import sys4import cv25import argparse6import numpy as np7import matplotlib.pyplot as plt8 9import torch10import torch.nn as nn11import torch.nn.functional as F12from torchvision import transforms13from PIL import Image14import rembg15 16class BLIP2():17    def __init__(self, device='cuda'):18        self.device = device19        from transformers import AutoProcessor, Blip2ForConditionalGeneration20        self.processor = AutoProcessor.from_pretrained("Salesforce/blip2-opt-2.7b")21        self.model = Blip2ForConditionalGeneration.from_pretrained("Salesforce/blip2-opt-2.7b", torch_dtype=torch.float16).to(device)22 23    @torch.no_grad()24    def __call__(self, image):25        image = Image.fromarray(image)26        inputs = self.processor(image, return_tensors="pt").to(self.device, torch.float16)27 28        generated_ids = self.model.generate(**inputs, max_new_tokens=20)29        generated_text = self.processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()30 31        return generated_text32 33 34if __name__ == '__main__':35 36    parser = argparse.ArgumentParser()37    parser.add_argument('path', type=str, help="path to image (png, jpeg, etc.)")38    parser.add_argument('--model', default='u2net', type=str, help="rembg model, see https://github.com/danielgatis/rembg#models")39    parser.add_argument('--size', default=256, type=int, help="output resolution")40    parser.add_argument('--border_ratio', default=0.2, type=float, help="output border ratio")41    parser.add_argument('--recenter', type=bool, default=True, help="recenter, potentially not helpful for multiview zero123")    42    opt = parser.parse_args()43 44    session = rembg.new_session(model_name=opt.model)45 46    if os.path.isdir(opt.path):47        print(f'[INFO] processing directory {opt.path}...')48        files = glob.glob(f'{opt.path}/*')49        out_dir = opt.path50    else: # isfile51        files = [opt.path]52        out_dir = os.path.dirname(opt.path)53    54    for file in files:55 56        out_base = os.path.basename(file).split('.')[0]57        out_rgba = os.path.join(out_dir, out_base + '_rgba.png')58 59        # load image60        print(f'[INFO] loading image {file}...')61        image = cv2.imread(file, cv2.IMREAD_UNCHANGED)62        63        # carve background64        print(f'[INFO] background removal...')65        carved_image = rembg.remove(image, session=session) # [H, W, 4]66        mask = carved_image[..., -1] > 067 68        # recenter69        if opt.recenter:70            print(f'[INFO] recenter...')71            final_rgba = np.zeros((opt.size, opt.size, 4), dtype=np.uint8)72            73            coords = np.nonzero(mask)74            x_min, x_max = coords[0].min(), coords[0].max()75            y_min, y_max = coords[1].min(), coords[1].max()76            h = x_max - x_min77            w = y_max - y_min78            desired_size = int(opt.size * (1 - opt.border_ratio))79            scale = desired_size / max(h, w)80            h2 = int(h * scale)81            w2 = int(w * scale)82            x2_min = (opt.size - h2) // 283            x2_max = x2_min + h284            y2_min = (opt.size - w2) // 285            y2_max = y2_min + w286            final_rgba[x2_min:x2_max, y2_min:y2_max] = cv2.resize(carved_image[x_min:x_max, y_min:y_max], (w2, h2), interpolation=cv2.INTER_AREA)87            88        else:89            final_rgba = carved_image90        91        # write image92        cv2.imwrite(out_rgba, final_rgba)