utils.py

import numpy
import sahi.predict
import sahi.utils
from PIL import Image

TEMP_DIR = "temp"


def sahi_yolov8m_inference(
    image,
    detection_model,
    slice_height=1240,
    slice_width=1240,
    overlap_height_ratio=0.1,
    overlap_width_ratio=0.1,
    image_size=4960,
    postprocess_match_threshold=0.8,
):
    # sliced inference
    detection_model.image_size = image_size
    prediction_result = sahi.predict.get_sliced_prediction(
        image=image,
        detection_model=detection_model,
        slice_height=slice_height,
        slice_width=slice_width,
        overlap_height_ratio=overlap_height_ratio,
        overlap_width_ratio=overlap_width_ratio,
        postprocess_match_threshold=postprocess_match_threshold,
    )
    visual_result = sahi.utils.cv.visualize_object_predictions(
        image=numpy.array(image),
        object_prediction_list=prediction_result.object_prediction_list,
        rect_th=3,
        text_size=3
    )

    output = Image.fromarray(visual_result["image"])

    return output