Inference step.
(img_file: str, model: tf_keras.layers.Layer)
| 92 | |
| 93 | |
| 94 | def inference(img_file: str, model: tf_keras.layers.Layer) -> Dict[str, Any]: |
| 95 | """Inference step.""" |
| 96 | img = cv2.cvtColor(cv2.imread(img_file), cv2.COLOR_BGR2RGB) |
| 97 | img_ndarray, ratio = _preprocess(img) |
| 98 | |
| 99 | output_dict = model.serve(img_ndarray) |
| 100 | class_tensor = output_dict['classes'].numpy() |
| 101 | mask_tensor = output_dict['masks'].numpy() |
| 102 | group_tensor = output_dict['groups'].numpy() |
| 103 | |
| 104 | indices = np.where(class_tensor[0])[0].tolist() # indices of positive slots. |
| 105 | mask_list = [ |
| 106 | mask_tensor[0, :, :, index] for index in indices] # List of mask ndarray. |
| 107 | |
| 108 | # Form lines and words |
| 109 | lines = [] |
| 110 | line_indices = [] |
| 111 | for index, mask in tqdm.tqdm(zip(indices, mask_list)): |
| 112 | line = { |
| 113 | 'words': [], |
| 114 | 'text': '', |
| 115 | } |
| 116 | |
| 117 | contours, _ = cv2.findContours( |
| 118 | (mask > 0.).astype(np.uint8), |
| 119 | cv2.RETR_TREE, |
| 120 | cv2.CHAIN_APPROX_SIMPLE)[-2:] |
| 121 | for contour in contours: |
| 122 | if (isinstance(contour, np.ndarray) and |
| 123 | len(contour.shape) == 3 and |
| 124 | contour.shape[0] > 2 and |
| 125 | contour.shape[1] == 1 and |
| 126 | contour.shape[2] == 2): |
| 127 | cnt_list = (contour[:, 0] * ratio).astype(np.int32).tolist() |
| 128 | line['words'].append({'text': '', 'vertices': cnt_list}) |
| 129 | else: |
| 130 | logging.error('Invalid contour: %s, discarded', str(contour)) |
| 131 | if line['words']: |
| 132 | lines.append(line) |
| 133 | line_indices.append(index) |
| 134 | |
| 135 | # Form paragraphs |
| 136 | line_grouping = utilities.DisjointSet(len(line_indices)) |
| 137 | affinity = group_tensor[0][line_indices][:, line_indices] |
| 138 | for i1, i2 in zip(*np.where(affinity > _PARA_GROUP_THR)): |
| 139 | line_grouping.union(i1, i2) |
| 140 | |
| 141 | line_groups = line_grouping.to_group() |
| 142 | paragraphs = [] |
| 143 | for line_group in line_groups: |
| 144 | paragraph = {'lines': []} |
| 145 | for id_ in line_group: |
| 146 | paragraph['lines'].append(lines[id_]) |
| 147 | if paragraph: |
| 148 | paragraphs.append(paragraph) |
| 149 | |
| 150 | return paragraphs |
| 151 |
no test coverage detected