endoreg-db 0.3.6__py3-none-any.whl → 0.3.8__py3-none-any.whl

endoreg_db/utils/legacy_ocr.py

@@ -1,201 +1,201 @@
-import pytesseract
-# import cv2
-from endoreg_db.models import EndoscopyProcessor
-import os
-from collections import Counter
-from tempfile import TemporaryDirectory
-import re
-from datetime import datetime
-from typing import Dict, List
-from icecream import ic
-import numpy as np
-
-N_FRAMES_MEAN_OCR = 2
-
-# Helper function to process date strings
-def process_date_text(date_text):
-    """
-    Processes a string of text that represents a date and returns a datetime.date object.
-
-    Args:
-        date_text (str): A string of text that represents a date.
-
-    Returns:
-        datetime.date: A datetime.date object representing the parsed date, or None if the text cannot be parsed.
-    """
-    try:
-        # Remove any non-digit characters
-        date_text_clean = re.sub(r'\D', '', date_text)
-        # Reformat to 'ddmmyyyy' if necessary
-        if len(date_text_clean) == 8:
-            return datetime.strptime(date_text_clean, "%d%m%Y").date()
-        elif len(date_text_clean) == 14:
-            return datetime.strptime(date_text_clean, "%d%m%Y%H%M%S").date()
-    except ValueError:
-        # Return None if the text cannot be parsed into a date
-        # set date to 1/1/1900
-        return datetime.strptime("01011900", "%d%m%Y").date()
-
-# Helper function to process patient names
-def process_name_text(name_text):
-    """
-    Remove all numbers, punctuation, and whitespace from a string of text and return the result.
-    """
-    name = re.sub(r'[0-9!"#$%&\'()*+,-./:;<=>?@[\\]^_`{|}~\s]+', '', name_text).strip()
-    # capitalize first letter of each word
-    name = ' '.join([word.capitalize() for word in name.split()])
-    return name
-
-
-# Helper function to process endoscope type text
-def process_general_text(endoscope_text):
-    """
-    This function takes in a string of text from an endoscope and returns a cleaned version of the text.
-    """
-    return ' '.join(endoscope_text.split())
-
-def roi_values_valid(roi):
-    """
-    Check if all values in an ROI dictionary are valid (>=0).
-    """
-    return all([value >= 0 for value in roi.values()])
-
-# Function to extract text from ROIs
-def extract_text_from_rois(image_path, processor:EndoscopyProcessor):
-    """
-    Extracts text from regions of interest (ROIs) in an image using OCR.
-
-    Args:
-        image_path (str): The path to the image file.
-        processor (EndoscopyProcessor): An instance of the EndoscopyProcessor class.
-
-    Returns:
-        dict: A dictionary containing the extracted text for each ROI.
-    """
-    # Read the image using OpenCV
-    image = cv2.imread(image_path)
-
-    # Initialize the dictionary to hold the extracted text
-    extracted_texts = {}
-
-    # Define your ROIs and their corresponding post-processing functions in tuples
-    rois_with_postprocessing = [
-        ('examination_date', processor.get_roi_examination_date, process_date_text),
-        ("patient_first_name", processor.get_roi_patient_first_name, process_name_text),
-        ('patient_last_name', processor.get_roi_patient_last_name, process_name_text),
-        ('patient_dob', processor.get_roi_patient_dob, process_date_text),
-        ('endoscope_type', processor.get_roi_endoscope_type, process_general_text),
-        ('endoscope_sn', processor.get_roi_endoscopy_sn, process_general_text),
-    ]
-
-    # Extract and post-process text for each ROI
-    for roi_name, roi_function, post_process in rois_with_postprocessing:
-        # Get the ROI dictionary
-        roi = roi_function()
-        
-        # Check if the ROI has values
-        
-        if roi_values_valid(roi):
-            # Crop the image to the ROI
-            x, y, w, h = roi['x'], roi['y'], roi['width'], roi['height']
-            roi_cropped = image[y:y+h, x:x+w]
-            # Convert to grayscale
-            gray = cv2.cvtColor(roi_cropped, cv2.COLOR_BGR2GRAY)
-
-            # Invert colors for white text on black background
-            gray = cv2.bitwise_not(gray)
-
-            # Binarize the image - using Otsu's method
-            _, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
-
-            # Dilate the image to improve the contour of the pixelated text
-            kernel = np.ones((2,2), np.uint8)
-            dilation = cv2.dilate(binary, kernel, iterations=1)
-
-            # OCR configuration: Recognize white text on black background without corrections
-            config = '--psm 10 -c tessedit_char_whitelist=0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ-üöäÜÖÄß'
-
-            # Use pytesseract to do OCR on the preprocessed ROI
-            text = pytesseract.image_to_string(dilation, config=config).strip()
-
-            # Post-process extracted text
-            processed_text = post_process(text)
-            # processed_text = text
-            
-            # Store the processed text in the dictionary
-            extracted_texts[roi_name] = processed_text
-
-        else:
-            ic(roi_name)
-            ic(roi)
-            ic("No values for this ROI")
-
-    return extracted_texts
-
-
-def get_most_frequent_values(rois_texts: Dict[str, List[str]]) -> Dict[str, str]:
-    """
-    Given a dictionary of ROIs and their corresponding texts, returns a dictionary of the most frequent text for each ROI.
-
-    Args:
-        rois_texts: A dictionary where the keys are the names of the ROIs and the values are lists of texts.
-
-    Returns:
-        A dictionary where the keys are the names of the ROIs and the values are the most frequent text for each ROI.
-    """
-    most_frequent = {}
-    for key in rois_texts.keys():
-        counter = Counter([text for text in rois_texts[key] if text])
-        ic(key)
-        ic(counter)
-        most_frequent[key], _ = counter.most_common(1)[0] if counter else (None, None)
-    return most_frequent
-
-def process_video(video_path, processor):
-    """
-    Processes a video file by extracting text from regions of interest (ROIs) in each frame.
-
-    Args:
-        video_path (str): The path to the video file to process.
-        processor (OCRProcessor): An instance of the OCRProcessor class that defines the ROIs to extract text from.
-
-    Returns:
-        dict: A dictionary containing the most frequent text values extracted from each ROI.
-    """
-    # Create a temporary directory to store frames
-    with TemporaryDirectory() as temp_dir:
-        ic(temp_dir)
-        # Capture the video
-        video = cv2.VideoCapture(video_path)
-        success, frame_number = True, 0
-        rois_texts = {roi_name: [] for roi_name in processor.get_rois().keys()}
-        frames_for_mean_extraction = 0
-
-        while success:
-            success, frame = video.read()
-
-            # Check if this is the 200th frame
-            if frame_number % 1000 == 0 and success:
-                frame_path = os.path.join(temp_dir, f"frame_{frame_number}.jpg")
-                cv2.imwrite(frame_path, frame)  # Save the frame as a JPEG file
-                # cv2.imwrite(f"_tmp/frame_{frame_number}.jpg", frame)
-                
-                # Extract text from ROIs
-                extracted_texts = extract_text_from_rois(frame_path, processor)
-                ic(extracted_texts)
-                
-                # Store the extracted text from each ROI
-                for key, text in extracted_texts.items():
-                    rois_texts[key].append(text)
-                frames_for_mean_extraction += 1
-
-            frame_number += 1
-
-            if frames_for_mean_extraction >= N_FRAMES_MEAN_OCR: break
-
-        # Release the video capture object
-        video.release()
-
-        # Get the most frequent values for each ROI
-        return get_most_frequent_values(rois_texts)
-
+import pytesseract
+# import cv2
+from endoreg_db.models import EndoscopyProcessor
+import os
+from collections import Counter
+from tempfile import TemporaryDirectory
+import re
+from datetime import datetime
+from typing import Dict, List
+from icecream import ic
+import numpy as np
+
+N_FRAMES_MEAN_OCR = 2
+
+# Helper function to process date strings
+def process_date_text(date_text):
+    """
+    Processes a string of text that represents a date and returns a datetime.date object.
+
+    Args:
+        date_text (str): A string of text that represents a date.
+
+    Returns:
+        datetime.date: A datetime.date object representing the parsed date, or None if the text cannot be parsed.
+    """
+    try:
+        # Remove any non-digit characters
+        date_text_clean = re.sub(r'\D', '', date_text)
+        # Reformat to 'ddmmyyyy' if necessary
+        if len(date_text_clean) == 8:
+            return datetime.strptime(date_text_clean, "%d%m%Y").date()
+        elif len(date_text_clean) == 14:
+            return datetime.strptime(date_text_clean, "%d%m%Y%H%M%S").date()
+    except ValueError:
+        # Return None if the text cannot be parsed into a date
+        # set date to 1/1/1900
+        return datetime.strptime("01011900", "%d%m%Y").date()
+
+# Helper function to process patient names
+def process_name_text(name_text):
+    """
+    Remove all numbers, punctuation, and whitespace from a string of text and return the result.
+    """
+    name = re.sub(r'[0-9!"#$%&\'()*+,-./:;<=>?@[\\]^_`{|}~\s]+', '', name_text).strip()
+    # capitalize first letter of each word
+    name = ' '.join([word.capitalize() for word in name.split()])
+    return name
+
+
+# Helper function to process endoscope type text
+def process_general_text(endoscope_text):
+    """
+    This function takes in a string of text from an endoscope and returns a cleaned version of the text.
+    """
+    return ' '.join(endoscope_text.split())
+
+def roi_values_valid(roi):
+    """
+    Check if all values in an ROI dictionary are valid (>=0).
+    """
+    return all([value >= 0 for value in roi.values()])
+
+# Function to extract text from ROIs
+def extract_text_from_rois(image_path, processor:EndoscopyProcessor):
+    """
+    Extracts text from regions of interest (ROIs) in an image using OCR.
+
+    Args:
+        image_path (str): The path to the image file.
+        processor (EndoscopyProcessor): An instance of the EndoscopyProcessor class.
+
+    Returns:
+        dict: A dictionary containing the extracted text for each ROI.
+    """
+    # Read the image using OpenCV
+    image = cv2.imread(image_path)
+
+    # Initialize the dictionary to hold the extracted text
+    extracted_texts = {}
+
+    # Define your ROIs and their corresponding post-processing functions in tuples
+    rois_with_postprocessing = [
+        ('examination_date', processor.get_roi_examination_date, process_date_text),
+        ("patient_first_name", processor.get_roi_patient_first_name, process_name_text),
+        ('patient_last_name', processor.get_roi_patient_last_name, process_name_text),
+        ('patient_dob', processor.get_roi_patient_dob, process_date_text),
+        ('endoscope_type', processor.get_roi_endoscope_type, process_general_text),
+        ('endoscope_sn', processor.get_roi_endoscopy_sn, process_general_text),
+    ]
+
+    # Extract and post-process text for each ROI
+    for roi_name, roi_function, post_process in rois_with_postprocessing:
+        # Get the ROI dictionary
+        roi = roi_function()
+        
+        # Check if the ROI has values
+        
+        if roi_values_valid(roi):
+            # Crop the image to the ROI
+            x, y, w, h = roi['x'], roi['y'], roi['width'], roi['height']
+            roi_cropped = image[y:y+h, x:x+w]
+            # Convert to grayscale
+            gray = cv2.cvtColor(roi_cropped, cv2.COLOR_BGR2GRAY)
+
+            # Invert colors for white text on black background
+            gray = cv2.bitwise_not(gray)
+
+            # Binarize the image - using Otsu's method
+            _, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
+
+            # Dilate the image to improve the contour of the pixelated text
+            kernel = np.ones((2,2), np.uint8)
+            dilation = cv2.dilate(binary, kernel, iterations=1)
+
+            # OCR configuration: Recognize white text on black background without corrections
+            config = '--psm 10 -c tessedit_char_whitelist=0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ-üöäÜÖÄß'
+
+            # Use pytesseract to do OCR on the preprocessed ROI
+            text = pytesseract.image_to_string(dilation, config=config).strip()
+
+            # Post-process extracted text
+            processed_text = post_process(text)
+            # processed_text = text
+            
+            # Store the processed text in the dictionary
+            extracted_texts[roi_name] = processed_text
+
+        else:
+            ic(roi_name)
+            ic(roi)
+            ic("No values for this ROI")
+
+    return extracted_texts
+
+
+def get_most_frequent_values(rois_texts: Dict[str, List[str]]) -> Dict[str, str]:
+    """
+    Given a dictionary of ROIs and their corresponding texts, returns a dictionary of the most frequent text for each ROI.
+
+    Args:
+        rois_texts: A dictionary where the keys are the names of the ROIs and the values are lists of texts.
+
+    Returns:
+        A dictionary where the keys are the names of the ROIs and the values are the most frequent text for each ROI.
+    """
+    most_frequent = {}
+    for key in rois_texts.keys():
+        counter = Counter([text for text in rois_texts[key] if text])
+        ic(key)
+        ic(counter)
+        most_frequent[key], _ = counter.most_common(1)[0] if counter else (None, None)
+    return most_frequent
+
+def process_video(video_path, processor):
+    """
+    Processes a video file by extracting text from regions of interest (ROIs) in each frame.
+
+    Args:
+        video_path (str): The path to the video file to process.
+        processor (OCRProcessor): An instance of the OCRProcessor class that defines the ROIs to extract text from.
+
+    Returns:
+        dict: A dictionary containing the most frequent text values extracted from each ROI.
+    """
+    # Create a temporary directory to store frames
+    with TemporaryDirectory() as temp_dir:
+        ic(temp_dir)
+        # Capture the video
+        video = cv2.VideoCapture(video_path)
+        success, frame_number = True, 0
+        rois_texts = {roi_name: [] for roi_name in processor.get_rois().keys()}
+        frames_for_mean_extraction = 0
+
+        while success:
+            success, frame = video.read()
+
+            # Check if this is the 200th frame
+            if frame_number % 1000 == 0 and success:
+                frame_path = os.path.join(temp_dir, f"frame_{frame_number}.jpg")
+                cv2.imwrite(frame_path, frame)  # Save the frame as a JPEG file
+                # cv2.imwrite(f"_tmp/frame_{frame_number}.jpg", frame)
+                
+                # Extract text from ROIs
+                extracted_texts = extract_text_from_rois(frame_path, processor)
+                ic(extracted_texts)
+                
+                # Store the extracted text from each ROI
+                for key, text in extracted_texts.items():
+                    rois_texts[key].append(text)
+                frames_for_mean_extraction += 1
+
+            frame_number += 1
+
+            if frames_for_mean_extraction >= N_FRAMES_MEAN_OCR: break
+
+        # Release the video capture object
+        video.release()
+
+        # Get the most frequent values for each ROI
+        return get_most_frequent_values(rois_texts)
+