Create app.py

app.py

@@ -0,0 +1,75 @@
+import torch
+from torch.nn import functional as F
+import torchvision.models as models
+from torchvision import transforms
+import requests
+
+preprocess = transforms.Compose([
+    transforms.Resize(256),
+    transforms.CenterCrop(224),
+    transforms.ToTensor(),
+    transforms.Normalize(
+        mean=[0.485, 0.456 , 0.406],
+        std=[0.229, 0.224, 0.225]
+    )
+])
+
+response = requests.get("https://git.io/JJkYN")
+labels = response.text.split("\n")
+
+image_prediction_models = {
+    'resnet': models.resnet50,
+    'alexnet': models.alexnet,
+    'vgg': models.vgg16,
+    'squeezenet': models.squeezenet1_0,
+    'densenet': models.densenet161,
+    'inception': models.inception_v3,
+    'googlenet': models.googlenet,
+    'shufflenet': models.shufflenet_v2_x1_0,
+    'mobilenet': models.mobilenet_v2,
+    'resnext': models.resnext50_32x4d,
+    'wide_resnet': models.wide_resnet50_2,
+    'mnasnet': models.mnasnet1_0,
+    'efficientnet': models.efficientnet_b0,
+    'regnet': models.regnet_y_400mf,
+    'vit': models.vit_b_16,
+    'convnext': models.convnext_tiny
+}
+
+def load_pretrained_model(model_name):
+    model_name_lower = model_name.lower()
+    if model_name_lower in image_prediction_models:
+        model_class = image_prediction_models[model_name_lower]
+        model = model_class(pretrained=True)
+        return model
+    else:
+        raise ValueError(f"Model {model_name} is not available for image prediction in torchvision.models")
+    
+def get_model_names(models_dict):
+    return [name.capitalize() for name in models_dict.keys()]
+
+model_list = get_model_names(image_prediction_models)
+
+def classify_image(input_image, selected_model):
+        input_tensor = preprocess(input_image)
+        input_batch  = input_tensor.unsqueeze(0)
+        model = load_pretrained_model(selected_model)
+        if torch.cuda.is_available():
+            input_batch = input_batch.to('cuda')
+
+        with torch.no_grad():
+            output = model(input_batch)
+            
+        probabilities = F.softmax(input = output[0] , dim = 0)
+        top_prob, top_catid = torch.topk(probabilities, 5)
+        confidences = {labels[top_catid[i].item()]: top_prob[i].item() for i in range(top_prob.size(0))}
+        return confidences
+
+import gradio as gr
+
+interface = gr.Interface(
+    fn=classify_image,
+    inputs= [gr.Image(type='pil'),
+             gr.Dropdown(model_list)],
+    outputs=gr.Label(num_top_classes=5))
+interface.launch()