Spaces:

WildOjisan
/

python_insightface_hf

Build error

6273393 about 1 year ago

5.11 kB

import os
import cv2
import numpy as np
import pickle
import faiss
from insightface.app import FaceAnalysis
from fastapi import FastAPI, UploadFile, File, HTTPException
from io import BytesIO

# 1. FastAPI 애플리케이션 인스턴스 생성
app = FastAPI(
    title="InsightFace Face Recognition API",
    description="InsightFace (buffalo_l)를 사용한 얼굴 인식 및 FAISS 인덱스 검색 API"
)

# 모델 및 인덱스 전역 변수
model = None
index = None
labels = None

# 🚀 환경 설정: Docker 환경에서는 모델/인덱스 파일을 프로젝트 디렉토리에 배치합니다.
# 인덱스 파일은 /app/data 폴더에 있다고 가정합니다.
LOAD_DIR = "data"
FAISS_INDEX_FILE = os.path.join(LOAD_DIR, "faiss_index_v2.index")
LABELS_FILE = os.path.join(LOAD_DIR, "faiss_labels_v2.pkl")

# ✅ 1. 모델 및 인덱스 준비 (서버 시작 시 한 번만 실행)
@app.on_event("startup")
async def startup_event():
    global model, index, labels
    print("🚀 서버 시작: InsightFace 모델 및 FAISS 인덱스 로딩 중...")

    # InsightFace ArcFace 모델 준비
    try:
        # CPUExecutionProvider 사용 (GPU가 없는 환경/Docker에 적합)
        model = FaceAnalysis(name='buffalo_l', providers=['CPUExecutionProvider'])
        model.prepare(ctx_id=0)
        print("✅ InsightFace model (buffalo_l) 로딩 완료.")
    except Exception as e:
        print(f"❌ InsightFace 모델 로딩 실패: {e}")
        raise HTTPException(status_code=500, detail=f"모델 로딩 실패: {e}")

    # 인덱스 & 라벨 로딩
    try:
        if not os.path.exists(FAISS_INDEX_FILE) or not os.path.exists(LABELS_FILE):
             print(f"❌ FAISS 파일이 없습니다. 경로 확인: {LOAD_DIR}")
             raise FileNotFoundError(f"필요한 FAISS 인덱스 파일 혹은 라벨 파일이 없습니다.")
        
        index = faiss.read_index(FAISS_INDEX_FILE)
        with open(LABELS_FILE, "rb") as f:
            labels = pickle.load(f)
        print(f"✅ FAISS 인덱스 로딩 완료. 총 {index.ntotal}개의 임베딩 로드.")
    except Exception as e:
        print(f"❌ FAISS 인덱스 로딩 실패: {e}")
        raise HTTPException(status_code=500, detail=f"FAISS 로딩 실패: {e}")
    
# 🚀 얼굴 임베딩 추출 함수 (원본 임베딩만 추출)
def get_face_embedding(img_np):
    """
    Numpy 배열 형태의 이미지에서 얼굴 임베딩을 추출합니다.
    """
    global model
    if model is None:
        raise HTTPException(status_code=500, detail="모델이 초기화되지 않았습니다.")

    # BGR을 RGB로 변환 (InsightFace 모델은 RGB를 선호)
    img = cv2.cvtColor(img_np, cv2.COLOR_BGR2RGB)
    faces = model.get(img)
    
    if faces:
        return faces[0].embedding
    else:
        return None
    


# --- API 엔드포인트 정의 ---
# 3. 루트 엔드포인트 (GET /)
@app.get("/")
def read_root():
    result={"success":True,"data":None,"msg":""}
    try:
        result["data"]="ok"
        return result
    except Exception as e:
        result["success"] = False
        result["msg"]=f"server error. {e!r}"
        return result
    
# 🚀 얼굴 예측 API 엔드포인트
@app.post("/predict_person/")
async def predict_person(
    image: UploadFile = File(..., description="분석할 얼굴 이미지 파일"),
    top_k: int = 1
):
    """
    업로드된 이미지에서 얼굴을 인식하고, FAISS 인덱스를 사용하여 가장 유사한 인물을 예측합니다.
    """
    global index, labels

    # 1. 이미지 파일 읽기
    content = await image.read()
    np_array = np.frombuffer(content, np.uint8)
    img_np = cv2.imdecode(np_array, cv2.IMREAD_COLOR)

    if img_np is None:
        raise HTTPException(status_code=400, detail="이미지 파일을 읽을 수 없습니다.")

    # 2. 얼굴 임베딩 추출
    embedding = get_face_embedding(img_np)
    
    if embedding is None:
        raise HTTPException(status_code=404, detail="이미지에서 얼굴을 인식할 수 없습니다.")

    # 3. 임베딩 정규화
    embedding = embedding.astype('float32')
    embedding /= np.linalg.norm(embedding)

    # 쿼리 형식에 맞게 [1, D] 형태로 변환
    query_vector = np.array([embedding])

    # 4. 유사도 검색
    # top_k는 최대 인덱스 크기(index.ntotal)를 초과할 수 없습니다.
    k = min(top_k, index.ntotal)
    scores, indices = index.search(query_vector, k)

    # 5. 결과 포맷팅
    results = []
    for idx, score in zip(indices[0], scores[0]):
        # labels[idx]는 인덱스에 저장된 임베딩 중 가장 유사한 임베딩의 라벨
        results.append({
            "rank": len(results) + 1,
            "person_id": labels[idx],
            "similarity_score": float(f"{score:.4f}")
        })
        
    return {"filename": image.filename, "predictions": results}


if __name__ == "__main__":
    # --reload 옵션을 추가하여 코드가 변경될 때마다 자동 재시작되게 설정합니다.
    uvicorn.run("app:app", host="0.0.0.0", port=8000, reload=True)

	import os
	import cv2
	import numpy as np
	import pickle
	import faiss
	from insightface.app import FaceAnalysis
	from fastapi import FastAPI, UploadFile, File, HTTPException
	from io import BytesIO

	# 1. FastAPI 애플리케이션 인스턴스 생성
	app = FastAPI(
	title="InsightFace Face Recognition API",
	description="InsightFace (buffalo_l)를 사용한 얼굴 인식 및 FAISS 인덱스 검색 API"
	)

	# 모델 및 인덱스 전역 변수
	model = None
	index = None
	labels = None

	# 🚀 환경 설정: Docker 환경에서는 모델/인덱스 파일을 프로젝트 디렉토리에 배치합니다.
	# 인덱스 파일은 /app/data 폴더에 있다고 가정합니다.
	LOAD_DIR = "data"
	FAISS_INDEX_FILE = os.path.join(LOAD_DIR, "faiss_index_v2.index")
	LABELS_FILE = os.path.join(LOAD_DIR, "faiss_labels_v2.pkl")

	# ✅ 1. 모델 및 인덱스 준비 (서버 시작 시 한 번만 실행)
	@app.on_event("startup")
	async def startup_event():
	global model, index, labels
	print("🚀 서버 시작: InsightFace 모델 및 FAISS 인덱스 로딩 중...")

	# InsightFace ArcFace 모델 준비
	try:
	# CPUExecutionProvider 사용 (GPU가 없는 환경/Docker에 적합)
	model = FaceAnalysis(name='buffalo_l', providers=['CPUExecutionProvider'])
	model.prepare(ctx_id=0)
	print("✅ InsightFace model (buffalo_l) 로딩 완료.")
	except Exception as e:
	print(f"❌ InsightFace 모델 로딩 실패: {e}")
	raise HTTPException(status_code=500, detail=f"모델 로딩 실패: {e}")

	# 인덱스 & 라벨 로딩
	try:
	if not os.path.exists(FAISS_INDEX_FILE) or not os.path.exists(LABELS_FILE):
	print(f"❌ FAISS 파일이 없습니다. 경로 확인: {LOAD_DIR}")
	raise FileNotFoundError(f"필요한 FAISS 인덱스 파일 혹은 라벨 파일이 없습니다.")

	index = faiss.read_index(FAISS_INDEX_FILE)
	with open(LABELS_FILE, "rb") as f:
	labels = pickle.load(f)
	print(f"✅ FAISS 인덱스 로딩 완료. 총 {index.ntotal}개의 임베딩 로드.")
	except Exception as e:
	print(f"❌ FAISS 인덱스 로딩 실패: {e}")
	raise HTTPException(status_code=500, detail=f"FAISS 로딩 실패: {e}")

	# 🚀 얼굴 임베딩 추출 함수 (원본 임베딩만 추출)
	def get_face_embedding(img_np):
	"""
	Numpy 배열 형태의 이미지에서 얼굴 임베딩을 추출합니다.
	"""
	global model
	if model is None:
	raise HTTPException(status_code=500, detail="모델이 초기화되지 않았습니다.")

	# BGR을 RGB로 변환 (InsightFace 모델은 RGB를 선호)
	img = cv2.cvtColor(img_np, cv2.COLOR_BGR2RGB)
	faces = model.get(img)

	if faces:
	return faces[0].embedding
	else:
	return None



	# --- API 엔드포인트 정의 ---
	# 3. 루트 엔드포인트 (GET /)
	@app.get("/")
	def read_root():
	result={"success":True,"data":None,"msg":""}
	try:
	result["data"]="ok"
	return result
	except Exception as e:
	result["success"] = False
	result["msg"]=f"server error. {e!r}"
	return result

	# 🚀 얼굴 예측 API 엔드포인트
	@app.post("/predict_person/")
	async def predict_person(
	image: UploadFile = File(..., description="분석할 얼굴 이미지 파일"),
	top_k: int = 1
	):
	"""
	업로드된 이미지에서 얼굴을 인식하고, FAISS 인덱스를 사용하여 가장 유사한 인물을 예측합니다.
	"""
	global index, labels

	# 1. 이미지 파일 읽기
	content = await image.read()
	np_array = np.frombuffer(content, np.uint8)
	img_np = cv2.imdecode(np_array, cv2.IMREAD_COLOR)

	if img_np is None:
	raise HTTPException(status_code=400, detail="이미지 파일을 읽을 수 없습니다.")

	# 2. 얼굴 임베딩 추출
	embedding = get_face_embedding(img_np)

	if embedding is None:
	raise HTTPException(status_code=404, detail="이미지에서 얼굴을 인식할 수 없습니다.")

	# 3. 임베딩 정규화
	embedding = embedding.astype('float32')
	embedding /= np.linalg.norm(embedding)

	# 쿼리 형식에 맞게 [1, D] 형태로 변환
	query_vector = np.array([embedding])

	# 4. 유사도 검색
	# top_k는 최대 인덱스 크기(index.ntotal)를 초과할 수 없습니다.
	k = min(top_k, index.ntotal)
	scores, indices = index.search(query_vector, k)

	# 5. 결과 포맷팅
	results = []
	for idx, score in zip(indices[0], scores[0]):
	# labels[idx]는 인덱스에 저장된 임베딩 중 가장 유사한 임베딩의 라벨
	results.append({
	"rank": len(results) + 1,
	"person_id": labels[idx],
	"similarity_score": float(f"{score:.4f}")
	})

	return {"filename": image.filename, "predictions": results}


	if __name__ == "__main__":
	# --reload 옵션을 추가하여 코드가 변경될 때마다 자동 재시작되게 설정합니다.
	uvicorn.run("app:app", host="0.0.0.0", port=8000, reload=True)