[Refactoring] Python server

main.py에 합쳐져있던 코드를 분리했습니다.
https://velog.io/@halfgene/FastAPI-AI-separate-main-into-modules

1. core/vision.py (비전 및 OCR 핵심 로직)

이미지를 처리하고 텍스트를 추출하는 엔진입니다. 무거운 ONNX 세션을 전역으로 한 번만 띄우고, 8GB 램 최적화 로직을 함수 하나로 깔끔하게 묶었습니다.

import io
import numpy as np
import onnxruntime as ort
from PIL import Image

# 전역 초기화: 모델을 한 번만 로드하여 메모리 파편화 방지
ocr_session = ort.InferenceSession("lightweight_ocr_model.onnx")

def extract_text_from_buffer(buffer: io.BytesIO) -> str:
    # 커서 초기화 및 포맷 정규화
    buffer.seek(0)
    original_img = Image.open(buffer).convert("RGB")
    
    # 8GB RAM 최적화 전처리
    resized_img = original_img.resize((224, 224))
    img_array = np.array(resized_img, dtype=np.float32)
    
    img_array /= 255.0
    img_array = np.transpose(img_array, (2, 0, 1))
    input_tensor = np.expand_dims(img_array, axis=0)
    
    # 초경량 ONNX 추론
    ocr_result = ocr_session.run(None, {"input": input_tensor})
    return str(ocr_result[0])

3. core/rag.py (벡터 DB 및 유사도 검색 로직)

ChromaDB 데이터베이스와 임베딩 모델을 담당합니다. 텍스트를 던지면 가장 유사한 판례 1건을 찾아오는 로직을 캡슐화했습니다.

import chromadb
from chromadb.utils import embedding_functions

# 전역 초기화: 경량 임베딩 모델 및 인메모리 DB 연결
sentence_transformer_ef = embedding_functions.SentenceTransformerEmbeddingFunction(model_name="all-MiniLM-L6-v2")
chroma_client = chromadb.Client()
fraud_collection = chroma_client.get_or_create_collection(
    name="fraud_cases", 
    embedding_function=sentence_transformer_ef
)

def find_similar_case(combined_text: str) -> str:
    # 1. 벡터 DB 유사도 검색 (메모리 방어를 위해 n_results=1 강제)
    rag_results = fraud_collection.query(
        query_texts=[combined_text],
        n_results=1
    )
    
    # 검색 결과 방어적 추출
    if rag_results['documents'] and rag_results['documents'][0]:
        return rag_results['documents'][0][0]
    return ""

4. main.py (FastAPI 컨트롤러 및 진입점)

이제 main.py는 지저분한 로직을 직접 처리하지 않습니다. 클라이언트의 요청을 받아 동시성(Semaphore)만 통제한 뒤, 앞서 만든 core 폴더의 도구들을 조립하여 결과를 반환하는 본연의 컨트롤러 역할만 수행합니다.

import io
import asyncio
from typing import List
from fastapi import FastAPI, File, Form, UploadFile

# 분리한 코어 모듈과 스키마를 불러오기
from models.schemas import FraudResponse
from core.vision import extract_text_from_buffer
from core.rag import find_similar_case

app = FastAPI()

concurrency_gate = asyncio.Semaphore(5)

@app.post("/analyze", response_model=FraudResponse)
async def analyze_images(
    files: List[UploadFile] = File(...),
    scamType: str = Form(...),
    imageType: str = Form(...)
):
    async with concurrency_gate:
        # 1. 디스크 I/O 배제 (메모리 버퍼링)
        image_buffers = []
        for file in files:
            image_bytes = await file.read()
            image_buffers.append(io.BytesIO(image_bytes))
            
        # 2. Vision 모듈 호출: 텍스트 추출
        extracted_texts = []
        for buffer in image_buffers:
            text = extract_text_from_buffer(buffer)
            extracted_texts.append(text)
            
        # 3. RAG 모듈 호출: 추출된 텍스트 병합 후 유사 판례 찾기
        combined_text = " ".join(extracted_texts)
        reference_case = find_similar_case(combined_text)
        
        # 추출된 combined_text와 reference_case를 LLM에 넘겨 점수 판단하는 로직은 여기에 나중에 추가
        
        # 5. 최종 규격화된 Pydantic 응답 반환
        return FraudResponse(
            status="SUSPICIOUS", 
            fraudScore=75.0, 
            description=f"유사 사례 참조: {reference_case[:20]}" if reference_case else "분석 완료"
        )