Upload 2 files

app.py

@@ -0,0 +1,142 @@
+import gradio as gr
+import torch
+from datasets import load_dataset, ClassLabel
+import os
+from transformers import LayoutLMv3ForTokenClassification, LayoutLMv3Processor,LayoutLMv3FeatureExtractor
+import pytesseract
+import numpy as np
+from PIL import ImageDraw, ImageFont
+
+os.system('pip3 install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cpu')
+os.system('sudo apt-get install tesseract-ocr')
+os.system('pip install -q pytesseract')
+print("pytesseract:",pytesseract.__version__)
+
+examples = [['./examples/example1.png'],['./examples/example2.png'],['./examples/example3.png']]
+dataset = load_dataset("nielsr/cord-layoutlmv3")['train']
+
+def get_label_list(labels):
+    unique_labels = set()
+    for label in labels:
+        unique_labels = unique_labels | set(label)
+    label_list = list(unique_labels)
+    label_list.sort()
+    return label_list
+
+def convert_l2n_n2l(dataset):
+    features = dataset.features
+    label_column_name = "ner_tags"
+
+    label_list = features[label_column_name].feature.names
+    if isinstance(features[label_column_name].feature, ClassLabel):
+        id2label = {k:v for k,v in enumerate(label_list)}
+        label2id = {v:k for k,v in enumerate(label_list)}
+    else:
+        label_list = get_label_list(dataset[label_column_name])
+        id2label = {k:v for k,v in enumerate(label_list)}
+        label2id = {v:k for k,v in enumerate(label_list)}
+
+    return label_list, id2label, label2id, len(label_list)
+
+def label_colour(label):
+    label2color = {'MENU.PRICE':'blue', 'MENU.NM':'green', 'other':'green','MENU.TOTAL_PRICE':'red'}
+    if label in label2color:
+        colour = label2color.get(label)
+    else:
+        colour = None
+    return colour
+
+def iob_to_label(label):
+    label = label[2:]
+    if not label:
+        return 'other'
+    return label
+
+def convert_results(words,tags):
+    ents = set()
+    completeword = ""
+    for word, tag in zip(words, tags):
+        if tag != "O":
+            ent_position, ent_type = tag.split("-")
+            if ent_position == "S":
+                ents.add((word,ent_type))
+            else:
+                if ent_position == "B":
+                    completeword = completeword+ " "+ word
+                elif ent_position == "I":
+                    completeword= completeword+ " " + word
+                elif ent_position == "E":
+                    completeword =completeword+" " + word
+                
+                ents.add((completeword,ent_type))
+                completeword= ""
+    return ents
+
+def unnormalize_box(bbox, width, height):
+     return [
+         width * (bbox[0] / 1000),
+         height * (bbox[1] / 1000),
+         width * (bbox[2] / 1000),
+         height * (bbox[3] / 1000),
+     ]
+
+def predict(image):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = LayoutLMv3ForTokenClassification.from_pretrained("keldrenloy/layoutlmv3cordfinetuned").to(device) #add your model directory here
+    processor = LayoutLMv3Processor.from_pretrained("microsoft/layoutlmv3-base")
+    label_list,id2label,label2id, num_labels = convert_l2n_n2l(dataset)
+    width, height = image.size
+    
+    encoding_inputs = processor(image,return_offsets_mapping=True, return_tensors="pt",truncation = True)
+    offset_mapping = encoding_inputs.pop('offset_mapping')
+    for k,v in encoding_inputs.items():
+        encoding_inputs[k] = v.to(device)
+
+    with torch.no_grad():
+        outputs = model(**encoding_inputs)
+    
+    predictions = outputs.logits.argmax(-1).squeeze().tolist()
+    token_boxes = encoding_inputs.bbox.squeeze().tolist()
+    
+    is_subword = np.array(offset_mapping.squeeze().tolist())[:,0] != 0
+    true_predictions = [id2label[pred] for idx, pred in enumerate(predictions) if not is_subword[idx]]
+    true_boxes = [unnormalize_box(box, width, height) for idx, box in enumerate(token_boxes) if not is_subword[idx]]
+    
+    return true_boxes, true_predictions
+    
+def text_extraction(image):
+    feature_extractor = LayoutLMv3FeatureExtractor()
+    encoding = feature_extractor(image, return_tensors="pt")
+    return encoding['words'][0]
+
+def image_render(image):
+    draw = ImageDraw.Draw(image)
+    font = ImageFont.load_default()
+    true_boxes,true_predictions = predict(image)
+        
+    for prediction, box in zip(true_predictions, true_boxes):
+        predicted_label = iob_to_label(prediction)
+        draw.rectangle(box, outline=label_colour(predicted_label))
+        draw.text((box[0]+10, box[1]-10), text=predicted_label, fill=label_colour(predicted_label), font=font)
+    
+    words = text_extraction(image)
+    print(words)
+    extracted_words = convert_results(words,true_predictions)
+    
+    return image,extracted_words
+
+css = """.output_image, .input_image {height: 600px !important}"""
+
+demo = gr.Interface(fn = image_render,
+                    inputs = gr.inputs.Image(type="pil"),
+                    outputs = [gr.outputs.Image(type="pil", label="annotated image"),'text'],
+                    css = css,
+                    examples = examples,
+                    allow_flagging=True,
+                    flagging_options=["incorrect", "correct"],
+                    flagging_callback = gr.CSVLogger(),             
+                    flagging_dir = "flagged"
+                   )
+
+if __name__ == "__main__":
+    demo.launch()

requirement.txt

@@ -0,0 +1,263 @@
+absl-py==1.2.0
+accelerate==0.12.0
+aiohttp==3.8.1
+aiosignal==1.2.0
+alembic==1.8.1
+analytics-python==1.4.0
+anyio==3.6.1
+appdirs==1.4.4
+argon2-cffi==21.3.0
+argon2-cffi-bindings==21.2.0
+asgiref==3.5.2
+asttokens==2.0.8
+async-timeout==4.0.2
+attr==0.3.1
+attrs==22.1.0
+azure-core==1.25.1
+azure-storage-blob==12.13.1
+backcall==0.2.0
+backoff==1.10.0
+bcrypt==4.0.0
+beautifulsoup4==4.11.1
+bleach==5.0.1
+boto==2.49.0
+boto3==1.16.63
+botocore==1.19.63
+boxing==0.1.4
+cachetools==4.2.4
+certifi==2022.6.15.1
+cffi==1.15.1
+charset-normalizer==2.0.12
+click==8.1.3
+cloudpickle==2.2.0
+colorama==0.4.5
+contourpy==1.0.5
+coreapi==2.3.3
+coreschema==0.0.4
+cryptography==38.0.1
+cuda-python==11.7.1
+cycler==0.11.0
+Cython==0.29.32
+databricks-cli==0.17.3
+datasets==2.4.0
+debugpy==1.6.3
+decorator==5.1.1
+defusedxml==0.7.1
+Deprecated==1.2.13
+dill==0.3.5.1
+Django==3.1.14
+django-annoying==0.10.6
+django-cors-headers==3.6.0
+django-debug-toolbar==3.2.1
+django-extensions==3.1.0
+django-filter==2.4.0
+django-model-utils==4.1.1
+django-ranged-fileresponse==0.1.2
+django-rest-swagger==2.2.0
+django-rq==2.5.1
+django-user-agents==0.4.0
+djangorestframework==3.13.1
+docker==6.0.0
+docker-pycreds==0.4.0
+docopt==0.6.2
+drf-dynamic-fields==0.3.0
+drf-flex-fields==0.9.5
+drf-generators==0.3.0
+drf-yasg==1.20.0
+entrypoints==0.4
+executing==1.0.0
+expiringdict==1.1.4
+fastapi==0.85.0
+fastjsonschema==2.16.2
+ffmpy==0.3.0
+filelock==3.8.0
+Flask==2.2.2
+fonttools==4.37.3
+frozenlist==1.3.1
+fsspec==2022.8.2
+gitdb==4.0.9
+GitPython==3.1.27
+google-api-core==1.31.5
+google-auth==1.35.0
+google-auth-oauthlib==0.4.6
+google-cloud-appengine-logging==1.1.0
+google-cloud-audit-log==0.2.0
+google-cloud-core==1.5.0
+google-cloud-logging==2.7.2
+google-cloud-storage==1.29.0
+google-resumable-media==0.5.1
+googleapis-common-protos==1.52.0
+gradio==3.3.1
+greenlet==1.1.3
+grpc-google-iam-v1==0.12.3
+grpcio==1.48.1
+h11==0.12.0
+htmlmin==0.1.12
+httpcore==0.15.0
+httpx==0.23.0
+huggingface-hub==0.9.1
+idna==3.3
+importlib-metadata==4.12.0
+inflection==0.5.1
+ipykernel==6.15.2
+ipython==8.5.0
+ipython-genutils==0.2.0
+ipywidgets==8.0.2
+isodate==0.6.1
+itsdangerous==2.1.2
+itypes==1.2.0
+jedi==0.18.1
+Jinja2==3.1.2
+jmespath==0.10.0
+joblib==1.2.0
+jsonschema==3.2.0
+jupyter-core==4.11.1
+jupyter_client==7.3.5
+jupyterlab-pygments==0.2.2
+jupyterlab-widgets==3.0.3
+kiwisolver==1.4.4
+label-studio==1.5.0.post0
+label-studio-converter==0.0.40
+label-studio-tools==0.0.0.dev14
+launchdarkly-server-sdk==7.3.0
+linkify-it-py==1.0.3
+lockfile==0.12.2
+lxml==4.9.1
+Mako==1.2.3
+Markdown==3.4.1
+markdown-it-py==2.1.0
+MarkupSafe==2.1.1
+matplotlib==3.6.0
+matplotlib-inline==0.1.6
+mdit-py-plugins==0.3.0
+mdurl==0.1.2
+mistune==2.0.4
+mlflow==1.29.0
+monotonic==1.6
+msrest==0.7.1
+multidict==6.0.2
+multiprocess==0.70.13
+nbclient==0.6.8
+nbconvert==7.0.0
+nbformat==5.6.0
+nest-asyncio==1.5.5
+nltk==3.6.7
+notebook==6.4.12
+numpy==1.23.3
+nvidia-ml-py3==7.352.0
+oauthlib==3.2.1
+openapi-codec==1.3.2
+ordered-set==4.0.2
+orjson==3.8.0
+packaging==21.3
+pandas==1.3.5
+pandocfilters==1.5.0
+paramiko==2.11.0
+parso==0.8.3
+pathtools==0.1.2
+pickleshare==0.7.5
+Pillow==9.0.1
+pipreqs==0.4.11
+prometheus-client==0.14.1
+prometheus-flask-exporter==0.20.3
+promise==2.3
+prompt-toolkit==3.0.31
+proto-plus==1.22.1
+protobuf==3.19.4
+psutil==5.9.2
+psycopg2-binary==2.9.1
+pure-eval==0.2.2
+pyarrow==9.0.0
+pyasn1==0.4.8
+pyasn1-modules==0.2.8
+pycparser==2.21
+pycryptodome==3.15.0
+pydantic==1.8.2
+pyDeprecate==0.3.2
+pydub==0.25.1
+Pygments==2.13.0
+PyJWT==2.5.0
+PyNaCl==1.5.0
+pyngrok==5.1.0
+pyparsing==3.0.9
+pyRFC3339==1.1
+pyrsistent==0.18.1
+pytesseract==0.3.10
+python-dateutil==2.8.2
+python-multipart==0.0.5
+pytorch-lightning==1.7.5
+pytz==2019.3
+pywin32==304
+pywinpty==2.0.8
+PyYAML==6.0
+pyzmq==23.2.1
+querystring-parser==1.2.4
+redis==4.3.4
+regex==2022.9.11
+requests==2.27.1
+requests-oauthlib==1.3.1
+responses==0.18.0
+rfc3986==1.5.0
+rq==1.10.1
+rsa==4.9
+ruamel.yaml==0.17.21
+ruamel.yaml.clib==0.2.6
+rules==2.2
+s3transfer==0.3.7
+scikit-learn==1.1.2
+scipy==1.9.1
+semver==2.13.0
+Send2Trash==1.8.0
+sentry-sdk==1.9.8
+seqeval==1.2.2
+setproctitle==1.3.2
+shortuuid==1.0.9
+simplejson==3.17.6
+six==1.16.0
+smmap==5.0.0
+sniffio==1.3.0
+soupsieve==2.3.2.post1
+SQLAlchemy==1.4.41
+sqlparse==0.4.2
+stack-data==0.5.0
+starlette==0.20.4
+tabulate==0.8.10
+tensorboard==2.10.0
+tensorboard-data-server==0.6.1
+tensorboard-plugin-wit==1.8.1
+terminado==0.15.0
+tesseract==0.1.3
+threadpoolctl==3.1.0
+tinycss2==1.1.1
+tokenizers==0.12.1
+torch==1.12.1+cu113
+torchaudio==0.12.1+cu113
+torchmetrics==0.9.3
+torchvision==0.13.1+cu113
+tornado==6.2
+tqdm==4.64.1
+traitlets==5.3.0
+transformers==4.21.3
+typing_extensions==4.3.0
+tzdata==2022.2
+ua-parser==0.16.1
+uc-micro-py==1.0.1
+ujson==5.5.0
+uritemplate==4.1.1
+urllib3==1.26.12
+user-agents==2.2.0
+uvicorn==0.18.3
+waitress==2.1.2
+wandb==0.13.3
+wcwidth==0.2.5
+webencodings==0.5.1
+websocket-client==1.4.1
+websockets==10.3
+Werkzeug==2.2.2
+widgetsnbextension==4.0.3
+wrapt==1.14.1
+xmljson==0.2.0
+xxhash==3.0.0
+yarg==0.1.9
+yarl==1.8.1
+zipp==3.8.1