Spaces:
Sleeping
Sleeping
import torch | |
from torchvision import models, transforms | |
from PIL import Image | |
import gradio as gr | |
# Mendefinisikan nama kelas | |
class_names = [ | |
"calculus", | |
"caries", | |
"gingivitis", | |
"hypodontia", | |
"mouth_ulcer", | |
"tooth_discoloration" | |
] | |
# Mengatur jumlah kelas | |
num_classes = len(class_names) | |
# Membuat dan mengkonfigurasi model | |
model = models.resnet50(pretrained=False) | |
model.fc = torch.nn.Linear(model.fc.in_features, num_classes) | |
# Memuat bobot model (sesuaikan path jika diperlukan) | |
model.load_state_dict(torch.load('best_model.pth', map_location=torch.device('cpu'))) | |
model.eval() | |
# Mengatur transformasi preprocessing | |
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]), | |
]) | |
# Fungsi untuk melakukan prediksi | |
def predict_image(image, model, preprocess, class_names): | |
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 | |
# Membuat interface Gradio | |
iface = gr.Interface( | |
fn=predict_image, | |
inputs=[gr.inputs.Image(type='pil')], | |
outputs=gr.outputs.Label(num_top_classes=1), | |
title="Klasifikasi Gambar Medis", | |
description="Upload gambar untuk memprediksi kelasnya." | |
) | |
# Menjalankan aplikasi Gradio | |
iface.launch() | |