用labelme标注的目标检测数据并用albumentations进行数据增强

思路：先将labelme格式改成albumentations能识别的coco格式，然后使用albumentations进行处理，再将处理后的albumentations格式改成labelme的格式保存

import os
import cv2
import json
import albumentations as A

def data_aug(source_annotated, data_aug_annotated, times=3):
    transform = A.Compose([
            A.RandomBrightnessContrast(p=0.3),
            A.OneOf([A.GaussNoise(), ], p=0.5),
            A.OneOf([
                A.MotionBlur(p=0.2),
                A.MedianBlur(blur_limit=3, p=0.5),
                A.Blur(blur_limit=3, p=0.2),
            ], p=0.5),

            A.RandomBrightnessContrast(p=0.5),   
            A.RandomCrop(300, 300, p=0.5),
            A.RandomRotate90(p=0.5),
            A.HorizontalFlip(p=0.5),
            A.VerticalFlip(p=0.5),
            # to do
        ], bbox_params=A.BboxParams(format='coco', min_area=0., min_visibility=0., label_fields=['class_labels']))

    for js_file in os.listdir(source_annotated):
        if js_file.endswith(".json"):
            for nth in range(times):
                with open(source_annotated + str(js_file), 'r', encoding='utf8') as js:
                    json_base_name = os.path.splitext(js_file)[0]

                    coco_bboxes = []
                    class_labels = []
                    points = []
                    json_data = {}
                    json_data = json.load(js)
                    json_data_new = {
                        "version": json.__version__,
                        "flags": {},
                        "shapes": [],
                        "imagePath": " ",
                        "imageData": None,
                        "imageHeight": None,
                        "imageWidth": None
                    }

                    image_path = json_data["imagePath"]
                    image_base_name = os.path.splitext(image_path)[0]
                    image_extension = os.path.splitext(image_path)[1]  #

                    if json_base_name == image_base_name:
                        image = cv2.imread(source_annotated + image_path, cv2.IMREAD_UNCHANGED)
                        for jd in range(len(json_data["shapes"]) ):
                            coco_bboxes_inner = []
                            label = json_data.get("shapes")[jd].get("label")
                            points = json_data.get("shapes")[jd].get("points")

                            #两坐标
                            X1 = points[0][0]
                            Y1 = points[0][1]
                            X2 = points[1][0]
                            Y2 = points[1][1]

                            # 转化为coco格式的坐标
                            # 获取左上角坐标
                            coco_bboxes_inner.append(X2 if X1 > X2 else X1)
                            coco_bboxes_inner.append(Y2 if Y1 > Y2 else Y1)
                            #width、height
                            coco_bboxes_inner.append(abs(X1 - X2))
                            coco_bboxes_inner.append(abs(Y1 - Y2))
                            coco_bboxes.append(coco_bboxes_inner)
                            class_labels.append(label)

                        transformed = transform(image=image, bboxes=coco_bboxes, class_labels=class_labels)
                        transformed_image = transformed['image']
                        transformed_bboxes = transformed['bboxes']

                        # 将经过Albumentations变化过的bbox转成labelMe格式的坐标
                        transed_bbox = []
                        if len(transformed_bboxes) == 0:
                            continue
                        for tb in transformed_bboxes:
                            #组合左上角坐标(X小Y小)
                            transed_bbox.append([float(tb[0]), float(tb[1])])
                            transed_bbox.append([float(tb[0]+tb[2]), float(tb[1] + tb[3])])

                        new_labelme_bbox = []
                        for i in range(len(transed_bbox)//2):
                            ii = i*2
                            new_labelme_bbox.append([transed_bbox[ii], transed_bbox[ii+1]])

                        save_base_name = json_base_name + "_" + str(nth)
                        cv2.imwrite(data_aug_annotated + "/" + save_base_name + image_extension, transformed_image)

                        for bbox, idx in zip(new_labelme_bbox, transformed['class_labels']):
                            dict_info = {'label': idx,
                                         'points': bbox,
                                         "group_id": None,
                                         "shape_type": "rectangle",
                                         "flags": {},
                                         }
                            json_data_new['shapes'].append(dict_info)
                            json_data_new["imagePath"] = save_base_name + image_extension
                            json_data_new["imageData"] = None
                            json_data_new["imageHeight"] = transformed_image.shape[0]
                            json_data_new["imageWidth"] = transformed_image.shape[1]

                        #写入新json文件
                        js_file_new = open(data_aug_annotated + "/" + save_base_name + '.json', 'w')
                        json.dump(json_data_new, js_file_new, indent=2)



if __name__ == "__main__":
    source_annotated = "data_annotated/"  # 原始的图片、标签路径
    data_aug_annotated = "data_aug_annotated" # 处理后的图片、标签路径
    times = 3  
    data_aug(source_annotated, data_aug_annotated, times)