Create app.py

app.py

@@ -0,0 +1,52 @@
+import torch
+import torch.nn as nn
+import torchvision.transforms as transforms
+from PIL import Image
+import gradio as gr
+
+# Load your trained model
+with torch.no_grad():
+    model = torch.load('classifier.pt')
+
+# Define the preprocessing function for the input image
+def preprocess(image):
+    transform = transforms.Compose([
+        transforms.Resize((224, 224)),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+    ])
+    image = Image.fromarray(image.astype('uint8'), 'RGB')
+    image = transform(image)
+    return image.unsqueeze(0)
+
+# Define the predict function
+def predict(image):
+    # Preprocess the image
+    input_tensor = preprocess(image)
+    
+    # Make a prediction
+    with torch.no_grad():
+        output = model(input_tensor)
+    
+    # Perform post-processing if needed (e.g., softmax for probabilities)
+    # Replace this with your actual post-processing logic
+    probabilities = torch.softmax(output, dim=1).squeeze().tolist()
+    
+    # Map the class indices to class labels
+    class_labels = ["Class1", "Class2", "Class3", "Class4"]
+    
+    # Create a dictionary with class labels and probabilities
+    predictions = {label: prob for label, prob in zip(class_labels, probabilities)}
+    
+    return predictions
+
+# Create the Gradio interface
+iface = gr.Interface(
+    fn=predict,
+    inputs=gr.Image(),
+    outputs=gr.Label(num_top_classes=4),
+    live=True
+)
+
+# Launch the Gradio app
+iface.launch()