Upload 6 files

app.py

@@ -0,0 +1,29 @@
+import gradio as gr
+import torch
+import torchvision.transforms.functional as TF
+from model import NeuralNetwork
+import json
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+def pokemon_classifier(inp):
+    model = NeuralNetwork()
+    model.load_state_dict(torch.load('model_best.pt', map_location=torch.device(device)))
+    model.eval()
+
+    with open('labels.json') as f:
+        labels = json.load(f)
+
+    x = TF.to_tensor(inp)
+    x = TF.resize(x, 64, antialias=True)
+    x = x.to(device)
+    x = x.unsqueeze(0)
+    
+    with torch.no_grad():
+        y_pred = model(x)
+    pokemon = torch.argmax(y_pred, dim=1).item()
+    
+    return labels[str(pokemon)]
+
+demo = gr.Interface(fn=pokemon_classifier, inputs=gr.Image(type="pil"), outputs="text")
+demo.launch()

labels.json

@@ -0,0 +1,154 @@
+{
+    "0": "Abra",
+    "1": "Aerodactyl",
+    "2": "Alakazam",
+    "3": "Arbok",
+    "4": "Arcanine",
+    "5": "Articuno",
+    "6": "Beedrill",
+    "7": "Bellsprout",
+    "8": "Blastoise",
+    "9": "Bulbasaur",
+    "10": "Butterfree",
+    "11": "Caterpie",
+    "12": "Chansey",
+    "13": "Charizard",
+    "14": "Charmander",
+    "15": "Charmeleon",
+    "16": "Clefable",
+    "17": "Clefairy",
+    "18": "Cloyster",
+    "19": "Cubone",
+    "20": "Dewgong",
+    "21": "Diglett",
+    "22": "Ditto",
+    "23": "Dodrio",
+    "24": "Doduo",
+    "25": "Dragonair",
+    "26": "Dragonite",
+    "27": "Dratini",
+    "28": "Drowzee",
+    "29": "Dugtrio",
+    "30": "Eevee",
+    "31": "Ekans",
+    "32": "Electabuzz",
+    "33": "Electrode",
+    "34": "Exeggcute",
+    "35": "Exeggutor",
+    "36": "Farfetchd",
+    "37": "Fearow",
+    "38": "Flareon",
+    "39": "Gastly",
+    "40": "Gengar",
+    "41": "Geodude",
+    "42": "Gloom",
+    "43": "Golbat",
+    "44": "Goldeen",
+    "45": "Golduck",
+    "46": "Golem",
+    "47": "Graveler",
+    "48": "Grimer",
+    "49": "Growlithe",
+    "50": "Gyarados",
+    "51": "Haunter",
+    "52": "Hitmonchan",
+    "53": "Hitmonlee",
+    "54": "Horsea",
+    "55": "Hypno",
+    "56": "Ivysaur",
+    "57": "Jigglypuff",
+    "58": "Jolteon",
+    "59": "Jynx",
+    "60": "Kabuto",
+    "61": "Kabutops",
+    "62": "Kadabra",
+    "63": "Kakuna",
+    "64": "Kangaskhan",
+    "65": "Kingler",
+    "66": "Koffing",
+    "67": "Krabby",
+    "68": "Lapras",
+    "69": "Lickitung",
+    "70": "Machamp",
+    "71": "Machoke",
+    "72": "Machop",
+    "73": "Magikarp",
+    "74": "Magmar",
+    "75": "Magnemite",
+    "76": "Magneton",
+    "77": "Mankey",
+    "78": "Marowak",
+    "79": "Meowth",
+    "80": "Metapod",
+    "81": "Mew",
+    "82": "Mewtwo",
+    "83": "Moltres",
+    "84": "Mr.Mime",
+    "85": "Muk",
+    "86": "Nidoking",
+    "87": "Nidoqueen",
+    "88": "Nidoran-f",
+    "89": "Nidoran-m",
+    "90": "Nidorina",
+    "91": "Nidorino",
+    "92": "Ninetales",
+    "93": "Oddish",
+    "94": "Omanyte",
+    "95": "Omastar",
+    "96": "Onix",
+    "97": "Paras",
+    "98": "Parasect",
+    "99": "Persian",
+    "100": "Pidgeot",
+    "101": "Pidgeotto",
+    "102": "Pidgey",
+    "103": "Pikachu",
+    "104": "Pinsir",
+    "105": "Poliwag",
+    "106": "Poliwhirl",
+    "107": "Poliwrath",
+    "108": "Ponyta",
+    "109": "Porygon",
+    "110": "Primeape",
+    "111": "Psyduck",
+    "112": "Raichu",
+    "113": "Rapidash",
+    "114": "Raticate",
+    "115": "Rattata",
+    "116": "Rhydon",
+    "117": "Rhyhorn",
+    "118": "Sandshrew",
+    "119": "Sandslash",
+    "120": "Scyther",
+    "121": "Seadra",
+    "122": "Seaking",
+    "123": "Seel",
+    "124": "Shellder",
+    "125": "Slowbro",
+    "126": "Slowpoke",
+    "127": "Snorlax",
+    "128": "Spearow",
+    "129": "Squirtle",
+    "130": "Starmie",
+    "131": "Staryu",
+    "132": "Tangela",
+    "133": "Tauros",
+    "134": "Tentacool",
+    "135": "Tentacruel",
+    "136": "Vaporeon",
+    "137": "Venomoth",
+    "138": "Venonat",
+    "139": "Venusaur",
+    "140": "Victreebel",
+    "141": "Vileplume",
+    "142": "Voltorb",
+    "143": "Vulpix",
+    "144": "Wartortle",
+    "145": "Weedle",
+    "146": "Weepinbell",
+    "147": "Weezing",
+    "148": "Wigglytuff",
+    "149": "Zapdos",
+    "150": "Zubat"
+  }
+  

model.py

@@ -0,0 +1,46 @@
+import torch
+from torch import nn
+
+class NeuralNetwork(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv = nn.Sequential(
+            nn.Conv2d(in_channels=3, out_channels=64, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(64),
+            nn.ReLU(),
+            nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=2, padding=1),
+            nn.Dropout2d(p=0.3),
+
+            nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, stride=1, padding=1), 
+            nn.BatchNorm2d(128),
+            nn.ReLU(),
+            nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3, stride=2, padding=1),
+            nn.Dropout2d(p=0.3),
+
+            nn.Conv2d(in_channels=128, out_channels=256, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(256),
+            nn.ReLU(),
+            nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, stride=2, padding=1),
+            nn.Dropout2d(p=0.4),
+
+            nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(256),
+            nn.ReLU(),
+            nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, stride=2, padding=1),
+        )
+        self.flatten = nn.Flatten()
+        self.fc = nn.Sequential(
+            nn.Linear(256*4*4, 256), 
+            nn.BatchNorm1d(256),
+            nn.ReLU(),
+            nn.Dropout(p=0.5),
+            nn.Linear(256, 151),
+        )
+
+    def forward(self, x):
+        out = self.conv(x)
+        out = self.flatten(out)
+        out = self.fc(out)
+        
+        return out
+    

requirements.txt

@@ -0,0 +1,2 @@
+torch
+torchvision