import os
import cv2
import streamlit as st
from ultralytics import YOLO
from streamlit_image_comparison import image_comparison
from utils import save_uploadedfile, apply_masks

model = YOLO('models/last.pt')

st.title("Image Segmentation with YOLOv8: A Web Integration")
st.subheader("Implementing for Image Segmentation and Object Detection")
st.write("Image segmentation is a critical task in computer vision that involves dividing an image into multiple "
         "segments or regions. YOLOv8 is a state-of-the-art deep learning model that can be used for "
         "image segmentation and object detection. In this web app, we will implement an interesting example "
         "using YOLOv8 for image segmentation. We can simply drop an image, View the identified segment "
         "a piece, a whole and the distribution of the identified segments. The essence of this application "
         "is to build a practical understanding and implementation of the powerful and light YOLOv8 "
         "for image segmentation and object detection")

url = "https://github.com/AkanimohOD19A/img-segmentation"
link = f'<a href="{url}">This sample app was heavily based on code shared by Akan Daniel. Huge credits to him.</a>'
st.markdown(link, unsafe_allow_html=True)



st.divider()

st.markdown('')
st.markdown('##### Segmented Pieces')

## Placeholder Image
parent_media_path = "media-directory"
img_file = './media-directory/example_input.png'

## Application States
APPLICATION_MODE = st.sidebar.selectbox("Our Options",
                                        ["Take Picture", "Upload Picture"]
                                        )

## captured_picture Image
if APPLICATION_MODE == "Take Picture":
    st.sidebar.write(
        """
            A computer aided application that segments your input image, built on 
            the powerful YOLOv8 instance segmentation algorithm developed by *ultralytics*.

            Simply take a captured_picture and it gets segmentated in real time.
        """
    )
    picture = st.camera_input("Take a picture")
    st.markdown('')
    if picture:
        st.sidebar.divider()
        st.sidebar.image(picture, caption="captured_picture")
        if st.button("Segment!"):
            img_file = os.path.join(parent_media_path, "captured_picture.jpg")
            save_uploadedfile(picture, save_path=img_file)
            st.sidebar.success("Saved File")
        st.write("Click on **Clear photo** to retake picture")

    st.divider()

elif APPLICATION_MODE == "Upload Picture":
    st.sidebar.write(
        """
            A computer aided application that segments your input image, built on 
            the powerful YOLOv8 object detection algorithm developed by *ultralytics*.

            Simply drop your image and it gets segmentated in real time.
        """
    )
    st.sidebar.divider()
    uploaded_file = st.sidebar.file_uploader("Drop a JPG/PNG file", accept_multiple_files=False, type=['jpg', 'png'])
    if uploaded_file is not None:
        img_file = os.path.join(parent_media_path, "uploaded_image.jpg")
        save_uploadedfile(uploaded_file, save_path=os.path.join(parent_media_path, "uploaded_image.jpg"))
        file_details = {"FileName": uploaded_file.name, "FileType": uploaded_file.type}

        st.sidebar.success("File saved successfully")
        print(f"File saved successfully to {os.path.abspath(img_file)}")
    else:
        st.sidebar.write("You are using a placeholder image, Upload your Image (.jpg for now) to explore")

results = model(img_file)
img = cv2.imread(img_file)
for result in results:
    # segmentation
    masks = [] if not result.masks else result.masks.data.cpu().numpy()
    numCols = len(masks)
    if numCols > 0:
        cols = st.columns(numCols)
        print(f"Number of instances found: {numCols}")
    else:
        st.warning("Unable to id Distinct items - Please retry with a clearer Image")
    img = apply_masks(img, masks)

    st.markdown('')
    st.markdown('##### Slider of Uploaded Image and Segments')
    image_comparison(
        img1=img_file,
        img2=img,
        label1="Actual Image",
        label2="Segmented Image",
        width=700,
        starting_position=50,
        show_labels=True,
        make_responsive=True,
        in_memory=True
    )

st.sidebar.divider()
st.sidebar.markdown('')
st.sidebar.markdown('#### Distribution of identified items')


st.markdown('')
st.markdown('')
st.markdown('')
st.markdown('')
st.markdown('')
st.markdown('')
st.sidebar.divider()