sahi小目标检测通过读取文件夹检测图片

学习记录

网上现有的sahi检测只能检测单张图片，加上循环之后，可以实现读取文件夹内的图片，循环检测。

import os
import ultralytics
from sahi.models.yolov8 import Yolov8DetectionModel
from sahi.predict import get_sliced_prediction

model_path = "D:/DeepL/ultralytics-main/runs/detect/train8/weights/best.pt"

# 使用的YOLOv5检测模型，使用gpu加速，置信度0.25
detection_model = Yolov8DetectionModel(
    model_path=model_path,
    confidence_threshold=0.5,
    device="cpu"
)

train_path = "D:/DeepL/ultralytics-main/data_fushi/testimages"
allpath = []

imglist = ultralytics.data.utils.getFileList(train_path, [], 'jpg')
print('本次执行检索到 ' + str(len(imglist)) + ' 张图像\n')

def get_lableandwav(path, dir):
    dirs = os.listdir(path)
    for a in dirs:
        imgname = os.path.splitext(os.path.basename(a))[0]
        print('获取图片: ' + imgname + '.jpg' )

        if os.path.isfile(path + "/" + a):
            #print("D:/DeepL/ultralytics-main/data_fushi/testimages/" + a + ".jpg")
            imgpath="D:/DeepL/ultralytics-main/data_fushi/testimages/" + a

            result = get_sliced_prediction(
                imgpath,
                detection_model,
                slice_height=160,  # 切分子图高度
                slice_width=160,  # 切分子图宽度
                overlap_height_ratio=0.2,  # 子图高度之间的重叠率
                overlap_width_ratio=0.2   # 子图宽度之间的重叠率
            )
            # 保存检测图片
            result.export_visuals(export_dir="D:/DeepL/ultralytics-main/sahi-result/", file_name=imgname)

            print('完成检测 ' + imgpath + '\n')
            allpath.append(a)
            # if dir != "":
            #     lllables.append(dir)
        else:
            get_lableandwav(str(path) + "/" + str(a), a)
        ##循环遍历这个文件夹

    return allpath

allpath = get_lableandwav(train_path, "")
print(allpath)
print('SAHI检测结束\n')