Spaces:
Build error
Build error
File size: 6,008 Bytes
4dbe5d1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 |
import os
import cv2
import numpy as np
from pdf2image import convert_from_path
from main import RapidOCR
from image_enhancement import enhance_image
# Initialize OCR engine once.
ocr_engine = RapidOCR()
def adaptive_threshold_to_rgb(image_rgb):
"""
Apply adaptive thresholding on the L channel of LAB color space
and reconstruct the thresholded image as RGB.
Parameters:
image_rgb (numpy.ndarray): Input RGB image.
Returns:
thresholded_rgb (numpy.ndarray): RGB image after thresholding the L channel.
"""
# Convert RGB to LAB color space and split channels.
image_lab = cv2.cvtColor(image_rgb, cv2.COLOR_RGB2LAB)
l_channel, a_channel, b_channel = cv2.split(image_lab)
# Apply adaptive thresholding to the L channel.
thresholded_l = cv2.adaptiveThreshold(
l_channel,
maxValue=255,
adaptiveMethod=cv2.ADAPTIVE_THRESH_GAUSSIAN_C, # or ADAPTIVE_THRESH_MEAN_C
thresholdType=cv2.THRESH_BINARY,
blockSize=11,
C=2
)
# Merge the thresholded L channel back with A and B channels.
updated_lab = cv2.merge((thresholded_l, a_channel, b_channel))
thresholded_rgb = cv2.cvtColor(updated_lab, cv2.COLOR_LAB2RGB)
return thresholded_rgb
def ocr_detect(image, ocr_engine):
"""
Run OCR on the image using the provided ocr_engine and check if consecutive
rows containing the '<' character are detected.
Parameters:
image (numpy.ndarray): Input image.
ocr_engine: The OCR engine instance.
Returns:
detected (bool): True if the desired pattern is detected, False otherwise.
"""
result, _ = ocr_engine(image, use_det=True, use_cls=False, use_rec=True)
if result:
test_list = [r[1] for r in result]
for j in range(len(test_list) - 1):
count1 = test_list[j].count("<")
count2 = test_list[j + 1].count("<")
if count1 > 1 and count2 > 1:
return True
return False
def rotate_until_detect(image, ocr_engine, max_attempts=4):
"""
Rotate the image 90° clockwise, up to max_attempts times, until the OCR
conveys the expected result.
Parameters:
image (numpy.ndarray): Input image.
ocr_engine: The OCR engine instance.
max_attempts (int): Maximum number of rotations.
Returns:
image (numpy.ndarray): Rotated image with detection (or final rotation if undetected).
detected (bool): Whether the expected OCR pattern was detected.
"""
attempt = 0
detected = False
while attempt < max_attempts:
if ocr_detect(image, ocr_engine):
detected = True
break
image = cv2.rotate(image, cv2.ROTATE_90_CLOCKWISE)
attempt += 1
return image, detected
def process_pdf(pdf_f, ocr_engine, enhance_params):
"""
Process a single PDF file by converting pages, enhancing images,
running OCR with rotation, and using adaptive thresholding as a fallback.
Parameters:
pdf_f (str): PDF file path.
ocr_engine: The OCR engine instance.
enhance_params (dict): Parameters for the image enhancement.
"""
# Convert specified pages of PDF into images.
images = convert_from_path(pdf_f, dpi=300, first_page=1, last_page=3)
bs_name = os.path.basename(pdf_f)
bs_name_0 = os.path.splitext(bs_name)[0]
for i, pil_image in enumerate(images):
# Convert PIL image to a NumPy array.
img = np.array(pil_image)
print("Original image shape:", img.shape)
# Enhance the image.
img = enhance_image(img, enhance_params, verbose=False)
img = np.uint8(img * 255.)
# Save the enhanced image as a reference.
enhanced_img_bgr = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
output_filename = f'{bs_name_0}_{i + 1}.jpg'
cv2.imwrite(output_filename, enhanced_img_bgr)
print(f"Saved enhanced image: {output_filename}")
# First: Try OCR on the enhanced image with rotation.
processed_img, detected = rotate_until_detect(img, ocr_engine)
if detected:
print(f"OCR success on {output_filename} with enhanced image rotation.")
else:
# Second: Apply adaptive thresholding and re-run OCR with rotation.
print(f"No OCR detection from enhanced image. Applying adaptive thresholding for {output_filename}.")
adaptive_img = adaptive_threshold_to_rgb(img)
processed_img, detected = rotate_until_detect(adaptive_img, ocr_engine)
if detected:
print(f"OCR success on {output_filename} with adaptive thresholding and rotation.")
else:
print(f"OCR detection failed for {output_filename} after fallback.")
def main():
# Set the data path and gather list of PDF files.
dataPath = '/home/tung/Tung_Works/OCR_code/OCR-20250423T073748Z-001/OCR/OCR辨識失敗-部分樣本'
list_pdf = [
os.path.join(root, file)
for root, _, files in os.walk(dataPath)
for file in files if file.endswith('.pdf')
]
# Optionally, sort the list.
# list_pdf = sorted(list_pdf)
# Define image enhancement parameters (applied to every image).
enhance_params = {
'local_contrast': 1.2, # 1.2x increase in details
'mid_tones': 0.5, # middle of range
'tonal_width': 0.5, # middle of range
'areas_dark': 0.7, # 70% improvement in dark areas
'areas_bright': 0.5, # 50% improvement in bright areas
'brightness': 0.1, # slight increase in overall brightness
'saturation_degree': 1.2, # 1.2x increase in color saturation
'preserve_tones': True,
'color_correction': True,
}
# Process each PDF.
for pdf_f in list_pdf:
process_pdf(pdf_f, ocr_engine, enhance_params)
if __name__ == '__main__':
main()
|