Merge commit 'bfbc270b655b1db87dfd9e3afaf79d9968042da5' as 'utils_v2'
This commit is contained in:
yatmesh
@@ -0,0 +1,410 @@
|
||||
"""
|
||||
|
||||
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")
|
||||
@@ -0,0 +1,379 @@
|
||||
"""
|
||||
|
||||
AUTHOR:
|
||||
|
||||
Khushal P Soonderji
|
||||
|
||||
DATE:
|
||||
|
||||
Tuesday, 24th 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/
|
||||
03. https://github.com/criistian14/flutter_document_scanner/blob/master/flutter_document_scanner_android/android/src/main/kotlin/com/christian/flutterDocumentScanner/OpenCVPlugin.kt
|
||||
04. https://www.geeksforgeeks.org/python-bilateral-filtering/
|
||||
|
||||
DOWNLOADS:
|
||||
|
||||
N/A
|
||||
|
||||
"""
|
||||
|
||||
|
||||
# *****************************************************************************************************************
|
||||
# ***** ****
|
||||
# *** IMPORT ***
|
||||
# ***** ****
|
||||
# *****************************************************************************************************************
|
||||
|
||||
|
||||
# To make sibling directories accessible for imports:
|
||||
import sys
|
||||
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 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 get_edges(image):
|
||||
|
||||
"""
|
||||
Apply all the preprocessing filters on the image before sending it off for contour-finding.
|
||||
:param image: The image to pre-process.
|
||||
:return: The pre-processed image.
|
||||
"""
|
||||
|
||||
# Denoise the image:
|
||||
processed_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
|
||||
processed_image = cv2.bilateralFilter(
|
||||
processed_image,
|
||||
d = 9,
|
||||
sigmaColor = 75,
|
||||
sigmaSpace = 75
|
||||
)
|
||||
|
||||
# Detect edges in the image:
|
||||
processed_image = cv2.Canny(
|
||||
processed_image,
|
||||
threshold1 = 75,
|
||||
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
|
||||
|
||||
# image = self.enhance(image)
|
||||
processed_image = self.get_edges(image)
|
||||
# self.show(processed_image)
|
||||
|
||||
# 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/sushmita_card.jpg")
|
||||
# image = my_scanner.open_image(r"/home/developer/Downloads/niranjan_card.jpg")
|
||||
# image = my_scanner.open_image(r"/home/developer/Downloads/niranjan_card_2.jpg")
|
||||
# image = my_scanner.open_image(r"/home/developer/Downloads/card_square.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")
|
||||
Reference in New Issue
Block a user