Create app.py

app.py

@@ -0,0 +1,299 @@
+import gradio as gr
+from transformers import AutoProcessor, Pix2StructForConditionalGeneration, T5Tokenizer, T5ForConditionalGeneration, Pix2StructProcessor, BartConfig,ViTConfig,VisionEncoderDecoderConfig, DonutProcessor, VisionEncoderDecoderModel, AutoTokenizer, AutoModel
+from PIL import Image  
+import torch
+import warnings
+import re
+import json
+import os
+import numpy as np
+import pandas as pd
+from tqdm import tqdm
+import argparse
+from scipy import optimize
+from typing import Optional
+import dataclasses
+import editdistance
+import itertools
+import sys
+import time
+import logging
+import subprocess
+import spaces
+import openai
+import base64
+from io import StringIO
+
+# Git LFS pull 명령어 실행
+result = subprocess.run(['git', 'lfs', 'pull'], capture_output=True, text=True)
+
+# 명령어 실행 결과 출력 (선택 사항)
+if result.returncode == 0:
+    print("LFS 파일이 성공적으로 다운로드되었습니다.")
+else:
+    print(f"오류 발생: {result.stderr}")
+
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+logger = logging.getLogger()
+
+warnings.filterwarnings('ignore')
+MAX_PATCHES = 512
+# Load the models and processor
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# Paths to the models
+ko_deplot_model_path = './deplot_model_ver_24.11.21_korean_only(exclude NUUA)_epoch1.bin'
+
+# Load first model ko-deplot
+def load_model1():
+    processor1 = Pix2StructProcessor.from_pretrained('nuua/ko-deplot')
+    model1 = Pix2StructForConditionalGeneration.from_pretrained('nuua/ko-deplot')
+    model1.load_state_dict(torch.load(ko_deplot_model_path, map_location="cpu"))
+    model1.to(torch.device("cuda"))
+    return processor1, model1
+
+processor1, model1 = load_model1()
+
+# Function to format output
+def format_output(prediction):
+    return prediction.replace('<0x0A>', '\n')
+
+# First model prediction: ko-deplot
+def predict_model1(image):
+    images = [image]
+    inputs = processor1(images=images, text="What is the title of the chart", return_tensors="pt", padding=True)
+    inputs = {k: v.to(device) for k, v in inputs.items()}  # Move to GPU
+
+    model1.eval()
+    with torch.no_grad():
+        predictions = model1.generate(**inputs, max_new_tokens=4096)
+    outputs = [processor1.decode(pred, skip_special_tokens=True) for pred in predictions]
+
+    formatted_output = format_output(outputs[0])
+    return formatted_output
+
+# Set your OpenAI API key
+openai.api_key = "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"
+
+# Function to encode the image as base64
+def encode_image(image_path):
+    with open(image_path, "rb") as image_file:
+        return base64.b64encode(image_file.read()).decode("utf-8")
+
+# Second model prediction: gpt-4o-mini
+def predict_model2(image):
+    # Encode the uploaded image to base64
+    image_data = encode_image(image)
+
+    # Prepare the request content
+    response = openai.ChatCompletion.create(
+        model="gpt-4o-mini",
+        messages=[
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "text",
+                        "text": "please extract chart title and chart data manually and present them as a table. you should only provide title and table without adding any additional comments such as **Chart Title:** ."
+                    },
+                    {
+                        "type": "image_url",
+                        "image_url": {
+                            "url": f"data:image/jpeg;base64,{image_data}"
+                        }
+                    }
+                ]
+            }
+        ]
+    )
+
+    # Return the table data from the response
+    return response.choices[0]["message"]["content"]
+
+def ko_deplot_convert_to_dataframe(label_table_str): #function that converts text generated by ko-deplot to pandas dataframe
+    lines = label_table_str.strip().split("\n")
+    data=[]
+    title= lines[0].split(" | ")[1]
+
+    if(len(lines[1].split("|")) == len(lines[2].split("|"))):
+      headers=lines[1].split(" | ")
+      for line in lines[2:]:
+          data.append(line.split(" | "))
+      df = pd.DataFrame(data, columns=headers)
+      return df, title
+    else:
+      legend_row=lines[1].split("|")
+      legend_row.insert(0," ")
+      for line in lines[2:]:
+          data.append(line.split(" | "))
+      df = pd.DataFrame(data, columns=legend_row)
+      return df, title
+
+def gpt_convert_to_dataframe(table_text): #function that converts text generated by gpt to pandas dataframe
+    try:
+        # Split the text into lines
+        lines = table_text.strip().split("\n")
+        title=lines[0]
+        lines.pop(1)
+        lines.pop(2)
+        # Process the remaining lines to create the DataFrame
+        data = [line.split("|")[1:-1] for line in lines[1:]]  # Split by | and remove empty first/last items
+        dataframe = pd.DataFrame(data[1:], columns=[col.strip() for col in data[0]])  # Use the first row as headers
+        
+        return dataframe, title
+    except Exception as e:
+        return f"Error converting table to DataFrame: {e}"
+
+def real_time_check(image_file):
+    image = Image.open(image_file)
+    ko_deplot_generated_txt = predict_model1(image)
+    parts=ko_deplot_generated_txt.split("\n")
+    del parts[-1]
+    ko_deplot_generated_txt="\n".join(parts)
+    gpt_generated_txt=predict_model2(image_file)
+    try:
+        ko_deplot_generated_df, ko_deplot_generated_title=ko_deplot_convert_to_dataframe(ko_deplot_generated_txt)
+        gpt_generated_df, gpt_generated_title=gpt_convert_to_dataframe(gpt_generated_txt)
+        return gr.DataFrame(ko_deplot_generated_df, label= ko_deplot_generated_title), gr.DataFrame(gpt_generated_df, label= gpt_generated_title), None,None,0
+    except Exception as e:
+        return None,None,ko_deplot_generated_txt,gpt_generated_txt,1
+    
+flag = 0 #flag to check whether exception happens or not. if flag is 1, it means that exception(generated txt cannot be converted to pandas dataframe) happens. 
+def inference(image_uploader,mode_selector):
+    if(mode_selector=="파일 업로드"):
+        ko_deplot_generated_df, gpt_generated_df,ko_deplot_generated_txt, gpt_generated_txt, flag= real_time_check(image_uploader)
+        if flag==1:
+            return gr.update(visible=False), gr.update(visible=False), gr.Text(ko_deplot_generated_txt,visible=True),gr.Text(gpt_generated_txt,visible=True)
+        else:
+            return ko_deplot_generated_df, gpt_generated_df, gr.update(visible=False),gr.update(visible=False)
+    else:
+        ko_deplot_generated_df, gpt_generated_df,ko_deplot_generated_txt, gpt_generated_txt, flag= real_time_check(image_files[current_image_index])
+        if flag==1:
+            return gr.update(visible=False), gr.update(visible=False), gr.Text(ko_deplot_generated_txt,visible=True),gr.Text(gpt_generated_txt,visible=True)
+        else:
+            return ko_deplot_generated_df, gpt_generated_df, gr.update(visible=False),gr.update(visible=False)
+        
+def toggle_model(selected_models,flag):
+    # Create a visibility list initialized to False for all components
+    visibility = [False] * 6
+    # Update visibility based on the selected models
+    if "VAIV_DePlot" in selected_models:
+        visibility[4]= True
+        if flag:
+            visibility[2]=  True
+        else:
+            visibility[0]= True
+    if "gpt-4o-mini" in selected_models:
+        visibility[5]= True
+        if flag:
+            visibility[3]=  True
+        else:
+            visibility[1]= True
+    if "all" in selected_models:
+        visibility[4]=True
+        visibility[5]=True
+        if flag:
+            visibility[2]=  True
+            visibility[3]= True
+        else:
+            visibility[0]= True
+            visibility[1]= True
+    # Return gr.update for each component with the corresponding visibility status
+    return tuple(gr.update(visible=v) for v in visibility)
+
+def toggle_mode(mode):
+        if mode == "파일 업로드":
+            return gr.update(visible=True), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
+        else:
+            return gr.update(visible=False), gr.update(visible=True), gr.update(visible=True), gr.update(visible=True)
+
+def display_image(image_file):
+    image=Image.open(image_file)
+    return image, os.path.basename(image_file)
+
+# Function to display the images in the folder sequentially
+image_files = []
+current_image_index = 0
+image_files_cnt=0
+
+def display_folder_images(image_file_path_list):
+    global image_files, current_image_index,image_files_cnt
+    image_files = image_file_path_list
+    image_files_cnt=len(image_files)
+    current_image_index = 0
+    if image_files:
+        return Image.open(image_files[current_image_index]), os.path.basename(image_files[current_image_index]), gr.update(interactive=False), gr.update(interactive=True)
+    return None, "No images found"
+    
+
+def next_image():
+    global current_image_index
+    if image_files:
+        current_image_index = (current_image_index + 1)
+        prev_disabled = current_image_index == 0
+        next_disabled = current_image_index == (len(image_files) - 1)
+        return Image.open(image_files[current_image_index]), os.path.basename(image_files[current_image_index]), gr.update(interactive=not prev_disabled), gr.update(interactive= not next_disabled)
+    return None, "No images found"
+
+def prev_image():
+    global current_image_index
+    if image_files:
+        current_image_index = (current_image_index - 1)
+        prev_disabled = current_image_index == 0
+        next_disabled = current_image_index == (len(image_files) - 1)
+        return Image.open(image_files[current_image_index]), os.path.basename(image_files[current_image_index]), gr.update(interactive=not prev_disabled), gr.update(interactive= not next_disabled)
+    return None, "No images found"
+
+css = """
+.dataframe-class {
+    overflow-y: auto !important; /* 스크롤을 가능하게 */
+    height: 250px
+}
+"""
+
+with gr.Blocks(css=css) as iface:
+    with gr.Row():
+        gr.Markdown("<h1 style='text-align: center;'>SKKU-VAIV Automatic chart understanding evaluation tool</h1>")
+    gr.Markdown("<hr style='border: 1px solid #ddd;' />")
+    with gr.Row():
+        with gr.Column():
+            mode_selector = gr.Radio(["파일 업로드", "폴더 업로드"], label="Upload Mode", value="파일 업로드")
+            image_uploader = gr.File(file_count="single", file_types=["image"], visible=True)
+            folder_uploader = gr.File(file_count="directory", file_types=["image"], visible=False, height=50)
+            model_type=gr.Dropdown(["VAIV_DePlot","gpt-4o-mini","all"],value="VAIV_DePlot",label="model",multiselect=True)
+            image_displayer = gr.Image(visible=True)
+            image_name = gr.Text("", visible=True)
+            with gr.Row():
+                prev_button = gr.Button("이전", visible=False, interactive=False)
+                next_button = gr.Button("다음", visible=False, interactive=False)
+            inference_button = gr.Button("추론")
+        with gr.Column():
+            md1 = gr.Markdown("# VAIV_DePlot Inference Result")
+            ko_deplot_generated_df = gr.DataFrame(visible=True, elem_classes="dataframe-class")
+            ko_deplot_generated_txt = gr.Text(visible=False)
+        with gr.Column():    
+            md2 = gr.Markdown("# gpt-4o-mini Inference Result", visible=False)
+            gpt_generated_df = gr.DataFrame(visible=False, elem_classes="dataframe-class")
+            gpt_generated_txt = gr.Text(visible=False)
+            #label_df = gr.DataFrame(visible=False, label="Ground Truth Table", elem_classes="dataframe-class",scale=1)
+
+    model_type.change(
+                        toggle_model,
+                        inputs=[model_type, gr.State(flag)],
+                        outputs=[ko_deplot_generated_df,gpt_generated_df,ko_deplot_generated_txt,gpt_generated_txt,md1,md2]
+                        )
+
+    mode_selector.change(
+        toggle_mode,
+        inputs=[mode_selector],
+        outputs=[image_uploader, folder_uploader, prev_button, next_button]
+    )
+
+    image_uploader.upload(display_image,inputs=[image_uploader],outputs=[image_displayer,image_name])
+    folder_uploader.upload(display_folder_images, inputs=[folder_uploader], outputs=[image_displayer, image_name, prev_button, next_button])
+    prev_button.click(prev_image, outputs=[image_displayer, image_name, prev_button, next_button])
+    next_button.click(next_image, outputs=[image_displayer, image_name, prev_button, next_button])
+    inference_button.click(inference,inputs=[image_uploader,mode_selector],outputs=[ko_deplot_generated_df, gpt_generated_df, ko_deplot_generated_txt, gpt_generated_txt])
+
+    if __name__ == "__main__":
+        iface.launch(share=True)