Spaces:
Sleeping
Sleeping
File size: 1,845 Bytes
11d0e57 b7fab1c 11d0e57 66ed482 11d0e57 d62b128 7bb89fc 11d0e57 7bb89fc 11d0e57 7bb89fc a16a11f 66ed482 11d0e57 b7fab1c 11d0e57 b7fab1c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 |
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 |