handler.py

"""
HuggingFace Inference Endpoint용 InsightFace 핸들러
얼굴 임베딩 추출 및 유사도 계산용
"""
import base64
import io
from typing import Dict, Any, Optional
from PIL import Image
import numpy as np
import insightface


class EndpointHandler:
    def __init__(self, path=""):
        """모델 초기화"""
        self.face_analyzer = None
        self._load_models(path)
    
    def _load_models(self, path: str):
        """모델 로드"""
        try:
            # InsightFace FaceAnalysis 초기화 (얼굴 감지 및 임베딩 추출용)
            self.face_analyzer = insightface.app.FaceAnalysis(
                name='buffalo_l',
                root=path,  # 모델 경로
                providers=['CUDAExecutionProvider', 'CPUExecutionProvider']
            )
            self.face_analyzer.prepare(ctx_id=0, det_size=(640, 640))
            
            print("✅ InsightFace 모델 로드 완료 (얼굴 임베딩 추출용)")
        except Exception as e:
            print(f"❌ 모델 로드 실패: {e}")
            raise
    
    def _base64_to_image(self, base64_str: str) -> Image.Image:
        """base64 문자열을 이미지로 변환"""
        if "," in base64_str:
            base64_str = base64_str.split(",")[1]
        
        img_bytes = base64.b64decode(base64_str)
        return Image.open(io.BytesIO(img_bytes))
    
    def _image_to_base64(self, image: Image.Image) -> str:
        """이미지를 base64 문자열로 변환"""
        buffer = io.BytesIO()
        image.save(buffer, format="PNG")
        img_bytes = buffer.getvalue()
        return base64.b64encode(img_bytes).decode("utf-8")
    
    def _pil_to_numpy(self, image: Image.Image) -> np.ndarray:
        """PIL Image를 BGR numpy 배열로 변환"""
        return np.array(image.convert('RGB'))[:, :, ::-1]
    
    def _numpy_to_pil(self, image: np.ndarray) -> Image.Image:
        """BGR numpy 배열을 PIL Image로 변환"""
        rgb = image[:, :, ::-1]  # BGR -> RGB
        return Image.fromarray(rgb)
    
    def face_detect(self, data: Dict[str, Any]) -> Dict[str, Any]:
        """
        얼굴 감지
        
        Args:
            data: {"image": "data:image/png;base64,..."}
        
        Returns:
            {"faces": [{"bbox": [x1, y1, x2, y2], "embedding": [...], ...}]}
        """
        try:
            image_b64 = data.get("image", "")
            if not image_b64:
                return {"error": "이미지가 제공되지 않았습니다"}
            
            # 이미지 변환
            image = self._base64_to_image(image_b64)
            image_np = self._pil_to_numpy(image)
            
            # 얼굴 감지
            faces = self.face_analyzer.get(image_np)
            
            # 결과 변환
            faces_data = []
            for face in faces:
                face_data = {
                    "bbox": face.bbox.tolist(),
                    "kps": face.kps.tolist() if hasattr(face, 'kps') and face.kps is not None else None,
                    "embedding": face.embedding.tolist() if hasattr(face, 'embedding') and face.embedding is not None else None,
                    "det_score": float(face.det_score) if hasattr(face, 'det_score') else None,
                    "gender": int(face.gender) if hasattr(face, 'gender') else None,
                    "age": int(face.age) if hasattr(face, 'age') else None,
                }
                faces_data.append(face_data)
            
            return {"faces": faces_data}
            
        except Exception as e:
            return {"error": str(e)}
    
    def face_similarity(self, data: Dict[str, Any]) -> Dict[str, Any]:
        """
        두 얼굴 간 유사도 계산 (코사인 유사도)
        
        Args:
            data: {
                "image1": "data:image/png;base64,...",
                "image2": "data:image/png;base64,...",
                "face_index1": 0,  # image1에서 사용할 얼굴 인덱스
                "face_index2": 0   # image2에서 사용할 얼굴 인덱스
            }
        
        Returns:
            {
                "similarity": float,  # 0.0 ~ 1.0 (1.0이 가장 유사)
                "embedding1": [...],   # 첫 번째 얼굴 임베딩
                "embedding2": [...],   # 두 번째 얼굴 임베딩
                "face1_info": {...},   # 첫 번째 얼굴 정보
                "face2_info": {...}    # 두 번째 얼굴 정보
            }
        """
        try:
            image1_b64 = data.get("image1", "")
            image2_b64 = data.get("image2", "")
            face_index1 = data.get("face_index1", 0)
            face_index2 = data.get("face_index2", 0)
            
            if not image1_b64 or not image2_b64:
                return {"error": "이미지1 또는 이미지2가 제공되지 않았습니다"}
            
            # 이미지 변환
            image1 = self._base64_to_image(image1_b64)
            image2 = self._base64_to_image(image2_b64)
            
            image1_np = self._pil_to_numpy(image1)
            image2_np = self._pil_to_numpy(image2)
            
            # 얼굴 감지
            faces1 = self.face_analyzer.get(image1_np)
            faces2 = self.face_analyzer.get(image2_np)
            
            if len(faces1) == 0:
                return {"error": "이미지1에서 얼굴을 찾을 수 없습니다"}
            
            if len(faces2) == 0:
                return {"error": "이미지2에서 얼굴을 찾을 수 없습니다"}
            
            # 얼굴 선택
            if face_index1 >= len(faces1):
                face_index1 = 0
            if face_index2 >= len(faces2):
                face_index2 = 0
            
            face1 = faces1[face_index1]
            face2 = faces2[face_index2]
            
            # 임베딩 추출
            if not hasattr(face1, 'embedding') or face1.embedding is None:
                return {"error": "이미지1의 얼굴에서 임베딩을 추출할 수 없습니다"}
            
            if not hasattr(face2, 'embedding') or face2.embedding is None:
                return {"error": "이미지2의 얼굴에서 임베딩을 추출할 수 없습니다"}
            
            emb1 = face1.embedding
            emb2 = face2.embedding
            
            # 코사인 유사도 계산
            # similarity = dot(emb1, emb2) / (norm(emb1) * norm(emb2))
            emb1_norm = np.linalg.norm(emb1)
            emb2_norm = np.linalg.norm(emb2)
            
            if emb1_norm == 0 or emb2_norm == 0:
                return {"error": "임베딩 벡터의 크기가 0입니다"}
            
            # 코사인 유사도 계산
            # InsightFace 임베딩은 이미 L2 정규화되어 있으므로 내적만으로 유사도 계산 가능
            similarity = np.dot(emb1, emb2) / (emb1_norm * emb2_norm)
            # 코사인 유사도는 -1 ~ 1 범위이지만, 얼굴 임베딩은 보통 0 ~ 1 범위
            # 음수 값이 나올 수 있으므로 0 ~ 1로 정규화 (선택사항)
            # 실제로는 InsightFace 임베딩이 정규화되어 있어 0 ~ 1 범위가 일반적
            similarity = max(0.0, min(1.0, similarity))
            
            # 얼굴 정보 추출
            face1_info = {
                "bbox": face1.bbox.tolist(),
                "det_score": float(face1.det_score) if hasattr(face1, 'det_score') else None,
                "gender": int(face1.gender) if hasattr(face1, 'gender') else None,
                "age": int(face1.age) if hasattr(face1, 'age') else None,
            }
            
            face2_info = {
                "bbox": face2.bbox.tolist(),
                "det_score": float(face2.det_score) if hasattr(face2, 'det_score') else None,
                "gender": int(face2.gender) if hasattr(face2, 'gender') else None,
                "age": int(face2.age) if hasattr(face2, 'age') else None,
            }
            
            return {
                "similarity": float(similarity),
                "embedding1": emb1.tolist(),
                "embedding2": emb2.tolist(),
                "face1_info": face1_info,
                "face2_info": face2_info
            }
            
        except Exception as e:
            return {"error": str(e)}
    
    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
        """
        메인 엔드포인트
        
        Args:
            data: {
                "task": "face-detect" | "face-similarity",
                "image": "data:image/png;base64,..." (face-detect용),
                "image1": "data:image/png;base64,..." (face-similarity용),
                "image2": "data:image/png;base64,..." (face-similarity용),
                ...
            }
        
        Returns:
            작업 결과
        """
        task = data.get("task", "")
        
        if task == "face-detect":
            return self.face_detect(data)
        elif task == "face-similarity":
            return self.face_similarity(data)
        else:
            return {"error": f"알 수 없는 작업: {task}. 지원 작업: 'face-detect', 'face-similarity'"}