|
|
|
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)), |
|
transforms.ToTensor(), |
|
]) |
|
|
|
input_image = transform(img).unsqueeze(0) |
|
|
|
|
|
|
|
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)]}' |