import torch
from efficientnet_pytorch import EfficientNet
import torch.nn.functional as F

# 모델 클래스 정의 및 모델 로드
model_name = 'efficientnet-b0'  # 사용한 EfficientNet 모델의 이름에 맞게 수정
model = EfficientNet.from_name(model_name)

num_classes = 2
model._fc = torch.nn.Linear(model._fc.in_features, num_classes)

# 모델 가중치 로드
model.load_state_dict(torch.load('promotion_model.pt', map_location=torch.device('cpu')))

# 모델을 평가 모드로 설정
model.eval()

# 예제 데이터 (입력 데이터의 형태에 맞게 수정해야 합니다)
# EfficientNet 모델에 입력되는 데이터는 224x224 크기의 이미지입니다.
# 여기서는 임의의 데이터를 생성합니다.
from PIL import Image
from torchvision import transforms

# 이미지 전처리
preprocess = transforms.Compose([
    transforms.Resize(224),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

# 예제 이미지 로드
input_image = Image.open('fake_test.webp')  # 실제 이미지 경로로 수정
# input_image = Image.open('real_test.jpg')

input_tensor = preprocess(input_image)
input_batch = input_tensor.unsqueeze(0)  # 배치 차원 추가

# CUDA 사용 (GPU가 있는 경우)
if torch.cuda.is_available():
    model = model.to('cuda')
    input_batch = input_batch.to('cuda')

# 예측 수행
with torch.no_grad():
    output = model(input_batch)

# 예측 결과 출력
# print(output)

probabilities = F.softmax(output, dim=1)

# 예측 결과 출력
# print(probabilities)

# 가장 높은 확률을 가진 클래스 확인
_, predicted_class = torch.max(probabilities, 1)

# 예측 결과 출력
print('Predicted class:')
print('Fake' if predicted_class.item()==0 else 'Real')