Add application file
Browse files- app.py +45 -0
- best_model.pth +3 -0
app.py
ADDED
|
@@ -0,0 +1,45 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import gradio as gr
|
| 2 |
+
import torch
|
| 3 |
+
import torch.nn as nn
|
| 4 |
+
from torchvision import transforms
|
| 5 |
+
from torchvision.models import efficientnet_v2_s, EfficientNet_V2_S_Weights
|
| 6 |
+
from PIL import Image
|
| 7 |
+
import json
|
| 8 |
+
import os
|
| 9 |
+
|
| 10 |
+
MODEL_PATH = "best_model.pth"
|
| 11 |
+
CLASS_NAMES = ["Achillies", "Adolf", "Al Waldo", "Alfie", "Ally", "Anselmo", "Baby Belgian", "Balls", "Bambi", "Bangus", "Barry", "Bingo", "Bingus", "Borgee", "Bozo", "Bulgogi", "Butter", "Butter", "Button", "Callie", "Cara", "Carme", "Chichi Shopwise", "Chico", "Choweder", "Cleo", "Coconut", "Daisy", "Daphne", "Darlene", "DOG 2", "DOG 3", "Dorian", "Doris", "Dot", "Ebi", "Echo", "Elise", "Enji", "Ensaymada Cheese Roll", "Faith", "Fave", "Flora", "Floyd", "Francis", "Frankie", "Georgie", "Gigi", "Ginger", "Gruyere", "Hansel", "Hany", "Harith", "Hazel", "Helga", "Huey", "Huni", "Inu", "Jack", "Jackie", "Jerome", "Josie", "JR", "Juliet", "Julio", "Juno", "Kalbi", "Kenta", "Kimchi", "Kimono", "KittyPerry", "Kucing", "Kucingcing", "Lambing", "Lara", "Leeca", "Lennon", "Leslie", "Limper_Bulag", "Limpkin", "Lion", "Louis", "Lucky", "Lyka", "Maggie", "Mama Waffle", "Maple Oscar", "Marge", "Marikit", "Marty", "Maxie", "Meemon", "Meowming", "Miki", "Milo", "Mingkay", "Minty", "Mitzki", "Moji", "Munchie", "Nala", "Nemo", "Nero", "NO_NAME 3", "NO_NAME 4", "NO_NAME 5", "NO_NAME 6", "NO_NAME 7- Batgirl", "NO_NAME 11", "NO_NAME 13 - Jinu", "NO_NAME Security Cat", "not_faith", "Nyssa Nisa", "Paris", "Patches", "Patchot", "Peeta", "Penny", "Pepper", "Percy", "Prof", "Pumpkin", "Pumpkin", "Raj", "Remy", "Renz", "Ringo", "Rocky", "Rory", "Ross", "Salvi", "Sam", "Shaq", "Sheeba", "Shibuya", "Sniper", "Socks", "Spaghetti", "Sunkiss", "Suzy", "Tarsier", "Tasi", "Teddy Chi", "Toes", "Toffee", "Toothless", "Trex", "Tutti", "Twister Fries", "Val", "Void", "Waffle", "Weena"]
|
| 12 |
+
|
| 13 |
+
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
| 14 |
+
|
| 15 |
+
weights = EfficientNet_V2_S_Weights.IMAGENET1K_V1
|
| 16 |
+
model = efficientnet_v2_s(weights=weights)
|
| 17 |
+
model.classifier[1] = nn.Linear(model.classifier[1].in_features, len(CLASS_NAMES))
|
| 18 |
+
model.load_state_dict(torch.load(MODEL_PATH, map_location=device))
|
| 19 |
+
model.eval().to(device)
|
| 20 |
+
|
| 21 |
+
transform = transforms.Compose([
|
| 22 |
+
transforms.Resize((384, 384)),
|
| 23 |
+
transforms.ToTensor(),
|
| 24 |
+
transforms.Normalize(mean=weights.meta["mean"], std=weights.meta["std"]),
|
| 25 |
+
])
|
| 26 |
+
|
| 27 |
+
def predict(image):
|
| 28 |
+
image = transform(image).unsqueeze(0).to(device)
|
| 29 |
+
with torch.no_grad():
|
| 30 |
+
outputs = model(image)
|
| 31 |
+
probs = torch.nn.functional.softmax(outputs, dim=1)[0]
|
| 32 |
+
results = {CLASS_NAMES[i]: float(probs[i]) for i in range(len(CLASS_NAMES))}
|
| 33 |
+
predicted_label = CLASS_NAMES[probs.argmax().item()]
|
| 34 |
+
return predicted_label, results
|
| 35 |
+
|
| 36 |
+
demo = gr.Interface(
|
| 37 |
+
fn=predict,
|
| 38 |
+
inputs=gr.Image(type="pil"),
|
| 39 |
+
outputs=[gr.Label(label="Prediction"), gr.JSON(label="Confidence Scores")],
|
| 40 |
+
title="StrAI - Cat Identifier",
|
| 41 |
+
description="Upload an image to identify which cat it is."
|
| 42 |
+
)
|
| 43 |
+
|
| 44 |
+
if __name__ == "__main__":
|
| 45 |
+
demo.launch()
|
best_model.pth
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:22e835f26352710b5a90b96e6085666b16b3cc7eacdca68d49bb547499685989
|
| 3 |
+
size 82426399
|