import torch
import torchvision.transforms as transforms
from torchvision import models
import gradio as gr

# 전체 모델 로드 (이미 정의된 MyModel 클래스는 필요하지 않음)
model = torch.load('resnet18_finetuned_full.pth')  # 전체 모델을 불러옴
model.eval()  # 평가 모드로 전환

# 데이터 변환 정의
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# 클래스 이름 리스트 정의
class_names = ['disposablecup', 'envmark', 'label', 'mug', 'nonlabel', 'reusablecup']

def predict(image):
    image = transform(image).unsqueeze(0)
    with torch.no_grad():
        outputs = model(image)
        print("Raw outputs:", outputs)  # 원본 출력값 확인
        _, predicted = torch.max(outputs, 1)
    label = class_names[predicted.item()]
    print("Predicted label:", label)
    
    # 예측 결과와 원본 출력값을 모두 반환
    return label, outputs.numpy()  # 원본 출력값을 NumPy 배열로 변환하여 반환

# Gradio 인터페이스 설정 (두 개의 출력)
iface = gr.Interface(fn=predict, inputs=gr.Image(type='pil'), outputs=['label', 'numpy'], live=True)
iface.launch()