from tensorflow.keras.models import load_model
from huggingface_hub import from_pretrained_keras
import streamlit as st
import numpy as np
import cv2
from PIL import Image

st.markdown('<img src="https://static.wixstatic.com/media/c2240c_1edbbcded56c4dcbb9cc90a7922021b1~mv2.png/v1/fill/w_333,h_370,al_c,lg_1,q_85,enc_auto/pomai.png" alt="Image" style="width: 200px;">', unsafe_allow_html=True)

st.header("Sniffusion PomerAInian")
st.write("Human/AI Art Classifier.")
st.write("NOTE: PomerAInian is a small model and was only trained on LEXICA Stable Diffusion images, images generated by other models may not be classified correctly.")
upload= st.file_uploader('Insert image for detection:', type=['png','jpg'])
c1, c2= st.columns(2)
if upload is not None:
    im= Image.open(upload)
    img= np.asarray(im)
    img = cv2.resize(img, (224, 224))
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    img = img / 255.0
    img = np.expand_dims(img, axis=0)
    c1.header('Input Image')
    c1.image(im)
    c1.write(img.shape)
    model = from_pretrained_keras("RaidedCluster/Sniffusion-PomerAInian")
    prediction = model.predict(img)
    hf=str(prediction[0][0]*100)+'% Human Factor'
    c2.header('Output')
    c2.subheader('Estimation:')
    if prediction >=0.5:
        est="Estimated to be Human Art."
    else:
        est="Estimated to be AI Art."
    c2.write(est)
    c2.write(hf)