khushalps
cf354c722a
git-subtree-dir: utils_v2 git-subtree-split: 3be5145c7a4cfede04d753324dfae31ace913c98
411 lines
16 KiB
Python
411 lines
16 KiB
Python
"""
|
|
|
|
AUTHOR:
|
|
|
|
Khushal P Soonderji
|
|
|
|
DATE:
|
|
|
|
Thursday, 19th Sept., 2024
|
|
|
|
OBJECTIVE:
|
|
|
|
To provide a way to easily scan documents like printouts and visiting cards.
|
|
|
|
REFERENCES:
|
|
|
|
01. https://medium.com/@victorolufemi/build-a-document-scanner-with-opencv-ff9f645a4085
|
|
02. https://github.com/JaidedAI/EasyOCR
|
|
|
|
DOWNLOADS:
|
|
|
|
N/A
|
|
|
|
"""
|
|
|
|
|
|
# *****************************************************************************************************************
|
|
# ***** ****
|
|
# *** IMPORT ***
|
|
# ***** ****
|
|
# *****************************************************************************************************************
|
|
|
|
|
|
# To make sibling directories accessible for imports:
|
|
import sys
|
|
|
|
from diffusers.utils.import_utils import candidates
|
|
|
|
sys.path.append(".")
|
|
sys.path.append("..")
|
|
|
|
# System-level activities:
|
|
import io
|
|
|
|
# To work with image processing:
|
|
import cv2
|
|
import imutils
|
|
import numpy as np
|
|
|
|
# To run OCR:
|
|
import easyocr
|
|
|
|
# To download images from the web:
|
|
import requests
|
|
|
|
# My utils:
|
|
from utils_v2.string import json
|
|
from utils_v2.ai.object_detection.yolo import YoloDetect
|
|
|
|
|
|
# *****************************************************************************************************************
|
|
# ***** ****
|
|
# *** MACROS / ONE-TIME INIT ***
|
|
# ***** ****
|
|
# *****************************************************************************************************************
|
|
|
|
|
|
# --- Nothing Yet
|
|
|
|
|
|
# *****************************************************************************************************************
|
|
# ***** ****
|
|
# *** VARIABLES ***
|
|
# ***** ****
|
|
# *****************************************************************************************************************
|
|
|
|
|
|
# --- Nothing Yet
|
|
|
|
|
|
# *****************************************************************************************************************
|
|
# ***** ****
|
|
# *** CLASSES ***
|
|
# ***** ****
|
|
# *****************************************************************************************************************
|
|
|
|
|
|
class DocumentScanner:
|
|
|
|
def __init__(
|
|
self,
|
|
layout_detection_yolo = None,
|
|
whitelisted_yolo_classes = None,
|
|
ocr_languages = None
|
|
):
|
|
|
|
self.__yolo = YoloDetect(
|
|
model_file = layout_detection_yolo,
|
|
debug = False
|
|
)
|
|
self.__whitelisted_yolo_classes = whitelisted_yolo_classes or []
|
|
self.__ocr_engine = easyocr.Reader(ocr_languages or ["en"])
|
|
pass
|
|
|
|
@staticmethod
|
|
def open_image(image):
|
|
|
|
"""
|
|
Opens an image from various sources.
|
|
:param image: The image either as a path to a file on the local disk, or a URL, or an io.BytesIO buffer.
|
|
:return: The image as a CV2 object.
|
|
"""
|
|
|
|
# When the image is provided as a string,
|
|
# it could be either as a URL or a path to a local file:
|
|
if isinstance(image, str):
|
|
if image.startswith("https://") or image.startswith("http://"):
|
|
image = io.BytesIO(requests.get(image).content)
|
|
else: image = cv2.imread(image)
|
|
|
|
# If the image is provided as a io.BytesIO buffer:
|
|
if isinstance(image, io.BytesIO):
|
|
image.seek(0)
|
|
image = np.asarray(bytearray(image.read()), dtype = np.uint8)
|
|
image = cv2.imdecode(image, cv2.IMREAD_COLOR)
|
|
|
|
# Done here:
|
|
return image
|
|
|
|
@staticmethod
|
|
def polygon_area(points):
|
|
|
|
"""
|
|
Computes the area occupied by a shape described by the array of points.
|
|
Example input: [[x1, y1], [x2, y2], [x3, y3], [x4, y4]]
|
|
:return: The area of the shape.
|
|
"""
|
|
|
|
n = len(points)
|
|
area = 0.0
|
|
for i in range(n):
|
|
j = (i + 1) % n # next vertex
|
|
area += points[i][0] * points[j][1]
|
|
area -= points[j][0] * points[i][1]
|
|
return abs(area) / 2.0
|
|
|
|
@staticmethod
|
|
def show(image, title = "Preview", wait = True):
|
|
|
|
"""
|
|
Just a quick wrapper to show the image in a window.
|
|
:param image: The image that you want to show.
|
|
:param title: The title of the window.
|
|
:param wait: Set this to True when you want to show the window(s). This is useful when you want to show many
|
|
windows at once. Suppose you want to show a lot of windows, you set this to False for all calls except the
|
|
very last one.
|
|
:return: None.
|
|
"""
|
|
|
|
cv2.imshow(title, image)
|
|
if wait:
|
|
cv2.waitKey()
|
|
cv2.destroyAllWindows()
|
|
|
|
@staticmethod
|
|
def enhance(image):
|
|
|
|
"""
|
|
Applies a form of contrast boost to make the edges more easily visible.
|
|
:param image: The OpenCV image that needs to be enhanced.
|
|
:return: The enhanced CV2 image.
|
|
"""
|
|
|
|
return cv2.detailEnhance(
|
|
src = image,
|
|
sigma_s = 20,
|
|
sigma_r = 0.15
|
|
)
|
|
|
|
@staticmethod
|
|
def get_edges(image):
|
|
|
|
"""
|
|
Gets the edges in an image.
|
|
:param image: A CV2 image.
|
|
:return: The image with the edges detected.
|
|
"""
|
|
|
|
# Denoise the image:
|
|
processed_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
|
|
processed_image = cv2.GaussianBlur(
|
|
processed_image,
|
|
ksize = (5, 5),
|
|
sigmaX = 0
|
|
)
|
|
|
|
# Detect edges in the image:
|
|
processed_image = cv2.Canny(
|
|
processed_image,
|
|
threshold1 = 50,
|
|
threshold2 = 200
|
|
)
|
|
|
|
# Close small gaps in the edges:
|
|
kernel = np.ones((5, 5), np.uint8)
|
|
processed_image = cv2.dilate(processed_image, kernel, iterations = 1)
|
|
# processed_image = cv2.morphologyEx(processed_image, cv2.MORPH_CLOSE, kernel)
|
|
# processed_image = cv2.erode(processed_image, kernel, iterations = 1)
|
|
|
|
# Done here:
|
|
return processed_image
|
|
|
|
def scan(self, image, min_area = 0.125):
|
|
|
|
"""
|
|
Looks for a rectangle in an image and flattens it out. No colour filters are applied here.
|
|
:param image: The image either as a path to a file on the local disk, or a URL, or an io.BytesIO buffer.
|
|
:param min_area: The minimum area occupied by the document in the image.
|
|
:return: The flattened image if any, else None.
|
|
"""
|
|
|
|
# Open the image and read the data:
|
|
image = self.open_image(image)
|
|
if image is None: return None
|
|
height, width = image.shape[:2]
|
|
image_area = height * width
|
|
|
|
# Enhance the image to have better visibility of edges:
|
|
# processed_image = self.enhance(image.copy())
|
|
|
|
# Detect edges in the image:
|
|
processed_image = self.get_edges(image.copy())
|
|
|
|
# Find the contours in the edges,
|
|
# and sort them in ascending order:
|
|
contours = cv2.findContours(processed_image, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
|
|
contours = imutils.grab_contours(contours)
|
|
contours = sorted(contours, key = cv2.contourArea, reverse = True)
|
|
|
|
# We try to get clean and passable contours.
|
|
# For our purposes, 'clean' is when it can be approximated to exactly 4 sides,
|
|
# and 'passable' is when the approximation comes close to it:
|
|
clean_rects = []
|
|
passable_rects = []
|
|
for contour in contours:
|
|
perimeter = cv2.arcLength(contour, closed = True)
|
|
approximation = cv2.approxPolyDP(contour, 0.025 * perimeter, closed = True)
|
|
if len(approximation) == 4: clean_rects.append(approximation)
|
|
elif len(approximation) <= 32:
|
|
candidate_contour = cv2.minAreaRect(approximation)
|
|
candidate_contour = cv2.boxPoints(candidate_contour)
|
|
candidate_contour = np.intp(candidate_contour)
|
|
passable_rects.append(candidate_contour)
|
|
clean_rects = sorted(clean_rects, key = cv2.contourArea, reverse = True)
|
|
passable_rects = sorted(passable_rects, key = cv2.contourArea, reverse = True)
|
|
|
|
# Now we decide between the best candidate:
|
|
document_outline = None
|
|
if len(clean_rects) > 0:
|
|
candidate_rect = np.array([p[0] for p in clean_rects[0]])
|
|
rectangle_area = cv2.contourArea(candidate_rect)
|
|
if rectangle_area / image_area >= min_area: document_outline = candidate_rect
|
|
if document_outline is None and len(passable_rects) > 0:
|
|
candidate_rect = cv2.minAreaRect(passable_rects[0])
|
|
candidate_rect = cv2.boxPoints(candidate_rect)
|
|
candidate_rect = np.intp(candidate_rect)
|
|
rectangle_area = cv2.contourArea(candidate_rect)
|
|
if rectangle_area / image_area >= min_area: document_outline = candidate_rect
|
|
|
|
# If there is no rectangular contour,
|
|
# we exit with failure:
|
|
if document_outline is None: return None
|
|
|
|
# Since we did get the best candidate for the document,
|
|
# we figure out which point n the contour is which corner of the document:
|
|
document_rectangle = np.zeros((4, 2), dtype = "float32")
|
|
sum_points = document_outline.sum(axis = 1)
|
|
document_rectangle[0] = document_outline[np.argmin(sum_points)]
|
|
document_rectangle[2] = document_outline[np.argmax(sum_points)]
|
|
diff_points = np.diff(document_outline, axis = 1)
|
|
document_rectangle[1] = document_outline[np.argmin(diff_points)]
|
|
document_rectangle[3] = document_outline[np.argmax(diff_points)]
|
|
(top_left, top_right, bottom_right, bottom_left) = document_rectangle
|
|
|
|
# let's compute the dimensions of the document in the image:
|
|
left_height = np.sqrt(((top_left[0] - bottom_left[0]) ** 2) + ((top_left[1] - bottom_left[1]) ** 2))
|
|
right_height = np.sqrt(((top_right[0] - bottom_right[0]) ** 2) + ((top_right[1] - bottom_right[1]) ** 2))
|
|
top_width = np.sqrt(((top_right[0] - top_left[0]) ** 2) + ((top_right[1] - top_left[1]) ** 2))
|
|
bottom_width = np.sqrt(((bottom_right[0] - bottom_left[0]) ** 2) + ((bottom_right[1] - bottom_left[1]) ** 2))
|
|
max_height = max(int(left_height), int(right_height))
|
|
max_width = max(int(top_width), int(bottom_width))
|
|
|
|
# We compute the destination of the transform:
|
|
destination = np.array(
|
|
object = [
|
|
[0, 0],
|
|
[max_width - 1, 0],
|
|
[max_width - 1, max_height - 1],
|
|
[0, max_height - 1]
|
|
],
|
|
dtype = "float32"
|
|
)
|
|
|
|
# We apply the transform that flattens out the document:
|
|
matrix = cv2.getPerspectiveTransform(document_rectangle, destination)
|
|
flattened_image = cv2.warpPerspective(image, matrix, dsize = (max_width, max_height))
|
|
|
|
# Done here:
|
|
return flattened_image
|
|
|
|
def extract_text(
|
|
self,
|
|
image,
|
|
min_confidence = 0.5,
|
|
margin = 0.05
|
|
):
|
|
|
|
"""
|
|
To run OCR on the input image.
|
|
:param image: The image either as a path to a file on the local disk, or a URL, or an io.BytesIO buffer.
|
|
:param min_confidence: The minimum amount of confidence for detected text to be considered.
|
|
:param margin: The extra pixels to include when cropping into a section for OCR.
|
|
:return: The extracted text.
|
|
"""
|
|
|
|
# Make a variable that will hold the final result:
|
|
extracted_text = {}
|
|
|
|
# Open the image and read the data:
|
|
image = self.open_image(image)
|
|
height, width = image.shape[:2]
|
|
|
|
# Run OCR on the entire canvas:
|
|
ocr_result = [r[1] for r in self.__ocr_engine.readtext(image) if r[-1] >= min_confidence]
|
|
extracted_text["fullDoc"] = " ".join(ocr_result)
|
|
|
|
# Get the doc layout in the image:
|
|
document_sections = self.__yolo.predict(image)
|
|
document_sections = document_sections["boxes"]
|
|
|
|
# For each section, we run the OCR process individually:
|
|
section_results = []
|
|
for section in document_sections:
|
|
if section["class"] in self.__whitelisted_yolo_classes:
|
|
section_width = section["x2"] - section["x1"]
|
|
section_height = section["y2"] - section["y1"]
|
|
x1 = max(int(section["x1"] - (margin * section_width)), 0)
|
|
y1 = max(int(section["y1"] - (margin * section_height)), 0)
|
|
x2 = min(int(section["x2"] + (margin * section_width)), width)
|
|
y2 = min(int(section["y2"] + (margin * section_height)), height)
|
|
sub_image = image[y1:y2, x1:x2]
|
|
ocr_result = [r[1] for r in self.__ocr_engine.readtext(sub_image) if r[-1] >= min_confidence]
|
|
section_results.append({
|
|
"text": " ".join(ocr_result),
|
|
"class": section["class"],
|
|
"className": section["className"],
|
|
"x1": x1,
|
|
"y1": y1,
|
|
"x2": x2,
|
|
"y2": y2
|
|
})
|
|
extracted_text["bySection"] = section_results
|
|
|
|
# Done here:
|
|
return extracted_text
|
|
|
|
|
|
# *****************************************************************************************************************
|
|
# ***** ****
|
|
# *** FUNCTIONS ***
|
|
# ***** ****
|
|
# *****************************************************************************************************************
|
|
|
|
|
|
# --- Nothing Yet
|
|
|
|
|
|
# *****************************************************************************************************************
|
|
# ***** ****
|
|
# *** MAIN PROGRAM ***
|
|
# ***** ****
|
|
# *****************************************************************************************************************
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
|
import time
|
|
|
|
my_scanner = DocumentScanner(
|
|
layout_detection_yolo = r"/home/developer/PycharmProjects/utils/data/ai/behaviour_models/hugging_face/object_detection/YOLOv10-Document-Layout-Analysis/yolov10x_best.pt",
|
|
whitelisted_yolo_classes = [0, 1, 3, 4, 5, 7, 9, 10],
|
|
# whitelisted_yolo_classes = [0, 1, 3, 4, 5, 6, 7, 8, 9, 10],
|
|
ocr_languages = ["en"]
|
|
)
|
|
|
|
image = my_scanner.open_image(r"/home/developer/Downloads/talkaholics_card.jpg")
|
|
document_image = my_scanner.scan(image)
|
|
if document_image is not None:
|
|
cv2.imwrite(r"/home/developer/Downloads/flattened_image.jpg", document_image)
|
|
start_time = time.time()
|
|
document_text = my_scanner.extract_text(
|
|
document_image,
|
|
min_confidence = 0.5
|
|
)
|
|
print("OCR RESULT:")
|
|
print(document_text)
|
|
|
|
print(f"FINISHED IN {time.time() - start_time} SECONDS!")
|
|
|
|
else: print("No image")
|