YOLOv5 小目标检测数据预处理——批量数据裁剪并切分json标注文件

# -*- coding:utf-8 -*-
import os
import matplotlib.pyplot as plt
import cv2
import numpy as np
import json
 
def divide_img(img_path, json_path, name, save_imgpath, save_jsonpath):

    with open(json_path, 'r', encoding='utf-8') as f:
        data = json.load(f)

    img = cv2.imread(img_path)
    # img = cv2.cvtColor(img,cv2.COLOR_BGR2RGB)
    h = img.shape[0]
    w = img.shape[1]
    n = int(np.floor(h * 1.0 / 1000)) + 1
    m = int(np.floor(w * 1.0 / 1000)) + 1
    print('h={},w={},n={},m={}'.format(h, w, n, m))
    dis_h = int(np.floor(h / n))
    dis_w = int(np.floor(w / m))
    num = 0
    for i in range(n):
        for j in range(m):
            num += 1            
            h0, h1, w0, w1 = dis_h*i, dis_h*(i+1)+30, dis_w*j, dis_w*(j+1)+30
            print('i={}, j={}'.format(i, j), h1-h0, w1-w0)
            sub = img[h0:h1, w0:w1, :]
            cv2.imwrite(os.path.join(save_imgpath,'{}_{}.jpg'.format(name, num)), sub)

            new_shapes = []
            for shape in data['shapes']:
                xmin, ymin, xmax, ymax = shape['points'][0][0], shape['points'][0][1], shape['points'][1][0], shape['points'][1][1]
                # print(xmin, ymin, xmax, ymax)
                if xmin<xmax and ymin<ymax:
                    if ymin>=h0 and ymax<=h1 and xmin>=w0 and xmax<=w1:
                        newshape={
                            "label": shape["label"],
                            "points": [[xmin-w0,ymin-h0],[xmax-w0,ymax-h0]],
                            "group_id": None,
                            "shape_type": "rectangle",
                            "flags": {}
                            },
                        new_shapes.extend(newshape)
                else:
                    print(f"{name}标注存在错误", shape)
            if len(new_shapes)>0:
                sub_data = {
                    "version": "4.5.7",
                    "flags": {},
                    "shapes": new_shapes,
                    "imagePath": '{}_{}.jpg'.format(name, num),
                    "imageData": None,
                    "imageHeight": h1-h0,
                    "imageWidth": w1-w0
                    }
                jsonPath = os.path.join(save_jsonpath, '{}_{}.json'.format(name, num))
                with open(jsonPath, "w", encoding="utf-8") as f:
                    json.dump(sub_data, f, indent=4, ensure_ascii=False)


if __name__ == '__main__':
    rootpath = '/TrainDatas/'
    img_rootpath = os.path.join(rootpath,'JPEGImages')
    json_rootpath = os.path.join(rootpath,'Annotations_LED')
    save_imgpath = os.path.join(rootpath,'crop_img','JPEGImages')
    os.makedirs(save_imgpath, exist_ok=True)
    save_jsonpath = os.path.join(rootpath,'crop_img','Annotations_LED')
    os.makedirs(save_jsonpath, exist_ok=True)

    for filename in os.listdir(json_rootpath):
        name, typ = filename.split('.')
        if typ=='json' :
            img_path = os.path.join(img_rootpath, name+'.jpg')
            json_path = os.path.join(json_rootpath, name+'.json')
            if os.path.exists(img_path) and os.path.exists(json_path):
                divide_img(img_path, json_path, name, save_imgpath, save_jsonpath)
                print(f"{name} >>>>>>>>>>> finish >>>>>>>>>>>")
            else:
                print(f"{name}图像不存在")