import streamlit as st from io import BytesIO from PIL import Image from transformers import ViltProcessor, ViltForQuestionAnswering import requests import torch import torchvision from langchain.prompts import PromptTemplate from langchain.chains import LLMChain from langchain.chat_models import ChatOpenAI from transformers import AutoProcessor, AutoModelForCausalLM from huggingface_hub import hf_hub_download from transformers import Pix2StructForConditionalGeneration, Pix2StructProcessor from transformers import BlipProcessor, BlipForConditionalGeneration import os os.environ["OPENAI_API_KEY"] = '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' llm = ChatOpenAI(temperature=0.2, model_name="gpt-3.5-turbo") prompt = PromptTemplate( input_variables=["question", "elements"], template="""You are a helpful assistant that can answer question related to an image. You have the ability to see the image and answer questions about it. I will give you a question and element about the image and you will answer the question. \n\n #Question: {question} #Elements: {elements} \n\n Your structured response:""", ) def convert_png_to_jpg(image): rgb_image = image.convert('RGB') byte_arr = BytesIO() rgb_image.save(byte_arr, format='JPEG') byte_arr.seek(0) return Image.open(byte_arr) def vilt(image, query): processor = ViltProcessor.from_pretrained("dandelin/vilt-b32-finetuned-vqa") model = ViltForQuestionAnswering.from_pretrained("dandelin/vilt-b32-finetuned-vqa") encoding = processor(image, query, return_tensors="pt") outputs = model(**encoding) logits = outputs.logits idx = logits.argmax(-1).item() sol = model.config.id2label[idx] return sol def blip(image, query): processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-large") model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-large") # unconditional image captioning inputs = processor(image, return_tensors="pt") out = model.generate(**inputs) sol = processor.decode(out[0], skip_special_tokens=True) return sol def GIT(image, query): processor = AutoProcessor.from_pretrained("microsoft/git-base-textvqa") model = AutoModelForCausalLM.from_pretrained("microsoft/git-base-textvqa") # file_path = hf_hub_download(repo_id="nielsr/textvqa-sample", filename="bus.png", repo_type="dataset") # image = Image.open(file_path).convert("RGB") pixel_values = processor(images=image, return_tensors="pt").pixel_values question = query input_ids = processor(text=question, add_special_tokens=False).input_ids input_ids = [processor.tokenizer.cls_token_id] + input_ids input_ids = torch.tensor(input_ids).unsqueeze(0) generated_ids = model.generate(pixel_values=pixel_values, input_ids=input_ids, max_length=50) response = processor.batch_decode(generated_ids, skip_special_tokens=True) generated_ids_1 = model.generate(pixel_values=pixel_values, max_length=50) generated_caption = processor.batch_decode(generated_ids_1, skip_special_tokens=True)[0] return response[0] + " " + generated_caption @st.cache_data(show_spinner="Processing image...") def generate_table(uploaded_file): image = Image.open(uploaded_file) print("graph start") model = Pix2StructForConditionalGeneration.from_pretrained('google/deplot') processor = Pix2StructProcessor.from_pretrained('google/deplot') print("graph start 1") inputs = processor(images=image, text="Generate underlying data table of the figure below and give the text as well:", return_tensors="pt") predictions = model.generate(**inputs, max_new_tokens=512) print("end") table = processor.decode(predictions[0], skip_special_tokens=True) print(table) return table def process_query(image, query): blip_sol = blip(image, query) vilt_sol = vilt(image, query) GIT_sol = GIT(image, query) llm_sol = blip_sol + " " + vilt_sol + " " + GIT_sol print(llm_sol) chain = LLMChain(llm=llm, prompt=prompt) response = chain.run(question=query, elements=llm_sol) return response def process_query_graph(data_table, query): prompt = PromptTemplate( input_variables=["question", "elements"], template="""You are a helpful assistant capable of answering questions related to graph images. You possess the ability to view the graph image and respond to inquiries about it. I will provide you with a question and the associated data table of the graph, and you will answer the question \n\n #Question: {question} #Elements: {elements} \n\n Your structured response:""", ) chain = LLMChain(llm=llm, prompt=prompt) response = chain.run(question=query, elements=data_table) return response def chart_with_Image(): st.header("Chat with Image", divider='rainbow') uploaded_file = st.file_uploader('Upload your IMAGE', type=['png', 'jpeg', 'jpg'], key="imageUploader") if uploaded_file is not None: image = Image.open(uploaded_file) # ViLT model only supports JPG images if image.format == 'PNG': image = convert_png_to_jpg(image) st.image(image, caption='Uploaded Image.', width=300) cancel_button = st.button('Cancel') query = st.text_input('Ask a question to the IMAGE') if query: with st.spinner('Processing...'): answer = process_query(image, query) st.write(answer) if cancel_button: st.stop() def chat_with_graph(): st.header("Chat with Graph", divider='rainbow') uploaded_file = st.file_uploader('Upload your GRAPH', type=['png', 'jpeg', 'jpg'], key="graphUploader") if uploaded_file is not None: image = Image.open(uploaded_file) # if image.format == 'PNG': # image = convert_png_to_jpg(image) # data_table = generate_table(uploaded_file) st.image(image, caption='Uploaded Image.') data_table = generate_table(uploaded_file) cancel_button = st.button('Cancel') query = st.text_input('Ask a question to the IMAGE') if query: with st.spinner('Processing...'): answer = process_query_graph(data_table, query) st.write(answer) if cancel_button: st.stop() st.title("Image Querying App ") option = st.selectbox( "Who would you like to chart with?", ("Image", "Graph"), index=None, placeholder="Select contact method...", ) st.write('You selected:', option) if option == "Image": chart_with_Image() elif option == "Graph": chat_with_graph()