图片增强组件实现

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

• 图片数据增强，包括且不限于：图片旋转、比例增强、高斯噪声、饱和度变换等
• 若图片包含对应标注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_noteselif 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_noteselif mode == 'img':return noisy_image

完整代码

source