import io
from flask import Flask, jsonify, request
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
from flask_ngrok import run_with_ngrok


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.RandomHorizontalFlip(), # 데이터 증진(augmentation)
    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)

   

def get_prediction(image_bytes):
    image = Image.open(io.BytesIO(image_bytes))
    image = transforms_test(image).unsqueeze(0).to(device)

    with torch.no_grad():
        outputs = model(image)
        _, preds = torch.max(outputs, 1)

    return class_names[preds[0]]


app = Flask(__name__)


@app.route('/', methods=['POST'])
def predict():
    if request.method == 'POST':
        # 이미지 바이트 데이터 받아오기 
        file = request.files['file']
        image_bytes = file.read()

        # 분류 결과 확인 및 클라이언트에게 결과 반환
        class_name = get_prediction(image_bytes=image_bytes)
        print("결과:", {'class_name': class_name})
        return jsonify({'class_name': class_name})
    

run_with_ngrok(app)
app.run()

#curl -X POST -F file=@c.jpg http://fb9b-34-146-6-117.ngrok.io
