sahi切片辅助训练推理

本文主要是介绍sahi切片辅助训练推理，希望对大家解决编程问题提供一定的参考价值，需要的开发者们随着小编来一起学习吧！

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

代码实现步骤如下

把yolov5格式转换成coco格式标签；
切片图片和coco标签；
把切片出来的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_jsondef 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.jsonsave_json(coco_json, "coco_dataset.json")return coco_jsondef convert2xywh(l,h,w): # 把(class cx xy w h)转换成左上角whnew_l = np.zeros_like(l)l[:,1] = l[:,1]*wl[:,3] = l[:,3]*wl[:,2] = l[:,2]*hl[:,4] = l[:,4]*hnew_l[:,0] = l[:,0]new_l[:,1] = l[:,1] - l[:,3]/2new_l[:,2] = l[:,2] - l[:,4]/2new_l[:,3] = l[:,3] new_l[:,4] = l[:,4]return new_ldef 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')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 splitcoco.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:passl = np.loadtxt(anno_dir+label,delimiter=' ') # class cx xy w himg_path = img_dir+label.replace('txt','jpg')img = cv2.imread(img_path)h,w,_ = img.shapenew_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'):continueres = []ball_bbox = []data = cv2.imread(root+img)ih,iw,c = data.shapetry:anno = open(root.replace('images','labels')+img[:-4]+'.txt').read().splitlines()except:continuefor 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])passelse: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/' + rootos.makedirs(save_dir,exist_ok=True)cv2.imwrite(os.path.join(save_dir,img),data)