Spaces:
Running
Running
| 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() |