Merge branch 'master' of git+ssh://github.com/foamliu/Deep-Image-Matting

demo.py

@@ -13,7 +13,7 @@
     img_rows, img_cols = 320, 320
     channel = 4
 
-    model_weights_path = 'models/model.62-0.0524.hdf5'
+    model_weights_path = 'models/model.98-0.0459.hdf5'
     model = build_encoder_decoder()
     model.load_weights(model_weights_path)
     print(model.summary())

model.py

@@ -1,5 +1,5 @@
 import keras.backend as K
-from keras.layers import Input, Conv2D, UpSampling2D, BatchNormalization, ZeroPadding2D, MaxPooling2D, Concatenate
+from keras.layers import Input, Conv2D, UpSampling2D, BatchNormalization, ZeroPadding2D, MaxPooling2D, Concatenate, Lambda
 from keras.models import Model
 from keras.utils import plot_model
 
@@ -99,8 +99,9 @@ def build_encoder_decoder():
 def build_refinement(encoder_decoder):
     input_tensor = encoder_decoder.input
 
-    x = encoder_decoder.output
-    x = Concatenate(axis=3)([input_tensor[:, :, :, 0:3], x])
+    input = Lambda(lambda i: i[:, :, :, 0:3])(input_tensor)
+
+    x = Concatenate(axis=3)([input, encoder_decoder.output])
     x = Conv2D(64, (3, 3), activation='relu', padding='same', kernel_initializer='he_normal',
                bias_initializer='zeros')(x)
     x = BatchNormalization()(x)

train_final.py

@@ -12,23 +12,23 @@
 
     # Callbacks
     tensor_board = keras.callbacks.TensorBoard(log_dir='./logs', histogram_freq=0, write_graph=True, write_images=True)
-    model_names = checkpoint_models_path + 'refinement.{epoch:02d}-{val_loss:.4f}.hdf5'
+    model_names = checkpoint_models_path + 'final.{epoch:02d}-{val_loss:.4f}.hdf5'
     model_checkpoint = ModelCheckpoint(model_names, monitor='val_loss', verbose=1, save_best_only=True)
     early_stop = EarlyStopping('val_loss', patience=patience)
     reduce_lr = ReduceLROnPlateau('val_loss', factor=0.1, patience=int(patience / 4), verbose=1)
 
     pretrained_path = 'models/'
     encoder_decoder = build_encoder_decoder()
-    refinement = build_refinement(encoder_decoder)
-    refinement.load_weights(pretrained_path)
+    final = build_refinement(encoder_decoder)
+    final.load_weights(pretrained_path)
     # finetune the whole network together.
-    for layer in refinement.layers:
+    for layer in final.layers:
         layer.trainable = True
 
     sgd = SGD(lr=1e-5, decay=1e-6, momentum=0.9, nesterov=True)
-    refinement.compile(optimizer=sgd, loss=custom_loss_wrapper(refinement.input))
+    final.compile(optimizer=sgd, loss=custom_loss_wrapper(final.input))
 
-    print(refinement.summary())
+    print(final.summary())
 
     # Summarize then go!
     num_cpu = get_available_cpus()
@@ -38,13 +38,13 @@
     callbacks = [tensor_board, model_checkpoint, early_stop, reduce_lr]
 
     # Start Fine-tuning
-    refinement.fit_generator(train_gen(),
-                             steps_per_epoch=num_train_samples // batch_size,
-                             validation_data=valid_gen(),
-                             validation_steps=num_valid_samples // batch_size,
-                             epochs=epochs,
-                             verbose=1,
-                             callbacks=callbacks,
-                             use_multiprocessing=True,
-                             workers=workers
-                             )
+    final.fit_generator(train_gen(),
+                        steps_per_epoch=num_train_samples // batch_size,
+                        validation_data=valid_gen(),
+                        validation_steps=num_valid_samples // batch_size,
+                        epochs=epochs,
+                        verbose=1,
+                        callbacks=callbacks,
+                        use_multiprocessing=True,
+                        workers=workers
+                        )

train_refinement.py

@@ -17,7 +17,7 @@
     early_stop = EarlyStopping('val_loss', patience=patience)
     reduce_lr = ReduceLROnPlateau('val_loss', factor=0.1, patience=int(patience / 4), verbose=1)
 
-    pretrained_path = 'models/'
+    pretrained_path = 'models/model.98-0.0459.hdf5'
     encoder_decoder = build_encoder_decoder()
     encoder_decoder.load_weights(pretrained_path)
     # fix encoder-decoder part parameters and then update the refinement part.

unit_tests.py

@@ -3,9 +3,10 @@
 
 import cv2 as cv
 import numpy as np
-
+import os
 from config import unknown
 from data_generator import generate_trimap
+from data_generator import get_alpha_test
 from data_generator import random_choice
 from utils import safe_crop
 
@@ -64,6 +65,28 @@ def test_different_sizes(self):
         crop_size = random.choice(different_sizes)
         print('crop_size=' + str(crop_size))
 
+    def test_resize(self):
+        with open('Combined_Dataset/Test_set/test_bg_names.txt') as f:
+            bg_test_files = f.read().splitlines()
+        name = '35_716.png'
+        filename = os.path.join('merged_test', name)
+        image = cv.imread(filename)
+        bg_h, bg_w = image.shape[:2]
+        a = get_alpha_test(name)
+        a_h, a_w = a.shape[:2]
+        alpha = np.zeros((bg_h, bg_w), np.float32)
+        alpha[0:a_h, 0:a_w] = a
+        trimap = generate_trimap(alpha)
+        # 剪切尺寸 320:640:480 = 3:1:1
+        crop_size = (480, 480)
+        x, y = random_choice(trimap, crop_size)
+        image = safe_crop(image, x, y, crop_size)
+        trimap = safe_crop(trimap, x, y, crop_size)
+        alpha = safe_crop(alpha, x, y, crop_size)
+        cv.imwrite('temp/test_resize_image.png', image)
+        cv.imwrite('temp/test_resize_trimap.png', trimap)
+        cv.imwrite('temp/test_resize_alpha.png', alpha)
+
 
 if __name__ == '__main__':
     unittest.main()