297 lines
13 KiB
Python
297 lines
13 KiB
Python
"""
|
|
Gemini API 연동 모듈 (좌표 추출 기능 추가)
|
|
Google Gemini API를 사용하여 이미지와 텍스트 좌표를 함께 분석합니다.
|
|
"""
|
|
|
|
import base64
|
|
import logging
|
|
import json
|
|
from google import genai
|
|
from google.genai import types
|
|
from typing import Optional, Dict, Any, List
|
|
|
|
from config import Config
|
|
|
|
# 로깅 설정
|
|
logging.basicConfig(level=logging.INFO)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# --- 새로운 스키마 정의 ---
|
|
|
|
# 좌표를 포함하는 값을 위한 재사용 가능한 스키마
|
|
ValueWithCoords = types.Schema(
|
|
type=types.Type.OBJECT,
|
|
properties={
|
|
"value": types.Schema(type=types.Type.STRING, description="추출된 텍스트 값"),
|
|
"x": types.Schema(type=types.Type.NUMBER, description="텍스트의 시작 x 좌표"),
|
|
"y": types.Schema(type=types.Type.NUMBER, description="텍스트의 시작 y 좌표"),
|
|
},
|
|
required=["value", "x", "y"]
|
|
)
|
|
|
|
# 범용 표의 한 행을 위한 스키마
|
|
GenericTableRow = types.Schema(
|
|
type=types.Type.ARRAY,
|
|
items=ValueWithCoords,
|
|
description="표의 한 행. 각 셀의 값과 좌표를 포함합니다."
|
|
)
|
|
|
|
# 페이지에서 추출된 범용 표를 위한 스키마
|
|
GenericTable = types.Schema(
|
|
type=types.Type.OBJECT,
|
|
properties={
|
|
"table_title": types.Schema(type=types.Type.STRING, description="추출된 표의 내용을 설명하는 제목 (예: '범례', 'IP 정보', '개정 이력')."),
|
|
"table_data": types.Schema(
|
|
type=types.Type.ARRAY,
|
|
items=GenericTableRow,
|
|
description="표의 데이터. 각 내부 리스트가 하나의 행을 나타냅니다."
|
|
)
|
|
},
|
|
description="도면에서 발견된 구조화된 정보 블록이나 표."
|
|
)
|
|
|
|
|
|
# 모든 필드가 ValueWithCoords를 사용하도록 스키마 업데이트
|
|
SCHEMA_EXPRESSWAY = types.Schema(
|
|
type=types.Type.OBJECT,
|
|
properties={
|
|
"도면명_line0": ValueWithCoords,
|
|
"도면명_line1": ValueWithCoords,
|
|
"도면명_line2": ValueWithCoords,
|
|
"편철번호": ValueWithCoords,
|
|
"도면번호": ValueWithCoords,
|
|
"Main_Title": ValueWithCoords,
|
|
"Sub_Title": ValueWithCoords,
|
|
"수평_도면_축척": ValueWithCoords,
|
|
"수직_도면_축척": ValueWithCoords,
|
|
"적용표준버전": ValueWithCoords,
|
|
"사업명_top": ValueWithCoords,
|
|
"사업명_bot": ValueWithCoords,
|
|
"시설_공구": ValueWithCoords,
|
|
"설계공구_공구명": ValueWithCoords,
|
|
"설계공구_범위": ValueWithCoords,
|
|
"시공공구_공구명": ValueWithCoords,
|
|
"시공공구_범위": ValueWithCoords,
|
|
"건설분야": ValueWithCoords,
|
|
"건설단계": ValueWithCoords,
|
|
"설계사": ValueWithCoords,
|
|
"시공사": ValueWithCoords,
|
|
"노선이정": ValueWithCoords,
|
|
"개정번호_1": ValueWithCoords,
|
|
"개정날짜_1": ValueWithCoords,
|
|
"개정내용_1": ValueWithCoords,
|
|
"작성자_1": ValueWithCoords,
|
|
"검토자_1": ValueWithCoords,
|
|
"확인자_1": ValueWithCoords,
|
|
"additional_tables": types.Schema(
|
|
type=types.Type.ARRAY,
|
|
items=GenericTable,
|
|
description="도면에서 발견된 추가적인 표나 정보 블록 목록."
|
|
)
|
|
},
|
|
)
|
|
|
|
SCHEMA_TRANSPORTATION = types.Schema(
|
|
type=types.Type.OBJECT,
|
|
properties={
|
|
"도면명": ValueWithCoords,
|
|
"편철번호": ValueWithCoords,
|
|
"도면번호": ValueWithCoords,
|
|
"Main Title": ValueWithCoords,
|
|
"Sub Title": ValueWithCoords,
|
|
"수평축척": ValueWithCoords,
|
|
"수직축척": ValueWithCoords,
|
|
"적용표준": ValueWithCoords,
|
|
"사업명": ValueWithCoords,
|
|
"시설_공구": ValueWithCoords,
|
|
"건설분야": ValueWithCoords,
|
|
"건설단계": ValueWithCoords,
|
|
"개정차수": ValueWithCoords,
|
|
"개정일자": ValueWithCoords,
|
|
"과업책임자": ValueWithCoords,
|
|
"분야별책임자": ValueWithCoords,
|
|
"설계자": ValueWithCoords,
|
|
"위치정보": ValueWithCoords,
|
|
"additional_tables": types.Schema(
|
|
type=types.Type.ARRAY,
|
|
items=GenericTable,
|
|
description="도면에서 발견된 추가적인 표나 정보 블록 목록."
|
|
)
|
|
},
|
|
)
|
|
|
|
|
|
class GeminiAnalyzer:
|
|
"""Gemini API 이미지 및 텍스트 분석 클래스"""
|
|
|
|
def __init__(self, api_key: Optional[str] = None, model: Optional[str] = None):
|
|
self.api_key = api_key or Config.GEMINI_API_KEY
|
|
self.model = model or Config.GEMINI_MODEL
|
|
self.default_prompt = Config.DEFAULT_PROMPT
|
|
|
|
if not self.api_key:
|
|
raise ValueError("Gemini API 키가 설정되지 않았습니다.")
|
|
|
|
try:
|
|
self.client = genai.Client(api_key=self.api_key)
|
|
logger.info(f"Gemini 클라이언트 초기화 완료 (모델: {self.model})")
|
|
except Exception as e:
|
|
logger.error(f"Gemini 클라이언트 초기화 실패: {e}")
|
|
raise
|
|
|
|
def _get_schema(self, organization_type: str) -> types.Schema:
|
|
"""조직 유형에 따른 스키마를 반환합니다."""
|
|
return SCHEMA_EXPRESSWAY if organization_type == "한국도로공사" else SCHEMA_TRANSPORTATION
|
|
|
|
def analyze_pdf_page(
|
|
self,
|
|
base64_data: str,
|
|
text_blocks: List[Dict[str, Any]],
|
|
prompt: Optional[str] = None,
|
|
mime_type: str = "image/png",
|
|
organization_type: str = "transportation"
|
|
) -> Optional[str]:
|
|
"""
|
|
Base64 이미지와 추출된 텍스트 좌표를 함께 분석합니다.
|
|
|
|
Args:
|
|
base64_data: Base64로 인코딩된 이미지 데이터.
|
|
text_blocks: PDF에서 추출된 텍스트와 좌표 정보 리스트.
|
|
prompt: 분석 요청 텍스트 (None인 경우 기본값 사용).
|
|
mime_type: 이미지 MIME 타입.
|
|
organization_type: 조직 유형 ("transportation" 또는 "expressway").
|
|
|
|
Returns:
|
|
분석 결과 JSON 문자열 또는 None (실패 시).
|
|
"""
|
|
try:
|
|
# 텍스트 블록 정보를 JSON 문자열로 변환하여 프롬프트에 추가
|
|
text_context = "\n".join([
|
|
f"- text: '{block['text']}', bbox: ({block['bbox'][0]:.0f}, {block['bbox'][1]:.0f})"
|
|
for block in text_blocks
|
|
])
|
|
|
|
analysis_prompt = (
|
|
(prompt or self.default_prompt) +
|
|
"\n\n--- 추출된 텍스트와 좌표 정보 ---\n" +
|
|
text_context +
|
|
"\n\n--- 지시사항 ---\n"
|
|
"1. 위 텍스트와 좌표 정보를 바탕으로, 이미지의 내용을 분석하여 JSON 스키마의 기본 필드를 채워주세요.\n"
|
|
"2. **(중요)** 도면 내에 표나 사각형으로 구분된 정보 블록이 있다면, `additional_tables` 필드에 추가해주세요. 예를 들어 'IP 정보 표', '범례(Legend)', '구조물 설명', '개정 이력' 등이 해당됩니다.\n"
|
|
" - 각 표/블록에 대해 `table_title`에 적절한 제목을 붙여주세요.\n"
|
|
" - `table_data`에는 표의 모든 행을 추출하여 리스트 형태로 넣어주세요. 각 행은 셀들의 리스트입니다.\n"
|
|
"3. 각 필드에 해당하는 텍스트를 찾고, 해당 텍스트의 'value'와 시작 'x', 'y' 좌표를 JSON에 기입하세요.\n"
|
|
"4. 해당하는 값이 없으면 빈 문자열이나 빈 리스트를 사용하세요.\n"
|
|
"\n--- 필드 설명 ---\n"
|
|
"- `{ }_Title`: 중앙 상단의 비교적 큰 폰트입니다.\n"
|
|
"- `사업명_top`에 해당하는 텍스트 아랫줄은 '시설_공구' 항목입니다.\n"
|
|
"- `도면명_line{n}`: 도면명에 해당하는 여러 줄의 텍스트를 위에서부터 0, 1, 2 순서로 기입합니다. 만약 두 줄이라면 line0은 비워두고 line1, line2를 채웁니다.\n"
|
|
"- `설계공구`/`시공공구`: '공구명'과 '범위'로 나뉘어 기입될 수 있습니다. (예: '설계공구 | 제2-1공구 | 12780.00-15860.00' -> `설계공구_공구명`: '제2-1공구', `설계공구_범위`: '12780.00-15860.00')\n"
|
|
)
|
|
|
|
contents = [
|
|
types.Content(
|
|
role="user",
|
|
parts=[
|
|
types.Part.from_bytes(
|
|
mime_type=mime_type,
|
|
data=base64.b64decode(base64_data),
|
|
),
|
|
types.Part.from_text(text=analysis_prompt),
|
|
],
|
|
)
|
|
]
|
|
|
|
selected_schema = self._get_schema(organization_type)
|
|
|
|
generate_content_config = types.GenerateContentConfig(
|
|
temperature=0,
|
|
top_p=0.05,
|
|
response_mime_type="application/json",
|
|
response_schema=selected_schema
|
|
)
|
|
|
|
logger.info("Gemini API 분석 요청 시작 (텍스트 좌표 포함)...")
|
|
|
|
response = self.client.models.generate_content(
|
|
model=self.model,
|
|
contents=contents,
|
|
config=generate_content_config,
|
|
)
|
|
|
|
if response and hasattr(response, 'text'):
|
|
result = response.text
|
|
# JSON 응답을 파싱하여 다시 직렬화 (일관된 포맷팅)
|
|
parsed_json = json.loads(result)
|
|
|
|
# 디버깅: Gemini 응답 내용 로깅
|
|
logger.info(f"=== Gemini 응답 디버깅 ===")
|
|
logger.info(f"조직 유형: {organization_type}")
|
|
logger.info(f"응답 필드 수: {len(parsed_json) if isinstance(parsed_json, dict) else 'N/A'}")
|
|
|
|
if isinstance(parsed_json, dict):
|
|
# 새로운 필드들이 응답에 포함되었는지 확인
|
|
new_fields = ["설계공구_Station_col1", "설계공구_Station_col2", "시공공구_Station_col1", "시공공구_Station_col2"]
|
|
old_fields = ["설계공구_Station", "시공공구_Station"]
|
|
|
|
logger.info("=== 새 필드 확인 ===")
|
|
for field in new_fields:
|
|
if field in parsed_json:
|
|
field_data = parsed_json[field]
|
|
if isinstance(field_data, dict) and field_data.get('value'):
|
|
logger.info(f"✅ {field}: '{field_data.get('value')}' at ({field_data.get('x', 'N/A')}, {field_data.get('y', 'N/A')})")
|
|
else:
|
|
logger.info(f"⚠️ {field}: 빈 값 또는 잘못된 형식 - {field_data}")
|
|
else:
|
|
logger.info(f"❌ {field}: 응답에 없음")
|
|
|
|
logger.info("=== 기존 필드 확인 ===")
|
|
for field in old_fields:
|
|
if field in parsed_json:
|
|
field_data = parsed_json[field]
|
|
if isinstance(field_data, dict) and field_data.get('value'):
|
|
logger.info(f"⚠️ {field}: '{field_data.get('value')}' (기존 필드가 여전히 존재)")
|
|
else:
|
|
logger.info(f"⚠️ {field}: 빈 값 - {field_data}")
|
|
else:
|
|
logger.info(f"✅ {field}: 응답에 없음 (예상됨)")
|
|
|
|
logger.info("=== 전체 응답 필드 목록 ===")
|
|
for key in parsed_json.keys():
|
|
value = parsed_json[key]
|
|
if isinstance(value, dict) and 'value' in value:
|
|
logger.info(f"필드: {key} = '{value.get('value', '')}' at ({value.get('x', 'N/A')}, {value.get('y', 'N/A')})")
|
|
else:
|
|
logger.info(f"필드: {key} = {type(value).__name__}")
|
|
|
|
logger.info("=== 디버깅 끝 ===")
|
|
|
|
pretty_result = json.dumps(parsed_json, ensure_ascii=False, indent=2)
|
|
logger.info(f"분석 완료: {len(pretty_result)} 문자")
|
|
return pretty_result
|
|
else:
|
|
logger.error("API 응답에서 텍스트를 찾을 수 없습니다.")
|
|
return None
|
|
|
|
except Exception as e:
|
|
logger.error(f"이미지 및 텍스트 분석 중 오류 발생: {e}")
|
|
return None
|
|
|
|
# --- 기존 다른 메서드들은 필요에 따라 수정 또는 유지 ---
|
|
# analyze_image_stream_from_base64, analyze_pdf_images 등은
|
|
# 새로운 analyze_pdf_page 메서드와 호환되도록 수정 필요.
|
|
# 지금은 핵심 기능에 집중.
|
|
|
|
def validate_api_connection(self) -> bool:
|
|
"""API 연결 상태를 확인합니다."""
|
|
try:
|
|
test_response = self.client.models.generate_content("안녕하세요")
|
|
if test_response and hasattr(test_response, 'text'):
|
|
logger.info("Gemini API 연결 테스트 성공")
|
|
return True
|
|
else:
|
|
logger.error("Gemini API 연결 테스트 실패")
|
|
return False
|
|
except Exception as e:
|
|
logger.error(f"Gemini API 연결 테스트 중 오류: {e}")
|
|
return False |