from transformers import OwlViTProcessor, OwlViTForObjectDetection
import matplotlib.colors as mcolors
import matplotlib.pyplot as plt
import streamlit as st
from PIL import Image
import warnings
import torch
import os
import io


# setttings
os.environ['CUDA_VISIBLE_DEVICES'] = '1'
warnings.filterwarnings('ignore')
st.set_page_config()


class owl_vit:

    def __init__(self, image_path, text, threshold):
        self.image_path = image_path
        self.text = text
        self.threshold = threshold

    def process(self, processor, model):
        image = Image.open(self.image_path)
        if len(image.split()) == 1:
            image = image.convert("RGB")
        inputs = processor(text=[self.text], images=[image], return_tensors="pt")
        outputs = model(**inputs)
        target_sizes = torch.tensor([[image.height, image.width] for image in [image]])
        self.results = processor.post_process(outputs=outputs, target_sizes=target_sizes)
        self.image = image
        return self.result_image()

    def result_image(self):
        boxes, scores, labels = self.results[0]["boxes"], self.results[0]["scores"], self.results[0]["labels"]
        plt.imshow(self.image)
        ax = plt.gca()
        for box, score, label in zip(boxes, scores, labels):
            if score >= self.threshold:
                box = box.detach().numpy()
                color = list(mcolors.CSS4_COLORS.keys())[label]
                ax.add_patch(plt.Rectangle(box[:2], box[2] - box[0], box[3] - box[1], fill=False, color=color, linewidth=3,))
                ax.text(box[0], box[1], f"{self.text[label]}: {round(score.item(), 2)}", fontsize=15, color=color)
        plt.tight_layout()
        img_buf = io.BytesIO()
        plt.savefig(img_buf, format='png')
        image = Image.open(img_buf)
        return image


def load_model():
    with st.spinner('Getting Neruons in Order ...'):
        processor = OwlViTProcessor.from_pretrained("google/owlvit-base-patch16")
        model = OwlViTForObjectDetection.from_pretrained("google/owlvit-base-patch16")
        return processor, model


def show_detects(image):
    st.title("Results")
    st.image(image, use_column_width=True, caption="Object Detection Results", clamp=True)


def process(upload, text, threshold):

    # save upload to file
    filetype = upload.name.split('.')[-1]
    name = len(os.listdir("images")) + 1
    file_path = os.path.join('images', f'{name}.{filetype}')
    with open(file_path, "wb") as f:
        f.write(upload.getbuffer())

    # predict detections and show results
    detector = owl_vit(file_path, text, threshold)
    results = detector.process(processor, model)
    show_detects(results)

    # clean up - if over 1000 images in folder, delete oldest 1
    if len(os.listdir("images")) > 1000:
        oldest = min(os.listdir("images"), key=os.path.getctime)
        os.remove(os.path.join("images", oldest))


def main(processor, model):

    # splash image
    st.image(os.path.join('refs', 'baseball_labeled.png'), use_column_width=True)

    # title project descriptions
    st.title("OWL-ViT")
    st.markdown("**OWL-ViT** is a zero-shot text-conditioned object detection model. OWL-ViT uses CLIP as its multi-modal \
                backbone, with a ViT-like Transformer to get visual features and a causal language model to get the text features. \
                To use CLIP for detection, OWL-ViT removes the final token pooling layer of the vision model and attaches a \
                lightweight classification and box head to each transformer output token. Open-vocabulary classification \
                is enabled by replacing the fixed classification layer weights with the class-name embeddings obtained \
                from the text model. The authors first train CLIP from scratch and fine-tune it end-to-end with the classification \
                and box heads on standard detection datasets using a bipartite matching loss. One or multiple text queries per image \
                can be used to perform zero-shot text-conditioned object detection.", unsafe_allow_html=True)

    # example
    if st.button("Run the Example Image/Text"):
        with st.spinner('Detecting Objects and Comparing Vocab...'):
            info = owl_vit(os.path.join('refs', 'baseball.jpg'), ["batter", "umpire", "catcher"], 0.50)
            results = info.process(processor, model)
            show_detects(results)
            if st.button("Clear Example"):
                st.markdown("")

    # upload
    col1, col2 = st.columns(2)
    threshold = st.slider('Confidence Threshold', min_value=0.0, max_value=1.0, value=0.1)
    with col1:
        upload = st.file_uploader('Image:', type=['jpg', 'jpeg', 'png'])
    with col2:
        text = st.text_area('Objects to Detect: (comma, seperated)', "batter, umpire, catcher")
        text = [x.strip() for x in text.split(',')]

    # process
    if upload is not None and text is not None:
        filetype = upload.name.split('.')[-1]
        if filetype in ['jpg', 'jpeg', 'png']:
            with st.spinner('Detecting and Counting Single Image...'):
                process(upload, text, threshold)
        else:
            st.warning('Unsupported file type.')


if __name__ == '__main__':
    processor, model = load_model()
    main(processor, model)