add kosmos-2

app.py

@@ -3,7 +3,8 @@ import warnings
 warnings.filterwarnings('ignore')
 
 import subprocess, io, os, sys, time
-os.system("pip install gradio==3.36.1")
+# os.system("pip install gradio==3.36.1")
+os.system("pip install gradio==3.41.2")
 import gradio as gr
 
 from loguru import logger
@@ -50,7 +51,7 @@ from io import BytesIO
 from diffusers import StableDiffusionInpaintPipeline
 from huggingface_hub import hf_hub_download
 
-from  utils import computer_info
+from utils import computer_info
 # relate anything
 from ram_utils import iou, sort_and_deduplicate, relation_classes, MLP, show_anns, ram_show_mask
 from ram_train_eval import RamModel,RamPredictor
@@ -61,6 +62,10 @@ from lama_cleaner.helper import (
     resize_max_size,
 )
 
+# from transformers import AutoProcessor, AutoModelForVision2Seq
+import ast
+from kosmos_utils import *
+
 config_file = 'GroundingDINO/groundingdino/config/GroundingDINO_SwinT_OGC.py'
 ckpt_repo_id = "ShilongLiu/GroundingDINO"
 ckpt_filenmae = "groundingdino_swint_ogc.pth"
@@ -81,6 +86,8 @@ sd_model = None
 lama_cleaner_model= None
 lama_cleaner_model_device = device
 ram_model = None
+kosmos_model = None
+kosmos_processor = None
 
 def get_sam_vit_h_4b8939():
     if not os.path.exists('./sam_vit_h_4b8939.pth'):
@@ -254,6 +261,7 @@ def set_device():
         device = 'cuda' if torch.cuda.is_available() else 'cpu'
     else:
         device = 'cpu'
+    print(f'device={device}')
 
 def load_groundingdino_model():
     # initialize groundingdino model
@@ -366,6 +374,8 @@ class Ram_Predictor(RamPredictor):
 def load_ram_model():
     # load ram model
     global ram_model
+    if os.environ.get('IS_MY_DEBUG') is not None:
+        return
     model_path = "./checkpoints/ram_epoch12.pth"
     ram_config = dict(
         model=dict(
@@ -510,19 +520,23 @@ mask_source_draw = "draw a mask on input image"
 mask_source_segment = "type what to detect below"
 
 def run_anything_task(input_image, text_prompt, task_type, inpaint_prompt, box_threshold, text_threshold, 
-            iou_threshold, inpaint_mode, mask_source_radio, remove_mode, remove_mask_extend, num_relation, cleaner_size_limit=1080):
-    
+            iou_threshold, inpaint_mode, mask_source_radio, remove_mode, remove_mask_extend, num_relation, kosmos_input, cleaner_size_limit=1080):
+    if (task_type == 'Kosmos-2'):
+        global kosmos_model, kosmos_processor
+        kosmos_image, kosmos_text, kosmos_entities = kosmos_generate_predictions(input_image, kosmos_input, kosmos_model, kosmos_processor)
+        return None, None, kosmos_image, kosmos_text, kosmos_entities
+
     if (task_type == 'relate anything'):
         output_images = relate_anything(input_image['image'], num_relation)
-        return output_images, gr.Gallery.update(label='relate images')
+        return output_images, gr.Gallery.update(label='relate images'), None, None, None
 
     text_prompt = text_prompt.strip()
     if not ((task_type == 'inpainting' or task_type == 'remove') and mask_source_radio == mask_source_draw):
         if text_prompt == '':
-            return [], gr.Gallery.update(label='Detection prompt is not found!😂😂😂😂')
+            return [], gr.Gallery.update(label='Detection prompt is not found!😂😂😂😂'), None, None, None
 
     if input_image is None:
-            return [], gr.Gallery.update(label='Please upload a image!😂😂😂😂')
+            return [], gr.Gallery.update(label='Please upload a image!😂😂😂😂'), None, None, None
 
     file_temp = int(time.time())
     logger.info(f'run_anything_task_002/{device}_[{file_temp}]_{task_type}/{inpaint_mode}/[{mask_source_radio}]/{remove_mode}/{remove_mask_extend}_[{text_prompt}]/[{inpaint_prompt}]___1_')
@@ -562,7 +576,7 @@ def run_anything_task(input_image, text_prompt, task_type, inpaint_prompt, box_t
         )
         if boxes_filt.size(0) == 0:
             logger.info(f'run_anything_task_[{file_temp}]_{task_type}_[{text_prompt}]_1___{groundingdino_device}/[No objects detected, please try others.]_')
-            return [], gr.Gallery.update(label='No objects detected, please try others.😂😂😂😂')
+            return [], gr.Gallery.update(label='No objects detected, please try others.😂😂😂😂'), None, None, None
         boxes_filt_ori = copy.deepcopy(boxes_filt)
 
         pred_dict = {
@@ -613,7 +627,7 @@ def run_anything_task(input_image, text_prompt, task_type, inpaint_prompt, box_t
     logger.info(f'run_anything_task_[{file_temp}]_{task_type}_3_')
     if task_type == 'detection' or task_type == 'segment':
         logger.info(f'run_anything_task_[{file_temp}]_{task_type}_9_')
-        return output_images, gr.Gallery.update(label='result images')
+        return output_images, gr.Gallery.update(label='result images'), None, None, None
     elif task_type == 'inpainting' or task_type == 'remove':
         if inpaint_prompt.strip() == '' and mask_source_radio == mask_source_segment:
             task_type = 'remove'
@@ -678,27 +692,48 @@ def run_anything_task(input_image, text_prompt, task_type, inpaint_prompt, box_t
         image_inpainting = image_inpainting.resize((image_pil.size[0], image_pil.size[1]))
         output_images.append(image_inpainting)
         logger.info(f'run_anything_task_[{file_temp}]_{task_type}_9_')
-        return output_images, gr.Gallery.update(label='result images')        
+        return output_images, gr.Gallery.update(label='result images'), None, None, None       
     else:
         logger.info(f"task_type:{task_type} error!")
     logger.info(f'run_anything_task_[{file_temp}]_9_9_')
-    return output_images, gr.Gallery.update(label='result images')
+    return output_images, gr.Gallery.update(label='result images'), None, None, None
 
 def change_radio_display(task_type, mask_source_radio):
     text_prompt_visible = True
     inpaint_prompt_visible = False
     mask_source_radio_visible = False
     num_relation_visible = False
-    if task_type == "inpainting":
-        inpaint_prompt_visible = True
-    if task_type == "inpainting" or task_type == "remove":
-        mask_source_radio_visible = True   
-        if mask_source_radio == mask_source_draw:
-            text_prompt_visible = False
-    if task_type == "relate anything":
+
+    image_gallery_visible = True
+    kosmos_input_visible = False
+    kosmos_output_visible = False
+    kosmos_text_output_visible = False
+
+    if task_type == "Kosmos-2":
         text_prompt_visible = False
-        num_relation_visible = True
-    return  gr.Textbox.update(visible=text_prompt_visible), gr.Textbox.update(visible=inpaint_prompt_visible), gr.Radio.update(visible=mask_source_radio_visible), gr.Slider.update(visible=num_relation_visible) 
+        image_gallery_visible = False
+        kosmos_input_visible = True
+        kosmos_output_visible = True
+        kosmos_text_output_visible = True
+    else: 
+        if task_type == "inpainting":
+            inpaint_prompt_visible = True
+        if task_type == "inpainting" or task_type == "remove":
+            mask_source_radio_visible = True   
+            if mask_source_radio == mask_source_draw:
+                text_prompt_visible = False
+        if task_type == "relate anything":
+            text_prompt_visible = False
+            num_relation_visible = True
+
+    return  (gr.Textbox.update(visible=text_prompt_visible), 
+            gr.Textbox.update(visible=inpaint_prompt_visible), 
+            gr.Radio.update(visible=mask_source_radio_visible), 
+            gr.Slider.update(visible=num_relation_visible),
+            gr.Gallery.update(visible=image_gallery_visible),
+            gr.Radio.update(visible=kosmos_input_visible),
+            gr.Image.update(visible=kosmos_output_visible),
+            gr.HighlightedText.update(visible=kosmos_text_output_visible))
 
 def get_model_device(module):
     try:
@@ -723,12 +758,18 @@ if __name__ == "__main__":
     print(f'args = {args}')
 
     set_device()
-    get_sam_vit_h_4b8939()
+
     load_groundingdino_model()
-    load_sam_model()
+    if os.environ.get('IS_MY_DEBUG') is None:
+        get_sam_vit_h_4b8939()
+        load_sam_model()
+
     load_sd_model()
     load_lama_cleaner_model()
     load_ram_model()
+    
+    if os.environ.get('IS_MY_DEBUG') is None:
+        kosmos_model, kosmos_processor = load_kosmos_model(device)
 
     if os.environ.get('IS_MY_DEBUG') is None:
         os.system("pip list")
@@ -744,7 +785,7 @@ if __name__ == "__main__":
         with gr.Row():
             with gr.Column():
                 input_image = gr.Image(source='upload', elem_id="image_upload", tool='sketch', type='pil', label="Upload")    
-                task_type = gr.Radio(["detection", "segment", "inpainting", "remove", "relate anything"],  value="detection", 
+                task_type = gr.Radio(["detection", "segment", "inpainting", "remove", "relate anything", "Kosmos-2"],  value="detection", 
                                                 label='Task type', visible=True) 
                 mask_source_radio = gr.Radio([mask_source_draw, mask_source_segment], 
                                     value=mask_source_segment, label="Mask from",
@@ -752,6 +793,9 @@ if __name__ == "__main__":
                 text_prompt = gr.Textbox(label="Detection Prompt[To detect multiple objects, seperating each name with '.', like this: cat . dog . chair ]", placeholder="Cannot be empty")                                                
                 inpaint_prompt = gr.Textbox(label="Inpaint Prompt (if this is empty, then remove)", visible=False)
                 num_relation = gr.Slider(label="How many relations do you want to see", minimum=1, maximum=20, value=5, step=1, visible=False)
+                
+                kosmos_input = gr.Radio(["Brief", "Detailed"], label="Kosmos Description Type", value="Brief", visible=False)
+
                 run_button = gr.Button(label="Run", visible=True)
                 with gr.Accordion("Advanced options", open=False) as advanced_options:
                     box_threshold = gr.Slider(
@@ -773,16 +817,50 @@ if __name__ == "__main__":
             with gr.Column():
                 image_gallery = gr.Gallery(label="result images", show_label=True, elem_id="gallery", visible=True
                     ).style(preview=True, columns=[5], object_fit="scale-down", height="auto")          
+                kosmos_output = gr.Image(type="pil", label="result images", visible=False)
+                kosmos_text_output = gr.HighlightedText(
+                                    label="Generated Description",
+                                    combine_adjacent=False,
+                                    show_legend=True,
+                                    visible=False,
+                                ).style(color_map=color_map)
+                # record which text span (label) is selected
+                selected = gr.Number(-1, show_label=False, placeholder="Selected", visible=False)
+
+                # record the current `entities`
+                entity_output = gr.Textbox(visible=False)
+
+                # get the current selected span label
+                def get_text_span_label(evt: gr.SelectData):
+                    if evt.value[-1] is None:
+                        return -1
+                    return int(evt.value[-1])
+                # and set this information to `selected`
+                kosmos_text_output.select(get_text_span_label, None, selected)
+                
+                # update output image when we change the span (enity) selection
+                def update_output_image(img_input, image_output, entities, idx):
+                    entities = ast.literal_eval(entities)
+                    updated_image = draw_entity_boxes_on_image(img_input, entities, entity_index=idx)
+                    return updated_image
+                selected.change(update_output_image, [kosmos_output, kosmos_output, entity_output, selected], [kosmos_output])
 
             run_button.click(fn=run_anything_task, inputs=[
-                            input_image, text_prompt, task_type, inpaint_prompt, box_threshold, text_threshold, iou_threshold, inpaint_mode, mask_source_radio, remove_mode, remove_mask_extend, num_relation], outputs=[image_gallery, image_gallery], show_progress=True, queue=True)
+                            input_image, text_prompt, task_type, inpaint_prompt, box_threshold, text_threshold, 
+                            iou_threshold, inpaint_mode, mask_source_radio, remove_mode, remove_mask_extend, num_relation, kosmos_input], 
+                            outputs=[image_gallery, image_gallery, kosmos_output, kosmos_text_output, entity_output], show_progress=True, queue=True)
             
-            mask_source_radio.change(fn=change_radio_display, inputs=[task_type, mask_source_radio], outputs=[text_prompt, inpaint_prompt, mask_source_radio, num_relation])
-            task_type.change(fn=change_radio_display, inputs=[task_type, mask_source_radio], outputs=[text_prompt, inpaint_prompt, mask_source_radio, num_relation])
+            mask_source_radio.change(fn=change_radio_display, inputs=[task_type, mask_source_radio], 
+                            outputs=[text_prompt, inpaint_prompt, mask_source_radio, num_relation])
+            task_type.change(fn=change_radio_display, inputs=[task_type, mask_source_radio], 
+                            outputs=[text_prompt, inpaint_prompt, mask_source_radio, num_relation,
+                            image_gallery, kosmos_input, kosmos_output, kosmos_text_output
+                            ])
 
         DESCRIPTION = f'### This demo from [Grounded-Segment-Anything](https://github.com/IDEA-Research/Grounded-Segment-Anything). <br>'
         DESCRIPTION += f'RAM from [RelateAnything](https://github.com/Luodian/RelateAnything). <br>'
         DESCRIPTION += f'Remove(cleaner) from [lama-cleaner](https://github.com/Sanster/lama-cleaner). <br>'
+        DESCRIPTION += f'Kosmos-2 from [RelateAnything](https://huggingface.co/spaces/ydshieh/Kosmos-2). <br>'
         DESCRIPTION += f'Thanks for their excellent work.'
         DESCRIPTION += f'<p>For faster inference without waiting in queue, you may duplicate the space and upgrade to GPU in settings. \
                         <a href="https://huggingface.co/spaces/yizhangliu/Grounded-Segment-Anything?duplicate=true"><img style="display: inline; margin-top: 0em; margin-bottom: 0em" src="https://bit.ly/3gLdBN6" alt="Duplicate Space" /></a></p>'

kosmos_utils.py

@@ -0,0 +1,233 @@
+import random
+import numpy as np
+import os
+import requests
+import torch
+import torchvision.transforms as torchvision_T
+from PIL import Image
+from transformers import AutoProcessor, AutoModelForVision2Seq
+import cv2
+import ast
+
+colors = [
+    (0, 255, 0),
+    (0, 0, 255),
+    (255, 255, 0),
+    (255, 0, 255),
+    (0, 255, 255),
+    (114, 128, 250),
+    (0, 165, 255),
+    (0, 128, 0),
+    (144, 238, 144),
+    (238, 238, 175),
+    (255, 191, 0),
+    (0, 128, 0),
+    (226, 43, 138),
+    (255, 0, 255),
+    (0, 215, 255),
+    (255, 0, 0),    
+]
+
+color_map = {
+    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)
+}
+
+
+def is_overlapping(rect1, rect2):
+    x1, y1, x2, y2 = rect1
+    x3, y3, x4, y4 = rect2
+    return not (x2 < x3 or x1 > x4 or y2 < y3 or y1 > y4)
+
+
+def draw_entity_boxes_on_image(image, entities, show=False, save_path=None, entity_index=-1):
+    """_summary_
+    Args:
+        image (_type_): image or image path
+        collect_entity_location (_type_): _description_
+    """
+    if isinstance(image, Image.Image):
+        image_h = image.height
+        image_w = image.width
+        image = np.array(image)[:, :, [2, 1, 0]]
+    elif isinstance(image, str):
+        if os.path.exists(image):
+            pil_img = Image.open(image).convert("RGB")
+            image = np.array(pil_img)[:, :, [2, 1, 0]]
+            image_h = pil_img.height
+            image_w = pil_img.width
+        else:
+            raise ValueError(f"invaild image path, {image}")
+    elif isinstance(image, torch.Tensor):
+        # pdb.set_trace()
+        image_tensor = image.cpu()
+        reverse_norm_mean = torch.tensor([0.48145466, 0.4578275, 0.40821073])[:, None, None]
+        reverse_norm_std = torch.tensor([0.26862954, 0.26130258, 0.27577711])[:, None, None]
+        image_tensor = image_tensor * reverse_norm_std + reverse_norm_mean
+        pil_img = torchvision_T.ToPILImage()(image_tensor)
+        image_h = pil_img.height
+        image_w = pil_img.width
+        image = np.array(pil_img)[:, :, [2, 1, 0]]
+    else:
+        raise ValueError(f"invaild image format, {type(image)} for {image}")
+    
+    if len(entities) == 0:
+        return image
+
+    indices = list(range(len(entities)))
+    if entity_index >= 0:
+        indices = [entity_index]
+
+    # Not to show too many bboxes
+    entities = entities[:len(color_map)]
+    
+    new_image = image.copy()
+    previous_bboxes = []
+    # size of text
+    text_size = 1
+    # thickness of text
+    text_line = 1  # int(max(1 * min(image_h, image_w) / 512, 1))
+    box_line = 3
+    (c_width, text_height), _ = cv2.getTextSize("F", cv2.FONT_HERSHEY_COMPLEX, text_size, text_line)
+    base_height = int(text_height * 0.675)
+    text_offset_original = text_height - base_height
+    text_spaces = 3
+
+    # num_bboxes = sum(len(x[-1]) for x in entities)
+    used_colors = colors  # random.sample(colors, k=num_bboxes)
+
+    color_id = -1
+    for entity_idx, (entity_name, (start, end), bboxes) in enumerate(entities):
+        color_id += 1
+        if entity_idx not in indices:
+            continue
+        for bbox_id, (x1_norm, y1_norm, x2_norm, y2_norm) in enumerate(bboxes):
+            # if start is None and bbox_id > 0:
+            #     color_id += 1
+            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)
+
+            # draw bbox
+            # random color
+            color = used_colors[color_id]  # tuple(np.random.randint(0, 255, size=3).tolist())
+            new_image = cv2.rectangle(new_image, (orig_x1, orig_y1), (orig_x2, orig_y2), color, box_line)
+
+            l_o, r_o = box_line // 2 + box_line % 2, box_line // 2 + box_line % 2 + 1
+
+            x1 = orig_x1 - l_o
+            y1 = orig_y1 - l_o
+
+            if y1 < text_height + text_offset_original + 2 * text_spaces:
+                y1 = orig_y1 + r_o + text_height + text_offset_original + 2 * text_spaces
+                x1 = orig_x1 + r_o
+
+            # add text background
+            (text_width, text_height), _ = cv2.getTextSize(f"  {entity_name}", cv2.FONT_HERSHEY_COMPLEX, text_size, text_line)
+            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, y1
+
+            for prev_bbox in previous_bboxes:
+                while is_overlapping((text_bg_x1, text_bg_y1, text_bg_x2, text_bg_y2), prev_bbox):
+                    text_bg_y1 += (text_height + text_offset_original + 2 * text_spaces)
+                    text_bg_y2 += (text_height + text_offset_original + 2 * text_spaces)
+                    y1 += (text_height + text_offset_original + 2 * text_spaces)
+
+                    if text_bg_y2 >= image_h:
+                        text_bg_y1 = max(0, image_h - (text_height + text_offset_original + 2 * text_spaces))
+                        text_bg_y2 = image_h
+                        y1 = image_h
+                        break
+
+            alpha = 0.5
+            for i in range(text_bg_y1, text_bg_y2):
+                for j in range(text_bg_x1, text_bg_x2):
+                    if i < image_h and j < image_w:
+                        if j < text_bg_x1 + 1.35 * c_width:
+                            # original color
+                            bg_color = color
+                        else:
+                            # white
+                            bg_color = [255, 255, 255]
+                        new_image[i, j] = (alpha * new_image[i, j] + (1 - alpha) * np.array(bg_color)).astype(np.uint8)
+
+            cv2.putText(
+                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_AA
+            )
+            # previous_locations.append((x1, y1))
+            previous_bboxes.append((text_bg_x1, text_bg_y1, text_bg_x2, text_bg_y2))
+
+    pil_image = Image.fromarray(new_image[:, :, [2, 1, 0]])
+    if save_path:
+        pil_image.save(save_path)
+    if show:
+        pil_image.show()
+
+    return pil_image
+
+def load_kosmos_model(device):
+    ckpt = "ydshieh/kosmos-2-patch14-224"
+    kosmos_model = AutoModelForVision2Seq.from_pretrained(ckpt, trust_remote_code=True).to(device)
+    kosmos_processor = AutoProcessor.from_pretrained(ckpt, trust_remote_code=True)
+    return kosmos_model, kosmos_processor
+
+def kosmos_generate_predictions(image_input, text_input, kosmos_model, kosmos_processor):
+    if kosmos_model is None:
+        return None, None, None
+
+    # Save the image and load it again to match the original Kosmos-2 demo.
+    # (https://github.com/microsoft/unilm/blob/f4695ed0244a275201fff00bee495f76670fbe70/kosmos-2/demo/gradio_app.py#L345-L346)
+    user_image_path = "/tmp/user_input_test_image.jpg"
+    image_input.save(user_image_path)
+    # This might give different results from the original argument `image_input`
+    image_input = Image.open(user_image_path)
+
+    if text_input == "Brief":
+        text_input = "<grounding>An image of"
+    elif text_input == "Detailed":
+        text_input = "<grounding>Describe this image in detail:"
+    else:
+        text_input = f"<grounding>{text_input}"
+
+    inputs = kosmos_processor(text=text_input, images=image_input, return_tensors="pt")
+
+    generated_ids = kosmos_model.generate(
+        pixel_values=inputs["pixel_values"].to("cuda"),
+        input_ids=inputs["input_ids"][:, :-1].to("cuda"),
+        attention_mask=inputs["attention_mask"][:, :-1].to("cuda"),
+        img_features=None,
+        img_attn_mask=inputs["img_attn_mask"][:, :-1].to("cuda"),
+        use_cache=True,
+        max_new_tokens=128,
+    )
+    generated_text = kosmos_processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+
+    # By default, the generated  text is cleanup and the entities are extracted.
+    processed_text, entities = kosmos_processor.post_process_generation(generated_text)
+
+    annotated_image = draw_entity_boxes_on_image(image_input, entities, show=False)
+
+    color_id = -1
+    entity_info = []
+    filtered_entities = []
+    for entity in entities:
+        entity_name, (start, end), bboxes = entity
+        if start == end:
+            # skip bounding bbox without a `phrase` associated
+            continue
+        color_id += 1
+        # for bbox_id, _ in enumerate(bboxes):
+            # if start is None and bbox_id > 0:
+            #     color_id += 1
+        entity_info.append(((start, end), color_id))
+        filtered_entities.append(entity)
+
+    colored_text = []
+    prev_start = 0
+    end = 0
+    for idx, ((start, end), color_id) in enumerate(entity_info):
+        if start > prev_start:
+            colored_text.append((processed_text[prev_start:start], None))
+        colored_text.append((processed_text[start:end], f"{color_id}"))
+        prev_start = end
+
+    if end < len(processed_text):
+        colored_text.append((processed_text[end:len(processed_text)], None))
+
+    return annotated_image, colored_text, str(filtered_entities)

requirements.txt

@@ -23,6 +23,7 @@ numba
 scipy 
 safetensors
 pynvml
+sentencepiece
 
 lama-cleaner==1.1.2
 openmim==0.1.5