PizzaNotPizza / read.py
dron3flyv3r's picture
Update read.py
296a3ff
raw
history blame
609 Bytes
import torch
from PIL import Image
def classify(model, img, trans=None, classes=[], device=torch.device("cpu")):
try:
model = model.eval()
img = img.convert("RGB")
img = trans(img)
img = img.unsqueeze(0)
img = img.to(device)
output = model(img)
_, pred = torch.max(output, 1)
procent = torch.sigmoid(output)
return f"It {classes[pred.item()]} i'm {procent[0][pred[0]]*100:.2f}% sure"
except Exception:
return "Something went wrong😕, please notify the developer with the following message: " + str(Exception)