Upload app.py

app.py

@@ -0,0 +1,56 @@
+from PIL import Image
+from torchvision import transforms
+from transformers import AutoModelForImageClassification
+import gradio as gr
+import torch
+from model import vit
+
+def predict(inp):
+    inputs = data_transforms(inp)[None]
+    model.eval()
+    with torch.no_grad():
+        logits = model(inputs)
+        probs = torch.softmax(logits,dim=1)
+        confidences = {labels[i]: probs[0][i] for i in range(num_classes)}
+    return confidences
+
+"""height=28
+width=28
+batch_size=128
+n_channels=3
+patch_size=14
+dim=384
+n_head=12
+feed_forward=1024
+num_blocks=8"""
+height=224
+batch_size=128
+width=224
+n_channels=3
+patch_size=16
+dim=256
+n_head=8
+feed_forward=512
+num_blocks=12
+num_classes=2
+data_transforms = transforms.Compose([
+	    transforms.Resize((height,width)),  # Resize the images to a specific size
+	    transforms.ToTensor(),           # Convert images to tensors
+	    #transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))  # Normalize the image data
+	])
+
+model = vit(height,width,n_channels,patch_size,batch_size,dim,n_head,feed_forward,num_blocks,num_classes)# Load saved weights
+model.load_state_dict(
+    torch.load(f="vit_model.pt",
+               map_location=torch.device("cpu")) # load to CPU
+)
+print(model.state_dict())
+"""labels = [
+    'airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck'
+]"""
+labels = [
+    'cat','dog'
+]
+gr.Interface(fn=predict,
+             inputs=gr.Image(type="pil"),
+             outputs=gr.Label(num_top_classes=3)).launch()