import torch
import torch.nn as nn
import torch.optim as optim
import os
import torchvision
from torchvision import datasets, models, transforms
from PIL import Image
import numpy as np
import time

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

model = models.resnet34(pretrained=True)
num_features = model.fc.in_features
# 전이 학습(transfer learning): 모델의 출력 뉴런 수를 3개로 교체하여 마지막 레이어 다시 학습
model.fc = nn.Linear(num_features, 5)
model = model.to(device)

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

PATH = '/workspace/capstone/pill_train.pth'
model.load_state_dict(torch.load(PATH))

model.eval()
start_time = time.time()


transforms_train = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) # 정규화(normalization)
])

transforms_test = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
 
data_dir = '/workspace/capstone/dataset'
train_datasets = torchvision.datasets.ImageFolder(os.path.join(data_dir, 'train'), transforms_train)



class_names = train_datasets.classes
#print('클래스:', class_names)


image = Image.open('/workspace/capstone/image/cefare.jpg') 
images = transforms_test(image).unsqueeze(0).to(device)


with torch.no_grad():
    outputs = model(images) 
    _, preds = torch.max(outputs, 1) 
    print(class_names[preds[0]]) 
   