import streamlit as st
from transformers import pipeline
from PIL import Image
checkpoint = "openai/clip-vit-large-patch14"
classifier = pipeline(model=checkpoint, task="zero-shot-image-classification")
def get_best_label(predictions):
    max_score = 0
    label = None
    for p in predictions:
        if p['score'] > max_score:
            max_score = p['score']
            label = str(p['label'])
    return label, max_score

st.markdown('<h1 style="color:black;">Document Classifier</h1>', unsafe_allow_html=True)
st.markdown('<h2 style="color:gray;">This model can classify input image to the following categories:</h2>', unsafe_allow_html=True)
st.markdown('<h3 style="color:gray;"> <ul> <li>Invoice</li> <li>Bank statement</li> <li>Credit bureau</li> </ul> </h3>', unsafe_allow_html=True)



upload= st.file_uploader('Insert image for classification', type=['png','jpg'])
c1, c2= st.columns(2)
if upload is not None:
  image = Image.open(upload)
  c1.header('Input Image')
  c1.image(image)
  print("c1", c1)
  print("c2", c2)
  c2.header('Output')
  predictions = classifier(image, candidate_labels=["invoice, receipt", "bank statement, financial statement", "credit report"])
  label, score = get_best_label(predictions)
  c2.subheader('Predicted class: ' + str(label) + " with score: " + str(score))
  c2.subheader('Probabilites:')
  c2.write(str(predictions))