gihakkk/mobilevit_model

사용예시

import onnxruntime as ort
import numpy as np
from transformers import AutoFeatureExtractor
from PIL import Image

# ONNX 모델 경로
onnx_model_path = r'C:\mobilevit_model.onnx'

# ONNX 런타임 세션 초기화
ort_session = ort.InferenceSession(onnx_model_path)

# 새로운 이미지 예측 함수 정의
def predict_image(image_path):
    # MobileViT 모델에 맞는 특징 추출기 로드
    feature_extractor = AutoFeatureExtractor.from_pretrained("apple/mobilevit-small")
    
    # 이미지를 로드하고 RGB로 변환
    image = Image.open(image_path).convert("RGB")
    
    # 이미지를 특징 추출기로 전처리
    inputs = feature_extractor(images=image, return_tensors="np")
    input_array = inputs['pixel_values']  # ONNX는 Numpy 형식을 사용
    
    # ONNX 모델에 입력 전달 및 추론
    ort_inputs = {ort_session.get_inputs()[0].name: input_array}
    ort_outputs = ort_session.run(None, ort_inputs)
    
    # 결과 해석
    logits = ort_outputs[0]
    predicted_class = np.argmax(logits, axis=-1).item()
    
    return "그냥 사진" if predicted_class == 1 else "로맨스 스캠 사진"

# 예측 예시
image_path = r'C:\1234567.jpg'
result = predict_image(image_path)
print(result)