ImageDataGeneretorにRandom Erasingを導入したい

前提・実現したいこと

pythonを最近始めたものです。
kerasVGG16の転移学習を行い画像を5クラスに分類する過程においてImageDataGeneretorによるデータ拡張にrandom erasing画像の一部をランダムにマスクするデータ拡張を取り入れることを目指しています。

python
1from keras.models import Model
2from keras.layers import Dense, GlobalAveragePooling2D,Input
3from keras.applications.vgg16 import VGG16
4from keras.preprocessing.image import ImageDataGenerator
5from keras.optimizers import SGD
6from keras.callbacks import CSVLogger
7from random_eraser import get_random_eraser #added
8import numpy as np
9
10n_categories=5
11batch_size=32
12train_dir='ver1.1/train'
13validation_dir='ver1.1/validation'
14file_name='filename'
15random_erasing = True
16
17base_model=VGG16(weights='imagenet',include_top=False,
18                 input_tensor=Input(shape=(224,224,3)))
19
20#add new layers instead of FC networks
21x=base_model.output
22x=GlobalAveragePooling2D()(x)
23x=Dense(1024,activation='relu')(x)
24prediction=Dense(n_categories,activation='softmax')(x)
25model=Model(inputs=base_model.input,outputs=prediction)
26
27#fix weights before VGG16 14layers
28for layer in base_model.layers[:15]:
29    layer.trainable=False
30
31model.compile(optimizer=SGD(lr=0.0001,momentum=0.9),
32              loss='categorical_crossentropy',
33              metrics=['accuracy'])
34
35model.summary()
36
37#save model
38json_string=model.to_json()
39open(file_name+'.json','w').write(json_string)
40
41
42#学習
43train_datagen=ImageDataGenerator(
44    rescale=1.0/255,
45    shear_range=0.2,
46    zoom_range=0.2,
47    horizontal_flip=True,
48    preprocessing_function=get_random_eraser(v_l=0, v_h=1)) #added
49
50validation_datagen=ImageDataGenerator(rescale=1.0/255)
51
52train_generator=train_datagen.flow_from_directory(
53    train_dir,
54    target_size=(224,224),
55    batch_size=batch_size,
56    class_mode='categorical',
57    shuffle=True
58)
59
60validation_generator=validation_datagen.flow_from_directory(
61    validation_dir,
62    target_size=(224,224),
63    batch_size=batch_size,
64    class_mode='categorical',
65    shuffle=True
66)
67
68hist=model.fit_generator(train_generator,
69                         epochs=200,
70                         verbose=1,
71                         validation_data=validation_generator,
72                         callbacks=[CSVLogger(file_name+'.csv')])
73
74#save weights
75model.save(file_name+'.h5')
76

random erasingのコード

python
1import numpy as np
2
3
4def get_random_eraser(p=0.5, s_l=0.02, s_h=0.4, r_1=0.3, r_2=1/0.3, v_l=0, v_h=255, pixel_level=False):
5    def eraser(input_img):
6        if input_img.ndim == 3:
7            img_h, img_w, img_c = input_img.shape
8        elif input_img.ndim == 2:
9            img_h, img_w = input_img.shape
10
11        p_1 = np.random.rand()
12
13        if p_1 > p:
14            return input_img
15
16        while True:
17            s = np.random.uniform(s_l, s_h) * img_h * img_w
18            r = np.random.uniform(r_1, r_2)
19            w = int(np.sqrt(s / r))
20            h = int(np.sqrt(s * r))
21            left = np.random.randint(0, img_w)
22            top = np.random.randint(0, img_h)
23
24            if left + w <= img_w and top + h <= img_h:
25                break
26
27        if pixel_level:
28            if input_img.ndim == 3:
29                c = np.random.uniform(v_l, v_h, (h, w, img_c))
30            if input_img.ndim == 2:
31                c = np.random.uniform(v_l, v_h, (h, w))
32        else:
33            c = np.random.uniform(v_l, v_h)
34
35        input_img[top:top + h, left:left + w] = c
36
37        return input_img
38
39    return eraser
40

発生している問題・エラーメッセージ

エラーメッセージ
~\python\graduation_research\running_shoes\random_eraser.py in eraser(input_img)
      5     def eraser(input_img):
----> 6         if input_img.ndim == 3:
      7             img_h, img_w, img_c = input_img.shape

AttributeError: 'JpegImageFile' object has no attribute 'ndim'

入力画像がnumpyの型に変更する必要があるということでしょうか?
また, 具体的にどの部分で型を変更すればよいか教えていただきたいです。よろしくお願いします

補足情報（FW/ツールのバージョンなど）

random erasingを導入しない場合では正しく動いていました。(ImagedataGeneretor内の#addedの行を除いた場合)

行動規範の内容に同意します

回答1件

自己解決

random erasingの処理を行う際にget_random_eraserにPIL形式で画像が渡されているため, 入力画像をnumpy形式に変換し処理をしたのち, PIL形式の画像を返すようにコードを変更することで解決しました。

python
1import numpy as np
2from PIL import Image#added
3
4
5def get_random_eraser(p=0.5, s_l=0.02, s_h=0.4, r_1=0.3, r_2=1/0.3, v_l=0, v_h=255, pixel_level=False):
6    def eraser(input_img):
7        input_img = np.array(input_img) #added
8        if input_img.ndim == 3:
9            img_h, img_w, img_c = input_img.shape
10        elif input_img.ndim == 2:
11            img_h, img_w = input_img.shape
12
13        p_1 = np.random.rand()
14
15        if p_1 > p:
16            return  Image.fromarray(np.uint8(input_img)) #changed
17
18        while True:
19            s = np.random.uniform(s_l, s_h) * img_h * img_w
20            r = np.random.uniform(r_1, r_2)
21            w = int(np.sqrt(s / r))
22            h = int(np.sqrt(s * r))
23            left = np.random.randint(0, img_w)
24            top = np.random.randint(0, img_h)
25
26            if left + w <= img_w and top + h <= img_h:
27                break
28
29        if pixel_level:
30            if input_img.ndim == 3:
31                c = np.random.uniform(v_l, v_h, (h, w, img_c))
32            if input_img.ndim == 2:
33                c = np.random.uniform(v_l, v_h, (h, w))
34        else:
35            c = np.random.uniform(v_l, v_h)
36
37        input_img[top:top + h, left:left + w] = c
38
39        return Image.fromarray(np.uint8(input_img)) #changed
40
41    return eraser
42