Spaces:
Runtime error
Runtime error
| from operator import itemgetter | |
| from collections import OrderedDict | |
| from typing import Dict, List, Iterator, Union, Tuple | |
| import re | |
| class TextExtractor: | |
| def __init__(self) -> None: | |
| pass | |
| def get_font_info(doc: Iterator, granularity=False) -> Tuple[List[Tuple[str, int]], Dict[str, Dict]]: | |
| """ | |
| This function return the fonts information inside the pdf such as size and type. | |
| Args: | |
| doc (<class 'fitz.fitz.Document'>): A fitz type document of the pdf file. | |
| granularity (bool, optional): Also use 'font', 'flags' and 'color' to discriminate text. Defaults to False. | |
| Raises: | |
| ValueError: Raises Value Error if there are no font detected | |
| Returns: | |
| Tuple[List[Tuple[str, int]], Dict[str, Dict]]: _description_ | |
| """ | |
| styles = {} | |
| font_counts = {} | |
| for block in [s for page in doc for b in page.get_text('dict')['blocks'] if b['type'] == 0 for l in b['lines'] for s in l['spans'] if s['text'].strip()]: | |
| identifier = "{0}_{1}_{2}".format(block['size'], block['flags'], block['font']) if granularity else "{0}".format(block['size']) | |
| styles[identifier] = {'size': block['size'], 'flags': block['flags'], 'font': block['font'], 'color': block['color']} if granularity else {'size': block['size'], 'font': block['font']} | |
| font_counts[identifier] = font_counts.get(identifier, 0) + 1 | |
| font_counts = sorted(font_counts.items(), key=lambda x: x[1], reverse=True) | |
| if not font_counts: | |
| raise ValueError("Zero discriminating fonts found!") | |
| return font_counts, styles | |
| def get_font_tags(font_counts, styles): | |
| """ | |
| _summary_ | |
| Args: | |
| font_counts (_type_): _description_ | |
| styles (_type_): _description_ | |
| Returns: | |
| _type_: _description_ | |
| """ | |
| p_size = styles[font_counts[0][0]]['size'] | |
| # sorting the font sizes high to low, so that we can append the right integer to each tag | |
| font_sizes = sorted(set(float(font_size) for font_size, _ in font_counts), reverse=True) | |
| size_tag = {p_size: "<p>"} | |
| for i, size in enumerate(font_sizes): | |
| if size > p_size: | |
| size_tag[size] = f"<h{i+1}>" | |
| elif size < p_size: | |
| size_tag[size] = f"<s{i+1}>" | |
| return size_tag | |
| def assign_tags(doc, size_tag): | |
| """ | |
| Scrapes headers & paragraphs from PDF and return texts with element tags. | |
| Args: | |
| doc (<class 'fitz.fitz.Document'>): PDF document to iterate through. | |
| size_tag (dict): Textual element tags for each size. | |
| Returns: | |
| list: Texts with pre-prended element tags | |
| """ | |
| texts = [] | |
| previous_s = {} | |
| block_string = "" | |
| for b in [b for page in doc for b in page.get_text("dict")["blocks"] if b['type'] == 0]: | |
| block_string = "" | |
| for l in b["lines"]: | |
| for s in l["spans"]: | |
| text = re.sub(r"[^\w\s]", '', s["text"]).strip() | |
| if text: | |
| if not previous_s: # First Span | |
| previous_s = s | |
| block_string = size_tag[s['size']] + s['text'] | |
| elif s['size'] == previous_s['size']: | |
| if not block_string or (block_string and all((c == "|") for c in block_string)): # New block | |
| block_string = size_tag[s['size']] + s['text'] | |
| else: # in the same block, so concatenate strings | |
| block_string += f" {s['text']}" | |
| else: | |
| texts.append(block_string) | |
| block_string = size_tag[s['size']] + s['text'] | |
| previous_s = s | |
| if block_string: | |
| block_string += "|" | |
| # if block_string: | |
| texts.append(block_string) | |
| return texts | |
| def get_slides(texts): | |
| slides = {} | |
| section = [] | |
| page = 1 | |
| current_header = "" | |
| for text, next_text in zip(texts, texts[1:] + [None]): | |
| tag_match = re.search(r'(?<=<)(.*?)(?=>)', text) | |
| if tag_match: | |
| tag = tag_match.group() | |
| if tag == 'h1': | |
| section = [] | |
| section.append(('h1', re.sub(r'<.*?>|\|', '', text).strip())) | |
| elif tag.startswith('h'): # non h1 headers | |
| # Remove tag and pipes from the text | |
| section.append((tag, re.sub(r'<.*?>|\|', '', text).strip())) | |
| elif tag.startswith('p'): | |
| text = re.split("((\|){2,})", text) # If encounter more than 1 pipe than split that text into different paragraphs | |
| for paragraph in text: | |
| paragraph = re.sub(r'<.*?>|\|', '', paragraph).strip() # Remove any pipe | |
| paragraph = re.sub(' +', ' ', paragraph) # Remove any double or more spaces into single space | |
| if paragraph and paragraph[0].islower(): # If a pargraph in a different block is found and the first character isn't an uppercase then concanate with last paragraph | |
| section[-1][1] += f" {paragraph}" | |
| elif paragraph: | |
| section.append([tag, paragraph]) | |
| try: | |
| if tag_match.group() == 'h1': # Create new page when current text is a type 1 header or title | |
| slides[f"Page {page}"] = section | |
| page += 1 | |
| except: | |
| continue | |
| return slides |