55 lines
2.2 KiB
Python
55 lines
2.2 KiB
Python
import os
|
|
import sys
|
|
import numpy as np
|
|
import pytesseract
|
|
import cv2
|
|
from PIL import Image
|
|
|
|
def process(image):
|
|
# Read image
|
|
img = cv2.imread(image)
|
|
# Convert to single-channel grayscale
|
|
img_gs = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
|
|
# Binary threshold to get rid of blemishes
|
|
_, img_gs = cv2.threshold(img_gs, 50, 255, cv2.THRESH_BINARY)
|
|
# OTSU threshold for contour detection
|
|
_, thr = cv2.threshold(img_gs, 0, 255, cv2.THRESH_OTSU | cv2.THRESH_BINARY_INV)
|
|
# Rectangular kernel for contour detection
|
|
rect_k = cv2.getStructuringElement(cv2.MORPH_RECT, (32, 32))
|
|
# Dilation
|
|
dil = cv2.dilate(thr, rect_k, iterations = 1)
|
|
# Finding contours
|
|
conts, _ = cv2.findContours(dil, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
|
|
txt=""
|
|
# Iterate over found contours and attemt to ocr text within
|
|
for cnt in conts:
|
|
x, y, w, h = cv2.boundingRect(cnt) # location and size of found feature
|
|
if x == 0 or y == 0:
|
|
# Do not do anything if the feature is on image edges, expected text is more or less centered
|
|
pass
|
|
else:
|
|
cv2.rectangle(img, (x, y), (x + w, y + h), (0, 255, 0), 2)
|
|
# Crop only the text feature with some padding in horizontal direction to help with ocr
|
|
txt_ft = img_gs[y:y+h, x:x+w]
|
|
try:
|
|
# Write image with the detected text bounding box
|
|
cv2.imwrite("output/{}_tmp.png".format(image.split("/")[1].split(".")[0]),img)
|
|
# Apply blur to remove sharp edges and help with ocr
|
|
txt_ft = cv2.GaussianBlur(txt_ft,(5,5),cv2.BORDER_DEFAULT)
|
|
# Actual ocr in one call
|
|
txt += pytesseract.image_to_string(txt_ft)
|
|
except:
|
|
# Fail silently, only the emtpy string will be returned
|
|
pass
|
|
return txt
|
|
|
|
if __name__ == "__main__":
|
|
if not os.path.exists("input"):
|
|
sys.exit(1)
|
|
if not os.path.exists("output"):
|
|
os.mkdir("output")
|
|
for item in ["input/"+file for file in os.listdir("input") if file.split(".")[1] in ["jpg", "jpeg", "png"]]:
|
|
print(item)
|
|
txt = process(item)
|
|
print(txt.strip())
|