| import gradio as gr |
| import torch |
| import torch.nn as nn |
| from torchvision import transforms |
| from torchvision.models import efficientnet_v2_s, EfficientNet_V2_S_Weights |
| from PIL import Image |
| import json |
| import os |
|
|
| MODEL_PATH = "best_model.pth" |
| with open("class_names.json", "r") as f: |
| CLASS_NAMES = json.load(f) |
|
|
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
|
|
| weights = EfficientNet_V2_S_Weights.IMAGENET1K_V1 |
| model = efficientnet_v2_s(weights=weights) |
| model.classifier[1] = nn.Linear(model.classifier[1].in_features, len(CLASS_NAMES)) |
| model.load_state_dict(torch.load(MODEL_PATH, map_location=device)) |
| model.eval().to(device) |
|
|
| mean = getattr(weights, "meta", {}).get("mean", [0.485, 0.456, 0.406]) |
| std = getattr(weights, "meta", {}).get("std", [0.229, 0.224, 0.225]) |
|
|
| transform = transforms.Compose([ |
| transforms.Resize((384, 384)), |
| transforms.ToTensor(), |
| transforms.Normalize(mean=mean, std=std), |
| ]) |
|
|
| def predict(image): |
| image = transform(image).unsqueeze(0).to(device) |
| with torch.no_grad(): |
| outputs = model(image) |
| probs = torch.nn.functional.softmax(outputs, dim=1)[0] |
| results = {CLASS_NAMES[i]: float(probs[i]) for i in range(len(CLASS_NAMES))} |
| predicted_label = CLASS_NAMES[probs.argmax().item()] |
| return predicted_label, results |
|
|
| demo = gr.Interface( |
| fn=predict, |
| inputs=gr.Image(type="pil"), |
| outputs=[gr.Label(label="Prediction"), gr.JSON(label="Confidence Scores")], |
| title="StrAI - Cat Identifier", |
| description="Upload an image to identify which cat it is." |
| ) |
|
|
| if __name__ == "__main__": |
| demo.launch() |
