|
import torch |
|
from torchvision import models, transforms |
|
from PIL import Image |
|
import gradio as gr |
|
|
|
class_names = [ |
|
"calculus", |
|
"caries", |
|
"gingivitis", |
|
"hypodontia", |
|
"mouth_ulcer", |
|
"tooth_discoloration" |
|
] |
|
|
|
model = models.resnet50(weights=None) |
|
model.fc = torch.nn.Linear(model.fc.in_features, len(class_names)) |
|
|
|
model.load_state_dict(torch.load('best_model.pth', map_location=torch.device('cpu'))) |
|
model.eval() |
|
|
|
preprocess = transforms.Compose([ |
|
transforms.Resize((224, 224)), |
|
transforms.ToTensor(), |
|
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), |
|
]) |
|
|
|
def predict_image(image): |
|
processed_image = preprocess(image).unsqueeze(0) |
|
|
|
with torch.no_grad(): |
|
outputs = model(processed_image) |
|
_, predicted = torch.max(outputs, 1) |
|
predicted_class = class_names[predicted.item()] |
|
|
|
return predicted_class |
|
|
|
iface = gr.Interface( |
|
fn=predict_image, |
|
inputs=gr.Image(type="pil"), |
|
outputs="label", |
|
title="Medical Image Classification", |
|
description="Upload an image to predict its class." |
|
) |
|
|
|
iface.launch() |