언어모델 LLM/프로젝트 예제

Streamlit으로 PDF 추출기 만들기

General AI 2024. 7. 12. 06:26
728x90
반응형

Streamlit으로 PDF 추출기 만들기: 전체 코드 및 단계별 설명

Streamlit을 사용하여 PDF 파일의 내용을 추출하고 Markdown 형식으로 변환하는 웹 애플리케이션을 만드는 방법을 단계별로 알아보겠습니다. 이 도구는 PDF 파일을 업로드하고, 내용을 추출한 후 Markdown 형식으로 다운로드할 수 있게 해줍니다.

Streamlit으로 PDF 추출기 만들기

전체 코드

먼저 전체 코드를 살펴보겠습니다:

import streamlit as st
import os
from datetime import datetime
from pdf2image import convert_from_path
import pytesseract
from PIL import Image
import io

# Tesseract 경로 설정 (시스템에 따라 다를 수 있음)
pytesseract.pytesseract.tesseract_cmd = r'/usr/bin/tesseract'

# 필요한 디렉토리 설정
UPLOAD_DIR = "uploaded_pdfs"
MARKDOWN_DIR = "converted_markdown"

# 필요한 디렉토리가 없으면 생성
for dir_path in [UPLOAD_DIR, MARKDOWN_DIR]:
    if not os.path.exists(dir_path):
        os.makedirs(dir_path)

def pdf_to_markdown(pdf_path, markdown_path):
    try:
        # PDF를 이미지로 변환
        images = convert_from_path(pdf_path)

        # 각 이미지에서 텍스트 추출 및 Markdown 파일에 쓰기
        with open(markdown_path, 'w', encoding='utf-8') as md_file:
            for i, image in enumerate(images):
                text = pytesseract.image_to_string(image)
                md_file.write(f"# Page {i+1}\n\n{text}\n\n")
        return True
    except Exception as e:
        st.error(f"PDF 변환 중 오류 발생: {str(e)}")
        return False

st.title("PDF2MD 변환기")

# 파일 업로더 위젯
uploaded_file = st.file_uploader("PDF 파일을 선택하세요", type="pdf")

if uploaded_file is not None:
    # 파일 이름에 타임스탬프 추가
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    file_name = f"{timestamp}_{uploaded_file.name}"
    pdf_path = os.path.join(UPLOAD_DIR, file_name)

    # PDF 파일 저장
    with open(pdf_path, "wb") as f:
        f.write(uploaded_file.getbuffer())

    st.success(f"파일이 성공적으로 업로드되었습니다: {file_name}")

    # Markdown으로 변환
    markdown_file_name = file_name.rsplit('.', 1)[0] + '.md'
    markdown_path = os.path.join(MARKDOWN_DIR, markdown_file_name)

    if pdf_to_markdown(pdf_path, markdown_path):
        st.success(f"PDF가 Markdown으로 변환되었습니다: {markdown_file_name}")
    else:
        st.error("PDF를 Markdown으로 변환하는 데 실패했습니다.")

# 현재 세션의 변환된 Markdown 파일만 표시 및 다운로드 버튼 제공
if 'converted_files' not in st.session_state:
    st.session_state.converted_files = []

if uploaded_file is not None and markdown_file_name:
    st.session_state.converted_files.append(markdown_file_name)

st.subheader("변환된 Markdown 파일 목록")
if st.session_state.converted_files:
    for file in st.session_state.converted_files:
        col1, col2 = st.columns([3, 1])
        with col1:
            st.write(file)
        with col2:
            with open(os.path.join(MARKDOWN_DIR, file), "r", encoding="utf-8") as md_file:
                markdown_content = md_file.read()
            st.download_button(
                label="다운로드",
                data=markdown_content,
                file_name=file,
                mime="text/markdown",
                key=file  # 각 버튼에 고유한 키 부여
            )
else:
    st.write("현재 세션에서 변환된 Markdown 파일이 없습니다.")

코드 설명

이제 코드의 각 부분을 단계별로 자세히 살펴보겠습니다.

STEP 1. 필요한 라이브러리 임포트

import streamlit as st
import os
from datetime import datetime
from pdf2image import convert_from_path
import pytesseract
from PIL import Image
import io

이 단계에서는 필요한 모든 라이브러리를 임포트합니다. streamlit은 웹 애플리케이션 구축을 위해, pdf2image는 PDF를 이미지로 변환하기 위해, pytesseract는 OCR(광학 문자 인식)을 위해 사용됩니다.

STEP 2. 디렉토리 설정

UPLOAD_DIR = "uploaded_pdfs"
MARKDOWN_DIR = "converted_markdown"

for dir_path in [UPLOAD_DIR, MARKDOWN_DIR]:
    if not os.path.exists(dir_path):
        os.makedirs(dir_path)

이 단계에서는 업로드된 PDF 파일과 변환된 Markdown 파일을 저장할 디렉토리를 설정하고 생성합니다.

STEP 3. PDF를 Markdown으로 변환하는 함수 정의

def pdf_to_markdown(pdf_path, markdown_path):
    try:
        images = convert_from_path(pdf_path)
        with open(markdown_path, 'w', encoding='utf-8') as md_file:
            for i, image in enumerate(images):
                text = pytesseract.image_to_string(image)
                md_file.write(f"# Page {i+1}\n\n{text}\n\n")
        return True
    except Exception as e:
        st.error(f"PDF 변환 중 오류 발생: {str(e)}")
        return False

이 함수는 PDF 파일을 이미지로 변환한 후, 각 이미지에서 텍스트를 추출하여 Markdown 형식으로 저장합니다.

STEP 4. Streamlit 인터페이스 구성

st.title("PDF2MD 변환기")

uploaded_file = st.file_uploader("PDF 파일을 선택하세요", type="pdf")

이 단계에서는 Streamlit을 사용하여 웹 인터페이스를 구성합니다. 사용자가 PDF 파일을 업로드할 수 있는 파일 업로더를 제공합니다.

STEP 5. 파일 처리 및 변환

if uploaded_file is not None:
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    file_name = f"{timestamp}_{uploaded_file.name}"
    pdf_path = os.path.join(UPLOAD_DIR, file_name)

    with open(pdf_path, "wb") as f:
        f.write(uploaded_file.getbuffer())

    st.success(f"파일이 성공적으로 업로드되었습니다: {file_name}")

    markdown_file_name = file_name.rsplit('.', 1)[0] + '.md'
    markdown_path = os.path.join(MARKDOWN_DIR, markdown_file_name)

    if pdf_to_markdown(pdf_path, markdown_path):
        st.success(f"PDF가 Markdown으로 변환되었습니다: {markdown_file_name}")
    else:
        st.error("PDF를 Markdown으로 변환하는 데 실패했습니다.")

이 단계에서는 업로드된 PDF 파일을 처리하고 Markdown으로 변환합니다. 파일 이름에 타임스탬프를 추가하여 고유성을 보장합니다.

STEP 6. 변환된 파일 목록 표시 및 다운로드 버튼 제공

if 'converted_files' not in st.session_state:
    st.session_state.converted_files = []

if uploaded_file is not None and markdown_file_name:
    st.session_state.converted_files.append(markdown_file_name)

st.subheader("변환된 Markdown 파일 목록")
if st.session_state.converted_files:
    for file in st.session_state.converted_files:
        col1, col2 = st.columns([3, 1])
        with col1:
            st.write(file)
        with col2:
            with open(os.path.join(MARKDOWN_DIR, file), "r", encoding="utf-8") as md_file:
                markdown_content = md_file.read()
            st.download_button(
                label="다운로드",
                data=markdown_content,
                file_name=file,
                mime="text/markdown",
                key=file
            )
else:
    st.write("현재 세션에서 변환된 Markdown 파일이 없습니다.")

마지막으로, 이 단계에서는 변환된 Markdown 파일의 목록을 표시하고 각 파일에 대한 다운로드 버튼을 제공합니다. Streamlit의 세션 상태를 사용하여 현재 세션에서 변환된 파일들의 목록을 관리합니다.

사용 방법

  1. 필요한 라이브러리를 설치합니다: streamlit, pdf2image, pytesseract
  2. 코드를 app.py로 저장합니다.
  3. 터미널에서 streamlit run app.py 명령을 실행합니다.
  4. 웹 브라우저에서 애플리케이션에 접속합니다.
  5. PDF 파일을 업로드하고 변환된 Markdown 파일을 다운로드합니다.

Streamlit으로 PDF 추출기 만들기

PDF 추출기의 한계점과 향후 계획

이번 프로젝트를 통해 PDF 추출기를 구현해보았지만, 몇 가지 아쉬운 점이 있었습니다. 이 프로젝트는 Streamlit의 강력함과 사용 편의성을 잘 보여줍니다. 간단한 코드로 실용적인 PDF 추출 도구를 만들 수 있었습니다. 이 코드를 기반으로 더 많은 기능을 추가하거나, 다른 파일 형식을 지원하는 등의 확장도 가능할 것입니다.

 

첫째, 생성된 md 파일에서 제목 등에 wiki 문법이 적용되지 않았습니다. 이로 인해 변환된 Markdown 파일의 구조가 원본 PDF의 구조를 제대로 반영하지 못하는 문제가 있었습니다.

 

둘째, 그림과 표 등이 제대로 변환되지 않았습니다. 이는 현재 사용한 OCR 솔루션의 한계로, 복잡한 레이아웃이나 이미지 요소를 텍스트로 정확하게 변환하는 데 어려움이 있었습니다.

 

셋째, 변환된 텍스트 중간 중간에 이해하기 어려운 문자들이 포함되어 있었습니다. 이는 OCR 과정에서 발생한 오류로, 텍스트의 가독성과 정확성을 떨어뜨리는 요인이 되었습니다.

 

이러한 한계점들을 고려하여, 다음 포스팅에서는 좀 더 좋은 OCR 솔루션을 사용해보고자 합니다. 이를 통해 더 정확한 텍스트 추출, 더 나은 구조 보존, 그리고 그림과 표의 적절한 처리가 가능할 것으로 기대합니다. 또한, Markdown 형식의 적절한 적용을 위한 후처리 과정도 추가할 계획입니다. 이러한 개선을 통해 PDF에서 Markdown으로의 변환 품질을 한층 높일 수 있을 것입니다.

유용한 링크

  1. Streamlit 공식 문서
  2. pytesseract 사용 가이드
  3. pdf2image 라이브러리

이 프로젝트에 대해 궁금한 점이나 개선 아이디어가 있다면 댓글로 남겨주세요. 함께 더 좋은 도구를 만들어 갈 수 있을 것입니다!

728x90
반응형

'언어모델 LLM > 프로젝트 예제' 카테고리의 다른 글

Streamlit과 Langchain으로 TTS 지원 챗봇 만들기  (1) 2024.07.12
로컬 모델로 Streamlit 챗봇 만들기  (0) 2024.07.12
RAG vs Non-RAG: 지하철 정보 검색 실험  (0) 2024.07.10
ollama API로 간단한 애플리케이션 만들기  (0) 2024.07.06
OpenAI API를 활용한 디스코드 챗봇 만들기  (0) 2024.07.06

'언어모델 LLM/프로젝트 예제'의 다른글

  • 현재글Streamlit으로 PDF 추출기 만들기

관련글

티스토리툴바