import streamlit as st
from transformers import pipeline
from transformers import ViTFeatureExtractor, ViTForImageClassification
from PIL import Image
import requests
import numpy as np

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) 
    
    feature_extractor = ViTFeatureExtractor.from_pretrained('rizvandwiki/gender-classification')
    model = ViTForImageClassification.from_pretrained('rizvandwiki/gender-classification')

    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])