server.py

from flask import Flask, request, jsonify
import io
import numpy as np
from model.layout_detection import LayoutDetectionModel
from util.pdf import PDFUtil
from flask_cors import CORS

app = Flask(__name__)
# CORS 설정 추가
CORS(app, resources={r"/*": {"origins": "*"}})

# 레이아웃 감지 모델 초기화 (크기 S 사용)
layout_model = LayoutDetectionModel(model_size="S")
pdf_util = PDFUtil()


@app.route("/health", methods=["GET"])
def health_check():
    return jsonify({"status": "ok"})


@app.route("/analyze-pdf", methods=["POST"])
def analyze_pdf():
    # 파일이 요청에 포함되어 있는지 확인
    if "file" not in request.files:
        return jsonify({"error": "PDF 파일이 필요합니다"}), 400

    file = request.files["file"]

    # 파일 이름이 비어있는지 확인
    if file.filename == "":
        return jsonify({"error": "선택된 파일이 없습니다"}), 400

    # 파일 확장자 확인
    if not file.filename.lower().endswith(".pdf"):
        return jsonify({"error": "PDF 파일만 허용됩니다"}), 400

    try:
        # 파일을 메모리에 로드
        pdf_bytes = io.BytesIO(file.read())

        # PDF 페이지 수 확인 및 이미지로 변환 (메모리 상에서 처리)
        images, page_count = pdf_util.convert_pdf_to_images_from_bytes(pdf_bytes)

        results = []

        # 각 페이지에 대해 레이아웃 감지 수행
        for i, image in enumerate(images):
            page_result = {
                "page": i + 1,
                "layout": layout_model.predict(np.array(image))[0],
            }
            results.append(page_result)

        return jsonify({"total_pages": page_count, "results": results})

    except Exception as e:
        return jsonify({"error": str(e)}), 500


if __name__ == "__main__":
    app.run(debug=True)