(argv: Sequence[str])
| 151 | |
| 152 | |
| 153 | def main(argv: Sequence[str]) -> None: |
| 154 | if len(argv) > 1: |
| 155 | raise app.UsageError('Too many command-line arguments.') |
| 156 | |
| 157 | # Get list of images |
| 158 | img_lists = [] |
| 159 | img_lists.extend(_IMG_FILE.value) |
| 160 | for img_dir in _IMG_DIR.value: |
| 161 | img_lists.extend(tf.io.gfile.glob(os.path.join(img_dir, '*'))) |
| 162 | |
| 163 | logging.info('Total number of input images: %d', len(img_lists)) |
| 164 | |
| 165 | model = load_model() |
| 166 | |
| 167 | vis_dis = _VIS_DIR.value |
| 168 | |
| 169 | output = {'annotations': []} |
| 170 | for img_file in tqdm.tqdm(img_lists): |
| 171 | output['annotations'].append({ |
| 172 | 'image_id': img_file.split('/')[-1].split('.')[0], |
| 173 | 'paragraphs': inference(img_file, model), |
| 174 | }) |
| 175 | |
| 176 | if vis_dis: |
| 177 | key = output['annotations'][-1]['image_id'] |
| 178 | paragraphs = output['annotations'][-1]['paragraphs'] |
| 179 | img = cv2.cvtColor(cv2.imread(img_file), cv2.COLOR_BGR2RGB) |
| 180 | word_bnds = [] |
| 181 | line_bnds = [] |
| 182 | para_bnds = [] |
| 183 | for paragraph in paragraphs: |
| 184 | paragraph_points_list = [] |
| 185 | for line in paragraph['lines']: |
| 186 | line_points_list = [] |
| 187 | for word in line['words']: |
| 188 | word_bnds.append( |
| 189 | np.array(word['vertices'], np.int32).reshape((-1, 1, 2))) |
| 190 | line_points_list.extend(word['vertices']) |
| 191 | paragraph_points_list.extend(line_points_list) |
| 192 | |
| 193 | line_points = np.array(line_points_list, np.int32) # (N,2) |
| 194 | left = int(np.min(line_points[:, 0])) |
| 195 | top = int(np.min(line_points[:, 1])) |
| 196 | right = int(np.max(line_points[:, 0])) |
| 197 | bottom = int(np.max(line_points[:, 1])) |
| 198 | line_bnds.append( |
| 199 | np.array([[[left, top]], [[right, top]], [[right, bottom]], |
| 200 | [[left, bottom]]], np.int32)) |
| 201 | para_points = np.array(paragraph_points_list, np.int32) # (N,2) |
| 202 | left = int(np.min(para_points[:, 0])) |
| 203 | top = int(np.min(para_points[:, 1])) |
| 204 | right = int(np.max(para_points[:, 0])) |
| 205 | bottom = int(np.max(para_points[:, 1])) |
| 206 | para_bnds.append( |
| 207 | np.array([[[left, top]], [[right, top]], [[right, bottom]], |
| 208 | [[left, bottom]]], np.int32)) |
| 209 | |
| 210 | for name, bnds in zip(['paragraph', 'line', 'word'], |
nothing calls this directly
no test coverage detected