Update app.py

app.py

@@ -4,15 +4,17 @@ from transformers import ViTFeatureExtractor, ViTForImageClassification
 from PIL import Image
 import requests
 
-url = 'http://images.cocodataset.org/val2017/000000039769.jpg'
-image = Image.open(requests.get(url, stream=True).raw)
+img_file_buffer = st.file_uploader('Upload a PNG image', type='png')
+if img_file_buffer is not None:
+    image = Image.open(img_file_buffer)
+    img_array = np.array(image)
+    
+    feature_extractor = ViTFeatureExtractor.from_pretrained('rizvandwiki/gender-classification')
+    model = ViTForImageClassification.from_pretrained('rizvandwiki/gender-classification')
 
-feature_extractor = ViTFeatureExtractor.from_pretrained('google/vit-base-patch16-224')
-model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224')
+    inputs = feature_extractor(images=image, return_tensors="pt")
+    outputs = model(**inputs)
+    logits = outputs.logits
 
-inputs = feature_extractor(images=image, return_tensors="pt")
-outputs = model(**inputs)
-logits = outputs.logits
-
-predicted_class_idx = logits.argmax(-1).item()
-print("Predicted class:", model.config.id2label[predicted_class_idx])
+    predicted_class_idx = logits.argmax(-1).item()
+    print("Predicted class:", model.config.id2label[predicted_class_idx])