fadliaulawi's picture
Initial commit
fb4710e
raw
history blame
2.05 kB
import pandas as pd
import base64
from img2table.document import Image
from img2table.ocr import DocTR
from langchain.schema.messages import HumanMessage, AIMessage
from langchain_experimental.agents import create_pandas_dataframe_agent
from langchain_openai import ChatOpenAI
from dotenv import load_dotenv
load_dotenv()
def encode_image(image_path):
with open(image_path, "rb") as image_file:
return base64.b64encode(image_file.read()).decode('utf-8')
path = '../NutriGenMe-Testing/ukmss-1.png'
vision = ChatOpenAI(model="gpt-4-vision-preview", max_tokens=4096)
img = encode_image(path)
msg = vision.invoke(
[
AIMessage(content="You are an experienced doctor specializing in genomics and want to identify names of genes, SNPs, and their related diseases based on the tables given."),
HumanMessage(
content=[
{ "type": "text",
"text": 'You will be provided with the image of a table. Extract all genes / locus names with its respective rsID / SNP and potential diseases in curly brackets like this: {"Genes" : "", "SNPs" : "", "Diseases" : ""}.'
},
{
"type": "image_url",
"image_url": {
"url": f"data:image/jpeg;base64,{img}",
"detail": "low"
},
},
]
)
]
)
print(msg.content)
# exit()
image = Image(path)
ocr = DocTR()
extracted_tables = image.extract_tables(ocr=ocr,
implicit_rows=True,
borderless_tables=True,
min_confidence=0)
df = extracted_tables[0].df
for et in extracted_tables[1:]:
df = pd.concat([df, et.df]).reset_index(drop=True)
print(df)
llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0)
agent = create_pandas_dataframe_agent(llm, df, verbose=True)
agent_output = agent.invoke("Is this table contain Gene names?")
print(agent_output)