import streamlit as st
import numpy as np
import cv2
from ultralytics import YOLO
import supervision as sv
import time

model = YOLO("yolov8x.pt")

def callback(x: np.ndarray) -> sv.Detections:
    result = model(x, verbose=False, conf=0.25)[0]
    return sv.Detections.from_ultralytics(result)

def main():
    st.title("Small Object Detection with SAHI")
    st.write("Slicing Aided Hyper Inference (SAHI) implementaion with Supervsion for small object detection")

    example_image_loaded = st.checkbox("Load example image")
    uploaded_image = None

    if example_image_loaded:
        image = cv2.imread("example-image.jpg")
    else:
        uploaded_image = st.file_uploader("Upload an image", type=["jpg", "png", "jpeg"])
        if uploaded_image is not None:
            image = cv2.imdecode(np.fromstring(uploaded_image.read(), np.uint8), 1)

    if uploaded_image is not None or example_image_loaded:
        with st.spinner("Loading..."):

            start_time_sahi = time.time()
            slicer = sv.InferenceSlicer(callback=callback)
            sliced_detections = slicer(image=image)
            end_time_sahi = time.time()

            start_time_yolo = time.time()
            yolo_results = model(image, verbose=False, conf=0.25)
            end_time_yolo = time.time()

        st.header("Original Image")
        st.image(image, channels="BGR")

        st.header("SAHI-Processed Image")
        sliced_image = sv.BoxAnnotator().annotate(image.copy(), detections=sliced_detections)
        st.image(sliced_image, channels="BGR")

        st.header("YOLO-Detected Image (Without SAHI)")
        yolo_image = sv.BoxAnnotator().annotate(image.copy(), detections=sv.Detections.from_ultralytics(yolo_results[0]))
        st.image(yolo_image, channels="BGR")

        st.subheader("Method Comparison")
        st.write("SAHI Inference Time:", round(end_time_sahi - start_time_sahi, 2), "seconds")
        st.write("YOLOv8 Inference Time:", round(end_time_yolo - start_time_yolo, 2), "seconds")

        st.write("SAHI Detection Count:", len(sliced_detections))
        st.write("YOLOv8 Detection Count:", len(yolo_results[0]))

if __name__ == "__main__":
    main()