首先使用sahi结合你的模型进行检测,这里面wh还有重叠率决定了切片数量和检测用时。
先是定义模型,将权重传入,并设置置信度
detection_model = AutoDetectionModel.from_pretrained(
model_type="ultralytics",
model_path=model_path,
confidence_threshold=0.3,
device='cuda:0', # or 'cuda:0'
)
使用get_sliced_prediction函数得到推理结果
from sahi.predict import get_sliced_prediction
result = get_sliced_prediction(
imagepath,
detection_model,
slice_height=640,
slice_width=640,
overlap_height_ratio=0.2,
overlap_width_ratio=0.2,
)
将结果转化为coco所需的json格式,这里需要与验证集的一致,我理解是imge_id
一致,可以通过下面函数传给他image_id,比如验证集json使用的是文件名,那么这里就要传给他文件名
pre = result.to_coco_predictions(image_id=imagepath)
1.之后是生成json的完整代码
import warnings
warnings.filterwarnings('ignore')
from ultralytics import YOLO
from sahi import AutoDetectionModel
from sahi.predict import predict, get_sliced_prediction
import os
import json
from tqdm import tqdm
#模型路径
model_path = 'runs/train/exp9/weights/best.pt'
#验证集路径
image_dir = 'ultralytics/assets'
def getjson(imagepath, detection_model):
try:
predictions = []
result = get_sliced_prediction(
imagepath,
detection_model,
slice_height=640,
slice_width=640,
overlap_height_ratio=0.2,
overlap_width_ratio=0.2,
)
id = os.path.splitext(os.path.basename(imagepath))[0]
pre = result.to_coco_predictions(image_id=id)
for prediction in pre:
new_prediction = {
"image_id": prediction["image_id"],
"category_id": prediction["category_id"],
"bbox": prediction["bbox"],
"score": prediction["score"]
}
predictions.append(new_prediction)
return predictions
except Exception as e:
print(f"处理图片 {imagepath} 时出错: {e}")
return []
detection_model = AutoDetectionModel.from_pretrained(
model_type="ultralytics",
model_path=model_path,
confidence_threshold=0.3,
device='cuda:0', # or 'cuda:0'
)
all_predictions = []
# 获取所有 .jpg 图片路径
jpg_images = [os.path.join(image_dir, img) for img in os.listdir(image_dir) if img.endswith('.jpg')]
# 使用 tqdm 包装循环,显示进度条
for image_path in tqdm(jpg_images, desc="处理图片进度"):
result = getjson(image_path, detection_model)
all_predictions += result
with open('results5.json', 'w', encoding='utf-8') as f:
json.dump(all_predictions, f, ensure_ascii=False, indent=4)
2.然后是对测试集生成
import os
import cv2
import json
from tqdm import tqdm
from sklearn.model_selection import train_test_split
import argparse
# 输入测试集类别
classes = ["mini", "serious"]
parser = argparse.ArgumentParser()
# 在default或者--给出测试集图片和标签路径
parser.add_argument('--image_path', default='/root/autodl-tmp/traffiicaccident/accident/images2',type=str, help="path of images")
parser.add_argument('--label_path', default='/root/autodl-tmp/traffiicaccident/accident/labels2',type=str, help="path of labels .txt")
parser.add_argument('--save_path', default='data.json', type=str, help="if not split the dataset, give a path to a json file")
arg = parser.parse_args()
def yolo2coco(arg):
print("Loading data from ", arg.image_path, arg.label_path)
assert os.path.exists(arg.image_path)
assert os.path.exists(arg.label_path)
originImagesDir = arg.image_path
originLabelsDir = arg.label_path
# images dir name
indexes = os.listdir(originImagesDir)
dataset = {'categories': [], 'annotations': [], 'images': []}
for i, cls in enumerate(classes, 0):
dataset['categories'].append({'id': i, 'name': cls, 'supercategory': 'mark'})
# 标注的id
ann_id_cnt = 0
for k, index in enumerate(tqdm(indexes)):
# 支持 png jpg 格式的图片.
txtFile = f'{index[:index.rfind(".")]}.txt'
stem = index[:index.rfind(".")]
# 读取图像的宽和高
try:
im = cv2.imread(os.path.join(originImagesDir, index))
height, width, _ = im.shape
except Exception as e:
print(f'{os.path.join(originImagesDir, index)} read error.\nerror:{e}')
# 添加图像的信息
if not os.path.exists(os.path.join(originLabelsDir, txtFile)):
# 如没标签,跳过,只保留图片信息.
continue
dataset['images'].append({'file_name': index,
'id': stem,
'width': width,
'height': height})
with open(os.path.join(originLabelsDir, txtFile), 'r') as fr:
labelList = fr.readlines()
for label in labelList:
label = label.strip().split()
x = float(label[1])
y = float(label[2])
w = float(label[3])
h = float(label[4])
# convert x,y,w,h to x1,y1,x2,y2
H, W, _ = im.shape
x1 = (x - w / 2) * W
y1 = (y - h / 2) * H
x2 = (x + w / 2) * W
y2 = (y + h / 2) * H
# 标签序号从0开始计算, coco2017数据集标号混乱,不管它了。
cls_id = int(label[0])
width = max(0, x2 - x1)
height = max(0, y2 - y1)
dataset['annotations'].append({
'area': width * height,
'bbox': [x1, y1, width, height],
'category_id': cls_id,
'id': ann_id_cnt,
'image_id': stem,
'iscrowd': 0,
# mask, 矩形是从左上角点按顺时针的四个顶点
'segmentation': [[x1, y1, x2, y1, x2, y2, x1, y2]]
})
ann_id_cnt += 1
# 保存结果
with open(arg.save_path, 'w') as f:
json.dump(dataset, f)
print('Save annotation to {}'.format(arg.save_path))
if __name__ == "__main__":
yolo2coco(arg)
3.使用上述两段代码生成的json来计算coco指标
import warnings
warnings.filterwarnings('ignore')
import argparse
from pycocotools.coco import COCO
from pycocotools.cocoeval import COCOeval
from tidecv import TIDE, datasets
#在default或者--给出上面的anno路径和pre路径,分别是测试集的json和预测的json
def parse_opt():
parser = argparse.ArgumentParser()
parser.add_argument('--anno_json', type=str, default='/root/autodl-tmp/UAV_smoke_fire/YOLO_UAV/data.json', help='label coco json path')
parser.add_argument('--pred_json', type=str, default='/root/autodl-tmp/drone/ultralytics-yolo11-main/runs/val/exp8/predictions.json', help='pred coco json path')
#autodl-tmp/drone/ultralytics-yolo11-main/runs/val/exp8/predictions.json
#sahi,results.json
return parser.parse_known_args()[0]
if __name__ == '__main__':
opt = parse_opt()
anno_json = opt.anno_json
pred_json = opt.pred_json
anno = COCO(anno_json) # init annotations api
pred = anno.loadRes(pred_json) # init predictions api
eval = COCOeval(anno, pred, 'bbox')
eval.evaluate()
eval.accumulate()
eval.summarize()
tide = TIDE()
tide.evaluate_range(datasets.COCO(anno_json), datasets.COCOResult(pred_json), mode=TIDE.BOX)
tide.summarize()
tide.plot(out_dir='result')
最后得到结果