similarity display

analysis.ipynb

@@ -11,9 +11,18 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 1,
+      "execution_count": 2,
       "metadata": {},
-      "outputs": [],
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "The autoreload extension is already loaded. To reload it, use:\n",
+            "  %reload_ext autoreload\n"
+          ]
+        }
+      ],
       "source": [
         "%load_ext autoreload\n",
         "%autoreload 2\n",
@@ -329,6 +338,10 @@
           "name": "stdout",
           "output_type": "stream",
           "text": [
+            "tensor(0)\n",
+            "tensor(0.)\n",
+            "['Text', 'Picture', 'Text', 'Picture', 'Title', 'Text', 'Picture', 'Text', 'Text', 'Picture', 'Picture', 'Text', 'Text', 'Title', 'Section-header', 'Picture', 'Title', 'Picture', 'Picture', 'Picture', 'Section-header', 'Title', 'Picture', 'Caption', 'Title', 'Text', 'Text', 'Picture', 'Caption', 'Title', 'Text', 'Title', 'Text', 'Page-header', 'Section-header', 'Section-header', 'Caption', 'Title', 'Page-header', 'Section-header', 'Section-header', 'Page-header', 'Text', 'Picture', 'Caption', 'Text', 'Caption', 'Text', 'Picture', 'Page-header', 'Title', 'Picture', 'Picture', 'Text', 'Page-footer', 'Section-header', 'Caption', 'Section-header', 'Title', 'Text', 'Picture', 'Page-header', 'Picture', 'Caption', 'Caption', 'Section-header', 'Section-header', 'Picture', 'Section-header', 'Title', 'Picture', 'Page-footer', 'Caption', 'Title', 'Text', 'Picture', 'Title', 'Picture', 'Text', 'Text', 'Section-header', 'Picture', 'Picture', 'Section-header', 'Caption', 'Text']\n",
+            "{0: 0.17, 1: 0.24545454545454548, 2: 0.3209090909090909, 3: 0.3963636363636364, 4: 0.4718181818181819, 5: 0.5472727272727274, 6: 0.6227272727272728, 7: 0.6981818181818182, 8: 0.7736363636363638, 9: 0.8490909090909092, 10: 0.9245454545454547}\n",
             "Keyboard interruption in main thread... closing server.\n"
           ]
         },
@@ -342,30 +355,11 @@
         }
       ],
       "source": [
-        "from app import app\n",
-        "\n",
-        "app()"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "from utils.get_RGB_image import get_RGB_image\n",
-        "from pdf2image import convert_from_path\n",
+        "from main import app\n",
         "\n",
-        "pdf = convert_from_path('./sdfes.png', 140)"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "get_RGB_image(pdf[0]) "
+        "model_path = '../detectron2-layout-parser/model_final.pth'\n",
+        "config_path = '../detectron2-layout-parser/config.yaml'\n",
+        "app(model_path=model_path, config_path=config_path, debug=True)"
       ]
     }
   ],

app.py

@@ -1,89 +1,10 @@
-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 os
+os.system("apt-get install poppler-utils")
+os.system("python -m pip install torch")
+os.system("python -m pip install 'git+https://github.com/facebookresearch/detectron2.git'")
+os.system("python -m pip install layoutparser layoutparser[layoutmodels] layoutparser[ocr]")
+os.system("python -m pip install Pillow==9.4.0")
+os.system("python -m pip install imagehash")
 
-def similarity_fn(document_image_1, document_image_2):
-    xxx = 'BOTH' if document_image_1 and document_image_2 else 'ONE' if document_image_1 or document_image_2 else 'NONE'
-    return gr.HTML(f'<pre>Similarity between the two documents: {xxx}<pre>', visible=True)
-
-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'<h4>Load the {"first" if document_number == 1 else "second"} PDF or Document Image<h4>', 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)
-            document_error_message = gr.HTML(label="Error Message", visible=False)
-            document_preview = gr.UploadButton(
-                "Click to 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():
-    title = 'Document Similarity Search Using Visual Layout Features'
-    description = f"<h2>{title}<h2>"
-    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('<hr/>', elem_classes=['hr'])
-        with gr.Row(elem_classes=['center']):
-            with gr.Column():
-                submit = gr.Button(value="Similarity", variant="primary")
-                reset = gr.Button(value="Reset", variant="secondary")
-            with gr.Column():
-                similarity_output = gr.HTML(visible=False)
-            submit.click(
-                fn=similarity_fn,
-                inputs=[document_1_image, document_2_image],
-                outputs=[similarity_output])
-    return app.launch(debug=True)
+from main import app
+app()

main.py

@@ -0,0 +1,140 @@
+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
+
+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 = {
+    'document_image_1_hash': None,
+    'document_image_2_hash': None,
+    'document_image_1_features': None,
+    'document_image_2_features': None,
+}
+pre_message_style = 'overflow: auto;border: 2px solid pink;padding: 4px;'
+
+def similarity_fn(document_image_1: Image.Image, document_image_2: Image.Image, model: lp.Detectron2LayoutModel):
+    message = None
+    try:
+        document_image_1_hash = str(average_hash(document_image_1))
+        document_image_2_hash = str(average_hash(document_image_2))
+
+        if document_image_1_hash == cache['document_image_1_hash']:
+            document_image_1_features = cache['document_image_1_features']
+        else:
+            document_image_1_features = get_features(document_image_1, model, label_names)
+            cache['document_image_1_hash'] = document_image_1_hash
+            cache['document_image_1_features'] = document_image_1_features
+
+        if document_image_2_hash == cache['document_image_2_hash']:
+            document_image_2_features = cache['document_image_2_features']
+        else:
+            document_image_2_features = get_features(document_image_2, model, label_names)
+            cache['document_image_2_hash'] = document_image_2_hash
+            cache['document_image_2_features'] = document_image_2_features
+
+        [[similarity]] = cosine_similarity(
+            [
+                cache['document_image_1_features']['vectors']
+            ], 
+            [
+                cache['document_image_2_features']['vectors']
+            ])
+        message = f'<pre style="{pre_message_style}">Similarity between the two documents is: {similarity}<pre>'
+    except Exception as e:
+        message = f'<pre style="{pre_message_style}">{traceback.format_exc()}<pre>'
+    return gr.HTML(message, visible=True)
+    
+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'<h4>Load the {"first" if document_number == 1 else "second"} PDF or Document Image<h4>', 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)
+            document_error_message = gr.HTML(label="Error Message", visible=False)
+            document_preview = gr.UploadButton(
+                "Click to 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"<h2>{title}<h2>"
+    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('<hr/>', elem_classes=['hr'])
+        with gr.Row(elem_classes=['center']):
+            with gr.Column():
+                submit = gr.Button(value="Similarity", variant="primary")
+                reset = gr.Button(value="Reset", variant="secondary")
+            with gr.Column():
+                similarity_output = gr.HTML(visible=False)
+            submit.click(
+                fn=lambda document_1_image, document_2_image: similarity_fn(
+                    document_1_image, 
+                    document_2_image, 
+                    model),
+                inputs=[document_1_image, document_2_image],
+                outputs=[similarity_output])
+    return app.launch(debug=debug)

utils/get_features.py

@@ -0,0 +1,127 @@
+import layoutparser as lp
+from PIL import Image
+import tensorflow as tf
+import numpy as np
+import torch
+import torchvision.ops.boxes as box_ops
+from typing import List, Tuple
+from .split_image import split_image
+from .get_unique_values import get_unique_values
+
+def get_vectors(*,
+    predicted_bboxes: List[Tuple[int, int, int, int]],
+    predicted_scores: List[float],
+    predicted_labels: List[str],
+    label_names: List[str],
+    sub_images_bboxes: List[Tuple[int, int, int, int]],
+    index_start: int = 0.17,
+    index_end: int = 1,
+    weighted_jaccard_index = False):
+  bboxes_tensor: torch.Tensor = torch.tensor(predicted_bboxes)
+  labels_nonce = { value:key for key, value in zip(get_unique_values(start = index_start, end = index_end, count = len(label_names)), list(label_names)) }
+  
+  def get_vector(bbox: Tuple[int, int, int, int], region_nonce: int):
+    # bbox: Expected to be in ``(x1, y1, x2, y2)`` format with ``0 <= x1 < x2`` and ``0 <= y1 < y2``.
+    bbox_tensor: torch.Tensor = torch.tensor([bbox])
+    [jaccard_indexes] = box_ops.box_iou(bbox_tensor, bboxes_tensor)
+    '''
+      Either get the index of bounding box with largest jaccard_index (Intersection Over Union) or
+      get the index of bounding box with largest jaccard_index (Intersection Over Union) multiplied by the score.
+      By doing this we strike a balance between accuracy and relative position.
+    '''
+    index_of_jaccard_index = jaccard_indexes.argmax() if not weighted_jaccard_index else np.multiply(jaccard_indexes, predicted_scores).argmax()
+    jaccard_index = jaccard_indexes[index_of_jaccard_index]
+    print(index_of_jaccard_index)
+    print(jaccard_index)
+    print(predicted_labels)
+    print(labels_nonce)
+    jaccard_index_bbox_label__nonce = labels_nonce[predicted_labels[index_of_jaccard_index]]
+    jaccard_index_bbox_score = predicted_scores[index_of_jaccard_index]
+    vector = region_nonce * jaccard_index * jaccard_index_bbox_label__nonce * jaccard_index_bbox_score
+    return vector.item()
+  sub_images_nonces = get_unique_values(start = index_start, end = index_end, count = len(sub_images_bboxes))
+  for sub_image_bbox, region_nonce in zip(sub_images_bboxes, sub_images_nonces):
+      yield get_vector(sub_image_bbox, region_nonce)
+
+def get_predictions(
+    image: Image.Image,
+    model: lp.Detectron2LayoutModel,
+    predictions_reducer = lambda *args: args):
+  layout_predicted = model.detect(image)
+  if len(layout_predicted) > 0:
+    predicted_bboxes = [block.coordinates for block in layout_predicted]
+    predicted_scores = [block.score for block in layout_predicted]
+    predicted_labels = [block.type for block in layout_predicted]
+    [predicted_bboxes, predicted_scores, predicted_labels] = predictions_reducer(
+        predicted_bboxes,
+        predicted_scores,
+        predicted_labels)
+    return {
+        'predicted_bboxes': predicted_bboxes,
+        'predicted_scores': predicted_scores,
+        'predicted_labels': predicted_labels,
+    }
+  else:
+    return {
+        'predicted_bboxes': [],
+        'predicted_scores': [],
+        'predicted_labels': [],
+    }
+
+def predictions_reducer(
+    predicted_bboxes: List[Tuple[int, int, int, int]],
+    predicted_scores: List[float],
+    predicted_labels: List[str]):
+  selected_indices = tf.image.non_max_suppression(
+    boxes = predicted_bboxes,
+    scores = predicted_scores ,
+    max_output_size = len(predicted_bboxes),
+    iou_threshold = 0.01)
+  return {
+      'predicted_bboxes': tf.gather(predicted_bboxes, selected_indices).numpy().tolist(), # List[List[int, int, int, int]]
+      'predicted_scores': tf.gather(predicted_scores, selected_indices).numpy().astype(float).tolist(),
+      'predicted_labels': tf.gather(predicted_labels, selected_indices).numpy().astype(str).tolist()
+    }
+
+def get_features(image: Image.Image, model: lp.Detectron2LayoutModel, label_names: List[str], width_parts = 100, height_parts = 100):
+  predictions = get_predictions(image, model)
+  reduced_predictions = predictions_reducer(**predictions)
+  sub_images_bboxes = list(split_image(np.array(image), width_parts, height_parts, result = 'bboxes'))
+
+  vectors = get_vectors(
+    sub_images_bboxes = sub_images_bboxes,
+    label_names = label_names,
+    weighted_jaccard_index = False,
+    **predictions)
+  
+  weighted_vectors = get_vectors(
+    sub_images_bboxes = sub_images_bboxes,
+    label_names = label_names,
+    weighted_jaccard_index = True,
+    **predictions)
+  
+  reduced_vectors = get_vectors(
+    sub_images_bboxes = sub_images_bboxes,
+    label_names = label_names,
+    weighted_jaccard_index = False,
+    **reduced_predictions)
+  
+  reduced_weighted_vectors = get_vectors(
+    sub_images_bboxes = sub_images_bboxes,
+    label_names = label_names,
+    weighted_jaccard_index = True,
+    **reduced_predictions)
+
+  return {
+      'predicted_bboxes': predictions['predicted_bboxes'],
+      'predicted_scores': predictions['predicted_scores'],
+      'predicted_labels': predictions['predicted_labels'],
+      'vectors': list(vectors),
+      'weighted_vectors': list(weighted_vectors),
+
+      'reduced_predicted_bboxes': reduced_predictions['predicted_bboxes'],
+      'reduced_predicted_scores': reduced_predictions['predicted_scores'],
+      'reduced_predicted_labels': reduced_predictions['predicted_labels'],
+      'reduced_vectors': list(reduced_vectors),
+      'weighted_reduced_vectors': list(reduced_weighted_vectors),
+  }

utils/split_image.py

@@ -0,0 +1,19 @@
+import numpy as np
+from typing import Literal
+
+def split_image(image: np.ndarray, width_parts: int, height_parts: int, result: Literal['np.ndarray', 'bboxes'] = 'np.ndarray'):
+  tile_width = image.shape[0] // width_parts
+  tile_height = image.shape[1] // height_parts
+  for height in range(height_parts):
+    for width in range(width_parts):
+      width_start = width * tile_width
+      width_end = tile_width * (width + 1) if (width + 1) < width_parts else image.shape[0]
+      height_start = height * tile_height
+      height_end = tile_height * (height + 1) if (height + 1) < height_parts else image.shape[1]
+      if result == 'np.ndarray':
+        # np.ndarray(height, width, channels)
+        yield image[height_start:height_end, width_start:width_end]
+      else:
+        # Both sets of boxes are expected to be in ``(x1, y1, x2, y2)`` format with ``0 <= x1 < x2`` and ``0 <= y1 < y2``.
+        yield (width_start, height_start, width_end, height_end)
+