diff --git a/src/eynollah/eynollah_ocr.py b/src/eynollah/eynollah_ocr.py index 9ed02d4..9cc2c9a 100644 --- a/src/eynollah/eynollah_ocr.py +++ b/src/eynollah/eynollah_ocr.py @@ -296,13 +296,11 @@ class Eynollah_ocr(Eynollah): probs[ver_index > 0][flipped_ver_is_better] = probs_ver[flipped_ver_is_better] def nooov(x): - return x != b'[UNK]' + if x == b'[UNK]': + return b'' + return x for pred, prob in zip(preds, probs): - text = b''.join( - filter(nooov, - map(bytes, - (filter(None, char) - for char in pred.tolist())))).decode('utf-8') + text = b''.join(map(nooov, pred.tolist())).decode('utf-8') extracted_texts.append(text) extracted_confs.append(prob) del cropped_lines_rgb diff --git a/src/eynollah/predictor.py b/src/eynollah/predictor.py index 526223a..6b7fe96 100644 --- a/src/eynollah/predictor.py +++ b/src/eynollah/predictor.py @@ -191,13 +191,19 @@ class Predictor(mp.context.SpawnProcess): #result = self.model.predict(data, verbose=0) # faster, less VRAM result = self.model.predict_on_batch(data) - if isinstance(result, tuple): + def make_shareable(x): + # convert tf.string/np.object to fixed-length bytes + # (because object segfaults in shm) + if x.dtype is np.dtype(object): + return x.astype(bytes) + return x + if isinstance(result, (list, tuple)): multi_output = True - results = zip(*(np.split(result0, len(jobs)) + results = zip(*(np.split(make_shareable(result0), len(jobs)) for result0 in result)) else: multi_output = False - results = np.split(result, len(jobs)) + results = np.split(make_shareable(result), len(jobs)) #self.logger.debug("sharing result array for '%d'", jobid) with ExitStack() as stack: for jobid, result in zip(jobs, results): diff --git a/src/eynollah/training/models.py b/src/eynollah/training/models.py index 9111225..eda6b0f 100644 --- a/src/eynollah/training/models.py +++ b/src/eynollah/training/models.py @@ -82,18 +82,17 @@ class CTCDecoder(Layer): inputs = tf.math.log( tf.transpose(inputs, perm=[1, 0, 2]) + tf.keras.backend.epsilon() ) - # tf.nn.ctc_greedy_decoder() is not as precise # tf.compat.v1.nn.ctc_beam_search_decoder() also needs merge_repeated=False - decoded, logits = tf.nn.ctc_beam_search_decoder( + # tf.nn.ctc_beam_search_decoder() is not supported by ONNX, yet + # tf.nn.ctc_greedy_decoder() is not as precise, though: + decoded, logits = tf.nn.ctc_greedy_decoder( inputs, lengths, - beam_width=10, - top_paths=2, ) # get top path for all sequences in batch decoded = decoded[0] - logits = logits[:, 0] - logits[:, 1] - probs = tf.exp(-logits) + logits = logits[:, 0] + probs = tf.exp(-logits / n_steps) # convert to dense outputs = tf.SparseTensor(decoded.indices, decoded.values, (n_samples, n_steps)) @@ -555,8 +554,6 @@ def cnn_rnn_ocr_model(input_height=None, input_width=None, n_classes=None, max_l voc = char2num.get_vocabulary() num2char = StringLookup(vocabulary=voc, invert=True) output = num2char(out) - # avoid output tf.dtype=string → np.dtype=object (which cannot be shm-ed) - output = tf.io.decode_raw(output, tf.uint8, fixed_length=max(map(len, voc))) return Model((inputs, inputs_bin), (output, prob)) @@ -585,8 +582,6 @@ def cnn_rnn_ocr_model4inference(model, model_path): voc = char2num.get_vocabulary() num2char = StringLookup(vocabulary=voc, invert=True) output = num2char(output) - # avoid output tf.dtype=string → np.dtype=object (which cannot be shm-ed) - output = tf.io.decode_raw(output, tf.uint8, fixed_length=max(map(len, voc))) inputs = (inputs, inputs_bin) outputs = (output, prob) return Model(inputs, outputs) diff --git a/src/eynollah/training/reload-models-v0.8.mk b/src/eynollah/training/reload-models-v0.8.mk index ef74bdc..022603e 100644 --- a/src/eynollah/training/reload-models-v0.8.mk +++ b/src/eynollah/training/reload-models-v0.8.mk @@ -59,8 +59,6 @@ $(MODELS_DST)/%.h5: $(MODELS_SRC)/% $(MODELS_DST)/%.onnx: $(MODELS_SRC)/% if jq -e '.task == "segmentation" and .backbone_type == "transformer"' $/dev/null; then \ echo skipping $@: vision transformer architecture currently does not work with ONNX; \ - elif jq -e '.task == "cnn-rnn-ocr"' $/dev/null || test x$(findstring _ocr,$@) = x_ocr; then \ - echo skipping $@: OCR CTC decoder does not work with ONNX; \ else \ eynollah-training convert \ $(and $(wildcard $