assets/skin.py

import torch
from torchvision import transforms
from torchvision.models import VGG19_BN_Weights, vgg19_bn
DEVICE = 'cpu'
model =torch.load('assets/skin.pth', map_location=DEVICE)
skin_map={0:'Добро',1:'Зло'}
def get_evil(img):
    
    transform = transforms.Compose([
            transforms.Resize((224, 224)),  # Размер, ожидаемый VGG19_bn
            transforms.ToTensor(),
        ])
    
    input_image = transform(img).unsqueeze(0) # Добавьте размерность пакета (batch dimension)
    
    
    #torch.device("cuda" if torch.cuda.is_available() else 'cpu')
    model.to(DEVICE)
    model.eval()
    input_image = input_image.to(DEVICE)
    with torch.no_grad():
        res=model(input_image).item()
    return f'Степень злобы: {res}\n\n Т.е. опухоль: {skin_map[round(res)]}'