sahi切片辅助训练推理

本文的目的切片yolov5标注格式的数据，并保存图片和标注文件

代码实现步骤如下

1. 把yolov5格式转换成coco格式标签；
2. 切片图片和coco标签；
3. 把切片出来的coco标签转换回yolov5标签格式

# 1. 把yolov5格式转换成coco格式标签；
# 2. 切片图片和coco标签；
# 3. 把切片出来的coco标签转换回yolov5标签格式

import os
import numpy as np
import cv2
from sahi.utils.coco import Coco, CocoCategory, CocoImage, CocoAnnotation
from sahi.slicing import slice_coco
from sahi.utils.file import save_json

def convert2coco(img_path,h,w,yololabel):

    coco = Coco()
    maps = {
        0 : 'person',
        1 : 'soccer'
    }
    coco.add_category(CocoCategory(id=0, name='person')) # 两个类别
    coco.add_category(CocoCategory(id=1, name='soccer'))

    coco_image = CocoImage(file_name=img_path, height=h, width=w)
    
    for label in yololabel:

        coco_image.add_annotation(
        CocoAnnotation(
            bbox=[label[1], label[2], label[3], label[4]],
            category_id=int(label[0]),
            category_name=maps[label[0]]
        )
        )
    
    coco.add_image(coco_image)
    coco_json = coco.json
    save_json(coco_json, "coco_dataset.json")
    return coco_json

def convert2xywh(l,h,w): # 把(class cx xy w h)转换成左上角wh
    new_l = np.zeros_like(l)
    l[:,1] = l[:,1]*w
    l[:,3] = l[:,3]*w
    l[:,2] = l[:,2]*h
    l[:,4] = l[:,4]*h
    
    new_l[:,0] = l[:,0]
    new_l[:,1] = l[:,1] - l[:,3]/2
    new_l[:,2] = l[:,2] - l[:,4]/2
    new_l[:,3] = l[:,3] 
    new_l[:,4] = l[:,4]
    
    return new_l

def slice_img(save_img_dir):
    
    coco_dict, coco_path = slice_coco(
                coco_annotation_file_path="coco_dataset.json",
                image_dir='',
                slice_height=640,
                slice_width=640,
                overlap_height_ratio=0.2,
                overlap_width_ratio=0.2,
                output_dir = save_img_dir,
                output_coco_annotation_file_name = 'sliced',
                min_area_ratio = 0.2,
                ignore_negative_samples = True
            )
    return  

def convert2yolov5(coco_dir,save_img_dir,save_label_dir):
    
    coco = Coco.from_coco_dict_or_path(coco_dir, save_img_dir)
    # export converted YoloV5 formatted dataset into given output_dir with a 85% train/15% val split
    coco.export_as_yolov5(
    output_dir=save_label_dir,
    disable_symlink = True
    )
    
    return 

if __name__ == '__main__':
    file = 'SNMOT-061'
    img_dir = f'datasets/soccernet/tracking/images/train/{file}/img1/'
    anno_dir = f'datasets/soccernet/tracking/labels/train/{file}/img1/'
    save_img_dir = 'datasets/sliced_soccernet/images/train/' + f'{file}/' # 把切分的图片保存到这里
    save_label_dir = 'datasets/sliced_soccernet/labels/train/' + f'{file}/'
    os.makedirs(save_img_dir,exist_ok=True)
    os.makedirs(save_label_dir,exist_ok=True)
    labels = os.listdir(anno_dir)
    for label in labels:
        if 'old' not in label:
            try:
                os.remove('coco_dataset.json') # 删除中间文件
                os.remove(save_img_dir+'sliced_coco.json')
            except:
                pass
            l = np.loadtxt(anno_dir+label,delimiter=' ') # class cx xy w h
            img_path = img_dir+label.replace('txt','jpg')
            img = cv2.imread(img_path)
            h,w,_ = img.shape
            new_l = convert2xywh(l,h,w)
            coco_json = convert2coco(img_path,h,w,new_l)
            slice_img(save_img_dir)  # 切分图片并保存
            convert2yolov5(save_img_dir+'sliced_coco.json', save_img_dir,save_label_dir) # 把切分完的coco标签转换回yolo格式并保存
            
            
    
            # for ll in new_l: # 验证是否转换正确
            #     if int(ll[0]) == 0:
            #         cv2.rectangle(img,(int(ll[1]),int(ll[2])),(int(ll[1]+ll[3]),int(ll[2]+ll[4])),(255,0,255),2)
            #     else:
            #         cv2.rectangle(img,(int(ll[1]),int(ll[2])),(int(ll[1]+ll[3]),int(ll[2]+ll[4])),(255,0,2),2)         
            # cv2.imwrite('./test.jpg',img)
            
            

在上述的基础上需要修改一下sahi的源码，它默认会保存图片的，注释掉：

最终结果


转换完成后，可视化代码：

import os
import cv2


# train_lists = os.listdir('datasets/soccernet/tracking/images/train')

# for tl in train_lists:

root = 'datasets/sliced_soccernet/images/train/' + 'SNMOT-061/'
# root2 = 'datasets/soccernet/tracking/images2/train/'+'Vision_State.v4i.yolov8'+'/images/'
sum = 0
train_list = os.listdir(root)

for img in train_list:
    if not img.endswith('jpg'):
        continue
    res = []
    ball_bbox = []
    data = cv2.imread(root+img)
    ih,iw,c = data.shape
    try:
        anno = open(root.replace('images','labels')+img[:-4]+'.txt').read().splitlines()
    except:
        continue
    for an in anno:
        a = an.split(' ')
        cls,x,y,w,h = int(a[0]),float(a[1]),float(a[2]),float(a[3]),float(a[4])
        x = int(iw*x)
        y = int(ih*y)
        w = int(w*iw)
        h = int(h*ih)
        if int(cls)==1:
            cv2.rectangle(data,(x-w//2,y-h//2),(x+w//2,y+h//2),(255,255,0),2)
            ball_bbox.append([x-w//2,y-h//2,w,h])
            pass
        else:
            res.append([x-w//2,y-h//2,x+w//2,y+h//2])
            cv2.rectangle(data,(x-w//2,y-h//2),(x+w//2,y+h//2),(0,0,255),2)
            pass
    # if len(ball_bbox) > 0:
    #     crop_img = data[max(0,ball_bbox[0][1]-50):ball_bbox[0][1]+ball_bbox[0][3]+50,max(0,ball_bbox[0][0]-50):ball_bbox[0][0]+ball_bbox[0][2]+50]

        save_dir = 'outimg/' + root
        os.makedirs(save_dir,exist_ok=True)
        cv2.imwrite(os.path.join(save_dir,img),data)
    

训练完后的切片推理

from sahi import AutoDetectionModel
from sahi.utils.cv import read_image
from sahi.utils.file import download_from_url
from sahi.predict import get_prediction, get_sliced_prediction, predict


yolov8_model_path = 'runs/detect/train13-sahi-n/weights/last.pt'


# 以下为单张切片推理
detection_model = AutoDetectionModel.from_pretrained(
    model_type='yolov8',
    model_path=yolov8_model_path,
    confidence_threshold=0.6,
    device="cuda:7", # or 'cuda:0'
)

result = get_sliced_prediction(
    "20231023-134213.jpeg",
    detection_model,
    slice_height = 640,
    slice_width = 640,
    overlap_height_ratio = 0.2,
    overlap_width_ratio = 0.2
)
result.export_visuals(export_dir="sahi/",text_size=1,rect_th=1)
object_prediction_list = result.object_prediction_list


# 以下为batch 切片推理
source_image_dir = "sahi_in/1"
result = predict(
    model_type='yolov8',
    model_path=yolov8_model_path,
    model_device="cuda:7",
    model_confidence_threshold=0.6, 
    source=source_image_dir,
    image_size = (640,640),
    slice_height=960,
    slice_width=1280,
    overlap_height_ratio=0.2,
    overlap_width_ratio=0.2,
    visual_text_size = 1,
    return_dict = True
)

另外batch推理默认返回的值不包括预测结果，需修改源码：