File size: 2,443 Bytes
dc17282 | 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 | import cv2
import numpy as np
from PIL import Image
import io
def bytes_to_cv2(image_bytes: bytes) -> np.ndarray:
nparr = np.frombuffer(image_bytes, np.uint8)
img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
return img
def cv2_to_pil(img: np.ndarray) -> Image.Image:
rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
return Image.fromarray(rgb)
def crop_to_foreground(img: np.ndarray, padding_ratio: float = 0.18) -> np.ndarray:
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) if len(img.shape) == 3 else img
blurred = cv2.GaussianBlur(gray, (5, 5), 0)
_, mask = cv2.threshold(blurred, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
kernel = np.ones((3, 3), np.uint8)
mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)
mask = cv2.dilate(mask, kernel, iterations=1)
contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
if not contours: return img
h, w = gray.shape[:2]
min_area = max(12, int(h * w * 0.0001))
boxes = [cv2.boundingRect(contour) for contour in contours if cv2.contourArea(contour) >= min_area]
if not boxes: return img
x1, y1, x2, y2 = min(x for x,_,_,_ in boxes), min(y for _,y,_,_ in boxes), max(x+bw for x,_,bw,_ in boxes), max(y+bh for _,y,_,bh in boxes)
pad_x, pad_y = max(8, int((x2 - x1) * padding_ratio)), max(8, int((y2 - y1) * padding_ratio))
x1, y1, x2, y2 = max(0, x1 - pad_x), max(0, y1 - pad_y), min(w, x2 + pad_x), min(h, y2 + pad_y)
return img[y1:y2, x1:x2]
def normalize_for_model(img: np.ndarray, target_height: int = 384, target_width: int = 384) -> np.ndarray:
h, w = img.shape[:2]
scale = min(target_height / h, target_width / w)
new_h, new_w = int(h * scale), int(w * scale)
resized = cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_AREA)
canvas = np.ones((target_height, target_width, 3), dtype=np.uint8) * 255
y_offset, x_offset = (target_height - new_h) // 2, (target_width - new_w) // 2
canvas[y_offset:y_offset + new_h, x_offset:x_offset + new_w] = resized
return canvas
def preprocess_for_ocr(image_bytes: bytes) -> Image.Image:
img = bytes_to_cv2(image_bytes)
if img is None: return None
h, w = img.shape[:2]
aspect_ratio = w / float(h)
if aspect_ratio <= 1.55:
img = crop_to_foreground(img)
elif aspect_ratio > 2.2:
img = crop_to_foreground(img)
img = normalize_for_model(img)
return cv2_to_pil(img)
|