import traceback import gradio as gr from utils.get_RGB_image import get_RGB_image, is_online_file, steam_online_file from pdf2image import convert_from_path, convert_from_bytes import layoutparser as lp from PIL import Image from utils.get_features import get_features from imagehash import average_hash from sklearn.metrics.pairwise import cosine_similarity from utils.visualize_bboxes_on_image import visualize_bboxes_on_image label_map = {0: 'Caption', 1: 'Footnote', 2: 'Formula', 3: 'List-item', 4: 'Page-footer', 5: 'Page-header', 6: 'Picture', 7: 'Section-header', 8: 'Table', 9: 'Text', 10: 'Title'} label_names = list(label_map.values()) color_map = {'Caption': '#acc2d9', 'Footnote': '#56ae57', 'Formula': '#b2996e', 'List-item': '#a8ff04', 'Page-footer': '#69d84f', 'Page-header': '#894585', 'Picture': '#70b23f', 'Section-header': '#d4ffff', 'Table': '#65ab7c', 'Text': '#952e8f', 'Title': '#fcfc81'} cache = { 'output_document_image_1_hash': None, 'output_document_image_2_hash': None, 'document_image_1_features': None, 'document_image_2_features': None, 'original_document_image_1': None, 'original_document_image_2': None } pre_message_style = 'overflow:auto;border:2px solid pink;padding:4px;border-radius:4px;' visualize_bboxes_on_image_kwargs = { 'label_text_color': 'white', 'label_fill_color': 'black', 'label_text_size': 12, 'label_text_padding': 3, 'label_rectangle_left_margin': 0, 'label_rectangle_top_margin': 0 } vectors_types = ['vectors', 'weighted_vectors', 'reduced_vectors', 'weighted_reduced_vectors'] def similarity_fn(model: lp.Detectron2LayoutModel, document_image_1: Image.Image, document_image_2: Image.Image, vectors_type: str): message = None annotations = { 'predicted_bboxes': 'predicted_bboxes' if vectors_type in ['vectors', 'weighted_vectors'] else 'reduced_predicted_bboxes', 'predicted_scores': 'predicted_scores' if vectors_type in ['vectors', 'weighted_vectors'] else 'reduced_predicted_scores', 'predicted_labels': 'predicted_labels' if vectors_type in ['vectors', 'weighted_vectors'] else 'reduced_predicted_labels', } show_vectors_type = False try: if document_image_1 is None or document_image_2 is None: message = f'
Please load both the documents to compare.
'
        else:
            input_document_image_1_hash = str(average_hash(document_image_1))
            input_document_image_2_hash = str(average_hash(document_image_2))

            if input_document_image_1_hash == cache['output_document_image_1_hash']:
                document_image_1_features = cache['document_image_1_features']
                document_image_1 = cache['original_document_image_1']
            else:
                document_image_1_features = get_features(document_image_1, model, label_names)
                cache['document_image_1_features'] = document_image_1_features
                cache['original_document_image_1'] = document_image_1

            if input_document_image_2_hash == cache['output_document_image_2_hash']:
                document_image_2_features = cache['document_image_2_features']
                document_image_2 = cache['original_document_image_2']
            else:
                document_image_2_features = get_features(document_image_2, model, label_names)
                cache['document_image_2_features'] = document_image_2_features
                cache['original_document_image_2'] = document_image_2

            [[similarity]] = cosine_similarity(
                [
                    cache['document_image_1_features'][vectors_type]
                ], 
                [
                    cache['document_image_2_features'][vectors_type]
                ])
            message = f'
Similarity between the two documents is: {round(similarity, 4)}
'
            document_image_1 = visualize_bboxes_on_image(
                image = document_image_1,
                bboxes = cache['document_image_1_features'][annotations['predicted_bboxes']],
                labels = [f'{label}, score:{round(score, 2)}' for label, score in zip(
                    cache['document_image_1_features'][annotations['predicted_labels']], 
                    cache['document_image_1_features'][annotations['predicted_scores']])],
                bbox_outline_color = [color_map[label] for label in cache['document_image_1_features'][annotations['predicted_labels']]],
                **visualize_bboxes_on_image_kwargs)
            document_image_2 = visualize_bboxes_on_image(
                image = document_image_2,
                bboxes = cache['document_image_2_features'][annotations['predicted_bboxes']],
                labels = [f'{label}, score:{score}' for label, score in zip(
                    cache['document_image_2_features'][annotations['predicted_labels']], 
                    cache['document_image_2_features'][annotations['predicted_scores']])],
                bbox_outline_color = [color_map[label] for label in cache['document_image_2_features'][annotations['predicted_labels']]], 
                **visualize_bboxes_on_image_kwargs)
            
            cache['output_document_image_1_hash'] = str(average_hash(document_image_1))
            cache['output_document_image_2_hash'] = str(average_hash(document_image_2))

            show_vectors_type = True
    except Exception as e:
        message = f'
{traceback.format_exc()}
'
    return [
        gr.HTML(message, visible=True), 
        document_image_1, 
        document_image_2,
        gr.Dropdown(visible=show_vectors_type)
    ]
    
def load_image(filename, page = 0):
    try:
        image = None
        try:
            if (is_online_file(filename)):
                image = get_RGB_image(convert_from_bytes(steam_online_file(filename))[page])
            else:
                image = get_RGB_image(convert_from_path(filename)[page])
        except:
            image = get_RGB_image(filename)
        return [
            gr.Image(value=image, visible=True), 
            None
        ]
    except:
        error = traceback.format_exc()
        return [None, gr.HTML(value=error, visible=True)]
    
def preview_url(url, page = 0):
    [image, error] = load_image(url, page = page)
    if image:
        return [gr.Tabs(selected=0), image, error]
    else:
        return [gr.Tabs(selected=1), image, error] 

def document_view(document_number: int):
    gr.HTML(value=f'

Load the {"first" if document_number == 1 else "second"} PDF or Document Image

', elem_classes=['center']) with gr.Tabs() as document_tabs: with gr.Tab("From Image", id=0): document = gr.Image(type="pil", label=f"Document {document_number}", visible=False, interactive=False, show_download_button=True) document_error_message = gr.HTML(label="Error Message", visible=False) document_preview = gr.UploadButton( "Upload PDF or Document Image", file_types=["image", ".pdf"], file_count="single") with gr.Tab("From URL", id=1): document_url = gr.Textbox( label=f"Document {document_number} URL", info="Paste a Link/URL to PDF or Document Image", placeholder="https://datasets-server.huggingface.co/.../image.jpg") document_url_error_message = gr.HTML(label="Error Message", visible=False) document_url_preview = gr.Button(value="Preview", variant="primary") document_preview.upload( fn = lambda file: load_image(file.name), inputs = [document_preview], outputs = [document, document_error_message]) document_url_preview.click( fn = preview_url, inputs = [document_url], outputs = [document_tabs, document, document_url_error_message]) return document def app(*, model_path, config_path, debug = False): model: lp.Detectron2LayoutModel = lp.Detectron2LayoutModel( config_path = config_path, model_path = model_path, label_map = label_map) title = 'Document Similarity Search Using Visual Layout Features' description = f"

{title}

" css = ''' image { max-height="86vh" !important; } .center { display: flex; flex: 1 1 auto; align-items: center; align-content: center; justify-content: center; justify-items: center; } .hr { width: 100%; display: block; padding: 0; margin: 0; background: gray; height: 4px; border: none; } ''' with gr.Blocks(title=title, css=css) as app: with gr.Row(): gr.HTML(value=description, elem_classes=['center']) with gr.Row(equal_height = False): with gr.Column(): document_1_image = document_view(1) with gr.Column(): document_2_image = document_view(2) gr.HTML('
', elem_classes=['hr']) with gr.Row(elem_classes=['center']): with gr.Column(): submit = gr.Button(value="Get Similarity", variant="primary") with gr.Column(): vectors_type = gr.Dropdown( choices = vectors_types, value = vectors_types[0], visible = False, label = "Vectors Type", info = "Select the Vectors Type to use for Similarity Calculation") similarity_output = gr.HTML(label="Similarity Score", visible=False) reset = gr.Button(value="Reset", variant="secondary") kwargs = { 'fn': lambda document_1_image, document_2_image, vectors_type: similarity_fn( model, document_1_image, document_2_image, vectors_type), 'inputs': [document_1_image, document_2_image, vectors_type], 'outputs': [similarity_output, document_1_image, document_2_image, vectors_type] } submit.click(**kwargs) vectors_type.change(**kwargs) return app.launch(debug=debug)