图片增强组件实现

设计并实现了一个图片增强的组件，具体功能如下：

• 图片数据增强，包括且不限于：图片旋转、比例增强、高斯噪声、饱和度变换等
• 若图片包含对应标注boundingbox，也支持对应变换，保证圈选内容的不变性
• 实现多种方式random组合，增加增强效果

具体依赖

import cv2
import numpy as np
import copy
import random
from PIL import Image
import os

数据增强方法

@classmethod
def resize_image(cls, image, mode: str = "img", notes: [[dict]] = None, scale_percent: int = None):
    if scale_percent is None:
        scale_percent = random.randrange(30, 80)
    width = int(image.shape[1] * scale_percent / 100)
    height = int(image.shape[0] * scale_percent / 100)
    resized_image = cv2.resize(image, (width, height))
    if mode == 'notes' and notes is None:
        raise ValueError("When mode is set to 'notes', you must provide notes data.")
    elif mode == 'notes':
        points = [i.get('points') for i in notes]
        resized_points = []
        for box in points:
            resized_box = []
            for point in box:
                x = int(point[0] * scale_percent / 100)
                y = int(point[1] * scale_percent / 100)
                resized_box.append([x, y])
            resized_points.append(resized_box)

        resized_notes = copy.deepcopy(notes)
        for i in range(len(resized_notes)):
            resized_notes[i]['points'] = resized_points[i]

        return resized_image, resized_notes
    elif mode == 'img':
        return resized_image

@classmethod
def add_gaussian_noise(cls, image, mode: str = "img", notes: [[dict]] = None, mean=70, std=20):
    noise = np.random.normal(mean, std, image.shape).astype(np.uint8)
    noisy_image = cv2.add(image, noise)
    if mode == 'notes' and notes is None:
        raise ValueError("When mode is set to 'notes', you must provide notes data.")
    elif mode == 'notes':
        gaussian_notes = copy.deepcopy(notes)
        return noisy_image, gaussian_notes
    elif mode == 'img':
        return noisy_image

完整代码

source