| import torch | |
| import torchvision | |
| from model import efficient_transformer , efficient_model | |
| CLASS_NAMES = ['pizza', 'steak', 'sushi'] | |
| def predict_gradio(image): | |
| image = efficient_transformer(image) | |
| efficient_model.eval() | |
| with torch.no_grad(): | |
| pred = efficient_model(torch.unsqueeze(image , dim = 0)) | |
| prediction_per_labels = {CLASS_NAMES[i]: float(torch.sigmoid(pred[0][i])) for i in range(len(CLASS_NAMES))} | |
| prediction = CLASS_NAMES[torch.argmax(pred).item()] | |
| return prediction_per_labels , prediction | |