utils.py

import numpy
import pandas as pd
import sahi.predict
import sahi.utils
from PIL import Image

TEMP_DIR = "temp"


def sahi_yolov8m_inference(
    image,
    detection_model,
    slice_height,
    slice_width,
    overlap_height_ratio,
    overlap_width_ratio,
    image_size,
):
    # 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,
    )
    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=2
    )
    output_visual = Image.fromarray(visual_result["image"])
    
    # object prediction annotation
    coco_annotations = prediction_result.to_coco_annotations()
    # base DataFrame with predefined categories
    output_df = pd.DataFrame(
        {'category': ['ball-valve', 'butterfly-valve', 'centrifugal-pump', 'check-valve', 'gate-valve'],
        'count': [0, 0, 0, 0, 0]
        }
    )
    # extract relevant data into a new DataFrame
    coco_df = pd.DataFrame(
        [(item['category_name'], round(item['score'], 2)) for item in coco_annotations],
        columns=['category', 'score']
    )
    # count occurrences of each category
    category_counts = coco_df['category'].value_counts().reset_index()
    category_counts.columns = ['category', 'count']
    # update the `count` column in the base DataFrame
    output_df['count'] = output_df['category'].map(category_counts.set_index('category')['count']).fillna(0).astype(int)

    return output_visual,coco_df,output_df