from typing import * from pydantic import BaseModel from abc import * import random import numpy as np import matplotlib.pyplot as plt import seaborn as sns from ..Internal import * from ..MathEx.Core import * #from ..Str.Core import UnWrapper as Unwrapper2Str from ..File.Core import tool_file, Wrapper as Wrapper2File, tool_file_or_str, is_image_file, loss_file, static_loss_file_dir from ..Visual.OpenCV import ImageObject, tool_file_cvex, WrapperFile2CVEX, Wrapper as Wrapper2Image, get_new_noise from PIL.Image import ( Image as PILImage, fromarray as PILFromArray, open as PILOpen ) from PIL.ImageFile import ImageFile as PILImageFile import cv2 as cv2 from io import BytesIO class data_visual_generator: def __init__(self, file:tool_file_or_str): self._file:tool_file = Wrapper2File(file) self._file.load() def open(self, mode='r', is_refresh=False, encoding:str='utf-8', *args, **kwargs): self._file.open(mode, is_refresh, encoding, *args, **kwargs) def reload(self, file:Optional[tool_file_or_str]): if file is not None: self._file = Wrapper2File(file) self._file.load() def plot_line(self, x, y, df=None, title="折线图", x_label=None, y_label=None): plt.figure(figsize=(10, 6)) sns.lineplot(data=df if df is not None else self._file.data, x=x, y=y) plt.title(title) plt.xlabel(x_label if x_label is not None else str(x)) plt.ylabel(y_label if y_label is not None else str(y)) plt.grid(True) plt.show() def plot_bar(self, x, y, df=None, figsize=(10,6), title="柱状图", x_label=None, y_label=None): plt.figure(figsize=figsize) sns.barplot(data=df if df is not None else self._file.data, x=x, y=y) plt.title(title) plt.xlabel(x_label if x_label is not None else str(x)) plt.ylabel(y_label if y_label is not None else str(y)) plt.grid(True) plt.show() def plot_scatter(self, x, y, df=None, title="散点图", x_label=None, y_label=None): plt.figure(figsize=(10, 6)) sns.scatterplot(data=df if df is not None else self._file.data, x=x, y=y) plt.title(title) plt.xlabel(x_label if x_label is not None else str(x)) plt.ylabel(y_label if y_label is not None else str(y)) plt.grid(True) plt.show() def plot_histogram(self, column, df=None, title="直方图", x_label=None, y_label=None): plt.figure(figsize=(10, 6)) sns.histplot(data=df if df is not None else self._file.data, x=column) plt.title(title) plt.xlabel(x_label if x_label is not None else str(column)) plt.ylabel(y_label if y_label is not None else "value") plt.grid(True) plt.show() def plot_pairplot(self, df=None, title="成对关系图"): sns.pairplot(df if df is not None else self._file.data) plt.suptitle(title, y=1.02) plt.show() def plot_pie(self, column, figsize=(10,6), df=None, title="饼图"): plt.figure(figsize=figsize) if df is not None: df[column].value_counts().plot.pie(autopct='%1.1f%%') else: self._file.data[column].value_counts().plot.pie(autopct='%1.1f%%') plt.title(title) plt.ylabel('') # 移除y轴标签 plt.show() def plot_box(self, x, y, df=None, figsize=(10,6), title="箱线图", x_label=None, y_label=None): plt.figure(figsize=figsize) sns.boxplot(data=df if df is not None else self._file.data, x=x, y=y) plt.title(title) plt.xlabel(x_label if x_label is not None else str(x)) plt.ylabel(y_label if y_label is not None else str(y)) plt.grid(True) plt.show() def plot_heatmap(self, df=None, figsize=(10,6), title="热力图", cmap='coolwarm'): plt.figure(figsize=figsize) sns.heatmap(df.corr() if df is not None else self._file.data.corr(), annot=True, cmap=cmap) plt.title(title) plt.show() def plot_catplot(self, x, y, hue=None, df=None, kind='bar', figsize=(10,6), title="分类数据图", x_label=None, y_label=None): plt.figure(figsize=figsize) sns.catplot(data=df if df is not None else self._file.data, x=x, y=y, hue=hue, kind=kind) plt.title(title) plt.xlabel(x_label if x_label is not None else str(x)) plt.ylabel(y_label if y_label is not None else str(y)) plt.grid(True) plt.show() def plot_catplot_strip(self, x, y, hue=None, df=None, figsize=(10,6), title="分类数据图", x_label=None, y_label=None): self.plot_catplot(x, y, hue=hue, df=df, kind='strip', figsize=figsize, title=title, x_label=x_label, y_label=y_label) def plot_catplot_swarm(self, x, y, hue=None, df=None, figsize=(10,6), title="分类数据图", x_label=None, y_label=None): self.plot_catplot(x, y, hue=hue, df=df, kind='swarm', figsize=figsize, title=title, x_label=x_label, y_label=y_label) def plot_catplot_box(self, x, y, hue=None, df=None, figsize=(10,6), title="分类数据图", x_label=None, y_label=None): self.plot_catplot(x, y, hue=hue, df=df, kind='box', figsize=figsize, title=title, x_label=x_label, y_label=y_label) def plot_catplot_violin(self, x, y, hue=None, df=None, figsize=(10,6), title="分类数据图", x_label=None, y_label=None): self.plot_catplot(x, y, hue=hue, df=df, kind='violin', figsize=figsize, title=title, x_label=x_label, y_label=y_label) def plot_jointplot(self, x, y, kind="scatter", df=None, title="联合图", x_label=None, y_label=None): sns.jointplot(data=df if df is not None else self._file.data, x=x, y=y, kind=kind) plt.suptitle(title, y=1.02) plt.xlabel(x_label if x_label is not None else str(x)) plt.ylabel(y_label if y_label is not None else str(y)) plt.show() def plot_jointplot_scatter(self, x, y, df=None, title="联合图", x_label=None, y_label=None): self.plot_jointplot(x, y, kind="scatter", df=df, title=title, x_label=x_label, y_label=y_label) def plot_jointplot_kde(self, x, y, df=None, title="联合图", x_label=None, y_label=None): self.plot_jointplot(x, y, kind="kde", df=df, title=title, x_label=x_label, y_label=y_label) def plot_jointplot_hex(self, x, y, df=None, title="联合图", x_label=None, y_label=None): self.plot_jointplot(x, y, kind="hex", df=df, title=title, x_label=x_label, y_label=y_label) class data_math_virsual_generator(data_visual_generator): def drop_missing_values(self, axis): """删除缺失值""" self._file.data = self._file.data.dropna(axis=axis) def fill_missing_values(self, value): """填充缺失值""" self._file.data = self._file.data.fillna(value) def remove_duplicates(self): """删除重复值""" self._file.data = self._file.data.drop_duplicates() def standardize_data(self): """数据标准化""" self._file.data = (self._file.data - self._file.data.mean()) / self._file.data.std() def normalize_data(self): """数据归一化""" self._file.data = (self._file.data - self._file.data.min()) / (self._file.data.max() - self._file.data.min()) # region image augmentation NDARRAY_ANY = TypeVar("numpy.ndarray") class BasicAugmentConfig(BaseModel, ABC): name: str = "unknown" @abstractmethod def augment( self, origin: ImageObject ) -> Tuple[Dict[str, Any], ImageObject]: ''' result: (change config, image) ''' raise NotImplementedError() class ResizeAugmentConfig(BasicAugmentConfig): width: Optional[int] = None height: Optional[int] = None name: str = "resize" @override def augment( self, origin: ImageObject ) -> Tuple[Dict[str, Any], ImageObject]: width = self.width height = self.height if width is None and height is None: rangewidth = origin.width rangeheight = origin.height width = rangewidth + random.randint( (-rangewidth*(random.random()%1)).__floor__(), (rangewidth*(random.random()%1)).__floor__() ) height = rangeheight + random.randint( (-rangeheight*(random.random()%1)).__floor__(), (rangeheight*(random.random()%1)).__floor__() ) elif width is None: width = origin.width elif height is None: height = origin.height change_config = { "width":width, "height":height } return (change_config, ImageObject(origin.get_resize_image(abs(width), abs(height)))) class ClipAugmentConfig(BasicAugmentConfig): mini: Union[float, NDARRAY_ANY] = 0 maxi: Union[float, NDARRAY_ANY] = 255 name: str = "clip" @override def augment( self, origin: ImageObject ) -> Tuple[Dict[str, Any], ImageObject]: mini = self.mini maxi = self.maxi if isinstance(mini, ImageObject): mini = mini.get_array() if isinstance(maxi, ImageObject): maxi = maxi.get_array() change_config = { "mini":mini, "maxi":maxi } return (change_config, ImageObject(origin.clip(mini, maxi))) class NormalizeAugmentConfig(BasicAugmentConfig): name: str = "normalize" mini: Optional[Union[NDARRAY_ANY, float]] = 0 maxi: Optional[Union[NDARRAY_ANY, float]] = 255 @override def augment( self, origin: ImageObject ) -> Tuple[Dict[str, Any], ImageObject]: change_config = { "mini":self.mini, "maxi":self.maxi } return (change_config, ImageObject(origin.normalize(self.mini, self.maxi))) class StandardizeAugmentConfig(BasicAugmentConfig): name: str = "standardize" mean: Optional[Union[NDARRAY_ANY, float]] = 0 std: Optional[Union[NDARRAY_ANY, float]] = 1 @override def augment( self, origin: ImageObject ) -> Tuple[Dict[str, Any], ImageObject]: change_config = { "mean":origin.get_array().mean(), "std":origin.get_array().std() } return (change_config, ImageObject(origin.standardize(self.mean, self.std))) class FlipAugmentConfig(BasicAugmentConfig): name: str = "flip" axis: Literal[-1, 1, 0] = 1 ''' 1: vertical 0: horizontal -1: both ''' @override def augment( self, origin: ImageObject ) -> Tuple[Dict[str, Any], ImageObject]: change_config = { "axis":self.axis } return (change_config, ImageObject(origin.flip(self.axis))) class CropAugmentConfig(BasicAugmentConfig): name: str = "crop" lbx: Optional[int] = None lby: Optional[int] = None width: Optional[int] = None height: Optional[int] = None @override def augment( self, origin: ImageObject ) -> Tuple[Dict[str, Any], ImageObject]: lbx = self.lbx if self.lbx is not None else random.randint(0, origin.width) lby = self.lby if self.lby is not None else random.randint(0, origin.height) width = self.width if self.width is not None else random.randint(1, origin.width - lbx) height = self.height if self.height is not None else random.randint(0, origin.height - lby) change_config = { "lbx":lbx, "lby":lby, "width":width, "height":height } return (change_config, ImageObject(origin.sub_image_with_rect((lbx, lby, width, height)))) class FilterAugmentConfig(BasicAugmentConfig): name: str = "filter" ddepth: int = -1 kernal: NDARRAY_ANY = cv2.getGaussianKernel(3, 1) def get_gaussian_kernal(self, kernal_size: int, sigma: float): return cv2.getGaussianKernel(kernal_size, sigma) @override def augment( self, origin: ImageObject ) -> Tuple[Dict[str, Any], ImageObject]: change_config = { "filter":self.ddepth, "kernal":self.kernal } return (change_config, ImageObject(origin.filter(self.ddepth, self.kernal))) class ColorSpaceAugmentConfig(BasicAugmentConfig): name: str = "color_space" space: int = cv2.COLOR_BGR2GRAY @override def augment( self, origin: ImageObject ) -> Tuple[Dict[str, Any], ImageObject]: change_config = { "color_space":self.space } return (change_config, ImageObject(origin.convert_to(self.space))) class LightingAugmentConfig(BasicAugmentConfig): name: str = "lighting" lighting: Optional[int] = None @override def augment( self, origin: ImageObject ) -> Tuple[Dict[str, Any], ImageObject]: lighting = self.lighting if self.lighting is not None else random.randint(0, 50) change_config = { "lighting":lighting } return (change_config, ImageObject(cv2.add(origin.image, lighting))) class DarkingAugmentConfig(BasicAugmentConfig): name: str = "darking" darking: Optional[FloatBetween01] = None @override def augment( self, origin: ImageObject ) -> Tuple[Dict[str, Any], ImageObject]: darking = self.darking if self.darking is not None else (random.random()%0.9+0.1) change_config = { "darking":darking } return (change_config, ImageObject(origin*darking)) class ContrastAugmentConfig(BasicAugmentConfig): name: str = "contrast" contrast: Optional[FloatBetween01] = None @override def augment( self, origin: ImageObject ) -> Tuple[Dict[str, Any], ImageObject]: change_config = { "contrast":self.contrast } contrast = self.contrast if self.contrast is not None else (random.random()%0.9+0.1) contrast = int(contrast*255) result = origin.image*(contrast / 127 + 1) - contrast return (change_config, ImageObject(result)) class SeparateSceneAugmentConfig(BasicAugmentConfig): scene: str = "separate_scene" is_front: bool = True @override def augment( self, origin: ImageObject ) -> Tuple[Dict[str, Any], ImageObject]: change_config = { "is_front":self.is_front } front, back = origin.SeparateFrontBackScenes() target_0 = back if self.is_front else front image = origin.image.copy() image[target_0] = 0 return (change_config, ImageObject(image)) class NoiseAugmentConfig(BasicAugmentConfig): name: str = "noise" mean: float = 0 sigma: float = 25 @override def augment( self, origin: ImageObject ) -> Tuple[Dict[str, Any], ImageObject]: change_config = { "mean":self.mean, "sigma":self.sigma } return (change_config, ImageObject( origin + get_new_noise( None, origin.height, origin.width, mean=self.mean, sigma=self.sigma ) )) class VignettingAugmentConfig(BasicAugmentConfig): name: str = "vignetting" ratio_min_dist: float = 0.2 range_vignette: Tuple[float, float] = (0.2, 0.8) random_sign: bool = False @override def augment( self, origin: ImageObject ) -> Tuple[Dict[str, Any], ImageObject]: change_config = { "ratio_min_dist":self.ratio_min_dist, "range_vignette":self.range_vignette, "random_sign":self.random_sign } h, w = origin.shape[:2] min_dist = np.array([h, w]) / 2 * np.random.random() * self.ratio_min_dist # 创建距离中心点在两个轴上的距离矩阵 x, y = np.meshgrid(np.linspace(-w/2, w/2, w), np.linspace(-h/2, h/2, h)) x, y = np.abs(x), np.abs(y) # 在两个轴上创建晕影遮罩 x = (x - min_dist[0]) / (np.max(x) - min_dist[0]) x = np.clip(x, 0, 1) y = (y - min_dist[1]) / (np.max(y) - min_dist[1]) y = np.clip(y, 0, 1) # 获取随机晕影强度 vignette = (x + y) / 2 * np.random.uniform(self.range_vignette[0], self.range_vignette[1]) vignette = np.tile(vignette[..., None], [1, 1, 3]) sign = 2 * (np.random.random() < 0.5) * (self.random_sign) - 1 return (change_config, ImageObject(origin * (1 + sign * vignette))) class LensDistortionAugmentConfig(BasicAugmentConfig): name: str = "lens_distortion" d_coef: Tuple[ float, float, float, float, float ] = (0.15, 0.15, 0.1, 0.1, 0.05) @override def augment( self, origin: ImageObject ) -> Tuple[Dict[str, Any], ImageObject]: change_config = { "d_coef":self.d_coef } # 获取图像的高度和宽度 h, w = origin.shape[:2] # 计算对角线长度 f = (h ** 2 + w ** 2) ** 0.5 # 设置图像投影到笛卡尔坐标系的维度 K = np.array([[f, 0, w / 2], [0, f, h / 2], [0, 0, 1]]) d_coef = self.d_coef * np.random.random(5) # 值 d_coef = d_coef * (2 * (np.random.random(5) < 0.5) - 1) # 符号 # 从参数生成新的相机矩阵 M, _ = cv2.getOptimalNewCameraMatrix(K, d_coef, (w, h), 0) # 生成用于重映射相机图像的查找表 remap = cv2.initUndistortRectifyMap(K, d_coef, None, M, (w, h), 5) # 将原始图像重映射到新图像 return (change_config, ImageObject(cv2.remap(origin.image, *remap, cv2.INTER_LINEAR))) class RotationAugmentConfig(BasicAugmentConfig): name: str = "rotation" angle: Optional[float] = None scale: float = 1.0 center: Optional[Tuple[int, int]] = None @override def augment( self, origin: ImageObject ) -> Tuple[Dict[str, Any], ImageObject]: angle = self.angle if self.angle is not None else random.uniform(-30, 30) center = self.center if self.center is not None else (origin.width // 2, origin.height // 2) change_config = { "angle": angle, "scale": self.scale, "center": center } # 获取旋转矩阵 rotation_matrix = cv2.getRotationMatrix2D(center, angle, self.scale) # 应用旋转变换 rotated_image = cv2.warpAffine(origin.image, rotation_matrix, (origin.width, origin.height)) return (change_config, ImageObject(rotated_image)) class BlurAugmentConfig(BasicAugmentConfig): name: str = "blur" kernel_size: Tuple[int, int] = (5, 5) @override def augment( self, origin: ImageObject ) -> Tuple[Dict[str, Any], ImageObject]: change_config = { "kernel_size": self.kernel_size } # 应用均值模糊 blurred_image = cv2.blur(origin.image, self.kernel_size) return (change_config, ImageObject(blurred_image)) class MedianBlurAugmentConfig(BasicAugmentConfig): name: str = "median_blur" ksize: int = 5 @override def augment( self, origin: ImageObject ) -> Tuple[Dict[str, Any], ImageObject]: change_config = { "ksize": self.ksize } # 应用中值模糊 blurred_image = cv2.medianBlur(origin.image, self.ksize) return (change_config, ImageObject(blurred_image)) class SaturationAugmentConfig(BasicAugmentConfig): name: str = "saturation" factor: Optional[float] = None @override def augment( self, origin: ImageObject ) -> Tuple[Dict[str, Any], ImageObject]: factor = self.factor if self.factor is not None else random.uniform(0.5, 1.5) change_config = { "factor": factor } # 转换为HSV颜色空间 hsv = cv2.cvtColor(origin.image, cv2.COLOR_BGR2HSV).astype(np.float32) # 调整饱和度 hsv[:, :, 1] = hsv[:, :, 1] * factor hsv[:, :, 1] = np.clip(hsv[:, :, 1], 0, 255) # 转换回BGR颜色空间 result = cv2.cvtColor(hsv.astype(np.uint8), cv2.COLOR_HSV2BGR) return (change_config, ImageObject(result)) class HueAugmentConfig(BasicAugmentConfig): name: str = "hue" shift: Optional[int] = None @override def augment( self, origin: ImageObject ) -> Tuple[Dict[str, Any], ImageObject]: shift = self.shift if self.shift is not None else random.randint(-20, 20) change_config = { "shift": shift } # 转换为HSV颜色空间 hsv = cv2.cvtColor(origin.image, cv2.COLOR_BGR2HSV).astype(np.float32) # 调整色调 hsv[:, :, 0] = (hsv[:, :, 0] + shift) % 180 # 转换回BGR颜色空间 result = cv2.cvtColor(hsv.astype(np.uint8), cv2.COLOR_HSV2BGR) return (change_config, ImageObject(result)) class GammaAugmentConfig(BasicAugmentConfig): name: str = "gamma" gamma: Optional[float] = None @override def augment( self, origin: ImageObject ) -> Tuple[Dict[str, Any], ImageObject]: gamma = self.gamma if self.gamma is not None else random.uniform(0.5, 2.0) change_config = { "gamma": gamma } # 应用伽马校正 inv_gamma = 1.0 / gamma table = np.array([((i / 255.0) ** inv_gamma) * 255 for i in range(256)]).astype(np.uint8) result = cv2.LUT(origin.image, table) return (change_config, ImageObject(result)) class PerspectiveTransformAugmentConfig(BasicAugmentConfig): name: str = "perspective" intensity: Optional[float] = None @override def augment( self, origin: ImageObject ) -> Tuple[Dict[str, Any], ImageObject]: intensity = self.intensity if self.intensity is not None else random.uniform(0.05, 0.1) change_config = { "intensity": intensity } h, w = origin.shape[:2] # 定义源点和目标点 src_points = np.float32([[0, 0], [w, 0], [0, h], [w, h]]) # 随机扰动目标点 dst_points = np.float32([ [0 + random.uniform(-intensity * w, intensity * w), 0 + random.uniform(-intensity * h, intensity * h)], [w + random.uniform(-intensity * w, intensity * w), 0 + random.uniform(-intensity * h, intensity * h)], [0 + random.uniform(-intensity * w, intensity * w), h + random.uniform(-intensity * h, intensity * h)], [w + random.uniform(-intensity * w, intensity * w), h + random.uniform(-intensity * h, intensity * h)] ]) # 计算透视变换矩阵 M = cv2.getPerspectiveTransform(src_points, dst_points) # 应用透视变换 result = cv2.warpPerspective(origin.image, M, (w, h)) return (change_config, ImageObject(result)) class ElasticTransformAugmentConfig(BasicAugmentConfig): name: str = "elastic" alpha: float = 50 sigma: float = 5 @override def augment( self, origin: ImageObject ) -> Tuple[Dict[str, Any], ImageObject]: change_config = { "alpha": self.alpha, "sigma": self.sigma } h, w = origin.shape[:2] # 创建随机位移场 dx = np.random.rand(h, w) * 2 - 1 dy = np.random.rand(h, w) * 2 - 1 # 高斯模糊位移场 dx = cv2.GaussianBlur(dx, (0, 0), self.sigma) dy = cv2.GaussianBlur(dy, (0, 0), self.sigma) # 归一化并缩放 dx = dx * self.alpha dy = dy * self.alpha # 创建网格 x, y = np.meshgrid(np.arange(w), np.arange(h)) # 应用位移 indices_x = np.clip(x + dx, 0, w - 1).astype(np.float32) indices_y = np.clip(y + dy, 0, h - 1).astype(np.float32) # 重映射 result = cv2.remap(origin.image, indices_x, indices_y, interpolation=cv2.INTER_LINEAR) return (change_config, ImageObject(result)) class ChannelShuffleAugmentConfig(BasicAugmentConfig): name: str = "channel_shuffle" @override def augment( self, origin: ImageObject ) -> Tuple[Dict[str, Any], ImageObject]: # 获取随机通道顺序 channels = list(range(origin.image.shape[2])) random.shuffle(channels) change_config = { "channels": channels } # 重新排列通道 result = origin.image[:, :, channels] return (change_config, ImageObject(result)) class MotionBlurAugmentConfig(BasicAugmentConfig): name: str = "motion_blur" kernel_size: int = 15 angle: Optional[float] = None @override def augment( self, origin: ImageObject ) -> Tuple[Dict[str, Any], ImageObject]: angle = self.angle if self.angle is not None else random.uniform(0, 360) change_config = { "kernel_size": self.kernel_size, "angle": angle } # 创建运动模糊核 kernel = np.zeros((self.kernel_size, self.kernel_size)) center = self.kernel_size // 2 # 计算角度的弧度值 rad = np.deg2rad(angle) # 在核上绘制一条线 x = np.cos(rad) * center y = np.sin(rad) * center # 使用Bresenham算法绘制线 cv2.line(kernel, (center - int(np.round(x)), center - int(np.round(y))), (center + int(np.round(x)), center + int(np.round(y))), 1, thickness=1) # 归一化核 kernel = kernel / np.sum(kernel) # 应用卷积 result = cv2.filter2D(origin.image, -1, kernel) return (change_config, ImageObject(result)) class SolarizeAugmentConfig(BasicAugmentConfig): name: str = "solarize" threshold: Optional[int] = None @override def augment( self, origin: ImageObject ) -> Tuple[Dict[str, Any], ImageObject]: threshold = self.threshold if self.threshold is not None else random.randint(100, 200) change_config = { "threshold": threshold } # 应用曝光效果 result = origin.image.copy() mask = origin.image > threshold result[mask] = 255 - result[mask] return (change_config, ImageObject(result)) class PosterizeAugmentConfig(BasicAugmentConfig): name: str = "posterize" bits: Optional[int] = None @override def augment( self, origin: ImageObject ) -> Tuple[Dict[str, Any], ImageObject]: bits = self.bits if self.bits is not None else random.randint(3, 7) change_config = { "bits": bits } # 应用海报效果（减少颜色位数） mask = 255 - (1 << (8 - bits)) result = origin.image & mask return (change_config, ImageObject(result)) class InvertAugmentConfig(BasicAugmentConfig): name: str = "invert" @override def augment( self, origin: ImageObject ) -> Tuple[Dict[str, Any], ImageObject]: change_config = {} # 反转图像颜色 result = 255 - origin.image return (change_config, ImageObject(result)) class EqualizationAugmentConfig(BasicAugmentConfig): name: str = "equalize" @override def augment( self, origin: ImageObject ) -> Tuple[Dict[str, Any], ImageObject]: change_config = {} # 对每个通道进行直方图均衡化 result = origin.image.copy() if len(origin.shape) > 2 and origin.shape[2] > 1: for i in range(origin.shape[2]): result[:, :, i] = cv2.equalizeHist(origin.image[:, :, i]) else: result = cv2.equalizeHist(origin.image) return (change_config, ImageObject(result)) class CutoutAugmentConfig(BasicAugmentConfig): name: str = "cutout" n_holes: int = 1 length: Optional[int] = None @override def augment( self, origin: ImageObject ) -> Tuple[Dict[str, Any], ImageObject]: length = self.length if self.length is not None else min(origin.width, origin.height) // 4 change_config = { "n_holes": self.n_holes, "length": length } result = origin.image.copy() h, w = origin.shape[:2] for _ in range(self.n_holes): # 随机选择矩形的中心点 y = np.random.randint(h) x = np.random.randint(w) # 计算矩形的边界 y1 = np.clip(y - length // 2, 0, h) y2 = np.clip(y + length // 2, 0, h) x1 = np.clip(x - length // 2, 0, w) x2 = np.clip(x + length // 2, 0, w) # 将矩形区域填充为黑色 result[y1:y2, x1:x2] = 0 return (change_config, ImageObject(result)) # Config.name -> (field, value) type ChangeConfig = Dict[str, Dict[str, Any]] # (field, value) type ResultImageObjects = Dict[str, ImageObject] class ImageAugmentConfig(BaseModel): resize: Optional[ResizeAugmentConfig] = None clip: Optional[ClipAugmentConfig] = None normalize: Optional[NormalizeAugmentConfig] = None standardize:Optional[StandardizeAugmentConfig] = None flip: Optional[FlipAugmentConfig] = None crop: Optional[CropAugmentConfig] = None filters: Sequence[FilterAugmentConfig] = [] colorspace: Optional[ColorSpaceAugmentConfig] = None lighting: Optional[LightingAugmentConfig] = None darking: Optional[DarkingAugmentConfig] = None contrast: Optional[ContrastAugmentConfig] = None separate_scene: Literal[0, 1, 2, 3] = 0 noise: Optional[NoiseAugmentConfig] = None vignette: Optional[VignettingAugmentConfig] = None lens_distortion: Optional[LensDistortionAugmentConfig] = None rotation: Optional[RotationAugmentConfig] = None blur: Optional[BlurAugmentConfig] = None median_blur:Optional[MedianBlurAugmentConfig] = None saturation: Optional[SaturationAugmentConfig] = None hue: Optional[HueAugmentConfig] = None gamma: Optional[GammaAugmentConfig] = None perspective:Optional[PerspectiveTransformAugmentConfig] = None elastic: Optional[ElasticTransformAugmentConfig] = None channel_shuffle: Optional[ChannelShuffleAugmentConfig] = None motion_blur:Optional[MotionBlurAugmentConfig] = None solarize: Optional[SolarizeAugmentConfig] = None posterize: Optional[PosterizeAugmentConfig] = None invert: Optional[InvertAugmentConfig] = None ''' None: 0 front: 1 back: 2 both: 3 ''' log_call: Optional[Callable[[Union[str, Dict[str, Any]]], None]] = None def get_all_configs(self) -> List[BasicAugmentConfig]: result = [ self.resize, self.clip, self.normalize, self.standardize, self.flip, self.crop, self.colorspace, self.lighting, self.darking, self.contrast, self.noise, self.vignette, self.lens_distortion ] result.extend(self.filters) if self.separate_scene&1: result.append(SeparateSceneAugmentConfig(is_front=True, name="front_scene")) if self.separate_scene&2: result.append(SeparateSceneAugmentConfig(is_front=False, name="back_scene")) return result def _inject_log(self, *args, **kwargs): if self.log_call is not None: self.log_call(*args, **kwargs) def augment( self, origin: ImageObject ) -> Tuple[ChangeConfig, ResultImageObjects]: result: Dict[str, ImageObject] = {} result_change_config: Dict[str, Dict[str, Any]] = {} augment_configs: List[BasicAugmentConfig] = self.get_all_configs() for item in augment_configs: if item is not None: result_change_config[item.name], result[item.name] = item.augment(ImageObject(origin.image)) self._inject_log(f"augmentation<{item.name}> change config: {result_change_config[item.name]}") return (result_change_config, result) def augment_to( self, input: Union[tool_file, str, ImageObject, np.ndarray, PILImage, PILImageFile], output_dir: tool_file_or_str, *, # 如果输出目录不存在,将调用must_exist # 如果输出目录存在但不是目录,将返回父目录 must_output_dir_exist: bool = False, output_file_name: str = "output.png", callback: Optional[Action[ChangeConfig]] = None, ) -> ResultImageObjects: # 初始化环境和变量 origin_image: ImageObject = self.__init_origin_image(input) result_dir: tool_file = self.__init_result_dir(output_dir, must_output_dir_exist) # 增强 self._inject_log(f"输出<{output_file_name}>开始增强") change_config, result = self._inject_augment( origin_image=origin_image, result_dir=result_dir, output_file_name=output_file_name, ) # 结果 if callback is not None: callback(change_config) return result def augment_from_dir_to( self, input_dir: Union[tool_file, str], output_dir: tool_file_or_str, *, # 如果输出目录不存在,将调用must_exist # 如果输出目录存在但不是目录,将返回父目录 must_output_dir_exist: bool = False, callback: Optional[Action2[tool_file, ChangeConfig]] = None, ) -> Dict[str, List[ImageObject]]: # 初始化环境和变量 origin_images: tool_file = Wrapper2File(input_dir) result_dir: tool_file = self.__init_result_dir(output_dir, must_output_dir_exist) if origin_images.exists() is False or origin_images.is_dir() is False: raise FileExistsError(f"input_dir<{origin_images}> is not exist or not dir") # augment result: Dict[str, List[ImageObject]] = {} for image_file in origin_images.dir_tool_file_iter(): if is_image_file(Unwrapper2Str(image_file)) is False: continue change_config, curResult = self._inject_augment( origin_image=WrapperFile2CVEX(image_file).load(), result_dir=result_dir, output_file_name=image_file.get_filename(), ) # 添加单个结果 for key in curResult: if key in result: result[key].append(curResult[key]) else: result[key] = [curResult[key]] # 调用回调 if callback is not None: callback(image_file, change_config) # 结果 return result def augment_from_images_to( self, inputs: Sequence[ImageObject], output_dir: tool_file_or_str, *, # 如果输出目录不存在,将调用must_exist # 如果输出目录存在但不是目录,将返回父目录 must_output_dir_exist: bool = False, callback: Optional[Action2[ImageObject, ChangeConfig]] = None, fileformat: str = "{}.jpg", indexbuilder: type = int ) -> Dict[str, List[ImageObject]]: # Init env and vars result_dir: tool_file = self.__init_result_dir(output_dir, must_output_dir_exist) index: Any = indexbuilder() # augment result: Dict[str, List[ImageObject]] = {} for image in inputs: current_output_name = fileformat.format(index) change_config, curResult = self._inject_augment( origin_image=image, result_dir=result_dir, output_file_name=current_output_name, ) # append single result for key in curResult: if key in result: result[key].append(curResult[key]) else: result[key] = [curResult[key]] index += 1 # call feedback if callback is not None: callback(image, change_config) # result return result def __init_origin_image(self, input:Union[tool_file, str, ImageObject, np.ndarray, PILImage, PILImageFile]) -> ImageObject: origin_image: ImageObject = None # check if isinstance(input, (tool_file, str)): inputfile = WrapperFile2CVEX(input) if inputfile.data is not None: origin_image = inputfile.data else: origin_image = inputfile.load() elif isinstance(input, (ImageObject, np.ndarray, PILImage, PILImageFile)): origin_image = Wrapper2Image(input) else: raise TypeError(f"input<{input}> is not support type") return origin_image def __init_result_dir(self, output_dir:tool_file_or_str, must_output_dir_exist:bool) -> tool_file: if output_dir is None or isinstance(output_dir, loss_file): return static_loss_file_dir result_dir: tool_file = Wrapper2File(output_dir) # check exist stats: bool = True if result_dir.exists() is False: if must_output_dir_exist: result_dir.must_exists_path() else: stats = False if stats is False: raise FileExistsError(f"output_dir<{result_dir}> is not exist") # check dir stats if result_dir.is_dir() is False: if must_output_dir_exist: result_dir.back_to_parent_dir() else: raise FileExistsError(f"output_dir<{result_dir}> is not a dir") # result return result_dir def _inject_augment( self, origin_image: ImageObject, result_dir: tool_file, output_file_name: str ) -> Tuple[ChangeConfig, ResultImageObjects]: self._inject_log(f"output<{output_file_name}> is start augment") result_dict, result_images = self.augment(origin_image) if not (result_dir is None or isinstance(result_dir, loss_file)): for key, value in result_images.items(): current_dir = result_dir|key current_result_file = current_dir|output_file_name value.save_image(current_result_file, True) return result_dict, result_images def image_augent( config:ImageAugmentConfig, source,*args, **kwargs ): if isinstance(source, ImageObject): return config.augment(source, *args, **kwargs) def get_config_of_gaussian_blur( intensity: float, blur_level: int = 3 ) -> FilterAugmentConfig: return FilterAugmentConfig( name="gaussian_blur", kernal=cv2.getGaussianKernel(blur_level, intensity) ) def get_config_of_smooth_blur(blur_level:int = 3): result = get_config_of_gaussian_blur(0, blur_level) result.name = "smooth_blur" return result def get_config_of_sharpen( intensity: float, sharpen_level: float = 1 ) -> FilterAugmentConfig: return FilterAugmentConfig( name="sharpen", kernal=np.array([ [0, -sharpen_level, 0], [-sharpen_level, intensity-sharpen_level, -sharpen_level], [0, -sharpen_level, 0] ]) ) def get_config_of_edge_enhance( intensity: float, sharpen_level: float = 8 ) -> FilterAugmentConfig: return FilterAugmentConfig( name="edge_enhance", kernal=np.array([ [-intensity, -intensity, -intensity], [-intensity, sharpen_level+intensity, -intensity], [-intensity, -intensity, -intensity] ]) ) def get_config_of_convert_to_gray() -> ColorSpaceAugmentConfig: return ColorSpaceAugmentConfig( name="convert_to_gray", color_space=cv2.COLOR_BGR2GRAY ) # region end # region image convert class BasicConvertConfig(BaseModel, ABC): name: str = "unknown" @abstractmethod def convert( self, origin: ImageObject ) -> Tuple[Dict[str, Any], ImageObject]: ''' result: (change config, image) ''' raise NotImplementedError() class PNGConvertConfig(BasicConvertConfig): name: str = "png" compression_level: int = 6 # 0-9, 9为最高压缩率 @override def convert( self, origin: ImageObject ) -> Tuple[Dict[str, Any], ImageObject]: change_config = { "compression_level": self.compression_level } # 转换为PIL Image以使用其PNG保存功能 pil_image = PILFromArray(cv2.cvtColor(origin.image, cv2.COLOR_BGR2RGB)) # 创建内存文件对象 buffer = BytesIO() # 保存为PNG pil_image.save(buffer, format='PNG', optimize=True, compress_level=self.compression_level) # 从内存读取图像数据 buffer.seek(0) result = PILOpen(buffer) # 转换回OpenCV格式 result = cv2.cvtColor(np.array(result), cv2.COLOR_RGB2BGR) return (change_config, ImageObject(result)) class JPGConvertConfig(BasicConvertConfig): name: str = "jpg" quality: int = 95 # 0-100, 100为最高质量 @override def convert( self, origin: ImageObject ) -> Tuple[Dict[str, Any], ImageObject]: change_config = { "quality": self.quality } # 转换为PIL Image pil_image = PILFromArray(cv2.cvtColor(origin.image, cv2.COLOR_BGR2RGB)) # 创建内存文件对象 buffer = BytesIO() # 保存为JPG pil_image.save(buffer, format='JPEG', quality=self.quality) # 从内存读取图像数据 buffer.seek(0) result = PILOpen(buffer) # 转换回OpenCV格式 result = cv2.cvtColor(np.array(result), cv2.COLOR_RGB2BGR) return (change_config, ImageObject(result)) class ICOConvertConfig(BasicConvertConfig): name: str = "ico" size: Tuple[int, int] = (16, 16) @override def convert( self, origin: ImageObject ) -> Tuple[Dict[str, Any], ImageObject]: change_config = { "size": self.size } return (change_config, ImageObject(origin.get_resize_image(*self.size))) class BMPConvertConfig(BasicConvertConfig): name: str = "bmp" @override def convert( self, origin: ImageObject ) -> Tuple[Dict[str, Any], ImageObject]: change_config = {} # 直接使用OpenCV保存为BMP _, buffer = cv2.imencode('.bmp', origin.image) # 解码回图像 result = cv2.imdecode(buffer, cv2.IMREAD_COLOR) return (change_config, ImageObject(result)) class WebPConvertConfig(BasicConvertConfig): name: str = "webp" quality: int = 80 # 0-100, 100为最高质量 @override def convert( self, origin: ImageObject ) -> Tuple[Dict[str, Any], ImageObject]: change_config = { "quality": self.quality } # 转换为PIL Image pil_image = PILFromArray(cv2.cvtColor(origin.image, cv2.COLOR_BGR2RGB)) # 创建内存文件对象 buffer = BytesIO() # 保存为WebP pil_image.save(buffer, format='WEBP', quality=self.quality) # 从内存读取图像数据 buffer.seek(0) result = PILOpen(buffer) # 转换回OpenCV格式 result = cv2.cvtColor(np.array(result), cv2.COLOR_RGB2BGR) return (change_config, ImageObject(result)) class ImageConvertConfig(BaseModel): png: Optional[PNGConvertConfig] = None jpg: Optional[JPGConvertConfig] = None ico: Optional[ICOConvertConfig] = None bmp: Optional[BMPConvertConfig] = None webp: Optional[WebPConvertConfig] = None log_call: Optional[Callable[[Union[str, Dict[str, Any]]], None]] = None def get_all_configs(self) -> List[BasicConvertConfig]: return [ self.png, self.jpg, self.ico, self.bmp, self.webp ] def _inject_log(self, *args, **kwargs): if self.log_call is not None: self.log_call(*args, **kwargs) def convert( self, origin: ImageObject ) -> Tuple[Dict[str, Dict[str, Any]], Dict[str, ImageObject]]: result: Dict[str, ImageObject] = {} result_change_config: Dict[str, Dict[str, Any]] = {} convert_configs: List[BasicConvertConfig] = self.get_all_configs() for item in convert_configs: if item is not None: result_change_config[item.name], result[item.name] = item.convert(ImageObject(origin.image)) self._inject_log(f"conversion<{item.name}> change config: {result_change_config[item.name]}") return (result_change_config, result) def convert_to( self, input: Union[tool_file, str, ImageObject, np.ndarray, PILImage, PILImageFile], output_dir: tool_file_or_str, *, must_output_dir_exist: bool = False, output_file_name: str = "output.png", callback: Optional[Action[Dict[str, Dict[str, Any]]]] = None, ) -> Dict[str, ImageObject]: # 初始化环境和变量 origin_image: ImageObject = self.__init_origin_image(input) result_dir: tool_file = self.__init_result_dir(output_dir, must_output_dir_exist) # 转换 self._inject_log(f"输出<{output_file_name}>开始转换") change_config, result = self._inject_convert( origin_image=origin_image, result_dir=result_dir, output_file_name=output_file_name, ) # 结果 if callback is not None: callback(change_config) return result def __init_origin_image(self, input: Union[tool_file, str, ImageObject, np.ndarray, PILImage, PILImageFile]) -> ImageObject: origin_image: ImageObject = None # check if isinstance(input, (tool_file, str)): inputfile = WrapperFile2CVEX(input) if inputfile.data is not None: origin_image = inputfile.data else: origin_image = inputfile.load() elif isinstance(input, (ImageObject, np.ndarray, PILImage, PILImageFile)): origin_image = Wrapper2Image(input) else: raise TypeError(f"input<{input}> is not support type") return origin_image def __init_result_dir(self, output_dir: tool_file_or_str, must_output_dir_exist: bool) -> tool_file: if output_dir is None or isinstance(output_dir, loss_file): return static_loss_file_dir result_dir: tool_file = Wrapper2File(output_dir) # check exist stats: bool = True if result_dir.exists() is False: if must_output_dir_exist: result_dir.must_exists_path() else: stats = False if stats is False: raise FileExistsError(f"output_dir<{result_dir}> is not exist") # check dir stats if result_dir.is_dir() is False: if must_output_dir_exist: result_dir.back_to_parent_dir() else: raise FileExistsError(f"output_dir<{result_dir}> is not a dir") # result return result_dir def _inject_convert( self, origin_image: ImageObject, result_dir: tool_file, output_file_name: str ) -> Tuple[Dict[str, Dict[str, Any]], Dict[str, ImageObject]]: self._inject_log(f"output<{output_file_name}> is start convert") result_dict, result_images = self.convert(origin_image) if not (result_dir is None or isinstance(result_dir, loss_file)): for key, value in result_images.items(): current_dir = result_dir|key current_result_file = current_dir|output_file_name value.save_image(current_result_file, True) return result_dict, result_images def image_convert( config: ImageConvertConfig, source, *args, **kwargs ): if isinstance(source, ImageObject): return config.convert(source, *args, **kwargs) # region end