File size: 1,298 Bytes
21a8746
 
 
 
 
 
 
 
 
 
 
 
735544b
21a8746
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
735544b
21a8746
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
---
license: apache-2.0
---

사용예시
```python
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)

```