import streamlit as st
import matplotlib.pyplot as plt
import numpy as np
import cv2
import PIL

from model import get_model, predict, prepare_prediction

print('Creating the model')
model = get_model('checkpoint.ckpt')

def plot_img_no_mask(image, boxes):
    # Show image
    boxes = boxes.cpu().detach().numpy().astype(np.int32)
    fig, ax = plt.subplots(1, 1, figsize=(12, 6))

    for i, box in enumerate(boxes):            
        [x1, y1, x2, y2] = np.array(box).astype(int)
        # Si no se hace la copia da error en cv2.rectangle
        image = np.array(image).copy()

        pt1 = (x1, y1)
        pt2 = (x2, y2)
        cv2.rectangle(image, pt1, pt2, (220,0,0), thickness=5)
        
    plt.axis('off')
    ax.imshow(image)
    fig.savefig("img.png", bbox_inches='tight')

image_file = st.file_uploader("Upload Images", type=["png","jpg","jpeg"])

if image_file is not None:
    print(image_file)
    print('Getting predictions')
    pred_dict = predict(model, image_file)
    print('Fixing the preds')
    boxes, image = prepare_prediction(pred_dict)
    print('Plotting')
    plot_img_no_mask(image, boxes)

    img = PIL.Image.open('img.png')
    st.image(img,width=750)