gihakkk/vit_modle

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

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)