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

# 모델 클래스 정의 (ResNet18)
class MyModel(torch.nn.Module):
    def __init__(self, num_classes):
        super(MyModel, self).__init__()
        self.model = models.resnet18(weights=None)  # ResNet18 모델 사용
        self.model.fc = torch.nn.Linear(self.model.fc.in_features, num_classes)  # 마지막 레이어 수정

    def forward(self, x):
        return self.model(x)

# 모델 초기화
num_classes = 6  # 클래스 수에 맞게 설정
model = MyModel(num_classes)

# 모델 가중치 로드 (CPU로 매핑)
model.load_state_dict(torch.load('resnet18_finetuned_2.pth', map_location=torch.device('cpu')), strict=False)
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()