generated from daniil-berg/boilerplate-py
24 lines
1000 B
Python
24 lines
1000 B
Python
import tensorflow as tf
|
|
import numpy as np
|
|
|
|
from .keras.backend import ctc_batch_cost
|
|
from .keras.layers import Layer
|
|
|
|
|
|
class CTCLayer(Layer):
|
|
def __init__(self, name: str = None):
|
|
super().__init__(name=name)
|
|
self.loss_fn = ctc_batch_cost
|
|
|
|
def call(self, y_true: np.ndarray = None, y_pred: np.ndarray = None) -> np.ndarray:
|
|
# Compute the training-time loss value and add it
|
|
# to the layer using `self.add_loss()`.
|
|
batch_len = tf.cast(tf.shape(y_true)[0], dtype='int64')
|
|
input_length = tf.cast(tf.shape(y_pred)[1], dtype='int64')
|
|
label_length = tf.cast(tf.shape(y_true)[1], dtype='int64')
|
|
input_length = input_length * tf.ones(shape=(batch_len, 1), dtype='int64')
|
|
label_length = label_length * tf.ones(shape=(batch_len, 1), dtype='int64')
|
|
loss = self.loss_fn(y_true, y_pred, input_length, label_length)
|
|
self.add_loss(loss)
|
|
# At test time, just return the computed predictions
|
|
return y_pred |