""" 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 downloads 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")