File size: 2,954 Bytes
cebb0c1 0599d82 cebb0c1 4c41a36 cebb0c1 0599d82 cebb0c1 4c41a36 0599d82 7361420 0599d82 4c41a36 cebb0c1 efed173 0599d82 efed173 cebb0c1 0599d82 cebb0c1 0599d82 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 |
import asyncio
import string
from collections import Counter
from itertools import count, tee
import cv2
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import streamlit as st
import torch
from PIL import Image
from transformers import (DetrImageProcessor, TableTransformerForObjectDetection)
from vietocr.tool.config import Cfg
from vietocr.tool.predictor import Predictor
st.set_option('deprecation.showPyplotGlobalUse', False)
st.set_page_config(layout='wide')
st.title("Table Detection and Table Structure Recognition")
st.write(
"Implemented by MSFT team: https://github.com/microsoft/table-transformer")
# Config (optional, comment out if not using)
# config = Cfg.load_config_from_name('vgg_transformer')
# config = Cfg.load_config_from_name('vgg_seq2seq')
# config['cnn']['pretrained'] = False
# config['device'] = 'cpu'
# config['predictor']['beamsearch'] = False
# detector = Predictor(config)
table_detection_model = TableTransformerForObjectDetection.from_pretrained(
"microsoft/table-transformer-detection")
table_recognition_model = TableTransformerForObjectDetection.from_pretrained(
"microsoft/table-transformer-structure-recognition")
def PIL_to_cv(pil_img):
return cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)
def cv_to_PIL(cv_img):
return Image.fromarray(cv2.cvtColor(cv_img, cv2.COLOR_BGR2RGB))
async def pytess(cell_pil_img, threshold: float = 0.5):
text, prob = detector.predict(cell_pil_img, return_prob=True) # Assuming detector is defined
if prob < threshold:
return ""
return text.strip()
def sharpen_image(pil_img):
img = PIL_to_cv(pil_img)
sharpen_kernel = np.array([[-1, -1, -1], [-1, 9, -1], [-1, -1, -1]])
sharpen = cv2.filter2D(img, -1, sharpen_kernel)
pil_img = cv_to_PIL(sharpen)
return pil_img
def uniquify(seq, suffs=count(1)):
"""Make all the items unique by adding a suffix (1, 2, etc).
Credit: https://stackoverflow.com/questions/30650474/python-rename-duplicates-in-list-with-progressive-numbers-without-sorting-list
"""
not_unique = [k for k, v in Counter(seq).items() if v > 1]
suff_gens = dict(zip(not_unique, tee(suffs, len(not_unique))))
for idx, s in enumerate(seq):
try:
suffix = str(next(suff_gens[s]))
except KeyError:
continue
else:
seq[idx] += suffix
return seq
def binarizeBlur_image(pil_img):
image = PIL_to_cv(pil_img)
thresh = cv2.threshold(image, 150, 255, cv2.THRESH_BINARY_INV)[1]
result = cv2.GaussianBlur(thresh, (5, 5), 0)
result = 255 - result
return cv_to_PIL(result)
def td_postprocess(pil_img):
'''
Removes gray background from tables
'''
img = PIL_to_cv(pil_img)
hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
mask = cv2.inRange(hsv, (0, 0, 100), (255, 5, 255)) # (0, 0, 100), (255, 5, 255)
nzmask = cv2.inRange(hsv, (0, 0, 5), (255, 255, 255)) # (0, 0, |