jackkuo commited on
Commit
b26838d
·
verified ·
1 Parent(s): 1f32a6b

Upload ocr_mathpix.py

Browse files
Files changed (1) hide show
  1. ocr_mathpix.py +213 -0
ocr_mathpix.py ADDED
@@ -0,0 +1,213 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import fitz
3
+ import requests
4
+ import json
5
+ import time
6
+ import logging
7
+
8
+ # Configure logging
9
+ logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
10
+
11
+
12
+ def get_pdf_pages(pdf_folder_dir, pdf_dir):
13
+ """
14
+ Get the number of pages in a PDF file.
15
+
16
+ Parameters:
17
+ pdf_folder_dir: str - The directory of the PDF folder.
18
+ pdf_dir: str - The name of the PDF file.
19
+
20
+ Returns:
21
+ int - The total number of pages in the PDF file, or None if the PDF cannot be read.
22
+ """
23
+ # Construct the full path to the PDF file
24
+ path = pdf_folder_dir + "/" + pdf_dir
25
+
26
+ # Attempt to open the PDF file
27
+ try:
28
+ doc = fitz.open(path)
29
+ except:
30
+ # If the file cannot be opened, print an error message and return None
31
+ print("can not read pdf")
32
+ return None
33
+
34
+ # Get and return the number of pages in the PDF file
35
+ page_count = doc.page_count
36
+
37
+ return page_count
38
+
39
+
40
+ def get_api_credentials():
41
+ """Retrieve Mathpix API credentials from environment variables"""
42
+ APP_ID = os.getenv('MATHPIX_APP_ID')
43
+ print(APP_ID)
44
+ APP_KEY = os.getenv('MATHPIX_APP_KEY')
45
+ if not APP_ID or not APP_KEY:
46
+ raise ValueError("Please set MATHPIX_APP_ID and MATHPIX_APP_KEY environment variables")
47
+ return APP_ID, APP_KEY
48
+
49
+ def upload_pdf_to_mathpix(pdf_file_path, headers, options):
50
+ """Upload the PDF file to Mathpix API"""
51
+ url = 'https://api.mathpix.com/v3/pdf'
52
+ with open(pdf_file_path, 'rb') as pdf_file:
53
+ files = {
54
+ 'file': pdf_file,
55
+ 'options_json': (None, json.dumps(options))
56
+ }
57
+ response = requests.post(url, headers=headers, files=files)
58
+ return response
59
+
60
+
61
+ def check_conversion_status(pdf_id, headers, max_retries=30, retry_interval=5):
62
+ """Check the conversion status with a maximum number of retries and interval"""
63
+ status_url = f'https://api.mathpix.com/v3/pdf/{pdf_id}'
64
+ retries = 0
65
+
66
+ while retries < max_retries:
67
+ status_response = requests.get(status_url, headers=headers)
68
+ status_data = status_response.json()
69
+ conversion_status = status_data.get('status', 'unknown')
70
+ logging.info(f"conversion_status: {conversion_status}")
71
+
72
+ # Log the full response data for debugging purposes
73
+ logging.debug(f"Full conversion status response: {status_data}")
74
+
75
+ if conversion_status == 'completed':
76
+ break
77
+ elif conversion_status in ['loaded', 'split', 'processing']:
78
+ logging.info(f"Conversion is {conversion_status}, waiting for processing to complete.")
79
+ time.sleep(retry_interval)
80
+ retries += 1
81
+ continue
82
+ else:
83
+ raise ValueError(f"Conversion failed, status: {conversion_status}")
84
+
85
+ logging.info('Processing... Please wait.')
86
+ time.sleep(retry_interval)
87
+ retries += 1
88
+
89
+ if retries >= max_retries:
90
+ raise TimeoutError("Conversion did not complete within the allowed time.")
91
+
92
+
93
+ def download_md_file(pdf_id, headers, output_dir, output_filename):
94
+ """Download and save the Markdown file"""
95
+ md_url = f'https://api.mathpix.com/v3/pdf/{pdf_id}.md'
96
+ md_response = requests.get(md_url, headers=headers)
97
+ if md_response.status_code == 200:
98
+ os.makedirs(output_dir, exist_ok=True)
99
+ output_path = os.path.join(output_dir, output_filename)
100
+ with open(output_path, "w", encoding="utf-8") as fout:
101
+ fout.write(md_response.text)
102
+ logging.info(f"OCR result saved to: {output_path}")
103
+ return md_response.text
104
+ else:
105
+ logging.error('Failed to download Markdown file.')
106
+ return None
107
+
108
+
109
+ def extract_pdf_mathpix(pdf_folder_dir, pdf_dir, md_folder_dir):
110
+ """
111
+ Extract content from a PDF file and convert it to Markdown format
112
+ """
113
+ try:
114
+ # Retrieve API credentials
115
+ APP_ID, APP_KEY = get_api_credentials()
116
+
117
+ # Build the PDF file path
118
+ pdf_file_path = os.path.join(pdf_folder_dir, pdf_dir)
119
+ logging.info(f"pdf_file_path: {pdf_file_path}")
120
+
121
+ # Check if the file exists
122
+ if not os.path.exists(pdf_file_path):
123
+ raise FileNotFoundError(f"File {pdf_file_path} does not exist")
124
+
125
+ # Set request headers and options
126
+ headers = {
127
+ 'app_id': APP_ID,
128
+ 'app_key': APP_KEY,
129
+ }
130
+ options = {
131
+ "conversion_formats": {
132
+ "md": True
133
+ },
134
+ "math_inline_delimiters": ["$", "$"],
135
+ "rm_spaces": True
136
+ }
137
+
138
+ # Upload the PDF file
139
+ response = upload_pdf_to_mathpix(pdf_file_path, headers, options)
140
+ if response.status_code != 200:
141
+ logging.error(f'Failed to upload PDF. Status code: {response.status_code}')
142
+ return None
143
+
144
+ # Get the PDF ID
145
+ pdf_id = response.json().get('pdf_id')
146
+ logging.info(f"pdf_id: {pdf_id}")
147
+
148
+ # Check the conversion status
149
+ check_conversion_status(pdf_id, headers)
150
+
151
+ # Download and save the Markdown file
152
+ output_filename = os.path.splitext(pdf_dir)[0] + ".md"
153
+ return download_md_file(pdf_id, headers, md_folder_dir, output_filename)
154
+
155
+ except Exception as e:
156
+ logging.error(f"An error occurred: {e}")
157
+ return None
158
+
159
+
160
+ def get_done_papers(md_folder_dir):
161
+ done_paper = []
162
+ if os.path.exists(md_folder_dir):
163
+ try:
164
+ done_paper = [i.replace(".md", ".pdf") for i in os.listdir(md_folder_dir)]
165
+ except (FileNotFoundError, PermissionError) as e:
166
+ print(f"Error reading md folder: {e}")
167
+ return done_paper
168
+
169
+
170
+ def process_pdfs(pdf_folder_dir, done_paper, md_folder_dir):
171
+ no_response_paper = []
172
+ pages_more_50 = []
173
+
174
+ try:
175
+ pdf_files = [i for i in os.listdir(pdf_folder_dir) if i.endswith("pdf")]
176
+ except (FileNotFoundError, PermissionError) as e:
177
+ print(f"Error reading pdf folder: {e}")
178
+ return no_response_paper, pages_more_50, done_paper
179
+
180
+ for pdf_file in pdf_files:
181
+ if pdf_file not in done_paper + no_response_paper + pages_more_50:
182
+ try:
183
+ pages = get_pdf_pages(pdf_folder_dir, pdf_file)
184
+ print(f"\nstart: {pdf_file} have pages: {pages}")
185
+
186
+ if pages <= 50:
187
+ print(f"start convert pdf 2 md: {pdf_file}")
188
+ content = extract_pdf_mathpix(pdf_folder_dir, pdf_file, md_folder_dir)
189
+ if content:
190
+ done_paper.append(pdf_file)
191
+ else:
192
+ no_response_paper.append(pdf_file)
193
+ else:
194
+ pages_more_50.append(pdf_file)
195
+ print(f"pages_more_50: {pages_more_50}")
196
+ except Exception as e:
197
+ print(f"Error processing {pdf_file}: {e}")
198
+
199
+ return no_response_paper, pages_more_50, done_paper
200
+
201
+
202
+ if __name__ == '__main__':
203
+ data_folder_dir = "../data/"
204
+ pdf_folder_dir = os.path.join(data_folder_dir, "pdf")
205
+ md_folder_dir = os.path.join(data_folder_dir, "md")
206
+
207
+ done_paper = get_done_papers(md_folder_dir)
208
+ print("done_paper:", done_paper)
209
+
210
+ no_response_paper, pages_more_50, done_paper = process_pdfs(pdf_folder_dir, done_paper, md_folder_dir)
211
+ print("done_paper:", done_paper)
212
+ print("no_response_paper:", no_response_paper)
213
+ print("pages_more_50:", pages_more_50)