Пример #1
0
def train(trainX, trainY):
    m = SRCNN()
    m.compile(Adam(lr=0.0003), 'mse')
    count = 1
    while True:
        m.fit(trainX, trainY, batch_size=128, nb_epoch=5)
        print("Saving model " + str(count * 5))
        m.save(join('./model_' + str(count * 5) + '.h5'))
        count += 1
Пример #2
0
def train_batch(aug, trainX, trainY, load_weight=True):
    print("[INFO] compiling model...")
    model = SRCNN()

    if load_weight and os.path.exists('my_model_weights.h5'):
        print("[INFO] import model weights...")
        model.load_weights('my_model_weights.h5')
    try:
        opt = Adam(lr=INIT_LR, decay=INIT_LR / EPOCHS)
        # opt = Adam(lr=INIT_LR)
        model.compile(loss="mean_squared_error",
                      optimizer=opt,
                      metrics=[psnr1])
        current_time = datetime.now().strftime("%Y%m%d-%H:%M")
        filepath = 'train_model/model-ep{epoch:03d}-loss{loss:.3f}-psnr{psnr1:.3f}-' + current_time + '.h5'

        checkpoint = ModelCheckpoint(filepath,
                                     monitor='psnr1',
                                     verbose=1,
                                     save_best_only=True,
                                     mode='max',
                                     period=2)
        print("[INFO] training network...")
        H = model.fit_generator(aug.flow(trainX, trainY, batch_size=BS),
                                steps_per_epoch=len(trainX) // BS,
                                epochs=EPOCHS,
                                callbacks=[checkpoint])
        print("[INFO] training finish")

    except KeyboardInterrupt:
        print("[INFO] KeyboardInterrupted...")

    except Exception as e:
        print(e)

    else:
        plt.figure()
        N = EPOCHS
        # plt.plot(np.arange(0, N), H.history["loss"], label="train_loss")
        # plt.plot(np.arange(0, N), H.history["acc"], label="train_acc")
        plt.plot(np.arange(0, N), H.history["psnr1"], label="psnr")
        plt.xlabel("Epoch #")
        plt.ylabel("psnr")
        plt.legend(loc="upper left")
        plt.savefig("plot.png")

    finally:
        print("[INFO] serializing network...")
        model.save("./model.h5")
        model.save_weights('my_model_weights.h5')
Пример #3
0
image_size = args.Image_patch_size      #image_patch_size
stride = args.Stride
scale = args.Scale
learning_rate = args.lr
batch_size = args.BATCH_SIZE
epochs = args.Epochs
is_training =args.is_training
dirname_train = args.dirname_train
dirname_test = args.dirname_test

if is_training:
   
    X_train,Y_train = load_train(image_size = image_size,stride = stride,scale = scale,dirname =dirname_train)
    srcnn = SRCNN(  image_size = image_size, learning_rate=learning_rate)
    optimizer = Adam(lr = learning_rate)
    srcnn.compile(optimizer=optimizer, loss='mean_squared_error')
    history = srcnn.fit(X_train, Y_train, batch_size=batch_size, epochs=epochs, verbose=2,validation_split = 0.1 )
    draw_loss_plot(history = history)

    #Saving trained model_weights in the current workspace .
    #make a folder named srcnn in your current workspace you are working in.In that folder your weights will be stored.
    json_string = srcnn.to_json()
    open(os.path.join('./srcnn/','srcnn_model.json'),'w').write(json_string)
    srcnn.save_weights(os.path.join('./srcnn/','srcnn_weight.hdf5'))

else:

    dir_list = os.listdir(dirname_test)
    for img in dir_list:

       image = cv2.cvtColor(cv2.imread(os.path.join(dirname_test,img)),cv2.COLOR_BGR2GRAY)