python实现图像的随机增强变换

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

21

22

23

24

25

26

27

28

29

30

31

32

33

34

35

36

37

38

39

40

41

42

43

44

45

46

47

48

49

50

51

52

53

54

55

56

57

58

59

60

61

62

63

64

65

66

67

68

69

70

71

72

73

74

75

76

77

78

79

80

81

82

83

84

85

86

87

88

89

90

91

92

93

94

95

96

97

98

99

100

101

102

103

104

105

106

107

108

109

110

111

112

113

import os

import random

import cv2

import numpy as np

from PIL import Image, ImageEnhance, ImageOps

# 定义各种数据增强方法

def random_rotate(image, angle_range=(-30, 30)):

    angle = random.uniform(angle_range[0], angle_range[1])

    (h, w) = image.shape[:2]

    center = (w // 2, h // 2)

    M = cv2.getRotationMatrix2D(center, angle, 1.0)

    rotated = cv2.warpAffine(image, M, (w, h), borderMode=cv2.BORDER_REFLECT)

    return rotated

def random_translate(image, translate_range=(-50, 50)):

    tx = random.randint(translate_range[0], translate_range[1])

    ty = random.randint(translate_range[0], translate_range[1])

    (h, w) = image.shape[:2]

    M = np.float32([[1, 0, tx], [0, 1, ty]])

    translated = cv2.warpAffine(image, M, (w, h), borderMode=cv2.BORDER_REFLECT)

    return translated

def random_flip(image):

    flip_code = random.choice([-1, 0, 1])

    flipped = cv2.flip(image, flip_code)

    return flipped

def random_scale(image, scale_range=(0.8, 1.2)):

    scale = random.uniform(scale_range[0], scale_range[1])

    (h, w) = image.shape[:2]

    new_dim = (int(w * scale), int(h * scale))

    scaled = cv2.resize(image, new_dim, interpolation=cv2.INTER_LINEAR)

    return scaled

def random_crop(image, crop_size=(224, 224)):

    (h, w) = image.shape[:2]

    if crop_size[0] > h or crop_size[1] > w:

        # 当裁剪尺寸大于图像尺寸时，抛出异常或调整裁剪尺寸

        raise ValueError("Crop size is larger than image size.")

    top = random.randint(0, h - crop_size[0])

    left = random.randint(0, w - crop_size[1])

    cropped = image[top:top+crop_size[0], left:left+crop_size[1]]

    return cropped

def random_color_jitter(image):

    pil_image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))

    color_jitter = ImageEnhance.Color(pil_image).enhance(random.uniform(0.6, 1.4))

    contrast_jitter = ImageEnhance.Contrast(color_jitter).enhance(random.uniform(0.5, 1.5))

    brightness_jitter = ImageEnhance.Brightness(contrast_jitter).enhance(random.uniform(0.6, 1.4))

    sharpness_jitter = ImageEnhance.Sharpness(brightness_jitter).enhance(random.uniform(0.6, 1.4))

    jittered = cv2.cvtColor(np.array(sharpness_jitter), cv2.COLOR_RGB2BGR)

    return jittered

def random_add_noise(image):

    row, col, ch = image.shape

    mean = 0

    var = 0.1

    sigma = var ** 0.5

    gauss = np.random.normal(mean, sigma, (row, col, ch))

    gauss = gauss.reshape(row, col, ch)

    noisy = image + gauss

    return np.clip(noisy, 0, 255).astype(np.uint8)

# 数据增强主函数

def augment_random_images(src_folder, dst_folder, num_images_to_select, num_augmentations_per_image):

    if not os.path.exists(dst_folder):

        os.makedirs(dst_folder)

    # 获取所有图像文件名

    all_filenames = [f for f in os.listdir(src_folder) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]

    # 如果选择的图像数量大于总图像数量，则只处理全部图像

    num_images_to_process = min(num_images_to_select, len(all_filenames))

    # 随机选择图像

    selected_filenames = random.sample(all_filenames, num_images_to_process)

    # 创建一个增强方法列表

    augmentation_methods = [

        random_rotate,

        #random_translate,

        random_flip,

        random_scale,

        #random_crop,

        random_color_jitter,

        random_add_noise

    ]

    for filename in selected_filenames:

        img_path = os.path.join(src_folder, filename)

        image = cv2.imread(img_path)

        for i in range(num_augmentations_per_image):

            # 随机选择一种增强方法

            augmentation_method = random.choice(augmentation_methods)

            # 应用选中的增强方法

            augmented_img = augmentation_method(image)

            # 保存增强后的图像

            base_name, ext = os.path.splitext(filename)

            save_path = os.path.join(dst_folder, f"{base_name}_aug_{i}{ext}")

            cv2.imwrite(save_path, augmented_img)

if __name__ == "__main__":

    src_folder = 'path/to/source/folder'  # 替换为你的源文件夹路径

    dst_folder = 'path/to/destination/folder'  # 替换为你要保存增强图像的文件夹路径

    num_images_to_select = 10  # 从源文件夹中随机选择的图像数量

    num_augmentations_per_image = 5  # 每张图像生成的增强图像数量

    augment_random_images(src_folder, dst_folder, num_images_to_select, num_augmentations_per_image)

    print(f"图像增强完成，增强后的图像已保存到 {dst_folder}")