(args, conv_res)
| 126 | |
| 127 | |
| 128 | def gen_image_and_table(args, conv_res): |
| 129 | input_token, output_token = 0, 0 |
| 130 | raw_source = args.poster_path |
| 131 | |
| 132 | output_dir = Path(f'<{args.model_name_t}_{args.model_name_v}>_images_and_tables/{args.poster_name}') |
| 133 | |
| 134 | output_dir.mkdir(parents=True, exist_ok=True) |
| 135 | doc_filename = args.poster_name |
| 136 | |
| 137 | # Save page images |
| 138 | for page_no, page in conv_res.document.pages.items(): |
| 139 | page_no = page.page_no |
| 140 | page_image_filename = output_dir / f"{doc_filename}-{page_no}.png" |
| 141 | with page_image_filename.open("wb") as fp: |
| 142 | page.image.pil_image.save(fp, format="PNG") |
| 143 | |
| 144 | # Save images of figures and tables |
| 145 | table_counter = 0 |
| 146 | picture_counter = 0 |
| 147 | for element, _level in conv_res.document.iterate_items(): |
| 148 | if isinstance(element, TableItem): |
| 149 | table_counter += 1 |
| 150 | element_image_filename = ( |
| 151 | output_dir / f"{doc_filename}-table-{table_counter}.png" |
| 152 | ) |
| 153 | with element_image_filename.open("wb") as fp: |
| 154 | element.get_image(conv_res.document).save(fp, "PNG") |
| 155 | |
| 156 | if isinstance(element, PictureItem): |
| 157 | picture_counter += 1 |
| 158 | element_image_filename = ( |
| 159 | output_dir / f"{doc_filename}-picture-{picture_counter}.png" |
| 160 | ) |
| 161 | with element_image_filename.open("wb") as fp: |
| 162 | element.get_image(conv_res.document).save(fp, "PNG") |
| 163 | |
| 164 | # Save markdown with embedded pictures |
| 165 | md_filename = output_dir / f"{doc_filename}-with-images.md" |
| 166 | conv_res.document.save_as_markdown(md_filename, image_mode=ImageRefMode.EMBEDDED) |
| 167 | |
| 168 | # Save markdown with externally referenced pictures |
| 169 | md_filename = output_dir / f"{doc_filename}-with-image-refs.md" |
| 170 | conv_res.document.save_as_markdown(md_filename, image_mode=ImageRefMode.REFERENCED) |
| 171 | |
| 172 | # Save HTML with externally referenced pictures |
| 173 | html_filename = output_dir / f"{doc_filename}-with-image-refs.html" |
| 174 | conv_res.document.save_as_html(html_filename, image_mode=ImageRefMode.REFERENCED) |
| 175 | |
| 176 | tables = {} |
| 177 | |
| 178 | table_index = 1 |
| 179 | for table in conv_res.document.tables: |
| 180 | caption = table.caption_text(conv_res.document) |
| 181 | if len(caption) > 0: |
| 182 | table_img_path = f'<{args.model_name_t}_{args.model_name_v}>_images_and_tables/{args.poster_name}/{args.poster_name}-table-{table_index}.png' |
| 183 | table_img = PIL.Image.open(table_img_path) |
| 184 | tables[str(table_index)] = { |
| 185 | 'caption': caption, |
no test coverage detected