105 lines
3.4 KiB
Python
105 lines
3.4 KiB
Python
import json # JSON 모듈 임포트
|
|
import os
|
|
from typing import Optional
|
|
|
|
from google.api_core.client_options import ClientOptions
|
|
from google.cloud import documentai # type: ignore
|
|
|
|
os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = (
|
|
"../drawingpdfocr-461103-2441e0b34216.json" # 이 경로가 정확한지 다시 한번 확인해주세요!
|
|
)
|
|
|
|
|
|
def process_document_sample(
|
|
project_id: str,
|
|
location: str,
|
|
processor_id: str,
|
|
file_path: str,
|
|
mime_type: str,
|
|
field_mask: Optional[str] = None,
|
|
processor_version_id: Optional[str] = None,
|
|
) -> documentai.Document:
|
|
opts = ClientOptions(api_endpoint=f"{location}-documentai.googleapis.com")
|
|
client = documentai.DocumentProcessorServiceClient(client_options=opts)
|
|
|
|
if processor_version_id:
|
|
name = client.processor_version_path(
|
|
project_id, location, processor_id, processor_version_id
|
|
)
|
|
else:
|
|
name = client.processor_path(project_id, location, processor_id)
|
|
|
|
with open(file_path, "rb") as image:
|
|
image_content = image.read()
|
|
|
|
raw_document = documentai.RawDocument(content=image_content, mime_type=mime_type)
|
|
process_options = documentai.ProcessOptions(
|
|
individual_page_selector=documentai.ProcessOptions.IndividualPageSelector(
|
|
pages=[1]
|
|
)
|
|
)
|
|
request = documentai.ProcessRequest(
|
|
name=name,
|
|
raw_document=raw_document,
|
|
field_mask=field_mask,
|
|
process_options=process_options,
|
|
)
|
|
result = client.process_document(request=request)
|
|
document = result.document
|
|
return document
|
|
|
|
|
|
def extract_and_convert_to_json(
|
|
document: documentai.Document,
|
|
) -> str: # 반환 타입을 str (JSON 문자열)로 명시
|
|
"""
|
|
Document AI의 Document 객체에서 entities의 type과 mention_text를 추출하여
|
|
JSON 문자열로 반환합니다.
|
|
"""
|
|
extracted_entities = []
|
|
if document and document.entities:
|
|
for entity in document.entities:
|
|
if (
|
|
hasattr(entity, "type_")
|
|
and hasattr(
|
|
entity, "mention_text"
|
|
) # type_와 mention_text 속성이 있는지 확인
|
|
and entity.type_
|
|
and entity.mention_text
|
|
):
|
|
extracted_entities.append(
|
|
{"type": entity.type_, "mention_text": entity.mention_text}
|
|
)
|
|
|
|
return json.dumps(extracted_entities, ensure_ascii=False, indent=2)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
project_id = "drawingpdfocr-461103"
|
|
location = "us"
|
|
processor_id = "b838676d4e3b4758"
|
|
file_path = "../data/UPLOAD_DOCS/3공구-설계도1-004.pdf"
|
|
mime_type = "application/pdf"
|
|
|
|
try:
|
|
document_result = process_document_sample(
|
|
project_id=project_id,
|
|
location=location,
|
|
processor_id=processor_id,
|
|
file_path=file_path,
|
|
mime_type=mime_type,
|
|
field_mask="text,entities", # entities 정보를 받아오도록 field_mask 설정
|
|
)
|
|
|
|
if document_result:
|
|
json_output_string = extract_and_convert_to_json(document_result)
|
|
|
|
print(json_output_string) # 변환된 JSON 문자열을 출력
|
|
|
|
except FileNotFoundError:
|
|
print(
|
|
f"오류: 파일 경로 '{file_path}'에서 파일을 찾을 수 없습니다. 파일 경로를 확인해주세요."
|
|
)
|
|
except Exception as e:
|
|
print(f"함수 실행 중 오류 발생: {e}")
|