import gradio as gr
import requests
from PIL import Image
from pdf2image import convert_from_path
from typing import List, Union, Dict, Optional, Tuple
from io import BytesIO
import base64
import numpy as np
import json
prompt = """You are an advanced document parsing bot. Given the fixture schedule I provided, you need to parse out
1. the name of the fixture
2. the company that produces this fixture
3. the description of this fixture. This is a 20-word description which summarize the size, function and the mounting method of the fixture and mention any necessary accesories. For example: 1" x 1" recessed downlight.
4. the part number of this fixture. It is a series of specification codes connected with - , and you can get the info by reading the texts marked in a different color or reading the top bar. Include every specification code in a correct order in your answer.
5. the input wattage of this fixture, short answer. Please answer the wattage according to the part number you found in question 3
Please format your response in json format
{
"fixture_name": <fixture name>,
"manufacture_name": <company name>,
"fixture_description": <description>,
"mfr": <part number>,
"input wattage": <numerical input wattage>
}
---
For example
{
"fixture_name": "SW24/1.5 Led Strips - Static White",
"manufacture_name": "Q-Tran Inc.",
"fixture_description": "Surface mounted static white LED strip."
"mfr": "SW24-1.5-DRY-30-BW-BW-WH-CL2-535",
"input wattage": "1.5W"
}"""
def query_openai_api(messages, model, temperature=0, api_key=None, organization_key=None, json_mode=False):
try:
url = "https://api.openai.com/v1/chat/completions"
if organization_key is not None:
headers = {
"Content-Type": "application/json",
"Authorization": f"Bearer {api_key}",
"OpenAI-Organization": f"{organization_key}",
}
else:
headers = {
"Content-Type": "application/json",
"Authorization": f"Bearer {api_key}",
}
data = {"model": model, "messages": messages, "temperature": temperature}
if json_mode:
data["response_format"] = {"type": "json_object"}
response = requests.post(url, headers=headers, data=json.dumps(data)).json()
print(response)
return response["choices"][0]["message"]["content"].lstrip(), response
except Exception as e:
print(f"An error occurred: {e}")
return f"API_ERROR: {e}", None
class GPT4V_Client:
def __init__(self, api_key, organization_key, model_name="gpt-4o", max_tokens=512):
self.api_key = api_key
self.organization_key = organization_key
self.model_name = model_name
self.max_tokens = max_tokens
def chat(self, messages, json_mode):
return query_openai_api(messages, self.model_name, api_key=self.api_key, organization_key=self.organization_key, json_mode=json_mode)
def one_step_chat(
self,
text,
image: Union[Image.Image, np.ndarray],
system_msg: Optional[str] = None,
json_mode=False,
):
jpeg_buffer = BytesIO()
# Save the image as JPEG to the buffer
if isinstance(image, np.ndarray):
image = Image.fromarray(image)
image = image.convert("RGB")
image.save(jpeg_buffer, format="JPEG")
# Get the byte data from the buffer
jpeg_data = jpeg_buffer.getvalue()
# Encode the JPEG image data in base64
jpg_base64 = base64.b64encode(jpeg_data)
# If you need it in string format
jpg_base64_str = jpg_base64.decode("utf-8")
messages = []
if system_msg is not None:
messages.append({"role": "system", "content": system_msg})
messages += [
{
"role": "user",
"content": [
{"type": "text", "text": text},
{
"type": "image_url",
"image_url": {
"url": f"data:image/jpeg;base64,{jpg_base64_str}"
},
},
],
}
]
return self.chat(messages, json_mode=json_mode)
def one_step_multi_image_chat(
self,
text,
images: list[Union[Image.Image, np.ndarray]],
system_msg: Optional[str] = None,
json_mode=False,
):
details = [i["detail"] for i in images]
img_strs = []
for img_info in images:
image = img_info["image"]
jpeg_buffer = BytesIO()
if isinstance(image, np.ndarray):
image = Image.fromarray(image)
image = image.convert("RGB")
image.save(jpeg_buffer, format="JPEG")
jpeg_data = jpeg_buffer.getvalue()
jpg_base64 = base64.b64encode(jpeg_data)
jpg_base64_str = jpg_base64.decode("utf-8")
img_strs.append(f"data:image/jpeg;base64,{jpg_base64_str}")
messages = []
if system_msg is not None:
messages.append({"role": "system", "content": system_msg})
img_sub_msg = [
{
"type": "image_url",
"image_url": {"url": img_str, "detail": detail},
}
for img_str, detail in zip(img_strs, details)
]
messages += [
{
"role": "user",
"content": [{"type": "text", "text": text}] + img_sub_msg,
}
]
return self.chat(messages, json_mode=json_mode)
def markdown_json_to_table(markdown_json_string, iteration, thumbnail_md):
"""
Convert the GPT JSON string into a markdown row with the first column as the PDF thumbnail.
Args:
markdown_json_string: the raw markdown (JSON) string from GPT
iteration: which row # we are on
thumbnail_md: something like 
Returns:
A string with either:
- The header row + first data row, if iteration==0
- Additional data row, if iteration>0
"""
# Try to detect if the JSON is enclosed in triple-backticks
# so we can parse it out properly:
if markdown_json_string.strip().startswith("```"):
# Remove the backticks and possible extra notations
json_string = markdown_json_string.strip().strip("```").strip("json").strip()
else:
# If the model didn't wrap it in markdown
json_string = markdown_json_string.strip()
# Safely parse JSON
try:
json_obj = json.loads(json_string)
except Exception:
# If it can't parse, return empty
return ""
# Turn the JSON object into a list of values for easier table building
# e.g. [fixture_name, manufacture_name, mfr, input wattage]
keys = list(json_obj.keys())
values = list(json_obj.values())
# We want the first column to be the PDF thumbnail
# So the table columns become: [Thumbnail, key1, key2, key3, ...]
# This means we have one extra column in front compared to the JSON.
# If iteration == 0, produce header
# e.g. | Thumbnail | fixture_name | manufacture_name | mfr | input wattage |
if iteration == 0:
header = ["Thumbnail"] + keys
header_row = "| " + " | ".join(header) + " |\n"
sep_row = "|" + "|".join(["---"] * len(header)) + "|\n"
else:
header_row = ""
sep_row = ""
# Then produce the data row
# e.g. |  | "SW24..." | "Q-Tran Inc." | ...
str_values = [str(v) for v in values]
data_row = "| " + thumbnail_md + " | " + " | ".join(str_values) + " |\n"
return header_row + sep_row + data_row
def gptRead(cutsheets, api_key, organization_key):
fixtureInfo = ""
iteration = 0
client = GPT4V_Client(api_key=api_key, organization_key=organization_key)
for cutsheet in cutsheets:
# Convert the first page of the PDF into an image
source = (convert_from_path(cutsheet.name))[0]
# Create a smaller thumbnail
thumbnail_img = source.copy()
thumbnail_img.thumbnail((100, 100))
# Encode the thumbnail to base64 for embedding in Markdown
thumb_io = BytesIO()
thumbnail_img.save(thumb_io, format="JPEG")
base64_thumb = base64.b64encode(thumb_io.getvalue()).decode('utf-8')
thumbnail_md = f""
# Chat with GPT about the original (non-thumbnail) image
response_text, _ = client.one_step_chat(prompt, source)
# Convert the GPT JSON to a Markdown row, including the thumbnail in the first column
fixtureInfo += markdown_json_to_table(response_text, iteration, thumbnail_md)
iteration += 1
return fixtureInfo
if __name__ == "__main__":
with gr.Blocks() as demo:
api_key = gr.Textbox(label="Input your ChatGPT4 API Key: ")
organization_key = gr.Textbox(label="Input your ChatGPT4 API Organization Key: ", info="(optional)")
gr.Markdown("# Lighting Manufacture Cutsheet GPT Tool")
file_uploader = gr.UploadButton("Upload cutsheets", type="filepath", file_count="multiple")
form = gr.Markdown()
# When user uploads, call gptRead -> produce the final Markdown w/ table
file_uploader.upload(fn=gptRead, inputs=[file_uploader, api_key, organization_key], outputs=form)
demo.launch(share=True)