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

api = "https://23d4-188-130-155-153.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, width = 700)

    detection_model = st.selectbox("Select detection model", ("YOLO-stamp", ))
    embedding_model = st.selectbox("Select embedding model", ("vits8", ))

    if st.button("Get prediction"):
        with st.spinner("Loading..."):
            response = requests.post(os.path.join(api, f"bounding-boxes-{detection_model}/"),
                                    files = {"file": input_image.getvalue(), "model_id": detection_model})
            prediction = ast.literal_eval(response.text)
            response = requests.post(os.path.join(api, f"image-w-boxes-{detection_model}/"),
                                    files = {"file": input_image.getvalue(), "model_id": detection_model})
            
            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, f"embeddings-from-cropped-{embedding_model}/"),
                                                        files = {"file": output.getvalue(), "model_id": embedding_model}).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("Embeddings")
        col3.write(arr)