ArTrOCR-HTR / preprocessing.py

Upload preprocessing.py with huggingface_hub

03e312d verified about 1 month ago

3.31 kB

def preprocess_image(image):
    from PIL import Image
    import cv2
    import numpy as np
    from skimage.filters import threshold_sauvola
    
    ####### crop #######
    padding=40
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    # Strong Blur to merge text strokes and smooth paper ripples
    blurred = cv2.GaussianBlur(gray, (7, 7), 0)
    # Adaptive Threshold (Block size 21, Constant 10)
    # ignores large shadows and focuses on local contrast
    thresh = cv2.adaptiveThreshold(blurred, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,  cv2.THRESH_BINARY_INV, 21, 10)
    # clean up: Remove small dots (salt & pepper)
    kernel = np.ones((3,3), np.uint8)
    # Opening removes small white noise from the background
    thresh = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, kernel, iterations=1)
    # Dilation makes the text "thick" so we don't miss thin strokes
    thresh = cv2.dilate(thresh, kernel, iterations=2)
    # Find Contours and filter by area
    contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) 
    if contours:
        # Filter out very small noise (area < 100 pixels)
        valid_cnts = [cnt for cnt in contours if cv2.contourArea(cnt) > 100]
        if valid_cnts:
            # Get one big bounding box for all valid text pieces
            all_pts = np.concatenate(valid_cnts)
            x, y, w, h = cv2.boundingRect(all_pts)
            # Apply Padding
            h_img, w_img = image.shape[:2]
            x1 = max(0, x - padding)
            y1 = max(0, y - padding)
            x2 = min(w_img, x + w + padding)
            y2 = min(h_img, y + h + padding)
            image = image[y1:y2, x1:x2]
    
    ####### binarize #######
    image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) 
    thresh = threshold_sauvola(image, window_size=25, k=0.2)
    image = ((image > thresh) * 255).astype(np.uint8)
    
    ####### deskew #######
    gray = image.copy()
    gray = cv2.bitwise_not(gray)
    # Threshold to make text white on black background
    thresh = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)[1]
    # Find all coordinates where pixels are "text" (white)
    coords = np.column_stack(np.where(thresh > 0))
    # find the minimum area rectangle that encloses all points
    # This rectangle will be tilted at the same angle as the text
    angle = cv2.minAreaRect(coords)[-1]
    # Handle OpenCV's angle logic
    if angle < -45:
        angle = -(90 + angle)
    else:
        angle = -angle
    # Rotate the image to deskew
    (h, w) = image.shape[:2]
    center = (w // 2, h // 2)
    M = cv2.getRotationMatrix2D(center, angle, 1.0)
    image = cv2.warpAffine(image, M, (w, h), flags=cv2.INTER_CUBIC, borderMode=cv2.BORDER_REPLICATE) 
    
    ####### resize #######
    image = Image.fromarray(image).convert("RGB")
    target_h = 102
    target_w = 512
    w, h = image.size
    scale = target_h / h
    new_w = int(w * scale)
    new_h = target_h
    image = image.resize((new_w, new_h), Image.Resampling.LANCZOS)
    canvas = Image.new("RGB", (target_w, target_h), 255)
    if new_w <= target_w:
        canvas.paste(image, (0, 0))
    else:
        image = image.resize((target_w, new_h), Image.Resampling.LANCZOS)
        canvas.paste(image, (0, 0))
        
    return canvas.convert("L").convert("RGB")

	def preprocess_image(image):
	from PIL import Image
	import cv2
	import numpy as np
	from skimage.filters import threshold_sauvola

	####### crop #######
	padding=40
	gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
	# Strong Blur to merge text strokes and smooth paper ripples
	blurred = cv2.GaussianBlur(gray, (7, 7), 0)
	# Adaptive Threshold (Block size 21, Constant 10)
	# ignores large shadows and focuses on local contrast
	thresh = cv2.adaptiveThreshold(blurred, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY_INV, 21, 10)
	# clean up: Remove small dots (salt & pepper)
	kernel = np.ones((3,3), np.uint8)
	# Opening removes small white noise from the background
	thresh = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, kernel, iterations=1)
	# Dilation makes the text "thick" so we don't miss thin strokes
	thresh = cv2.dilate(thresh, kernel, iterations=2)
	# Find Contours and filter by area
	contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
	if contours:
	# Filter out very small noise (area < 100 pixels)
	valid_cnts = [cnt for cnt in contours if cv2.contourArea(cnt) > 100]
	if valid_cnts:
	# Get one big bounding box for all valid text pieces
	all_pts = np.concatenate(valid_cnts)
	x, y, w, h = cv2.boundingRect(all_pts)
	# Apply Padding
	h_img, w_img = image.shape[:2]
	x1 = max(0, x - padding)
	y1 = max(0, y - padding)
	x2 = min(w_img, x + w + padding)
	y2 = min(h_img, y + h + padding)
	image = image[y1:y2, x1:x2]

	####### binarize #######
	image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
	thresh = threshold_sauvola(image, window_size=25, k=0.2)
	image = ((image > thresh) * 255).astype(np.uint8)

	####### deskew #######
	gray = image.copy()
	gray = cv2.bitwise_not(gray)
	# Threshold to make text white on black background
	thresh = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY \| cv2.THRESH_OTSU)[1]
	# Find all coordinates where pixels are "text" (white)
	coords = np.column_stack(np.where(thresh > 0))
	# find the minimum area rectangle that encloses all points
	# This rectangle will be tilted at the same angle as the text
	angle = cv2.minAreaRect(coords)[-1]
	# Handle OpenCV's angle logic
	if angle < -45:
	angle = -(90 + angle)
	else:
	angle = -angle
	# Rotate the image to deskew
	(h, w) = image.shape[:2]
	center = (w // 2, h // 2)
	M = cv2.getRotationMatrix2D(center, angle, 1.0)
	image = cv2.warpAffine(image, M, (w, h), flags=cv2.INTER_CUBIC, borderMode=cv2.BORDER_REPLICATE)

	####### resize #######
	image = Image.fromarray(image).convert("RGB")
	target_h = 102
	target_w = 512
	w, h = image.size
	scale = target_h / h
	new_w = int(w * scale)
	new_h = target_h
	image = image.resize((new_w, new_h), Image.Resampling.LANCZOS)
	canvas = Image.new("RGB", (target_w, target_h), 255)
	if new_w <= target_w:
	canvas.paste(image, (0, 0))
	else:
	image = image.resize((target_w, new_h), Image.Resampling.LANCZOS)
	canvas.paste(image, (0, 0))

	return canvas.convert("L").convert("RGB")