ethanrom/sv-small-object-detection

1import streamlit as st2import numpy as np3import cv24from ultralytics import YOLO5import supervision as sv6import time7 8model = YOLO("yolov8x.pt")9 10def callback(x: np.ndarray) -> sv.Detections:11    result = model(x, verbose=False, conf=0.25)[0]12    return sv.Detections.from_ultralytics(result)13 14def main():15    st.title("Small Object Detection with SAHI")16    st.write("Slicing Aided Hyper Inference (SAHI) implementaion with Supervsion for small object detection")17 18    example_image_loaded = st.checkbox("Load example image")19    uploaded_image = None20 21    if example_image_loaded:22        image = cv2.imread("example-image.jpg")23    else:24        uploaded_image = st.file_uploader("Upload an image", type=["jpg", "png", "jpeg"])25        if uploaded_image is not None:26            image = cv2.imdecode(np.fromstring(uploaded_image.read(), np.uint8), 1)27 28    if uploaded_image is not None or example_image_loaded:29        with st.spinner("Loading..."):30 31            start_time_sahi = time.time()32            slicer = sv.InferenceSlicer(callback=callback)33            sliced_detections = slicer(image=image)34            end_time_sahi = time.time()35 36            start_time_yolo = time.time()37            yolo_results = model(image, verbose=False, conf=0.25)38            end_time_yolo = time.time()39 40        st.header("Original Image")41        st.image(image, channels="BGR")42 43        st.header("SAHI-Processed Image")44        sliced_image = sv.BoxAnnotator().annotate(image.copy(), detections=sliced_detections)45        st.image(sliced_image, channels="BGR")46 47        st.header("YOLO-Detected Image (Without SAHI)")48        yolo_image = sv.BoxAnnotator().annotate(image.copy(), detections=sv.Detections.from_ultralytics(yolo_results[0]))49        st.image(yolo_image, channels="BGR")50 51        st.subheader("Method Comparison")52        st.write("SAHI Inference Time:", round(end_time_sahi - start_time_sahi, 2), "seconds")53        st.write("YOLOv8 Inference Time:", round(end_time_yolo - start_time_yolo, 2), "seconds")54 55        st.write("SAHI Detection Count:", len(sliced_detections))56        st.write("YOLOv8 Detection Count:", len(yolo_results[0]))57 58if __name__ == "__main__":59    main()