import streamlit as st
from segment_anything import sam_model_registry, SamAutomaticMaskGenerator, SamPredictor
# from IPython.display import display, Image
import cv2
import numpy as np
import matplotlib.pyplot as plt
from rembg import remove
from PIL import Image


# Content of detectObjects.py file 
# import detectObjects
import ultralytics
from ultralytics import YOLO

model = YOLO('yolov8n.pt')
sam_checkpoint = "sam_vit_b_01ec64.pth"
model_type = "vit_b"
sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
predictor = SamPredictor(sam)

def detected_objects(filename:str):
    results = model.predict(source=filename, conf=0.25)

    categories = results[0].names

    dc = []
    for i in range(len(results[0])):
        cat = results[0].boxes[i].cls
        dc.append(categories[int(cat)])

    print(dc)
    return results, dc

def show_mask(mask, ax, random_color=False):
    if random_color:
        color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)
    else:
        color = np.array([30/255, 144/255, 255/255, 0.6])
    h, w = mask.shape[-2:]
    mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
    ax.imshow(mask_image)

def show_points(coords, labels, ax, marker_size=375):
    pos_points = coords[labels==1]
    neg_points = coords[labels==0]
    ax.scatter(pos_points[:, 0], pos_points[:, 1], color='green', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)
    ax.scatter(neg_points[:, 0], neg_points[:, 1], color='red', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)

def show_box(box, ax):
    x0, y0 = box[0], box[1]
    w, h = box[2] - box[0], box[3] - box[1]
    ax.add_patch(plt.Rectangle((x0, y0), w, h, edgecolor='green', facecolor=(0,0,0,0), lw=2))

st.title('Extract Objects From Image')

uploaded_file = st.file_uploader('Upload an image')

if uploaded_file is not None:
    # To read file as bytes:
    bytes_data = uploaded_file.getvalue()
    with open('uploaded_file.png','wb') as file:
        file.write(uploaded_file.getvalue())

    # Detect objects in the uploaded image
    # results, dc = detectObjects.detected_objects('uploaded_file.png')
    results, dc = detected_objects('uploaded_file.png')

    st.write(dc)

    option = st.selectbox("Which object would you like to extract?", tuple(dc))
    # print(option)
    index_of_the_choosen_detected_object = tuple(dc).index(option)

    if st.button('Extract'):
        for result in results:
            boxes = result.boxes

        bbox=boxes.xyxy.tolist()[index_of_the_choosen_detected_object]
        # sam_checkpoint = "sam_vit_b_01ec64.pth"
        # model_type = "vit_b"
        # sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
        # predictor = SamPredictor(sam)

        image = cv2.cvtColor(cv2.imread('uploaded_file.png'), cv2.COLOR_BGR2RGB)
        predictor.set_image(image)

        input_box = np.array(bbox)

        masks, _, _ = predictor.predict(
            point_coords=None,
            point_labels=None,
            box=input_box[None, :],
            multimask_output=False,
        )

        # plt.figure(figsize=(10, 10))
        # st.image(image)
        # plt.imshow(image)
        # show_mask(masks[0], plt.gca())
        # show_box(input_box, plt.gca())
        # plt.axis('off')
        # plt.show()

        segmentation_mask = masks[0]
        binary_mask = np.where(segmentation_mask > 0.5, 1, 0)

        white_background = np.ones_like(image) * 255

        new_image = white_background * (1 - binary_mask[..., np.newaxis]) + image * binary_mask[..., np.newaxis]


        plt.imsave('extracted_image.jpg', new_image.astype(np.uint8))
        # st.image('extracted_image.jpg')

        # Store path of the image in the variable input_path
        input_path = 'extracted_image.jpg'

        # Store path of the output image in the variable output_path
        output_path = 'finalExtracted.png'

        # Processing the image
        input = Image.open(input_path)

        # Removing the background from the given Image
        output = remove(input)

        #Saving the image in the given path
        output.save(output_path)
        # st.image(output_path)

        with open("finalExtracted.png", "rb") as file:
            btn = st.download_button(
                label="Download final image",
                data=file,
                file_name="finalExtracted.png",
                mime="image/png",
            )

        # bbox=boxes.xyxy.tolist()[0]