2023-11-27-03-48-27

app.py

@@ -0,0 +1,146 @@
+
+import os
+
+os.system("pip install torch")
+os.system("pip install 'git+https://github.com/facebookresearch/detectron2.git'")
+os.system("pip install layoutparser")
+os.system("pip install layoutparser[layoutmodels]")
+os.system("pip install layoutparser[ocr]")
+os.system("pip install Pillow==9.4.0")
+os.system("pip install requests")
+
+import gradio as gr
+import layoutparser as lp
+from PIL import Image
+from urllib.parse import urlparse
+import requests
+
+def get_RGB_image(image_or_path: str | Image.Image) -> bytes:
+    if isinstance(image_or_path, str):
+        if urlparse(image_or_path).scheme in ["http", "https"]:  # Online
+            image_or_path = Image.open(
+                requests.get(image_or_path, stream=True).raw)
+        else:  # Local
+            image_or_path = Image.open(image_or_path)
+    return image_or_path.convert("RGB")
+
+def inference_factory(config_path: str, model_path: str, label_map: dict, color_map: dict, examples=[], launch=True):
+    import traceback
+    model: lp.elements.layout.Layout = lp.Detectron2LayoutModel(
+        config_path=config_path,
+        model_path=model_path,
+        # extra_config = ["MODEL.ROI_HEADS.SCORE_THRESH_TEST", 0.8],
+        label_map=label_map)
+    default_threshold = 0.8
+    cache = {
+        'annotated_image': None,
+        'message': None,
+        'threshold': default_threshold,
+        'image': None,
+        'predicted': None
+    }
+
+    def truncate(f, n):
+        return int(f * 10 ** n) / 10 ** n
+
+    def fn(image: Image.Image, threshold: float = default_threshold, just_image=True):
+        try:
+            nonlocal cache
+            if cache['image'] == image and cache['threshold'] == threshold and bool(cache['annotated_image']):
+                return [cache['annotated_image'], cache['message'], cache['threshold']]
+            layout_predicted = cache['predicted'] if cache['image'] == image else model.detect(
+                image)
+            threshold = truncate(
+                min([max([block.score for block in layout_predicted] + [0])] + [threshold]), 1)
+            blocks: List[lp.elements.layout_elements.TextBlock] = [block.set(
+                id=f'{block.type}/{block.score:.2f}') for block in layout_predicted if block.score >= threshold]
+            annotated_image = lp.draw_box(
+                image,
+                blocks,
+                color_map=color_map,
+                show_element_id=True,
+                id_font_size=14,
+                id_text_background_color='black',
+                id_text_color='white')
+            message = \
+                f'{len(blocks)} bounding boxes matched for {threshold} threshold, out of {len(layout_predicted)} total bounding boxes' if len(blocks) > 0 \
+                else f'No bounding boxesfor {threshold} threshold.'
+            cache = {
+                'annotated_image': annotated_image,
+                'message': message,
+                'threshold': threshold,
+                'image': image,
+                'predicted': layout_predicted
+            }
+            return annotated_image if just_image else [annotated_image, message, threshold]
+        except Exception as e:
+            error = traceback.format_exc()
+            return error if just_image else [None, error, threshold]
+    if not launch:
+        return fn
+    
+    ###########################################################
+    ################### Start of Gradio setup #################
+    ###########################################################
+    title = "Document Similarity Search using Detectron2"
+    description = "<h2>Document Similarity Search using Detectron2<h2>"
+    article = "<h4>More details, Links about this! - Document Similarity Search using Detectron2<h4>"
+    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; }
+  '''
+
+    def preview(image_url):
+        try:
+            return [gr.Tabs(selected=0), get_RGB_image(image_url), None]
+        except:
+            error = traceback.format_exc()
+            return [gr.Tabs(selected=1), None, gr.HTML(value=error, visible=True)]
+
+    with gr.Blocks(title=title, css=css) as app:
+        with gr.Row():
+            gr.HTML(value=description, elem_classes=['center'])
+        with gr.Row():
+            with gr.Column():
+                with gr.Tabs() as tabs:
+                    with gr.Tab("From Image", id=0):
+                        document_image = gr.Image(type="pil", label="Document Image")
+                        submit = gr.Button(value="Submit", variant="primary")
+                        if len(examples) > 0:
+                            gr.Examples(
+                                examples=examples,
+                                inputs=document_image,
+                                label='Select any of these test examples')
+                    with gr.Tab("From URL", id=1):
+                        image_url = gr.Textbox(
+                            label="Document Image Link",
+                            info="Paste a Link to Document Image",
+                            placeholder="https://datasets-server.huggingface.co/assets/ds4sd/icdar2023-doclaynet/--/2023.01/validation/6/image/image.jpg")
+                        error_message = gr.HTML(label="Error Message", visible=False)
+                        preview_btn = gr.Button(value="Preview", variant="primary")
+            with gr.Column():
+                with gr.Group():
+                    annotated_document_image = gr.Image(type="pil", label="Annotated Document Image")
+                    message = gr.HTML(label="Message")
+                    threshold = gr.Slider(0.0, 1.0, value=0.0, label="Threshold", info="Choose between 0.0 and 1.0")
+        with gr.Row():
+            gr.HTML(value=article, elem_classes=['center'])
+        preview_btn.click(preview, [image_url], [tabs, document_image, error_message])
+        submit.click(
+            fn=lambda image: fn(image, just_image=False),
+            inputs=document_image,
+            outputs=[annotated_document_image, message, threshold])
+        threshold.change(
+            fn=lambda image, threshold: fn(image, threshold, just_image=False),
+            inputs=[document_image, threshold],
+            outputs=[annotated_document_image, message])
+    return app.launch
+
+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'}
+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'}
+config_path = './config.yaml'
+model_path = './model_final.pth'
+examples = ['./example.1.jpg', './example.2.jpg', './example.3.jpg']
+
+infer = inference_factory(config_path, model_path, label_map, color_map, examples = examples)
+infer(debug=True)