mangalathkedar commited on
Commit
83a5b2f
·
verified ·
1 Parent(s): c077595

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +59 -102
app.py CHANGED
@@ -10,81 +10,64 @@ import tempfile
10
 
11
  # Load environment variables
12
  load_dotenv()
13
-
14
- # Initialize OpenAI client with error handling
15
- try:
16
- client = OpenAI(api_key=os.getenv('OPENAI_API_KEY'))
17
- if not os.getenv('OPENAI_API_KEY'):
18
- st.error("OpenAI API key not found. Please check your .env file.")
19
- st.stop()
20
- except Exception as e:
21
- st.error(f"Error initializing OpenAI client: {str(e)}")
22
- st.stop()
23
 
24
  def convert_pdf_to_images(pdf_file):
25
  """Convert PDF to list of images using PyMuPDF"""
26
  images = []
27
- try:
28
- with tempfile.NamedTemporaryFile(delete=False, suffix='.pdf') as tmp_file:
29
- tmp_file.write(pdf_file.getvalue())
30
- pdf_path = tmp_file.name
 
 
 
 
 
 
 
31
 
32
- pdf_document = fitz.open(pdf_path)
33
- for page_number in range(pdf_document.page_count):
34
- page = pdf_document[page_number]
35
- pix = page.get_pixmap()
36
- img_data = pix.tobytes("png")
37
- image = Image.open(io.BytesIO(img_data))
38
- images.append(image)
39
-
40
- pdf_document.close()
41
- os.unlink(pdf_path)
42
- return images
43
- except Exception as e:
44
- st.error(f"Error converting PDF to images: {str(e)}")
45
- return []
46
 
47
  def format_response(text):
48
  """Format the analysis response with clean styling"""
49
- try:
50
- formatted_text = ""
 
 
 
 
 
51
 
52
- # Split into pages
53
- pages = text.split("Page")
 
 
 
 
 
 
 
 
 
54
 
55
- for page_num, page_content in enumerate(pages[1:], 1):
56
- formatted_text += f'\n### Page {page_num}\n'
57
-
58
- # Process each line
59
- lines = page_content.split('\n')
60
- for line in lines:
61
- if line.strip() and not line.strip().startswith('*') and not line.strip().startswith('Here'):
62
- line = line.replace('**', '').replace('- ', '')
63
-
64
- if ':' in line:
65
- label, value = line.split(':', 1)
66
- formatted_text += f'- **{label.strip()}**: {value.strip()}\n'
67
-
68
- return formatted_text
69
- except Exception as e:
70
- st.error(f"Error formatting response: {str(e)}")
71
- return text
72
 
73
  def analyze_image(image):
74
  """Analyze image using OpenAI API"""
75
  try:
76
- # Add timeout parameter for API request
77
- timeout = 30 # seconds
78
-
79
  img_byte_arr = io.BytesIO()
80
  image.save(img_byte_arr, format='PNG')
81
  img_byte_arr = img_byte_arr.getvalue()
82
 
83
  base64_image = base64.b64encode(img_byte_arr).decode("utf-8")
84
 
85
- # Fix model name typo and add error handling for API response
86
  response = client.chat.completions.create(
87
- model="gpt-4o-mini", # Corrected model name
88
  messages=[
89
  {
90
  "role": "user",
@@ -110,69 +93,43 @@ def analyze_image(image):
110
  ],
111
  }
112
  ],
113
- max_tokens=1000,
114
- timeout=timeout
115
  )
116
 
117
  return response.choices[0].message.content
118
  except Exception as e:
119
- error_message = str(e)
120
- if "timeout" in error_message.lower():
121
- return "The request timed out. Please try again."
122
- elif "api key" in error_message.lower():
123
- return "Invalid or missing API key. Please check your configuration."
124
- else:
125
- return f"An error occurred during analysis: {error_message}"
126
 
127
  def main():
128
  st.set_page_config(page_title="Document Analysis App", layout="wide")
129
 
130
  st.title("Document Analysis App")
131
-
132
- # Add API key input if not in environment
133
- if not os.getenv('OPENAI_API_KEY'):
134
- api_key = st.text_input("Enter your OpenAI API key:", type="password")
135
- if api_key:
136
- os.environ['OPENAI_API_KEY'] = api_key
137
- st.success("API key set successfully!")
138
- else:
139
- st.warning("Please enter your OpenAI API key to continue.")
140
- st.stop()
141
-
142
  uploaded_file = st.file_uploader("Upload document (PDF/Image)", type=['pdf', 'png', 'jpg', 'jpeg'])
143
 
144
  if uploaded_file is not None:
145
- try:
146
- if uploaded_file.type == "application/pdf":
147
- # Handle PDF
148
- with st.spinner("Processing PDF..."):
149
- images = convert_pdf_to_images(uploaded_file)
150
-
151
- if not images:
152
- st.error("Failed to process PDF file.")
153
- st.stop()
154
 
155
- if st.button("Extract Information"):
156
- with st.spinner("Analyzing document..."):
157
- all_results = []
158
- for i, image in enumerate(images, 1):
159
- result = analyze_image(image)
160
- all_results.append(f"Page {i} Information:\n{result}")
161
 
162
- combined_results = "\n\n".join(all_results)
163
- st.markdown(format_response(combined_results))
164
 
165
- else:
166
- # Handle single image
167
- image = Image.open(uploaded_file)
168
 
169
- if st.button("Extract Information"):
170
- with st.spinner("Analyzing document..."):
171
- result = analyze_image(image)
172
- st.markdown(format_response(result))
173
-
174
- except Exception as e:
175
- st.error(f"An error occurred while processing the file: {str(e)}")
176
 
177
- if __name__ == "__main__":
178
  main()
 
10
 
11
  # Load environment variables
12
  load_dotenv()
13
+ # Initialize OpenAI client
14
+ client = OpenAI(api_key=os.getenv('OPENAI_API_KEY'))
 
 
 
 
 
 
 
 
15
 
16
  def convert_pdf_to_images(pdf_file):
17
  """Convert PDF to list of images using PyMuPDF"""
18
  images = []
19
+ with tempfile.NamedTemporaryFile(delete=False, suffix='.pdf') as tmp_file:
20
+ tmp_file.write(pdf_file.getvalue())
21
+ pdf_path = tmp_file.name
22
+
23
+ pdf_document = fitz.open(pdf_path)
24
+ for page_number in range(pdf_document.page_count):
25
+ page = pdf_document[page_number]
26
+ pix = page.get_pixmap()
27
+ img_data = pix.tobytes("png")
28
+ image = Image.open(io.BytesIO(img_data))
29
+ images.append(image)
30
 
31
+ pdf_document.close()
32
+ os.unlink(pdf_path)
33
+ return images
34
+
 
 
 
 
 
 
 
 
 
 
35
 
36
  def format_response(text):
37
  """Format the analysis response with clean styling"""
38
+ formatted_text = ""
39
+
40
+ # Split into pages
41
+ pages = text.split("Page")
42
+
43
+ for page_num, page_content in enumerate(pages[1:], 1): # Skip first empty split
44
+ formatted_text += f'\n### Page {page_num}\n'
45
 
46
+ # Process each line
47
+ lines = page_content.split('\n')
48
+ for line in lines:
49
+ # Skip empty lines and lines with asterisks
50
+ if line.strip() and not line.strip().startswith('*') and not line.strip().startswith('Here'):
51
+ # Remove asterisks and dashes
52
+ line = line.replace('**', '').replace('- ', '')
53
+
54
+ if ':' in line:
55
+ label, value = line.split(':', 1)
56
+ formatted_text += f'- *{label.strip()}*: {value.strip()}\n'
57
 
58
+ return formatted_text
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
59
 
60
  def analyze_image(image):
61
  """Analyze image using OpenAI API"""
62
  try:
 
 
 
63
  img_byte_arr = io.BytesIO()
64
  image.save(img_byte_arr, format='PNG')
65
  img_byte_arr = img_byte_arr.getvalue()
66
 
67
  base64_image = base64.b64encode(img_byte_arr).decode("utf-8")
68
 
 
69
  response = client.chat.completions.create(
70
+ model="gpt-4o-mini",
71
  messages=[
72
  {
73
  "role": "user",
 
93
  ],
94
  }
95
  ],
96
+ max_tokens=1000
 
97
  )
98
 
99
  return response.choices[0].message.content
100
  except Exception as e:
101
+ return f"An error occurred: {str(e)}"
 
 
 
 
 
 
102
 
103
  def main():
104
  st.set_page_config(page_title="Document Analysis App", layout="wide")
105
 
106
  st.title("Document Analysis App")
 
 
 
 
 
 
 
 
 
 
 
107
  uploaded_file = st.file_uploader("Upload document (PDF/Image)", type=['pdf', 'png', 'jpg', 'jpeg'])
108
 
109
  if uploaded_file is not None:
110
+ if uploaded_file.type == "application/pdf":
111
+ # Handle PDF
112
+ with st.spinner("Processing PDF..."):
113
+ images = convert_pdf_to_images(uploaded_file)
 
 
 
 
 
114
 
115
+ if st.button("Extract Information"):
116
+ with st.spinner("Analyzing document..."):
117
+ all_results = []
118
+ for i, image in enumerate(images, 1):
119
+ result = analyze_image(image)
120
+ all_results.append(f"Page {i} Information:\n{result}")
121
 
122
+ combined_results = "\n\n".join(all_results)
123
+ st.markdown(format_response(combined_results))
124
 
125
+ else:
126
+ # Handle single image
127
+ image = Image.open(uploaded_file)
128
 
129
+ if st.button("Extract Information"):
130
+ with st.spinner("Analyzing document..."):
131
+ result = analyze_image(image)
132
+ st.markdown(format_response(result))
 
 
 
133
 
134
+ if _name_ == "_main_":
135
  main()