main.py

# -*- coding: utf-8 -*-
"""Untitled0.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1GJPdFMjb049-ylHyng_M1MaLtF2R9gOa
"""

!apt install ghostscript python3-tk
!apt-get install libmagic-dev
!pip install camelot-py
!pip install python-magic
!pip install pytesseract
!pip install ghostscript
!pip install sqlalchemy
!pip install kraken
!pip install --ignore-installed Pillow==9.0.0
!pip install openpyxl

import magic
import pprint
import cv2
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import csv
import pytesseract
import camelot
import ghostscript
from PIL import Image
import warnings
from openpyxl import Workbook, load_workbook
from openpyxl.utils import get_column_interval
import re
import itertools
 
try:
    from PIL import Image
except ImportError:
    import Image
from sqlalchemy import create_engine, inspect
from kraken.lib.models import load_any
from kraken import rpred, binarization
import random
from subprocess import call
from imutils import contours
import argparse

def detect_file(path):
  f = magic.Magic(mime=True, uncompress=True)
  print(f.from_file(path))
  return f.from_file(path)

def load_workbook_range(range_string, ws):
    col_start, col_end = re.findall("[A-Z]+", range_string)

    data_rows = []
    for row in ws[range_string]:
        data_rows.append([cell.value for cell in row])

    return pd.DataFrame(data_rows, columns=get_column_interval(col_start, col_end))

warnings.filterwarnings("ignore", category=FutureWarning)

## ---Loading Kraken Model---
model = load_any("en-default.mlmodel")


def preprocessing_non_tabular(path):
    img = cv2.imread(path)

    ## ---Binarization of image---
    genrator_image = Image.fromarray(img)
    genrator_image = binarization.nlbin(genrator_image)

    # ----Grayscaling Image----
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    # --- performing Otsu threshold ---
    ret, thresh1 = cv2.threshold(gray, 0, 255, cv2.THRESH_OTSU | cv2.THRESH_BINARY_INV)
    # cv2.imwrite("processed_image/threshold.png", thresh1)
    # cv2.imshow('thresh1', thresh1)
    # cv2.waitKey(0)

    # ----Image dialation----
    rect_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (15, 3))
    dilation = cv2.dilate(thresh1, rect_kernel, iterations=1)
    # cv2.imwrite("processed_image/dilation.png", dilation)
    # cv2.imshow('dilation', dilation)
    # cv2.waitKey(0)

    # ---Finding contours ---
    contours, hierarchy = cv2.findContours(dilation, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
    return img, genrator_image, contours[::-1]


def preprocessing_tabular(path):
    # Load image
    img = cv2.imread(path)

    ## ---Binarization of image---
    genrator_image = Image.fromarray(img)
    genrator_image = binarization.nlbin(genrator_image)

    # ----Grayscaling Image----
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    # --- performing Otsu threshold ---
    thresh = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)[1]

    # Remove text characters with morph open and contour filtering
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
    opening = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, kernel, iterations=1)

    cnts = cv2.findContours(opening, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
    cnts = cnts[0] if len(cnts) == 2 else cnts[1]
    for c in cnts:
        area = cv2.contourArea(c)
        if area < 500:
            cv2.drawContours(opening, [c], -1, (0, 0, 0), -1)

    # Repair table lines, sort contours, and extract ROI
    close = 255 - cv2.morphologyEx(opening, cv2.MORPH_CLOSE, kernel, iterations=1)

    cnts = cv2.findContours(close, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
    cnts = cnts[0] if len(cnts) == 2 else cnts[1]
    cnts, _ = contours.sort_contours(cnts, method="top-to-bottom")
    return img, genrator_image, cnts


def imgtry(path):

    img, genrator_image, cnts = preprocessing_non_tabular(path)
    if len(cnts) < 8:
        img, genrator_image, cnts = preprocessing_tabular(path)

    row_list = list()
    old_y = 0
    single_row = list()
    for idx, cnt in enumerate(cnts):

        x, y, w, h = cv2.boundingRect(cnt)
        if idx == 0:
            single_row.append([x, y, w, h])
        else:
            if y - old_y < 5:
                single_row.append([x, y, w, h])
            else:
                row_list.append(single_row)
                single_row = list()
                single_row.append([x, y, w, h])
        old_y = y

    # color = (np.random.random(size=3) * 256)
    dummy_image = img.copy()

    all_text = list()
    row_count = list()

    for row_boxes in row_list:
        row_text = list()
        for one_box in row_boxes:
            x, y, w, h = one_box
            ## Different color for each row
            r = random.randint(0, 255)
            g = random.randint(0, 255)
            b = random.randint(0, 255)

            # Drawing box
            cv2.rectangle(dummy_image, (x, y), (x + w, y + h), (b, g, r), 2)
            #########################################################################################

            ## Kraken Text Extraction
            cord = [x, y, x + w, y + h]
            bound = {'boxes': [tuple(cord)], 'text_direction': 'horizontal-lr'}

            ## Using Kraken API
            generator = rpred.rpred(network=model, im=genrator_image, bounds=bound)
            nxt_gen = next(generator)
            box_text = nxt_gen.prediction
            print("Box_Text = {} | Y = {}".format(box_text, y))

            ##########################################################################################

            row_text.append(box_text)

        print(row_text)
        row_count.append(len(row_text))
        all_text.append(row_text)
        print("======================================================================")
    print(all_text)

    # cv2.imwrite("processed_image/show_box.png", dummy_image)
    # cv2.imshow('final', dummy_image)
    # cv2.waitKey(0)

    updated_text_rows = list()
    columns = max(set(row_count), key=row_count.count)

    for rows in all_text:
        diff = columns - len(rows)
        rows = rows + [" "] * diff
        updated_text_rows.append(rows)

    # Creating a dataframe of the generated OCR list
    arr = np.array(updated_text_rows)
    dataframe = pd.DataFrame(arr, columns=range(0, columns))
    dataframe.to_csv("output.csv", index=False)

def to_pandasDF(path):
  ft = detect_file(path)
  df = pd.DataFrame()

  if 'xml' in ft:
    # df = pd.read_excel(path)
    pp = pprint.PrettyPrinter(indent=4)
    workbook = load_workbook(filename=path)
    pp.pprint(workbook.sheetnames)
    ws = workbook[input('sheet-name: ')]
    df = load_workbook_range(input("Range: "), ws)
    
  elif 'plain' in ft:
    print('csv-file')
    df = pd.read_csv(path)

    
  elif 'image' in ft:
    # ---Image_Path---
    # path = "images/patient.png"
    imgtry(path)

  elif 'pdf' in ft:
    # pdf tables working
    p_DF = camelot.read_pdf(path, pages="1-end")
    df = p_DF[0].df

  else:
    print('Invalid file format.')

  disk_engine = create_engine('sqlite:///my_lite_store.db')
  # two types of storage sql and csv
  rnd = str(random.randint(0,100))+path.split('.')[0]
  df.to_sql(rnd, disk_engine, if_exists='replace')
  # df.to_csv(rnd+".csv")
  

  return df

# output format to csv or to excel

path = input("Enter filename: ")
DF = to_pandasDF(path)