Spaces:
Sleeping
Sleeping
File size: 6,187 Bytes
ed5ccb1 01133a9 6dc8e89 ed5ccb1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 |
import base64
import streamlit as st
from openai import OpenAI
import pdf2image
from PIL import Image
import io
import tempfile
# Load environment variables
# Initialize OpenAI client
client = OpenAI()
def convert_pdf_to_images(pdf_file):
"""Convert PDF to list of images"""
with tempfile.NamedTemporaryFile(delete=False, suffix='.pdf') as tmp_file:
tmp_file.write(pdf_file.getvalue())
pdf_path = tmp_file.name
images = pdf2image.convert_from_path(pdf_path)
os.unlink(pdf_path)
return images
def format_response(text):
"""Format the analysis response with clean styling"""
formatted_text = """
<div style="
background-color: white;
padding: 20px;
border-radius: 5px;
font-family: Arial, sans-serif;
box-shadow: 0 2px 4px rgba(0,0,0,0.1);
">
"""
# Split into pages
pages = text.split("Page")
for page_num, page_content in enumerate(pages[1:], 1): # Skip first empty split
# Add page header
formatted_text += f'<div style="margin-bottom: 30px;">'
formatted_text += f'<h3 style="color: #2c3e50; margin-bottom: 15px;">Page {page_num}</h3>'
# Process each line
lines = page_content.split('\n')
for line in lines:
# Skip empty lines and lines with asterisks
if line.strip() and not line.strip().startswith('*') and not line.strip().startswith('Here'):
# Remove asterisks and dashes
line = line.replace('**', '').replace('- ', '')
if ':' in line:
label, value = line.split(':', 1)
formatted_text += f'<div style="margin-bottom: 10px; display: flex;">'
formatted_text += f'<span style="font-weight: bold; color: #2c3e50; min-width: 200px;">{label.strip()}</span>'
formatted_text += f'<span style="color: #34495e; flex: 1;">{value.strip()}</span>'
formatted_text += '</div>'
formatted_text += '</div>'
# Add separator between pages except for the last page
if page_num < len(pages) - 1:
formatted_text += '<hr style="border: 1px solid #eee; margin: 20px 0;">'
formatted_text += '</div>'
return formatted_text
def analyze_image(image):
"""Analyze image using OpenAI API"""
try:
img_byte_arr = io.BytesIO()
image.save(img_byte_arr, format='PNG')
img_byte_arr = img_byte_arr.getvalue()
base64_image = base64.b64encode(img_byte_arr).decode("utf-8")
response = client.chat.completions.create(
model="gpt-4o-mini",
messages=[
{
"role": "user",
"content": [
{
"type": "text",
"text": """Please analyze the image and extract the following information:
- Sender information
- Recipient information
- Container details
- Weights and measurements
- Dates and reference numbers
- Cargo details
Format the response as 'Label: Value' pairs."""
},
{
"type": "image_url",
"image_url": {
"url": f"data:image/jpeg;base64,{base64_image}"
},
},
],
}
],
max_tokens=1000
)
return response.choices[0].message.content
except Exception as e:
return f"An error occurred: {str(e)}"
def main():
st.set_page_config(page_title="Document Analysis App", layout="wide")
# Custom CSS to set light background and improve button styling
st.markdown("""
<style>
.stApp {
background-color: white;
}
.stButton>button {
width: 100%;
background-color: #2c3e50;
color: white;
border: none;
padding: 10px 20px;
border-radius: 5px;
margin-top: 20px;
margin-bottom: 20px;
}
.stButton>button:hover {
background-color: #34495e;
}
.uploadedFile {
margin-bottom: 20px;
}
</style>
""", unsafe_allow_html=True)
col1, col2, col3 = st.columns([1,2,1])
with col2:
st.title("Document Analysis App")
uploaded_file = st.file_uploader("Upload document (PDF/Image)", type=['pdf', 'png', 'jpg', 'jpeg'])
if uploaded_file is not None:
if uploaded_file.type == "application/pdf":
# Handle PDF
with st.spinner("Processing PDF..."):
images = convert_pdf_to_images(uploaded_file)
if st.button("Extract Information"):
with st.spinner("Analyzing document..."):
all_results = []
for i, image in enumerate(images, 1):
result = analyze_image(image)
all_results.append(f"Page {i} Information:\n{result}")
combined_results = "\n\n".join(all_results)
st.markdown(format_response(combined_results), unsafe_allow_html=True)
else:
# Handle single image
image = Image.open(uploaded_file)
if st.button("Extract Information"):
with st.spinner("Analyzing document..."):
result = analyze_image(image)
st.markdown(format_response(result), unsafe_allow_html=True)
if __name__ == "__main__":
main() |