Create app.py

app.py

@@ -0,0 +1,37 @@
+import gradio as gr
+import torch
+from PIL import Image
+
+def predict(image):
+  """Predicts the class of the input image using the fine-tuned student model."""
+  # Convert the Gradio image to a PIL Image
+  image = Image.fromarray(image)
+  # Preprocess the image
+  inputs = processor(image, return_tensors="pt").to(device)
+  # Make prediction
+  with torch.no_grad():
+    outputs = student_model(**inputs)
+    predicted_class_idx = torch.argmax(outputs.logits, dim=1).item()
+  # Get predicted class label
+  predicted_class = id2label[predicted_class_idx]
+  return predicted_class
+
+# Get a few examples from the test dataset
+num_examples = 3  # Adjust the number of examples as needed
+examples = []
+for i in range(num_examples):
+    # Access the image data using 'bytes' key and convert to PIL Image
+    image = Image.open(io.BytesIO(train_test_valid_dataset["test"][i]["image"]['bytes'])) 
+    
+    # Convert the PIL Image to a NumPy array
+    image_np = np.array(image)  
+    examples.append(image_np)
+
+iface = gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="numpy"),
+    outputs="text",
+    title="Pets Image Classification",
+    description="Upload an image of a cat or dog to get its breed prediction.",
+    examples=examples, 
+).launch()