Spaces:
Running on Zero
Running on Zero
File size: 2,048 Bytes
98e1d9f | 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 | from __future__ import annotations
from PIL import Image
from ocr_studio.config import DESKEW_MAX_ANGLE
def _rotate(image: Image.Image, angle: float) -> Image.Image:
if abs(angle) < 0.25:
return image
fill = (255, 255, 255)
return image.convert("RGB").rotate(angle, expand=True, fillcolor=fill, resample=Image.Resampling.BICUBIC)
def enhance_scan(image: Image.Image) -> Image.Image:
try:
import cv2
import numpy as np
except Exception:
return image.convert("RGB")
rgb = np.array(image.convert("RGB"))
lab = cv2.cvtColor(rgb, cv2.COLOR_RGB2LAB)
lightness, axis_a, axis_b = cv2.split(lab)
clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
lightness = clahe.apply(lightness)
merged = cv2.merge((lightness, axis_a, axis_b))
enhanced = cv2.cvtColor(merged, cv2.COLOR_LAB2RGB)
return Image.fromarray(enhanced)
def estimate_deskew_angle(image: Image.Image) -> float:
try:
import cv2
import numpy as np
except Exception:
return 0.0
gray = np.array(image.convert("L"))
if gray.size < 400:
return 0.0
inverted = cv2.bitwise_not(gray)
_, binary = cv2.threshold(inverted, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
coords = np.column_stack(np.where(binary > 0))
if coords.shape[0] < 80:
return 0.0
_rect, _size, raw_angle = cv2.minAreaRect(coords)
angle = raw_angle
if angle < -45:
angle = 90.0 + angle
if abs(angle) > DESKEW_MAX_ANGLE:
return 0.0
return float(-angle)
def deskew_image(image: Image.Image) -> tuple[Image.Image, float]:
rgb = image.convert("RGB")
angle = estimate_deskew_angle(rgb)
straightened = _rotate(rgb, angle)
return enhance_scan(straightened), angle
def prepare_page(image: Image.Image, deskew: bool) -> tuple[Image.Image, float]:
if not deskew:
return enhance_scan(image.convert("RGB")), 0.0
return deskew_image(image)
|