alivegames/Grounded-Segment-Anything

1import random2import numpy as np3import os,sys4import requests5import torch6import torchvision.transforms as torchvision_T7from PIL import Image8 9from transformers import AutoProcessor, AutoModelForVision2Seq10# import subprocess, io, os, sys, time11# sys.path.insert(0, './transformers_4_35_0')12# from transformers_4_35_0 import AutoProcessor, AutoModelForVision2Seq13 14import cv215import ast16 17colors = [18    (0, 255, 0),19    (0, 0, 255),20    (255, 255, 0),21    (255, 0, 255),22    (0, 255, 255),23    (114, 128, 250),24    (0, 165, 255),25    (0, 128, 0),26    (144, 238, 144),27    (238, 238, 175),28    (255, 191, 0),29    (0, 128, 0),30    (226, 43, 138),31    (255, 0, 255),32    (0, 215, 255),33    (255, 0, 0),    34]35 36color_map = {37    f"{color_id}": f"#{hex(color[2])[2:].zfill(2)}{hex(color[1])[2:].zfill(2)}{hex(color[0])[2:].zfill(2)}" for color_id, color in enumerate(colors)38}39 40 41def is_overlapping(rect1, rect2):42    x1, y1, x2, y2 = rect143    x3, y3, x4, y4 = rect244    return not (x2 < x3 or x1 > x4 or y2 < y3 or y1 > y4)45 46 47def draw_entity_boxes_on_image(image, entities, show=False, save_path=None, entity_index=-1):48    """_summary_49    Args:50        image (_type_): image or image path51        collect_entity_location (_type_): _description_52    """53    if isinstance(image, Image.Image):54        image_h = image.height55        image_w = image.width56        image = np.array(image)[:, :, [2, 1, 0]]57    elif isinstance(image, str):58        if os.path.exists(image):59            pil_img = Image.open(image).convert("RGB")60            image = np.array(pil_img)[:, :, [2, 1, 0]]61            image_h = pil_img.height62            image_w = pil_img.width63        else:64            raise ValueError(f"invaild image path, {image}")65    elif isinstance(image, torch.Tensor):66        # pdb.set_trace()67        image_tensor = image.cpu()68        reverse_norm_mean = torch.tensor([0.48145466, 0.4578275, 0.40821073])[:, None, None]69        reverse_norm_std = torch.tensor([0.26862954, 0.26130258, 0.27577711])[:, None, None]70        image_tensor = image_tensor * reverse_norm_std + reverse_norm_mean71        pil_img = torchvision_T.ToPILImage()(image_tensor)72        image_h = pil_img.height73        image_w = pil_img.width74        image = np.array(pil_img)[:, :, [2, 1, 0]]75    else:76        raise ValueError(f"invaild image format, {type(image)} for {image}")77    78    if len(entities) == 0:79        return image80 81    indices = list(range(len(entities)))82    if entity_index >= 0:83        indices = [entity_index]84 85    # Not to show too many bboxes86    entities = entities[:len(color_map)]87    88    new_image = image.copy()89    previous_bboxes = []90    # size of text91    text_size = 192    # thickness of text93    text_line = 1  # int(max(1 * min(image_h, image_w) / 512, 1))94    box_line = 395    (c_width, text_height), _ = cv2.getTextSize("F", cv2.FONT_HERSHEY_COMPLEX, text_size, text_line)96    base_height = int(text_height * 0.675)97    text_offset_original = text_height - base_height98    text_spaces = 399 100    # num_bboxes = sum(len(x[-1]) for x in entities)101    used_colors = colors  # random.sample(colors, k=num_bboxes)102 103    color_id = -1104    for entity_idx, (entity_name, (start, end), bboxes) in enumerate(entities):105        color_id += 1106        if entity_idx not in indices:107            continue108        for bbox_id, (x1_norm, y1_norm, x2_norm, y2_norm) in enumerate(bboxes):109            # if start is None and bbox_id > 0:110            #     color_id += 1111            orig_x1, orig_y1, orig_x2, orig_y2 = int(x1_norm * image_w), int(y1_norm * image_h), int(x2_norm * image_w), int(y2_norm * image_h)112 113            # draw bbox114            # random color115            color = used_colors[color_id]  # tuple(np.random.randint(0, 255, size=3).tolist())116            new_image = cv2.rectangle(new_image, (orig_x1, orig_y1), (orig_x2, orig_y2), color, box_line)117 118            l_o, r_o = box_line // 2 + box_line % 2, box_line // 2 + box_line % 2 + 1119 120            x1 = orig_x1 - l_o121            y1 = orig_y1 - l_o122 123            if y1 < text_height + text_offset_original + 2 * text_spaces:124                y1 = orig_y1 + r_o + text_height + text_offset_original + 2 * text_spaces125                x1 = orig_x1 + r_o126 127            # add text background128            (text_width, text_height), _ = cv2.getTextSize(f"  {entity_name}", cv2.FONT_HERSHEY_COMPLEX, text_size, text_line)129            text_bg_x1, text_bg_y1, text_bg_x2, text_bg_y2 = x1, y1 - (text_height + text_offset_original + 2 * text_spaces), x1 + text_width, y1130 131            for prev_bbox in previous_bboxes:132                while is_overlapping((text_bg_x1, text_bg_y1, text_bg_x2, text_bg_y2), prev_bbox):133                    text_bg_y1 += (text_height + text_offset_original + 2 * text_spaces)134                    text_bg_y2 += (text_height + text_offset_original + 2 * text_spaces)135                    y1 += (text_height + text_offset_original + 2 * text_spaces)136 137                    if text_bg_y2 >= image_h:138                        text_bg_y1 = max(0, image_h - (text_height + text_offset_original + 2 * text_spaces))139                        text_bg_y2 = image_h140                        y1 = image_h141                        break142 143            alpha = 0.5144            for i in range(text_bg_y1, text_bg_y2):145                for j in range(text_bg_x1, text_bg_x2):146                    if i < image_h and j < image_w:147                        if j < text_bg_x1 + 1.35 * c_width:148                            # original color149                            bg_color = color150                        else:151                            # white152                            bg_color = [255, 255, 255]153                        new_image[i, j] = (alpha * new_image[i, j] + (1 - alpha) * np.array(bg_color)).astype(np.uint8)154 155            cv2.putText(156                new_image, f"  {entity_name}", (x1, y1 - text_offset_original - 1 * text_spaces), cv2.FONT_HERSHEY_COMPLEX, text_size, (0, 0, 0), text_line, cv2.LINE_AA157            )158            # previous_locations.append((x1, y1))159            previous_bboxes.append((text_bg_x1, text_bg_y1, text_bg_x2, text_bg_y2))160 161    pil_image = Image.fromarray(new_image[:, :, [2, 1, 0]])162    if save_path:163        pil_image.save(save_path)164    if show:165        pil_image.show()166 167    return pil_image168 169def load_kosmos_model(device):170    ckpt = "ydshieh/kosmos-2-patch14-224"171    kosmos_model = AutoModelForVision2Seq.from_pretrained(ckpt, trust_remote_code=True).to(device)172    kosmos_processor = AutoProcessor.from_pretrained(ckpt, trust_remote_code=True)173    return kosmos_model, kosmos_processor174 175def kosmos_generate_predictions(image_input, text_input, kosmos_model, kosmos_processor):176    if kosmos_model is None:177        return None, None, None178 179    # Save the image and load it again to match the original Kosmos-2 demo.180    # (https://github.com/microsoft/unilm/blob/f4695ed0244a275201fff00bee495f76670fbe70/kosmos-2/demo/gradio_app.py#L345-L346)181    user_image_path = "/tmp/user_input_test_image.jpg"182    image_input.save(user_image_path)183    # This might give different results from the original argument `image_input`184    image_input = Image.open(user_image_path)185 186    if text_input == "Brief":187        text_input = "<grounding>An image of"188    elif text_input == "Detailed":189        text_input = "<grounding>Describe this image in detail:"190    else:191        text_input = f"<grounding>{text_input}"192 193    inputs = kosmos_processor(text=text_input, images=image_input, return_tensors="pt")194 195    generated_ids = kosmos_model.generate(196        pixel_values=inputs["pixel_values"].to("cuda"),197        input_ids=inputs["input_ids"][:, :-1].to("cuda"),198        attention_mask=inputs["attention_mask"][:, :-1].to("cuda"),199        img_features=None,200        img_attn_mask=inputs["img_attn_mask"][:, :-1].to("cuda"),201        use_cache=True,202        max_new_tokens=128,203    )204    generated_text = kosmos_processor.batch_decode(generated_ids, skip_special_tokens=True)[0]205 206    # By default, the generated  text is cleanup and the entities are extracted.207    processed_text, entities = kosmos_processor.post_process_generation(generated_text)208 209    annotated_image = draw_entity_boxes_on_image(image_input, entities, show=False)210 211    color_id = -1212    entity_info = []213    filtered_entities = []214    for entity in entities:215        entity_name, (start, end), bboxes = entity216        if start == end:217            # skip bounding bbox without a `phrase` associated218            continue219        color_id += 1220        # for bbox_id, _ in enumerate(bboxes):221            # if start is None and bbox_id > 0:222            #     color_id += 1223        entity_info.append(((start, end), color_id))224        filtered_entities.append(entity)225 226    colored_text = []227    prev_start = 0228    end = 0229    for idx, ((start, end), color_id) in enumerate(entity_info):230        if start > prev_start:231            colored_text.append((processed_text[prev_start:start], None))232        colored_text.append((processed_text[start:end], f"{color_id}"))233        prev_start = end234 235    if end < len(processed_text):236        colored_text.append((processed_text[end:len(processed_text)], None))237 238    return annotated_image, colored_text, str(filtered_entities)239