Spaces:
Sleeping
Sleeping
File size: 1,696 Bytes
b3f257f 7ab069e b3f257f 280c7c8 b3f257f 280c7c8 b3f257f 7ab069e b3f257f da2622b b3f257f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 |
import gradio as gr
import torch
from torchvision import transforms
from torchvision.models import resnet18, ResNet18_Weights
from torch import nn
from PIL import Image # pip install pillow
labels = ['fractured','not fractured']
# Same data transformation that was used for inputs (except data augmentation)
imgSize = 128
data_transform = transforms.Compose([
transforms.Resize(size=(imgSize, imgSize)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
# https://pytorch.org/tutorials/beginner/saving_loading_models.html
# Loading Model for Inference with state_dict (recommended)
model = resnet18(weights=ResNet18_Weights.DEFAULT)
model.fc = nn.Linear(in_features=512, out_features=len(labels))
model.load_state_dict(torch.load("model.pth",map_location=torch.device('cpu')))
model.eval()
def predict(img):
X = data_transform(img).unsqueeze(0) # returns tensor
with torch.no_grad():
predictions = model(X).flatten()
predictions = torch.nn.functional.softmax(predictions)
confidences = {labels[i]: float(predictions[i]) for i in range(len(labels))}
return confidences
title = "Bone Fractures"
description = "Bone fractures classifier trained on the Kaggle dataset using Resnet18"
demo=gr.Interface(fn=predict,
inputs=gr.Image(type="pil"),
outputs=gr.Label(num_top_classes=len(labels)),
title=title,
description=description,
examples=["Elbow_Fractured.jpeg","Finger_Fractured.jpg","Shoulder_Fractured.jpeg","Foot_Not_Fractured.jpg","Leg_Not_Fractured.jpg"])
demo.launch('share=True') |