Merge pull request #44 from huridocs/text_extraction

Add text extraction

src/app.py

@@ -9,6 +9,7 @@
 from pdf_layout_analysis.get_xml import get_xml
 from pdf_layout_analysis.run_pdf_layout_analysis import analyze_pdf
 from pdf_layout_analysis.run_pdf_layout_analysis_fast import analyze_pdf_fast
+from text_extraction.get_text_extraction import get_text_extraction
 from toc.get_toc import get_toc
 
 service_logger.info(f"Is PyTorch using GPU: {torch.cuda.is_available()}")
@@ -51,3 +52,9 @@ async def get_xml_by_name(xml_file_name: str):
 @catch_exceptions
 async def get_toc_endpoint(file: UploadFile = File(...), fast: bool = Form(False)):
  return await run_in_threadpool(get_toc, file, fast)
+
+
+@app.post("/text")
+@catch_exceptions
+async def get_text_endpoint(file: UploadFile = File(...), fast: bool = Form(False), types: str = Form("all")):
+ return await run_in_threadpool(get_text_extraction, file, fast, types)

src/test_end_to_end.py

@@ -178,3 +178,28 @@ def test_toc_fast(self):
  self.assertEqual(response_json[0]["indentation"], 0)
  self.assertEqual(response_json[-1]["label"], "C. TITLE LONGER")
  self.assertEqual(response_json[-1]["indentation"], 2)
+
+ def test_text_extraction(self):
+ with open(f"{ROOT_PATH}/test_pdfs/test.pdf", "rb") as stream:
+ files = {"file": stream}
+
+ response = requests.post(f"{self.service_url}/text", files=files)
+
+ response_json = response.json()
+ self.assertEqual(response.status_code, 200)
+ self.assertEqual(response_json.split()[0], "Document")
+ self.assertEqual(response_json.split()[1], "Big")
+ self.assertEqual(response_json.split()[-1], "TEXT")
+
+ def test_text_extraction_fast(self):
+ with open(f"{ROOT_PATH}/test_pdfs/test.pdf", "rb") as stream:
+ files = {"file": stream}
+ data = {"fast": "True"}
+
+ response = requests.post(f"{self.service_url}/text", files=files, data=data)
+
+ response_json = response.json()
+ self.assertEqual(response.status_code, 200)
+ self.assertEqual(response_json.split()[0], "Document")
+ self.assertEqual(response_json.split()[1], "Big")
+ self.assertEqual(response_json.split()[-1], "TEXT")

src/text_extraction/extract_text.py

@@ -0,0 +1,14 @@
+from configuration import service_logger
+from pdf_token_type_labels.TokenType import TokenType
+
+
+def extract_text(segment_boxes: list[dict], types: list[TokenType]):
+ service_logger.info(f"Extracted types: {[t.name for t in types]}")
+ text = "\n".join(
+ [
+ segment_box["text"]
+ for segment_box in segment_boxes
+ if TokenType.from_text(segment_box["type"].replace(" ", "_")) in types
+ ]
+ )
+ return text

src/text_extraction/get_text_extraction.py

@@ -0,0 +1,16 @@
+from fastapi import UploadFile
+from pdf_token_type_labels.TokenType import TokenType
+from pdf_layout_analysis.run_pdf_layout_analysis import analyze_pdf
+from pdf_layout_analysis.run_pdf_layout_analysis_fast import analyze_pdf_fast
+from text_extraction.extract_text import extract_text
+
+
+def get_text_extraction(file: UploadFile, fast: bool, types: str):
+ file_content = file.file.read()
+ if types == "all":
+ token_types: list[TokenType] = [t for t in TokenType]
+ else:
+ token_types = list(set([TokenType.from_text(t.strip().replace(" ", "_")) for t in types.split(",")]))
+ if fast:
+ return extract_text(analyze_pdf_fast(file_content), token_types)
+ return extract_text(analyze_pdf(file_content, ""), token_types)