import torch
import os
import yaml
from PIL import Image
from torchvision import transforms

import yaml

import torch
import torch.nn as nn
import torch.nn.functional as F
import timm

class CustomResNext2(nn.Module):
    def __init__(self, model_name='resnext50d_32x4d', pretrained=True):
        super().__init__()
        self.status_size = 768
        self.type_size = 512
        self.num_worker = 2
        self.target_size_status = 2
        self.target_size_type = 8
        self.target_size_fault = 23
        self.model = timm.create_model(model_name, pretrained=pretrained)
        n_features = self.model.fc.in_features
        self.model.fc = nn.Linear(n_features, self.target_size_fault)

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

def get_transforms(cfg):
    return transforms.Compose([
        transforms.Resize((cfg['size'], cfg['size'])),
        transforms.ToTensor(),
        transforms.Normalize(
            mean=[0.485, 0.456, 0.406],
            std=[0.229, 0.224, 0.225],
        ),
    ])

def load_image(image_path, transform):
    image = Image.open(image_path)
    image = transform(image).unsqueeze(0)
    return image

def inference(model, state, image, device):
    model.to(device)
    model.load_state_dict(state['model'])
    model.eval()
    image = image.to(device)
    with torch.no_grad():
        y_preds = model(image)
    return y_preds.softmax(1).cpu().numpy()

if __name__ == "__main__":
    cuda = 'cuda:0'
    image_path = '/data/si/KEPCO/src/images/7815_2075S791_LP애자_설치불량.JPG'


    device = cuda
    MODEL_DIR = 'weights/resnext50d_32x4d_fold0_best.pth'

    transform = get_transforms(cfg)
    image = load_image(image_path, transform)
    
    modelC = CustomResNext2('resnext50d_32x4d', pretrained=True)
    state = torch.load(MODEL_DIR, map_location=device)
    
    # Move model and image to the specified device
    modelC = modelC.to(device)
    image = image.to(device)
    
    modelC.load_state_dict(state['model'])
    modelC.eval()
    with torch.no_grad():
        y_preds = modelC(image)  # use 'image' instead of 'img'
    
    print(f'Predictions: {y_preds}')
    
classes = ['정상', '부식', '석면노출', '크랙', '아크흔적', '파손', '휴즈링크 석면 노출', '박리',
       '휴즈링크 석면노출', '홀더불량', '설치불량', '균열', '침식', '탄화', '정상차측탈락',
       '설치불량 (정상)', '불순물', '탈락정상', '탈락', '커버없음', '박리 (정상)', '커버탈락',
       '스프리트핀유실']