Create app.py

app.py

@@ -0,0 +1,87 @@
+import os
+import torch
+import pickle
+import joblib
+import torch.nn.functional as F
+from PIL import Image
+import gradio as gr
+from transformers import AutoModelForImageClassification
+from torch import nn
+from torchvision import transforms
+from huggingface_hub import hf_hub_download
+
+# Paths in Hugging Face model repository
+MODEL_PATH = "DeiT_Model_Parameter.pth"
+ENCODER_PATH = "label_encoder.pkl"
+
+# Ensure device is set
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+def load_label_encoder():
+    # Load label encoder from Hugging Face repository
+    label_encoder_path = hf_hub_download(repo_id="bobs24/DeiT-Classification-Apparel", filename=ENCODER_PATH)
+    label_encoder = joblib.load(label_encoder_path)
+    return label_encoder
+
+# Define the model class
+class CustomModel(nn.Module):
+    def __init__(self, num_classes):
+        super(CustomModel, self).__init__()
+        self.base_model = AutoModelForImageClassification.from_pretrained(
+            "facebook/deit-base-patch16-224",
+            num_labels=num_classes,
+            ignore_mismatched_sizes=True
+        )
+
+    def forward(self, x):
+        return self.base_model(x).logits
+
+def load_model():
+    # Load the model from Hugging Face repository
+    model_path = hf_hub_download(repo_id="bobs24/DeiT-Classification-Apparel", filename=MODEL_PATH)
+    label_encoder = load_label_encoder()
+    model = CustomModel(num_classes=len(label_encoder.classes_)).to(device)
+    model.load_state_dict(torch.load(model_path, map_location=device))
+    model.device = device
+    model.eval()
+    
+    return model, label_encoder
+
+# Load the model and label encoder
+model, label_encoder = load_model()
+
+# Preprocessing as per your training setup
+preprocess = transforms.Compose([
+    transforms.Resize(256),  # Resize to 256x256 (a bit larger than 224)
+    transforms.CenterCrop(224),  # Crop the center to 224x224
+    transforms.ToTensor(),  # Convert to tensor
+    transforms.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225))  # Normalize as per DeiT
+])
+
+# Function to perform predictions and show probabilities
+def predict(image):
+    # Apply preprocessing to the input image
+    image = Image.fromarray(image).convert("RGB")
+    input_tensor = preprocess(image).unsqueeze(0).to(device)
+    
+    # Perform inference
+    with torch.no_grad():
+        output = model(input_tensor)
+        
+        # Apply softmax to get probabilities
+        probabilities = F.softmax(output, dim=1)
+        
+        # Get the predicted label and confidence
+        predicted_label = torch.argmax(probabilities, dim=1).item()
+        confidence = probabilities[0, predicted_label].item()
+        
+        # Get the class name using label encoder
+        class_name = label_encoder.inverse_transform([predicted_label])[0]
+    
+    return f"Predicted class: {class_name}, Confidence: {confidence:.4f}"
+
+# Create Gradio interface
+iface = gr.Interface(fn=predict, inputs=gr.Image(type="numpy"), outputs="text", live=True)
+
+# Launch the interface
+iface.launch()