工业质检-LabelMe-albumentations-BOX的图像数据增强

最新推荐文章于 2024-07-09 09:18:11 发布

Leagues

最新推荐文章于 2024-07-09 09:18:11 发布

阅读量267

点赞数

文章标签： python

本文链接：https://blog.csdn.net/Leagues/article/details/130986080

版权

格式转换需求: LabelMe进行打标，数据不足，通过albumentations进行图像数据量增加

import os

import albumentations as A
import cv2
import json
import tqdm

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.6),    # 中值滤波
        A.Blur(blur_limit=3, p=0.2),   # 使用随机大小的内核模糊输入图像。
    ], p=0.6),
        # 随机应用仿射变换：平移，缩放和旋转输入
        A.RandomBrightnessContrast(p=0.3),   # 随机明亮对比度



], bbox_params=A.BboxParams(format='coco', min_area=1000000, min_visibility=0.6, label_fields=['class_labels']))

source_path = r'/workpace/W8_data/2'
for i in tqdm(range(2000)):
    for file in os.listdir(source_path):
        if file.endswith(".json"):
            with open(source_path + str(file), 'r', encoding='utf8') as fp:
                bboxes_inner = []
                class_labels = []
                bboxes = []
                points = []
                json_data = {}
                json_data = json.load(fp)
                label = json_data.get("shapes")[0].get("label")
                points = json_data.get("shapes")[0].get("points")
                #两个X坐标
                X1 = points[0][0]
                X2 = points[1][0]

                #两个Y坐标
                Y1 = points[0][1]
                Y2 = points[1][1]

                #获取左上角坐标
                bboxes_inner.append(X2 if X1 > X2 else X1)
                bboxes_inner.append(Y2 if Y1 > Y2 else Y1)

                #width、height
                bboxes_inner.append(abs(int(points[0][0]) - int(points[1][0])))
                bboxes_inner.append(abs(int(points[0][1]) - int(points[1][1])))
                #bboxes_inner.append(label)
                bboxes.append(bboxes_inner)
                json_data["imageData"] = None
                class_labels.append(label)

                image = cv2.imread(source_path + file.split('.')[0] + '.jpg')
                image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

                print(file)
                transformed = transform(image=image, bboxes=bboxes,class_labels=class_labels)
                transformed_image = transformed['image']
                transformed_bboxes = transformed['bboxes']
                new_box = []
                print(transformed_bboxes)
                if len(transformed_bboxes) == 0:
                    continue
                #组合左上角坐标(X小Y小)
                new_box.append([float(transformed_bboxes[0][0]) ,float(transformed_bboxes[0][1])])

                new_box.append([float(transformed_bboxes[0][0]) + float(transformed_bboxes[0][2]), float(transformed_bboxes[0][1]) + float(transformed_bboxes[0][3])  ])

                #将经过Albumentations变化过的box转成labelMe能识别的坐标
                print('转变之前的box ',   bboxes    ,  '转变之后的transformed_bboxes' , transformed_bboxes , '转变之后的new_box',new_box )

                points = json_data.get("shapes")[0]["points"] = new_box
                transformed_class_labels = transformed['class_labels']
                #json_data.get("shapes")[0]["points"] = transformed_bboxes
                #写入json新文件
                file_new = open(r"/workpace/W8_data/1/"  + str(file), 'w')
                json.dump(json_data, file_new, indent=2)  # indent的目的是为了使得JSON格式样子更加好看
                transformed_image = cv2.cvtColor(transformed_image, cv2.COLOR_RGB2BGR)
                cv2.imwrite(r"/workpace/W8_data/1/" + file.split('.')[0] + '.jpg', transformed_image)