Create app.py

app.py

@@ -0,0 +1,76 @@
+import torch
+import torchvision
+from torch import nn
+from torchvision import transforms
+from torchvision.transforms import InterpolationMode
+from PIL import Image
+import gradio as gr
+import os
+import matplotlib.pyplot as plt
+import seaborn as sns
+
+os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
+
+# Device configuration
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Assuming 'class_names' is already defined in your script
+class_names = [line.strip() for line in open("classes.txt")]
+
+# Load the model
+model = torchvision.models.vit_b_16(weights=None)  # Initialize the model architecture
+model.heads = nn.Linear(in_features=768, out_features=len(class_names))  # Adjust the classifier head
+checkpoint = torch.load('08_pretrained_vit_feature_extractor_pizza_steak_sushi.pth')
+model.load_state_dict(checkpoint, strict=False)
+model = model.to(device)
+model.eval()
+
+# Define transformations
+transform = transforms.Compose([
+    transforms.Resize(256, interpolation=InterpolationMode.BILINEAR),
+    transforms.CenterCrop(224),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+])
+
+# Prediction function
+def predict(image):
+    img = Image.fromarray(image)
+    transformed_image = transform(img).unsqueeze(dim=0).to(device)
+    
+    with torch.inference_mode():
+        target_image_pred = model(transformed_image)
+    
+    target_image_pred_probs = torch.softmax(target_image_pred, dim=1)
+    top_probs, top_indices = torch.topk(target_image_pred_probs, k=5)
+    top_probs = top_probs.squeeze().cpu().numpy()
+    top_indices = top_indices.squeeze().cpu().numpy()
+    
+    top_classes = [class_names[i] for i in top_indices]
+
+    # Plotting the probabilities as a bar chart
+    fig, ax = plt.subplots(figsize=(10, 6))
+    sns.barplot(x=top_probs, y=top_classes, palette="viridis", ax=ax)
+    ax.set_xlabel('Probability')
+    ax.set_ylabel('Class')
+    ax.set_title('Top 5 Predictions')
+    ax.set_xlim(0, 1)
+    for i in ax.patches:
+        ax.text(i.get_width() + 0.02, i.get_y() + 0.55, f'{i.get_width():.2f}', 
+                ha='center', va='center', fontsize=10, color='black')
+    sns.despine(left=True, bottom=True)
+    
+    plt.tight_layout()
+    
+    return top_classes[0], fig
+
+# Create Gradio interface
+iface = gr.Interface(
+    fn=predict, 
+    inputs=gr.Image(type="numpy"), 
+    outputs=[gr.Textbox(label="Top Prediction"), gr.Plot()],  # Textbox for top prediction and Plot for the bar chart
+    examples=[r"C:\Users\Asus\Desktop\download (1).jpg", r"C:\Users\Asus\Desktop\download (3).jpg"]  # Optional: Add paths to example images
+)
+
+# Launch the Gradio app
+iface.launch()