add gpt box

app.py

@@ -1,6 +1,9 @@
 import os
 import gradio as gr
 import torch
+import base64
+import requests
+from io import BytesIO
 
 from API_LLaVA.functions import get_model as llava_get_model, get_preanswer as llava_get_preanswer, from_preanswer_to_mask as llava_from_preanswer_to_mask
 from API_LLaVA.hook import hook_logger as llava_hook_logger
@@ -23,8 +26,55 @@ MARKDOWN = """
 </div>
 """
 
+def get_base64_images(image):
+    image = image.convert('RGB')
+    buffer = BytesIO()
+    image.save(buffer, format='JPEG')
+    image_base64 = base64.b64encode(buffer.getvalue()).decode('utf-8')
+    return image_base64
+
+def vqa(image, question, api_key):
+    base64_image = get_base64_images(image)
+    headers = {
+    "Content-Type": "application/json",
+    "Authorization": f"Bearer {api_key}"
+    }
+
+    payload = {
+    "model": "gpt-4-turbo-2024-04-09",
+    "messages": [
+        {
+        "role": "user",
+        "content": [
+            {
+            "type": "text",
+            "text": question
+            },
+            {
+            "type": "image_url",
+            "image_url": {
+                "url": f"data:image/jpeg;base64,{base64_image}",
+                "detail":"low"
+            }
+            }
+        ]
+        }
+    ],
+    "max_tokens": 300
+    }
+
+    response = requests.post("https://api.openai.com/v1/chat/completions", headers=headers, json=payload)
+    return response.json()["choices"][0]["message"]["content"]
+
+def compare(input_image, output_image, query, api_key):
+    original_response = vqa(input_image, query, api_key)
+    api_response = vqa(output_image, query, api_key)
+    return original_response, api_response
+
 def init_clip():
-    clip_model, clip_prs, clip_preprocess, _, clip_tokenizer = clip_get_model(model_name = "ViT-L-14-336", layer_index = 22, device= DEVICE)
+    clip_model, clip_prs, clip_preprocess, _, clip_tokenizer = clip_get_model(
+        model_name = "ViT-L-14-336" if torch.cuda.is_available() else "ViT-L-14", 
+        layer_index = 22, device= DEVICE)
     return {"clip_model": clip_model, "clip_prs": clip_prs, "clip_preprocess": clip_preprocess, "clip_tokenizer": clip_tokenizer}
 
 def init_llava():
@@ -133,21 +183,38 @@ image_output = gr.Image(
 text_query = gr.Textbox(
     label="Query",
     placeholder="Enter a query about the image",
-    lines=4,
+    lines=2,
     type="text")
 text_pre_answer = gr.Textbox(
     label="LLaVA Response",
     info = 'Only used for LLaVA-Based API. Press "Pre-Answer" to generate the response.',
     placeholder="",
-    lines=4,
+    lines=2,
     interactive=False,
     type="text")
 text_highlight_text = gr.Textbox(
     label = "Hint Text.",
-    info = "The text based on which the mask will be generated. For CLIP-Based API, it should be a substring of the query. For LLaVA-Based API, it should be a substring of the pre-answer.",
+    info = "The text based on which the mask will be generated. For LLaVA-Based API, it should be a substring of the pre-answer.",
     placeholder="Enter the hint text",
     lines=1,
     type="text")
+text_api_token = gr.Textbox(
+    label = "OpenAI API Token",
+    placeholder="Input your OpenAI API token",
+    lines=1,
+    type="text")
+text_original_image_response = gr.Textbox(
+    label="GPT Response (Original Image)",
+    placeholder="",
+    lines=2,
+    interactive=False,
+    type="text")
+text_API_image_response = gr.Textbox(
+    label="GPT Response (API-maksed Image)",
+    placeholder="",
+    lines=2,
+    interactive=False,
+    type="text")
 
 radio_api_method = gr.Radio(
     ["CLIP_Based API", "LLaVA_Based API"] if torch.cuda.is_available() else ["CLIP_Based API"], 
@@ -187,38 +254,56 @@ radio_interpolate_method_name = gr.Radio(
 
 generate_llava_response_button = gr.Button("Pre-Answer", interactive=False)
 generate_mask_button = gr.Button("API Go!")
+ask_gpt_button = gr.Button("GPT Go!")
 
 with gr.Blocks() as demo:
     gr.Markdown(MARKDOWN)
     state_cache = gr.State({})
     state_model = gr.State(init_clip())
     with gr.Row():
-        with gr.Column():
-            image_input.render()
-        with gr.Column():
-            image_output.render()           
-    with gr.Row():
-        radio_api_method.render()
-    with gr.Row():
-        with gr.Column():
-            with gr.Row():
+        image_input.render()
+        image_output.render()    
+    with gr.Accordion("Query and API Processing"):       
+        with gr.Row():
+            radio_api_method.render()
+        with gr.Row(equal_height=True):
+            with gr.Column():
                 text_query.render()
-            with gr.Row():
                 generate_llava_response_button.render()
-            with gr.Row():
                 text_pre_answer.render()
-            with gr.Row():
                 text_highlight_text.render()
-        with gr.Column():
-            with gr.Row():
+            with gr.Column():
                 slider_enhance_coe.render()
-            with gr.Row():
                 slider_kernel_size.render()
-            with gr.Row():
                 radio_interpolate_method_name.render()
-            with gr.Row():
                 slider_mask_grayscale.render()
+        with gr.Row():
             generate_mask_button.render()
+    with gr.Accordion("GPT Response"):      
+        text_api_token.render()
+        ask_gpt_button.render()
+        with gr.Row():
+            text_original_image_response.render()
+            text_API_image_response.render()
+    with gr.Accordion("Examples"):  
+        examples_images_responses = gr.Examples(
+            [
+
+            ],
+            [
+                image_input, 
+                image_output, 
+                text_query, 
+                text_pre_answer, 
+                text_highlight_text, 
+                slider_enhance_coe,
+                slider_kernel_size,
+                radio_interpolate_method_name,
+                slider_mask_grayscale,
+                text_original_image_response,
+                text_API_image_response
+                ],
+            )
 
     radio_api_method.change(
         fn=change_api_method,
@@ -264,5 +349,10 @@ with gr.Blocks() as demo:
             ],
         outputs=[image_output, state_cache]
     )
+    ask_gpt_button.click(
+        fn=compare,
+        inputs=[image_input, image_output, text_query, text_api_token],
+        outputs=[text_original_image_response, text_API_image_response]
+    )
 
 demo.queue(max_size = 1).launch(show_error=True)