#!/usr/bin/env python3
"""
ONNX 모델을 사용한 이미지 유사도 추론 예제
"""

import onnxruntime as ort
import numpy as np
from PIL import Image
import torchvision.transforms as transforms

# 전처리 파이프라인
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                        std=[0.229, 0.224, 0.225])
])

def preprocess_image(image_path):
    """이미지 전처리"""
    image = Image.open(image_path).convert('RGB')
    tensor = transform(image)
    return tensor.unsqueeze(0).numpy()  # 배치 차원 추가

def predict_similarity(onnx_model_path, image1_path, image2_path):
    """이미지 쌍 유사도 예측"""

    # ONNX 세션 생성
    session = ort.InferenceSession(onnx_model_path)

    # 이미지 전처리
    img1 = preprocess_image(image1_path)
    img2 = preprocess_image(image2_path)

    # 추론 실행
    inputs = {'image1': img1, 'image2': img2}
    logits = session.run(None, inputs)[0]

    # 시그모이드로 확률 변환
    similarity = 1 / (1 + np.exp(-logits[0][0]))

    return similarity

# 사용 예시
if __name__ == "__main__":
    onnx_path = "room_image_comparator.onnx"
    img1_path = "room1.jpg"
    img2_path = "room2.jpg"

    similarity = predict_similarity(onnx_path, img1_path, img2_path)
    print(f"유사도: {similarity:.4f}")