import ast
import os
import requests
from PIL import Image
import streamlit as st
import io

api = "https://7bb4-188-130-155-168.ngrok-free.app"


st.set_page_config(layout="wide")
st.title("Stamp2vec")

input_image = st.file_uploader("insert image")
if(input_image):
    image = Image.open(input_image)
    st.header("Original")
    st.image(input_image, use_column_width=True)
    if st.button("Get prediction"):
        with st.spinner("Loading..."):
            response = requests.post(os.path.join(api, "bounding-boxes/"), files = {"file": input_image.getvalue()})
            prediction = ast.literal_eval(response.text)
            response = requests.post(os.path.join(api, "image-w-boxes/"), files = {"file": input_image.getvalue()})
            image_with_boxes = response.content
            arr = []
            for b in prediction["bboxes"]:
                stamp = image.crop((b["xmin"], b["ymin"], b["xmin"] + b["width"], b["ymin"] + b["height"]))
                output = io.BytesIO()
                stamp.save(output, format="BMP")
                response = ast.literal_eval(requests.post(os.path.join(api, "embeddings-from-cropped/"), files = {"file": output.getvalue()}).text)
                arr.extend(response["embedding"])

        col1, col2, col3 = st.columns(3)

        col1.subheader("Prediction")
        col1.write(prediction)

        col2.subheader("Image")
        col2.image(image_with_boxes, use_column_width=True)

        col3.subheader("Embedding")
        col3.write(arr)