006f86 commited on
Commit
1949817
·
1 Parent(s): d6e1fcb

Add application file

Browse files
Files changed (2) hide show
  1. app.py +45 -0
  2. 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