Resetting utils subtree.

This commit is contained in:
2025-01-08 18:56:05 +05:30
parent 81fc28295f
commit 278951e8de
189 changed files with 80 additions and 142029 deletions
-379
View File
@@ -1,379 +0,0 @@
"""
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")