killah-t-cell/EditAnything

1# Edit Anything trained with Stable Diffusion + ControlNet + SAM  + BLIP22# pip install mmcv3 4from torchvision.utils import save_image5from PIL import Image6import subprocess7from collections import OrderedDict8import numpy as np9import cv210import textwrap11import torch12import os13from annotator.util import resize_image, HWC314import mmcv15import random16 17# device = "cuda" if torch.cuda.is_available() else "cpu" # > 15GB GPU memory required18device = "cpu"19use_blip = True20use_gradio = True21 22if device == 'cpu':23    data_type = torch.float3224else:25    data_type = torch.float1626# Diffusion init using diffusers.27 28# diffusers==0.14.0 required.29from diffusers.utils import load_image30 31base_model_path = "stabilityai/stable-diffusion-2-inpainting"32config_dict = OrderedDict([('SAM Pretrained(v0-1): Good Natural Sense', 'shgao/edit-anything-v0-1-1'),33                        ('LAION Pretrained(v0-3): Good Face', 'shgao/edit-anything-v0-3'),34                        ('SD Inpainting: Not keep position', 'stabilityai/stable-diffusion-2-inpainting')35                        ])36 37# Segment-Anything init.38# pip install git+https://github.com/facebookresearch/segment-anything.git39try:40    from segment_anything import sam_model_registry, SamAutomaticMaskGenerator41except ImportError:42    print('segment_anything not installed')43    result = subprocess.run(['pip', 'install', 'git+https://github.com/facebookresearch/segment-anything.git'], check=True)44    print(f'Install segment_anything {result}')   45    from segment_anything import sam_model_registry, SamAutomaticMaskGenerator46if not os.path.exists('./models/sam_vit_h_4b8939.pth'):47    result = subprocess.run(['wget', 'https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth', '-P', 'models'], check=True)48    print(f'Download sam_vit_h_4b8939.pth {result}')   49sam_checkpoint = "models/sam_vit_h_4b8939.pth"50model_type = "default"51sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)52sam.to(device=device)53mask_generator = SamAutomaticMaskGenerator(sam)54 55 56# BLIP2 init.57if use_blip:58    # need the latest transformers59    # pip install git+https://github.com/huggingface/transformers.git60    from transformers import AutoProcessor, Blip2ForConditionalGeneration61    processor = AutoProcessor.from_pretrained("Salesforce/blip2-opt-2.7b")62    blip_model = Blip2ForConditionalGeneration.from_pretrained(63        "Salesforce/blip2-opt-2.7b", torch_dtype=data_type)64 65 66def region_classify_w_blip2(image):67    inputs = processor(image, return_tensors="pt").to(device, data_type)68    generated_ids = blip_model.generate(**inputs, max_new_tokens=15)69    generated_text = processor.batch_decode(70        generated_ids, skip_special_tokens=True)[0].strip()71    return generated_text72 73def region_level_semantic_api(image, topk=5):74    """75    rank regions by area, and classify each region with blip276    Args:77        image: numpy array78        topk: int79    Returns:80        topk_region_w_class_label: list of dict with key 'class_label'81    """82    topk_region_w_class_label = []83    anns = mask_generator.generate(image)84    if len(anns) == 0:85        return []86    sorted_anns = sorted(anns, key=(lambda x: x['area']), reverse=True)87    for i in range(min(topk, len(sorted_anns))):88        ann = anns[i]89        m = ann['segmentation']90        m_3c = m[:,:, np.newaxis]91        m_3c = np.concatenate((m_3c,m_3c,m_3c), axis=2)92        bbox = ann['bbox']93        region = mmcv.imcrop(image*m_3c, np.array([bbox[0], bbox[1], bbox[0] + bbox[2], bbox[1] + bbox[3]]), scale=1)94        region_class_label = region_classify_w_blip2(region)95        ann['class_label'] = region_class_label96        print(ann['class_label'], str(bbox))97        topk_region_w_class_label.append(ann)98    return topk_region_w_class_label99 100def show_semantic_image_label(anns):101    """102    show semantic image label for each region103    Args:104        anns: list of dict with key 'class_label'105    Returns:106        full_img: numpy array107    """108    full_img = None109    # generate mask image110    for i in range(len(anns)):111        m = anns[i]['segmentation']112        if full_img is None:113            full_img = np.zeros((m.shape[0], m.shape[1], 3))114        color_mask = np.random.random((1, 3)).tolist()[0]115        full_img[m != 0] = color_mask116    full_img = full_img*255117    # add text on this mask image118    for i in range(len(anns)):119        m = anns[i]['segmentation']120        class_label = anns[i]['class_label']121        # add text to region122        # Calculate the centroid of the region to place the text123        y, x = np.where(m != 0)124        x_center, y_center = int(np.mean(x)), int(np.mean(y))125 126        # Split the text into multiple lines127        max_width = 20  # Adjust this value based on your preferred maximum width128        wrapped_text = textwrap.wrap(class_label, width=max_width)129 130        # Add text to region131        font = cv2.FONT_HERSHEY_SIMPLEX132        font_scale = 1.2133        font_thickness = 2134        font_color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))  # red135        line_spacing = 40  # Adjust this value based on your preferred line136 137        for idx, line in enumerate(wrapped_text):138            y_offset = y_center - (len(wrapped_text) - 1) * line_spacing // 2 + idx * line_spacing139            text_size = cv2.getTextSize(line, font, font_scale, font_thickness)[0]140            x_offset = x_center - text_size[0] // 2141            # Draw the text multiple times with small offsets to create a bolder appearance142            offsets = [(-1, -1), (-1, 0), (-1, 1), (0, -1), (0, 1), (1, -1), (1, 0), (1, 1)]143            for off_x, off_y in offsets:144                cv2.putText(full_img, line, (x_offset + off_x, y_offset + off_y), font, font_scale, font_color, font_thickness, cv2.LINE_AA)145 146    return full_img147 148 149 150image_path = "images/sa_224577.jpg"151input_image = Image.open(image_path)152detect_resolution=1024153input_image = resize_image(np.array(input_image, dtype=np.uint8), detect_resolution)154region_level_annots = region_level_semantic_api(input_image, topk=5)155output = show_semantic_image_label(region_level_annots)156 157image_list = []158input_image = resize_image(input_image, 512)159output = resize_image(output, 512)160input_image = np.array(input_image, dtype=np.uint8)161output = np.array(output, dtype=np.uint8)162image_list.append(torch.tensor(input_image).float())163image_list.append(torch.tensor(output).float())164for each in image_list:165    print(each.shape, type(each))166    print(each.max(), each.min())167 168 169image_list = torch.stack(image_list).permute(0, 3, 1, 2)170print(image_list.shape)171 172save_image(image_list, "images/sample_semantic.jpg", nrow=2,173        normalize=True)174 175