README.md

---
license: unknown
---

로맨스 스캠 사진과, 그냥 사진을 구별할 수 있는 ViT 모델 입니다.  
기존의 CNN 모델에 비해 훨신 성능이 좋습니다.  
추후 데이터를 추가해 성능을 더욱 늘릴것 입니다.
사용 코드는 다음과 같습니다.  

```python
import torch
from transformers import ViTForImageClassification, ViTFeatureExtractor
from PIL import Image

# Hugging Face에서 모델 및 특징 추출기 불러오기
model = ViTForImageClassification.from_pretrained("gihakkk/vit_modle")
feature_extractor = ViTFeatureExtractor.from_pretrained("gihakkk/vit_modle")

# 새로운 이미지 예측 함수 정의
def predict_image(image_path):
    # 이미지를 로드하고 RGB로 변환
    image = Image.open(image_path).convert("RGB")
    
    # 이미지를 특징 추출기로 전처리하여 모델 입력 형식으로 변환
    inputs = feature_extractor(images=image, return_tensors="pt")
    
    # 예측 수행
    with torch.no_grad():
        outputs = model(**inputs).logits
    predicted_class = torch.argmax(outputs, dim=-1).item()

    return "그냥 사진" if predicted_class == 1 else "로맨스 스캠 사진"

# 예측 예시
image_path = r'path\to\your\img.jpg'
result = predict_image(image_path)
print(result)

```