from fastapi import FastAPI, Request, File, UploadFile, Form
from fastapi.responses import HTMLResponse
from pydantic import BaseModel
import yaml
import tempfile
import os
import traceback
from model.cer_module import calculate_korean_crr
from model.wav2vec2 import Wav2Vec2
# ---------------- 설정 로드 ----------------
with open("config/wav2vec2.yaml", "r") as f:
config = yaml.safe_load(f)
# ---------------- 모델 초기화 ----------------
wav2vec2_model = Wav2Vec2(config)
# ---------------- FastAPI 앱 ----------------
app = FastAPI(
title="Korean Speech Recognition API",
description="FastAPI + Wav2Vec2 기반 한국어 음성 인식 서버",
version="1.0.0"
)
# ---------------- 입력 모델 ----------------
class TranscriptionResponse(BaseModel):
transcription: str
status: str
crr: float = None # CRR 값, 선택적 필드
# ---------------- API: 파일 업로드 POST ----------------
@app.post("/transcribe", response_model=TranscriptionResponse)
async def transcribe_audio(file: UploadFile = File(...), reference: str = None):
# 파일 형식 검증
if not file.filename.lower().endswith(('.wav', '.mp3', '.flac', '.m4a')):
return TranscriptionResponse(
transcription="",
status="error: 지원되지 않는 파일 형식입니다. wav, mp3, flac, m4a 파일만 지원됩니다.",
crr=None
)
try:
audio_bytes = await file.read()
result = wav2vec2_model.transcribe_from_bytes(audio_bytes, file.filename)
# reference가 전달된 경우 CRR 계산
crr = None
if reference:
crr_result = calculate_korean_crr(reference, result)
crr = crr_result['crr']
return TranscriptionResponse(
transcription=result,
status="success",
crr=crr
)
except Exception as e:
return TranscriptionResponse(
transcription="",
status=f"error: {str(e)}",
crr=None
)
class CRRRequest(BaseModel):
original: str
corrected: str
class CRRResponse(BaseModel):
crr: float
@app.post("/calculate-crr", response_model=CRRResponse)
async def calculate_crr_api(data: CRRRequest):
"""
두 문장(original, corrected)을 받아 CRR(정확도)만 계산해서 반환
"""
result = calculate_korean_crr(data.original, data.corrected)
return CRRResponse(crr=result['crr'])
# ---------------- HTML UI ----------------
@app.get("/", response_class=HTMLResponse)
async def main_ui():
return """
Korean Speech Recognition
🎤 한국어 음성 인식
지원 형식: WAV, MP3, FLAC, M4A
모델: Wav2Vec2 Korean Fine-tuned
"""
# ---------------- 결과 렌더링 ----------------
@app.post("/submit", response_class=HTMLResponse)
async def handle_form(request: Request, audio_file: UploadFile = File(...)):
try:
# 파일 형식 검증
if not audio_file.filename.lower().endswith(('.wav', '.mp3', '.flac', '.m4a')):
return f"""
에러
❌ 파일 형식 오류
지원되지 않는 파일 형식입니다.
지원 형식: WAV, MP3, FLAC, M4A
← 돌아가기
"""
# 파일 내용 읽기
audio_bytes = await audio_file.read()
# 음성 인식 수행
result = wav2vec2_model.transcribe_from_bytes(audio_bytes, audio_file.filename)
except Exception as e:
error_details = traceback.format_exc()
return f"""
에러
❌ 서버 오류 발생
오류 메시지:
{str(e)}
에러 상세 (클릭하여 펼치기)
{error_details}
← 돌아가기
"""
return f"""
결과
✅ 음성 인식 결과
업로드된 파일: {audio_file.filename}
파일 크기: {len(audio_bytes):,} bytes
🎯 인식된 텍스트:
← 다시 시도하기
"""
# ---------------- 헬스 체크 ----------------
@app.get("/health")
async def health_check():
return {
"status": "ok",
"model": config["model"]["id"],
"device": config["model"]["device"],
"sampling_rate": config["model"]["sampling_rate"]
}
# ---------------- 모델 정보 ----------------
@app.get("/info")
async def model_info():
return {
"model_id": config["model"]["id"],
"device": config["model"]["device"],
"sampling_rate": config["model"]["sampling_rate"],
"supported_formats": ["wav", "mp3", "flac", "m4a"],
"description": "Korean Speech Recognition using Wav2Vec2"
}