update gradio ui

app.py

@@ -5,6 +5,32 @@ from PIL import Image
 import base64
 from io import BytesIO
 from models.image_text_transformation import ImageTextTransformation
+import argparse
+import torch
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--gpt_version', choices=['gpt-3.5-turbo', 'gpt4'], default='gpt-3.5-turbo')
+parser.add_argument('--image_caption', action='store_true', dest='image_caption', default=True, help='Set this flag to True if you want to use BLIP2 Image Caption')
+parser.add_argument('--dense_caption', action='store_true', dest='dense_caption', default=True, help='Set this flag to True if you want to use Dense Caption')
+parser.add_argument('--semantic_segment', action='store_true', dest='semantic_segment', default=False, help='Set this flag to True if you want to use semantic segmentation')
+parser.add_argument('--image_caption_device', choices=['cuda', 'cpu'], default='cpu', help='Select the device: cuda or cpu, gpu memory larger than 14G is recommended')
+parser.add_argument('--dense_caption_device', choices=['cuda', 'cpu'], default='cpu', help='Select the device: cuda or cpu, < 6G GPU is not recommended>')
+parser.add_argument('--semantic_segment_device', choices=['cuda', 'cpu'], default='cpu', help='Select the device: cuda or cpu, gpu memory larger than 14G is recommended')
+parser.add_argument('--contolnet_device', choices=['cuda', 'cpu'], default='cpu', help='Select the device: cuda or cpu, <6G GPU is not recommended>')
+
+args = parser.parse_args()
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+if device == "cuda":
+    args.image_caption_device = "cuda"
+    args.dense_caption_device = "cuda"
+    args.semantic_segment_device = "cuda"
+    args.contolnet_device = "cuda"
+else:
+    args.image_caption_device = "cpu"
+    args.dense_caption_device = "cpu"
+    args.semantic_segment_device = "cpu"
+    args.contolnet_device = "cpu"
 
 def pil_image_to_base64(image):
     buffered = BytesIO()
@@ -17,7 +43,8 @@ def add_logo():
         logo_base64 = base64.b64encode(f.read()).decode()
     return logo_base64
 
-def process_image(image_src, processor):
+def process_image(image_src, options, processor):
+    processor.args.semantic_segment = "Semantic Segment" in options
     gen_text = processor.image_to_text(image_src)
     gen_image = processor.text_to_image(gen_text)
     gen_image_str = pil_image_to_base64(gen_image)
@@ -38,10 +65,11 @@ def process_image(image_src, processor):
 
     return custom_output
 
-processor = ImageTextTransformation()
+processor = ImageTextTransformation(args)
 
 # Create Gradio input and output components
 image_input = gr.inputs.Image(type='filepath', label="Input Image")
+semantic_segment_checkbox = gr.inputs.Checkbox(label="Semantic Segment", default=False)
 
 logo_base64 = add_logo()
 # Create the title with the logo
@@ -49,12 +77,18 @@ title_with_logo = f'<img src="data:image/jpeg;base64,{logo_base64}" width="400"
 
 # Create Gradio interface
 interface = gr.Interface(
-    fn=lambda image: process_image(image, processor),  # Pass the processor object using a lambda function
-    inputs=image_input,
+    fn=lambda image, options, devices: process_image(image, options, devices, processor),
+    inputs=[image_input,        
+            gr.CheckboxGroup(
+            label="Options",
+            choices=["Semantic Segment"],
+            ),
+            ],
     outputs=gr.outputs.HTML(),
     title=title_with_logo,
     description="""
     This code support image to text transformation. Then the generated text can do retrieval, question answering et al to conduct zero-shot.
+    \n Semantic segment is very slow in cpu(~8m), best use on gpu or run local.
     """
 )
 

main.py

@@ -1,20 +0,0 @@
-import argparse
-from models.image_text_transformation import ImageTextTransformation
-from utils.util import display_images_and_text
-
-if __name__ == '__main__':
-    parser = argparse.ArgumentParser()
-    parser.add_argument('--image_src', default='examples/1.jpg')
-    parser.add_argument('--out_image_name', default='output/1_result.jpg')
-    args = parser.parse_args()
-
-    processor = ImageTextTransformation()
-    generated_text = processor.image_to_text(args.image_src)
-    generated_image = processor.text_to_image(generated_text)
-    ## then text to image
-    print("*" * 50)
-    print("Generated Text:")
-    print(generated_text)
-    print("*" * 50)
-
-    results = display_images_and_text(args.image_src, generated_image, generated_text, args.out_image_name)

main_gradio.py

@@ -1,84 +0,0 @@
-import gradio as gr
-import cv2
-import numpy as np
-from PIL import Image
-import base64
-from io import BytesIO
-from models.image_text_transformation import ImageTextTransformation
-
-def pil_image_to_base64(image):
-    buffered = BytesIO()
-    image.save(buffered, format="JPEG")
-    img_str = base64.b64encode(buffered.getvalue()).decode()
-    return img_str
-
-def add_logo():
-    with open("examples/logo.png", "rb") as f:
-        logo_base64 = base64.b64encode(f.read()).decode()
-    return logo_base64
-
-def process_image(image_src, processor):
-    gen_text = processor.image_to_text(image_src)
-    gen_image = processor.text_to_image(gen_text)
-    gen_image_str = pil_image_to_base64(gen_image)
-    # Combine the outputs into a single HTML output
-    custom_output = f'''
-    <h2>Image->Text->Image:</h2>
-    <div style="display: flex; flex-wrap: wrap;">
-        <div style="flex: 1;">
-            <h3>Image2Text</h3>
-            <p>{gen_text}</p>
-        </div>
-        <div style="flex: 1;">
-            <h3>Text2Image</h3>
-            <img src="data:image/jpeg;base64,{gen_image_str}" width="100%" />
-        </div>
-    </div>
-    <h2>Using Source Image to do Retrieval on COCO:</h2>
-    <div style="display: flex; flex-wrap: wrap;">
-        <div style="flex: 1;">
-            <h3>Retrieval Top-3 Text</h3>
-            <p>{gen_text}</p>
-        </div>
-        <div style="flex: 1;">
-            <h3>Retrieval Top-3 Image</h3>
-            <img src="data:image/jpeg;base64,{gen_image_str}" width="100%" />
-        </div>
-    </div>
-    <h2>Using Generated texts to do Retrieval on COCO:</h2>
-    <div style="display: flex; flex-wrap: wrap;">
-        <div style="flex: 1;">
-            <h3>Retrieval Top-3 Text</h3>
-            <p>{gen_text}</p>
-        </div>
-        <div style="flex: 1;">
-            <h3>Retrieval Top-3 Image</h3>
-            <img src="data:image/jpeg;base64,{gen_image_str}" width="100%" />
-        </div>
-    </div>
-    '''
-
-    return custom_output
-
-processor = ImageTextTransformation()
-
-# Create Gradio input and output components
-image_input = gr.inputs.Image(type='filepath', label="Input Image")
-
-logo_base64 = add_logo()
-# Create the title with the logo
-title_with_logo = f'<img src="data:image/jpeg;base64,{logo_base64}" width="400" style="vertical-align: middle;"> Understanding Image with Text'
-
-# Create Gradio interface
-interface = gr.Interface(
-    fn=lambda image: process_image(image, processor),  # Pass the processor object using a lambda function
-    inputs=image_input,
-    outputs=gr.outputs.HTML(),
-    title=title_with_logo,
-    description="""
-    This code support image to text transformation. Then the generated text can do retrieval, question answering et al to conduct zero-shot.
-    """
-)
-
-# Launch the interface
-interface.launch()

models/blip2_model.py

@@ -5,14 +5,11 @@ import torch
 
 
 class ImageCaptioning:
-    def __init__(self) -> None:
-        self.device = None
-        # self.processor, self.model = None, None
+    def __init__(self, device):
+        self.device = device
         self.processor, self.model = self.initialize_model()
 
     def initialize_model(self):
-        self.device = "cuda" if torch.cuda.is_available() else "cpu"
-        # self.device = "cpu" # for low gpu memory devices
         if self.device == 'cpu':
             self.data_type = torch.float32
         else:
@@ -29,9 +26,10 @@ class ImageCaptioning:
         inputs = self.processor(images=image, return_tensors="pt").to(self.device, self.data_type)
         generated_ids = self.model.generate(**inputs)
         generated_text = self.processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()
-        print('*'*100 + '\nStep1, BLIP2 caption:')
+        print('\033[1;35m' + '*' * 100 + '\033[0m')
+        print('\nStep1, BLIP2 caption:')
         print(generated_text)
-        print('\n' + '*'*100)
+        print('\033[1;35m' + '*' * 100 + '\033[0m')
         return generated_text
     
     def image_caption_debug(self, image_src):

models/controlnet_model.py

@@ -10,8 +10,8 @@ from diffusers import (
 
 
 class TextToImage:
-    def __init__(self):
-        # self.model = None
+    def __init__(self, device):
+        self.device = device
         self.model = self.initialize_model()
 
     def initialize_model(self):
@@ -29,6 +29,7 @@ class TextToImage:
             pipeline.scheduler.config
         )
         pipeline.enable_model_cpu_offload()
+        pipeline.to(self.device)
         return pipeline
 
     @staticmethod
@@ -42,8 +43,12 @@ class TextToImage:
         return image
 
     def text_to_image(self, text, image):
+        print('\033[1;35m' + '*' * 100 + '\033[0m')
+        print('\nStep5, Text to Image:')
         image = self.preprocess_image(image)
         generated_image = self.model(text, image, num_inference_steps=20).images[0]
+        print("Generated image has been svaed.")
+        print('\033[1;35m' + '*' * 100 + '\033[0m')
         return generated_image
     
     def text_to_image_debug(self, text, image):

models/gpt_model.py

@@ -1,9 +1,10 @@
 import openai
 
 class ImageToText:
-    def __init__(self, api_key):
+    def __init__(self, api_key, gpt_version="gpt-3.5-turbo"):
         self.template = self.initialize_template()
         openai.api_key = api_key
+        self.gpt_version = gpt_version
 
     def initialize_template(self):
         prompt_prefix_1 = """Generate only an informative and nature paragraph based on the given information(a,b,c,d):\n"""  
@@ -16,6 +17,7 @@ class ImageToText:
         Use nouns rather than coordinates to show position information of each object.
         No more than 7 sentences.
         Only use one paragraph.
+        Describe position detailedly.
         Do not appear number.
         """
         template = f"{prompt_prefix_1}{prompt_prefix_2}{{width}}X{{height}}{prompt_prefix_3}{{caption}}{prompt_prefix_4}{{dense_caption}}{prompt_prefix_5}{{region_semantic}}{prompt_suffix}"
@@ -23,15 +25,17 @@ class ImageToText:
     
     def paragraph_summary_with_gpt(self, caption, dense_caption, region_semantic, width, height):
         question = self.template.format(width=width, height=height, caption=caption, dense_caption=dense_caption, region_semantic=region_semantic)
-        print('*'*100)
-        print("question:", question)
+        print('\033[1;35m' + '*' * 100 + '\033[0m')
+        print('\nStep4, Paragraph Summary with GPT-3:')
+        print('\033[1;34m' + "Question:".ljust(10) + '\033[1;36m' + question + '\033[0m')
         completion = openai.ChatCompletion.create(
-            model="gpt-3.5-turbo", 
+            model=self.gpt_version, 
             messages = [
             {"role": "user", "content" : question}]
         )
-        print("chatgpt response:", completion['choices'][0]['message']['content'])
-        print('*'*100)
+
+        print('\033[1;34m' + "ChatGPT Response:".ljust(18) + '\033[1;32m' + completion['choices'][0]['message']['content'] + '\033[0m')
+        print('\033[1;35m' + '*' * 100 + '\033[0m')
         return completion['choices'][0]['message']['content']
 
     def paragraph_summary_with_gpt_debug(self, caption, dense_caption, width, height):

models/grit_model.py

@@ -2,8 +2,8 @@ import os
 from models.grit_src.image_dense_captions import image_caption_api
 
 class DenseCaptioning():
-    def __init__(self) -> None:
-        self.model = None
+    def __init__(self, device):
+        self.device = device
 
 
     def initialize_model(self):
@@ -18,9 +18,10 @@ class DenseCaptioning():
         return dense_caption
     
     def image_dense_caption(self, image_src):
-        dense_caption = image_caption_api(image_src)
+        dense_caption = image_caption_api(image_src, self.device)
+        print('\033[1;35m' + '*' * 100 + '\033[0m')
         print("Step2, Dense Caption:\n")
         print(dense_caption)
-        print('\n'+'*'*100)
+        print('\033[1;35m' + '*' * 100 + '\033[0m')
         return dense_caption
     

models/grit_src/image_dense_captions.py

@@ -50,12 +50,14 @@ def setup_cfg(args):
     return cfg
 
 
-def get_parser():
+def get_parser(device):
     arg_dict = {'config_file': "models/grit_src/configs/GRiT_B_DenseCap_ObjectDet.yaml", 'cpu': False, 'confidence_threshold': 0.5, 'test_task': 'DenseCap', 'opts': ["MODEL.WEIGHTS", "pretrained_models/grit_b_densecap_objectdet.pth"]}
+    if device == "cpu":
+        arg_dict["cpu"] = True
     return arg_dict
 
-def image_caption_api(image_src):
-    args2 = get_parser()
+def image_caption_api(image_src, device):
+    args2 = get_parser(device)
     cfg = setup_cfg(args2)
     demo = VisualizationDemo(cfg)
     if image_src:

models/image_text_transformation.py

@@ -18,27 +18,42 @@ def pil_image_to_base64(image):
 
 
 class ImageTextTransformation:
-    def __init__(self):
+    def __init__(self, args):
         # Load your big model here
+        self.args = args
         self.init_models()
         self.ref_image = None
     
     def init_models(self):
         openai_key = os.environ['OPENAI_KEY']
-        self.image_caption_model = ImageCaptioning()
-        self.dense_caption_model = DenseCaptioning()
+        print('\033[1;34m' + "Welcome to the Image2Paragraph toolbox...".center(50, '-') + '\033[0m')
+        print('\033[1;33m' + "Initializing models...".center(50, '-') + '\033[0m')
+        print('\033[1;31m' + "This is time-consuming, please wait...".center(50, '-') + '\033[0m')
+        self.image_caption_model = ImageCaptioning(device=self.args.image_caption_device)
+        self.dense_caption_model = DenseCaptioning(device=self.args.dense_caption_device)
         self.gpt_model = ImageToText(openai_key)
-        self.controlnet_model = TextToImage()
-        self.region_semantic_model = RegionSemantic()
+        self.controlnet_model = TextToImage(device=self.args.contolnet_device)
+        self.region_semantic_model = RegionSemantic(device=self.args.semantic_segment_device)
+        print('\033[1;32m' + "Model initialization finished!".center(50, '-') + '\033[0m')
 
     
     def image_to_text(self, img_src):
         # the information to generate paragraph based on the context
         self.ref_image = Image.open(img_src)
         width, height = read_image_width_height(img_src)
-        image_caption = self.image_caption_model.image_caption(img_src)
-        dense_caption = self.dense_caption_model.image_dense_caption(img_src)
-        region_semantic = self.region_semantic_model.region_semantic(img_src)
+        print(self.args)
+        if self.args.image_caption:
+            image_caption = self.image_caption_model.image_caption(img_src)
+        else:
+            image_caption = " "
+        if self.args.dense_caption:
+            dense_caption = self.dense_caption_model.image_dense_caption(img_src)
+        else:
+            dense_caption = " "
+        if self.args.semantic_segment:
+            region_semantic = self.region_semantic_model.region_semantic(img_src)
+        else:
+            region_semantic = " "
         generated_text = self.gpt_model.paragraph_summary_with_gpt(image_caption, dense_caption, region_semantic, width, height)
         return generated_text
 

models/region_semantic.py

@@ -3,12 +3,13 @@ from models.segment_models.semantic_segment_anything_model import SemanticSegmen
 
 
 class RegionSemantic():
-    def __init__(self) -> None:
+    def __init__(self, device):
+        self.device = device
         self.init_models()
 
     def init_models(self):
-        self.segment_model = SegmentAnything()
-        self.semantic_segment_model = SemanticSegment()
+        self.segment_model = SegmentAnything(self.device)
+        self.semantic_segment_model = SemanticSegment(self.device)
 
     def semantic_prompt_gen(self, anns):
         """
@@ -21,12 +22,12 @@ class RegionSemantic():
         # Select the top 10 largest regions
         top_10_largest_regions = sorted_annotations[:10]
         semantic_prompt = ""
-        print('*'*100)
+        print('\033[1;35m' + '*' * 100 + '\033[0m')
         print("\nStep3, Semantic Prompt:")
         for region in top_10_largest_regions:
             semantic_prompt += region['class_name'] + ': ' + str(region['bbox']) + "; "
         print(semantic_prompt)
-        print('*'*100)
+        print('\033[1;35m' + '*' * 100 + '\033[0m')
         return semantic_prompt
 
     def region_semantic(self, img_src):

models/segment_models/semantic_segment_anything_model.py

@@ -15,12 +15,11 @@ from models.segment_models.configs.coco_id2label import CONFIG as CONFIG_COCO_ID
 nlp = spacy.load('en_core_web_sm')
 
 class SemanticSegment():
-    def __init__(self):
+    def __init__(self, device):
+        self.device = device
         self.model_init()
 
     def model_init(self):
-        self.device = "cuda" if torch.cuda.is_available() else "cpu"
-        # self.device = 'cpu'
         self.init_clip()
         self.init_oneformer_ade20k()
         self.init_oneformer_coco()

models/segment_models/semgent_anything_model.py

@@ -3,14 +3,13 @@ from segment_anything import SamAutomaticMaskGenerator, sam_model_registry
 import torch
 
 class SegmentAnything:
-    def __init__(self, arch="vit_h", pretrained_weights="pretrained_models/sam_vit_h_4b8939.pth"):
-        # self.model = None
+    def __init__(self, device, arch="vit_h", pretrained_weights="pretrained_models/sam_vit_h_4b8939.pth"):
+        self.device = device
         self.model = self.initialize_model(arch, pretrained_weights)
     
     def initialize_model(self, arch, pretrained_weights):
-        device = "cuda" if torch.cuda.is_available() else "cpu"
         sam = sam_model_registry[arch](checkpoint=pretrained_weights)
-        sam.to(device=device)
+        sam.to(device=self.device)
         mask_generator = SamAutomaticMaskGenerator(sam)
         return mask_generator