kerasで正規化(Imagedatageneratorのいろいろな変換)＋Random Erasingの実装

前提

kerasのImagedatageneratorをカスタマイズしてRandom Erasingを実装できたのですが、正規化(255で割る)をすると真っ黒な画像が生成されてしまいます。
そのため、しっかりと正規化(Imagedatageneratorのいろいろな変換)＋Random Erasingができるようにしたいです。

実現したいこと

・kerasで正規化(Imagedatageneratorのいろいろな変換)＋Random Erasing

発生している問題

図1
正規化していない場合のImagedatageneratorの変換＋Random Erasing
図2
正規化した場合のImagedatageneratorの変換＋Random Erasing

該当のソースコード

python
1import numpy as np
2from PIL import Image
3import tensorflow as tf
4import time
5import matplotlib.pyplot as plt
6# import os
7# os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
8from tensorflow.keras.models import Model
9from tensorflow.keras.preprocessing import image
10from tensorflow.keras.preprocessing.image import ImageDataGenerator,array_to_img, img_to_array
11from tensorflow.keras.optimizers import Adam
12from tensorflow.keras.callbacks import EarlyStopping, LearningRateScheduler,ReduceLROnPlateau, ModelCheckpoint
13
14class MyImagedatagenerator(image.ImageDataGenerator):
15    
16    def __init__(self,
17                 random_erasing_probability = None,
18                 random_erasing_area_ratio = [0.02,0.4],
19                 random_erasing_aspect_ratio = [0.3,1/0.3],
20                 random_erasing_mask_value = [0,1],
21                 *args, **kwargs
22    ):
23        super().__init__(*args, **kwargs)
24        self.random_erasing_probability = random_erasing_probability
25        self.random_erasing_area_ratio = random_erasing_area_ratio
26        self.random_erasing_aspect_ratio = random_erasing_aspect_ratio
27        self.random_erasing_mask_value = random_erasing_mask_value
28
29    # def random_erasing(self, X, erasing_probability, erasing_area_ratio_range, 
30    #     erasing_aspect_ratio_range, random_erasing_mask_value):
31    #     # ...一旦省略...
32    def __init__(self,
33        random_erasing_probability = None,
34        random_erasing_area_ratio = [0.02, 0.4],
35        random_erasing_aspect_ratio = [0.3, 1/0.3],
36        random_erasing_mask_value = [0, 1],
37        *args, **kwargs
38    ):
39        super().__init__(*args, **kwargs)
40        #random erasingを発生させる確率
41        self.random_erasing_probability = random_erasing_probability
42        #マスクする面積の範囲(相対値で表現)
43        self.random_erasing_area_ratio = random_erasing_area_ratio
44        #アスペクト比の範囲
45        self.random_erasing_aspect_ratio = random_erasing_aspect_ratio
46        #マスクする値の取りうる範囲
47        self.random_erasing_mask_value = random_erasing_mask_value
48        
49    def random_erasing(self, X, erasing_probability, erasing_area_ratio_range, 
50        erasing_aspect_ratio_range, random_erasing_mask_value
51    ):
52        X_copy = X.copy()
53        batch_size, H, W, C = X_copy.shape
54        original_area = H * W
55
56        for batch_index in range(batch_size):
57
58            if erasing_probability < np.random.rand():
59                continue
60
61            # はみ出さないようリトライするループ
62            while True:
63
64                # マスクする面積をサンプリング
65                erasing_area = np.random.uniform(
66                    erasing_area_ratio_range[0], erasing_area_ratio_range[1]
67                ) * original_area
68
69                # マスクするアスペクト比をサンプリング
70                erasing_aspect_ratio = np.random.uniform(
71                    erasing_aspect_ratio_range[0], erasing_aspect_ratio_range[1]
72                )
73
74                # 面積とアスペクト比から高さと幅を計算
75                erasing_height = int(np.sqrt(erasing_area * erasing_aspect_ratio))
76                erasing_width  = int(np.sqrt(erasing_area / erasing_aspect_ratio))
77
78                # マスクを配置する端点をサンプリング
79                erasing_left_top_x = np.random.randint(0, W)
80                erasing_left_top_y = np.random.randint(0, H)
81
82                # マスクが元画像をはみ出すかどうかを計算
83                if erasing_left_top_x + erasing_width <= W \
84                    and erasing_left_top_y + erasing_height <= H:
85                    break
86
87            # マスクする値の生成
88            erasing_values = np.random.uniform(
89                random_erasing_mask_value[0], random_erasing_mask_value[1], 
90                (erasing_height, erasing_width, C)
91            )
92
93            X_copy[batch_index, 
94                erasing_left_top_y:erasing_left_top_y + erasing_height,
95                erasing_left_top_x:erasing_left_top_x + erasing_width, :] = erasing_values
96            
97            
98
99        return X_copy
100    
101    def flow(self, seed=None, *args, **kwargs):
102
103        batch_gen = super().flow(seed=seed, *args, **kwargs)
104
105        while True:
106
107            batch_x, batch_y = next(batch_gen)
108
109            # random erasing
110            if self.random_erasing_probability is not None:
111                batch_x = self.random_erasing(
112                    batch_x, 
113                    self.random_erasing_probability, 
114                    self.random_erasing_area_ratio, 
115                    self.random_erasing_aspect_ratio,
116                    self.random_erasing_mask_value)
117                
118                # batch_x = batch_x/255.
119                
120            yield (batch_x, batch_y)
121
122def plot_augmentation_image(train_image_sample, train_label_sample, params):
123
124    # 同じデータを16個複製する
125    train_image_samples = np.repeat(
126        train_image_sample.reshape((1, *train_image_sample.shape)), 16, axis=0)
127    train_label_sample = np.repeat(
128        train_label_sample.reshape((1, *train_label_sample.shape)), 16, axis=0)
129
130    # 16個に対してparamsで与えられた変換を実施
131    data_generator = MyImagedatagenerator(**params)
132    generator = data_generator.flow(
133        x=train_image_samples, y=train_label_sample, batch_size=16)
134
135    # 変換後のデータを取得
136    batch_x, batch_y = generator.__next__()
137
138    # 変換後はfloat32となっているため、uint8に変換
139    batch_x = batch_x.astype(np.uint8)
140
141    # 描画処理
142    plt.figure(figsize=(10,10))
143    for i in range(16):
144        plt.subplot(4,4,i+1)
145        plt.imshow(batch_x[i])
146        plt.tick_params(labelbottom='off')
147        plt.tick_params(labelleft='off')
148
149params = {
150    'random_erasing_probability': 0.5,
151    'random_erasing_area_ratio': [0.02, 0.4],
152    'random_erasing_aspect_ratio': [0.3, 1/0.3],
153    'random_erasing_mask_value': [0, 255],
154    'rescale' :1./255,
155    'rotation_range':180,
156    'fill_mode':"constant",
157    'vertical_flip': True,
158    'brightness_range' : [0.7, 1.117],
159    'horizontal_flip' : True,
160    'dtype' : np.float64
161}
162plot_augmentation_image(x_train[1], y_train[1], params)
163
164

試したこと

https://dev.classmethod.jp/articles/tensorflow-image-generator-custom/
こちらのサイトを参考(コピペ)させていだいてRandom Erasing を実装しました。こちらのサイトでは、正規化されていなかったため正規化したかったのですが真っ黒になってしまい行き詰っております。
paramsにrescaleの追加やdef random_erasingやdef flowで255で割ってみても真っ黒になってしまいます。

回答いただけると幸いです。