added vggcrop

app.py

@@ -8,6 +8,13 @@ from huggingface_hub import hf_hub_download
 from PIL import Image
 import numpy as np
 
+from datasets import load_dataset
+import random
+
+# Load the VGGFace2 dataset using Hugging Face's datasets library
+ds = load_dataset("chronopt-research/cropped-vggface2-224")
+
+
 # Load the model checkpoint from Hugging Face
 checkpoint_path = hf_hub_download(repo_id="ttoosi/resnet50_robust_face", filename="100_checkpoint.pt")
 
@@ -44,10 +51,24 @@ def predict(image):
     with torch.no_grad():
         output = model(image)  # Perform inference on CPU
     _, predicted_class = output.max(1)
-    return f"Predicted class: {predicted_class.item()}"
+    # Fetch 9 random samples from the predicted class
+    class_samples = ds.filter(lambda example: example['label'] == predicted_class.item())['train']
+    sample_images = random.sample(list(class_samples), min(len(class_samples), 9))
+    
+    sample_images_urls = [sample['image'] for sample in sample_images]
+    
+    return f"Predicted class: {predicted_class.item()}", sample_images_urls
+
+# # Create the Gradio interface
+# iface = gr.Interface(fn=predict, inputs=gr.Image(type="numpy"), outputs="text")  # Updated from gr.inputs.Image to gr.Image
 
 # Create the Gradio interface
-iface = gr.Interface(fn=predict, inputs=gr.Image(type="numpy"), outputs="text")  # Updated from gr.inputs.Image to gr.Image
+iface = gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="numpy"),
+    outputs=[gr.Textbox(label="Predicted Class"), gr.Gallery(label="Class Samples")],
+    title="ResNet-50 VGGFace2 Classifier"
+)
 
 # Launch the interface
 iface.launch()