Create app.py

app.py

@@ -0,0 +1,74 @@
+import torch
+
+
+model = torch.hub.load("facebookresearch/swag", model="vit_h14_in1k")
+
+# we also convert the model to eval mode
+model.eval()
+
+resolution = 518
+
+import os
+os.system("wget https://s3.amazonaws.com/deep-learning-models/image-models/imagenet_class_index.json -O in_cls_idx.json")
+
+import gradio as gr
+
+from PIL import Image
+from torchvision import transforms
+import matplotlib.pyplot as plt
+
+
+def load_image(image_path):
+    return Image.open(image_path).convert("RGB")
+
+
+
+def transform_image(image, resolution):
+    transform = transforms.Compose([
+        transforms.Resize(
+            resolution,
+            interpolation=transforms.InterpolationMode.BICUBIC,
+        ),
+        transforms.CenterCrop(resolution),
+        transforms.ToTensor(),
+        transforms.Normalize(
+            mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
+        ),
+    ])
+    image = transform(image)
+    # we also add a batch dimension to the image since that is what the model expects
+    image = image[None, :]
+    return image
+    
+def visualize_and_predict(model, resolution, image_path):
+    image = load_image(image_path)
+    image = transform_image(image, resolution)
+    
+    # we do not need to track gradients for inference
+    with torch.no_grad():
+        _, preds = model(image).topk(5)
+    # convert preds to a Python list and remove the batch dimension
+    preds = preds.tolist()[0]
+    return preds
+    
+os.system("wget https://github.com/pytorch/hub/raw/master/images/dog.jpg -O dog.jpg")
+
+
+
+def inference(img):
+    preds = visualize_and_predict(model, resolution, img)
+    
+    return preds
+
+inputs = gr.inputs.Image(type='pil')
+outputs = gr.outputs.Textbox(label="Output")
+
+title = "SWAG"
+
+description = "Gradio demo for Revisiting Weakly Supervised Pre-Training of Visual Perception Models. To use it, simply upload your image, or click one of the examples to load them. Read more at the links below."
+
+article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2201.08371' target='_blank'>Revisiting Weakly Supervised Pre-Training of Visual Perception Models</a> | <a href='https://github.com/facebookresearch/SWAG' target='_blank'>Github Repo</a></p>"
+
+examples = ['dog.jpg']
+
+gr.Interface(inference, inputs, outputs, title=title, description=description, article=article, analytics_enabled=False, examples=examples).launch(enable_queue=True)