Update README.md

README.md

@@ -1,3 +1,74 @@
 ---
 license: cc-by-4.0
 ---
+
+# This model doesn't inherit huggingface/transformers so it needs to be downloaded 
+```
+wget https://huggingface.co/Lancelot53/icon_classifier_maxvit/blob/main/id_2_class_89.json
+wget https://huggingface.co/Lancelot53/icon_classifier_maxvit/blob/main/best_model_89.pth
+```
+
+# Inference Code
+```
+import torch
+import torch.nn as nn
+from torchvision import transforms, models
+from PIL import Image
+import torch.nn.functional as F
+
+#load id_2_class.json
+import json
+
+with open('id_2_class_89.json') as json_file:
+    id_2_class = json.load(json_file)
+
+#make class_2_id dict
+
+class_2_id = {}
+for key, value in id_2_class.items():
+    class_2_id[value] = key
+
+test_transform = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.5,0.5,0.5], std=[0.5,0.5,0.5])
+])
+
+class MaxViT(nn.Module):
+    def __init__(self):
+        super(MaxViT, self).__init__()
+        model = models.maxvit_t(weights="DEFAULT")
+        num_ftrs = model.classifier[5].in_features
+        model.classifier[5] = nn.Linear(num_ftrs, len(class_2_id))
+        self.model = model
+    def forward(self, x):
+        return self.model(x)
+
+# Instantiate the model
+model = MaxViT()
+model.load_state_dict(torch.load('best_model_89.pth'))
+model.eval()
+
+def inference(image_path, CONFIDENT_THRESHOLD=None):
+    img = Image.open(image_path).convert("L").convert("RGB")
+    img = test_transform(img)
+    img = img.unsqueeze(0)
+
+    with torch.no_grad():
+        output = F.softmax(model(img), dim=1)
+        confidence, predicted = torch.max(output.data, 1)
+
+    if CONFIDENT_THRESHOLD is not None and confidence.item() < CONFIDENT_THRESHOLD:
+        return "UNKNOWN_CLASS", confidence.item()
+    
+    return id_2_class[str(predicted.item())], confidence.item()
+
+inference("images/7820.jpg", 0.9) #0.9 should be good enough
+```
+
+
+# Training
+Check the repo
+
+# Dataset
+Trained on 8K icons in 43 classes. The dataset is proprietary for now (Email me if you want it).