Exemple #1
0
    def predict_word(self, X):
        crnn = CRNN()
        crnn.build(dropout=False)
        crnn.model.load_weights(self.weight_name + ".h5")

        out = crnn.test_func([X])[0]
        ret = []

        for j in range(out.shape[0]):
            out_best = list(np.argmax(out[j, 2:], 1))
            out_best = [k for k, g in itertools.groupby(out_best)]
            outstr = ''
            for c in out_best:
                if 0 <= c <= 9:
                    outstr += chr(c + ord('0'))
                elif 10 <= c <= 35:
                    outstr += chr(c - 10 + ord('A'))
                elif 36 <= c <= 61:
                    outstr += chr(c - 36 + ord('a'))
            ret.append(outstr)
        return ret
def train(args):
    @tf.function
    def train_step(x, y):
        with tf.GradientTape() as tape:
            y_pred = model(x["the_input"])
            # loss = tf.reduce_mean(ctc_lambda_func((y_pred, x["the_labels"], x["input_length"].reshape((-1,1)), x["label_length"].reshape((-1,1)))))
            loss = tf.reduce_mean(ctc_lambda_func((y_pred, x["the_labels"], tf.reshape(x["input_length"], [-1, 1]), tf.reshape(x["label_length"], [-1, 1]))))
        
        # Compute gradients
        trainable_vars = model.trainable_variables
        gradients = tape.gradient(loss, trainable_vars)

        # Update weights
        model.optimizer.apply_gradients(zip(gradients, trainable_vars))
        return loss


    epochs = 1000
    iter_per_epoch = 100
    #model, test_func = get_CResRNN(weights=os.path.join("OUTPUT_DIR", "exp1", "weights06.h5"))
    #model, test_func = get_CResRNN(weights=os.path.join("OUTPUT_DIR", "weights0995.h5"))
    #model.load_weights(os.path.join("OUTPUT_DIR", "exp1", "weights15.h5"))
    #model.load_weights(os.path.join("OUTPUT_DIR", "weights0995.h5"))
    model2, test_func = CRNN_model()

    train_generator = FakeImageGenerator(args).next_gen()
    

    model = CRNN(ALPHABET)
    model.build()
    model.summary()

    # model = tf.keras.load_model('checkpoints/checkpoint')
    
    model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=0.0001, clipnorm=5))

    loss_train = []

    for epoch in range(1, epochs):
        print(f"Start of epoch {epoch}")

        pb = Progbar(iter_per_epoch, stateful_metrics="loss")

        for iter in range(iter_per_epoch):
            x, y = next(train_generator)
            with tf.GradientTape() as tape:
                y_pred = model(x["the_input"])
                # loss = tf.reduce_mean(ctc_lambda_func((y_pred, x["the_labels"], x["input_length"].reshape((-1,1)), x["label_length"].reshape((-1,1)))))
                loss = tf.reduce_mean(ctc_lambda_func((y_pred, x["the_labels"], tf.reshape(x["input_length"], [-1, 1]), tf.reshape(x["label_length"], [-1, 1]))))
            
            # Compute gradients
            trainable_vars = model.trainable_variables
            gradients = tape.gradient(loss, trainable_vars)

            # Update weights
            model.optimizer.apply_gradients(zip(gradients, trainable_vars))

            values = [('loss', loss)]
            pb.add(1, values=values)

        if epoch % 5 == 0:
            model.save("checkpoints/base_crnn.h5")


    
    

    # print("test2")
    # x, y = next(train_generator)
    # model.fit(x, y)
    # print("test1")
    
    x, y = next(train_generator)
    print(model(x["the_input"]))

    """