print('document_image_1.info.get(annotation_key) == True, end:', document_image_1.info.get(annotation_key) == True)

analysis.ipynb

@@ -11,18 +11,9 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 2,
+      "execution_count": 1,
       "metadata": {},
-      "outputs": [
-        {
-          "name": "stdout",
-          "output_type": "stream",
-          "text": [
-            "The autoreload extension is already loaded. To reload it, use:\n",
-            "  %reload_ext autoreload\n"
-          ]
-        }
-      ],
+      "outputs": [],
       "source": [
         "%load_ext autoreload\n",
         "%autoreload 2\n",
@@ -310,7 +301,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 5,
+      "execution_count": 33,
       "metadata": {},
       "outputs": [
         {
@@ -338,10 +329,14 @@
           "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",
+            "{'predicted_bboxes': 'predicted_bboxes', 'predicted_scores': 'predicted_scores', 'predicted_labels': 'predicted_labels'}\n",
+            "document_image_1.info.get(annotation_key) == True, start: False\n",
+            "document_image_1.info.get(annotation_key) == True, middle: False\n",
+            "document_image_1.info.get(annotation_key) == True, end: True\n",
+            "{'predicted_bboxes': 'reduced_predicted_bboxes', 'predicted_scores': 'reduced_predicted_scores', 'predicted_labels': 'reduced_predicted_labels'}\n",
+            "document_image_1.info.get(annotation_key) == True, start: False\n",
+            "document_image_1.info.get(annotation_key) == True, middle: False\n",
+            "document_image_1.info.get(annotation_key) == True, end: True\n",
             "Keyboard interruption in main thread... closing server.\n"
           ]
         },
@@ -349,7 +344,7 @@
           "data": {
             "text/plain": []
           },
-          "execution_count": 5,
+          "execution_count": 33,
           "metadata": {},
           "output_type": "execute_result"
         }

main.py

@@ -7,6 +7,7 @@ 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())
@@ -17,39 +18,96 @@ cache = {
     'document_image_1_features': None,
     'document_image_2_features': None,
 }
-pre_message_style = 'overflow: auto;border: 2px solid pink;padding: 4px;'
+pre_message_style = 'overflow:auto;border:2px solid pink;padding:4px;border-radius:4px;'
+visualize_bboxes_on_image_kwargs = {
+    'label_text_color': 'white',
+    'label_rectangle_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', 'reduced_weighted_vectors']
+
+annotation_key = 'is_annotated_document_image'
+annotation_original_image_key = 'original_image'
+def annotate_document_image(document_image: Image.Image, original_document_image: Image.Image):
+    document_image.info.update({
+        annotation_key: True,
+        annotation_original_image_key: original_document_image
+    })
+    return document_image
+
+def get_original_document_image(document_image: Image.Image):
+    if document_image.info.get(annotation_key) == True:
+        return document_image.info.get(annotation_original_image_key)
+    return document_image
 
-def similarity_fn(document_image_1: Image.Image, document_image_2: Image.Image, model: lp.Detectron2LayoutModel):
+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:
-        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']
+        if document_image_1 is None or document_image_2 is None:
+            message = f'<pre style="{pre_message_style}">Please load both the documents to compare.<pre>'
         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
+            document_image_1 = get_original_document_image(document_image_1)
+            document_image_2 = get_original_document_image(document_image_2)
 
-        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
+            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>'
+            [[similarity]] = cosine_similarity(
+                [
+                    cache['document_image_1_features'][vectors_type]
+                ], 
+                [
+                    cache['document_image_2_features'][vectors_type]
+                ])
+            message = f'<pre style="{pre_message_style}">Similarity between the two documents is: {round(similarity, 4)}<pre>'
+            annotated_document_image_1 = visualize_bboxes_on_image(
+                image = document_image_1,
+                bboxes = cache['document_image_1_features'][annotations['predicted_bboxes']],
+                titles = [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']])],
+                **visualize_bboxes_on_image_kwargs)
+            annotated_document_image_2 = visualize_bboxes_on_image(
+                image = document_image_2,
+                bboxes = cache['document_image_2_features'][annotations['predicted_bboxes']],
+                titles = [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']])],
+                **visualize_bboxes_on_image_kwargs)
+            show_vectors_type = True
+            document_image_1 = annotate_document_image(annotated_document_image_1, document_image_1)
+            document_image_2 = annotate_document_image(annotated_document_image_2, document_image_2)
     except Exception as e:
         message = f'<pre style="{pre_message_style}">{traceback.format_exc()}<pre>'
-    return gr.HTML(message, visible=True)
+    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:
@@ -83,7 +141,7 @@ def document_view(document_number: int):
             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", 
+                "Upload PDF or Document Image", 
                 file_types=["image", ".pdf"], 
                 file_count="single")
         with gr.Tab("From URL", id=1):
@@ -114,7 +172,7 @@ def app(*, model_path, config_path, debug = False):
     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'])
@@ -126,15 +184,25 @@ def app(*, model_path, config_path, debug = False):
         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")
+                submit = gr.Button(value="Get Similarity", variant="primary")
             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])
+                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)

utils/get_features.py

@@ -31,10 +31,6 @@ def get_vectors(*,
     '''
     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