Spaces:
Sleeping
Sleeping
| import torch | |
| import torch.nn as nn | |
| from torchvision import models, transforms | |
| from PIL import Image | |
| import gradio as gr | |
| DEVICE = torch.device("cpu") | |
| NUM_CLASSES = 64 | |
| MODEL_PATH = "resnet18_animals_64.pth" | |
| CLASS_NAMES_FILE = "classes.txt" | |
| with open(CLASS_NAMES_FILE, "r") as f: | |
| class_names = [line.strip() for line in f.readlines()] | |
| def load_model(): | |
| model = models.resnet18(weights=None) | |
| num_ftrs = model.fc.in_features | |
| model.fc = nn.Linear(num_ftrs, NUM_CLASSES) | |
| model.load_state_dict(torch.load(MODEL_PATH, map_location=DEVICE)) | |
| model.to(DEVICE) | |
| model.eval() | |
| return model | |
| model = load_model() | |
| image_transforms = transforms.Compose([ | |
| transforms.Resize((224, 224)), | |
| transforms.ToTensor(), | |
| transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) | |
| ]) | |
| def predict(image): | |
| img_tensor = image_transforms(image).unsqueeze(0) | |
| img_tensor = img_tensor.to(DEVICE) | |
| with torch.no_grad(): | |
| outputs = model(img_tensor) | |
| probabilities = torch.nn.functional.softmax(outputs[0], dim=0) | |
| top3_prob, top3_catid = torch.topk(probabilities, 3) | |
| confidences = {} | |
| for i in range(top3_prob.size(0)): | |
| confidences[class_names[top3_catid[i]]] = float(top3_prob[i]) | |
| return confidences | |
| title = "Animal Classifier 🐾" | |
| description = "Nhận diện 64 loài động vật sử dụng mô hình ResNet-18." | |
| examples = [] | |
| interface = gr.Interface( | |
| fn=predict, | |
| inputs=gr.Image(type="pil"), | |
| outputs=gr.Label(num_top_classes=3), | |
| title=title, | |
| description=description | |
| ) | |
| if __name__ == "__main__": | |
| interface.launch() |