110 lines
3.1 KiB
Python
110 lines
3.1 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""Image decoding and preprocessing helpers."""
|
|
import io
|
|
import numpy as np
|
|
from PIL import Image, ImageSequence
|
|
|
|
try:
|
|
import cv2
|
|
except ImportError: # pragma: no cover
|
|
cv2 = None
|
|
|
|
|
|
def decode_image(data):
|
|
"""Decode bytes to RGB numpy array. Returns None on failure."""
|
|
img = _decode_with_pillow(data)
|
|
if img is None and cv2 is not None:
|
|
img = _decode_with_cv2(data)
|
|
return img
|
|
|
|
|
|
def _decode_with_pillow(data):
|
|
try:
|
|
with Image.open(io.BytesIO(data)) as im:
|
|
# pick first frame for GIFs
|
|
frame = next(ImageSequence.Iterator(im))
|
|
rgb = frame.convert('RGB')
|
|
arr = np.array(rgb)
|
|
return arr
|
|
except Exception:
|
|
return None
|
|
|
|
|
|
def _decode_with_cv2(data):
|
|
try:
|
|
arr = np.frombuffer(data, dtype=np.uint8)
|
|
img = cv2.imdecode(arr, cv2.IMREAD_COLOR)
|
|
if img is None:
|
|
return None
|
|
rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
|
return rgb
|
|
except Exception:
|
|
return None
|
|
|
|
|
|
def prepare_for_ocr(img_rgb, use_threshold=True, upscale=1.0, clahe=False,
|
|
bilateral=False, blur=0, invert=False):
|
|
"""Convert RGB to grayscale and apply light preprocessing for OCR."""
|
|
if img_rgb is None:
|
|
return None
|
|
|
|
gray = _to_gray(img_rgb)
|
|
|
|
# optional upscale (cap at 3x to avoid memory blow-up)
|
|
if upscale and upscale > 1.0:
|
|
factor = min(float(upscale), 3.0)
|
|
h, w = gray.shape[:2]
|
|
new_w = int(w * factor)
|
|
new_h = int(h * factor)
|
|
if cv2 is not None:
|
|
gray = cv2.resize(gray, (new_w, new_h), interpolation=cv2.INTER_CUBIC)
|
|
else:
|
|
gray = np.array(Image.fromarray(gray).resize((new_w, new_h)))
|
|
|
|
if clahe and cv2 is not None:
|
|
try:
|
|
clahe_obj = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
|
|
gray = clahe_obj.apply(gray)
|
|
except Exception:
|
|
pass
|
|
|
|
if bilateral and cv2 is not None:
|
|
try:
|
|
gray = cv2.bilateralFilter(gray, d=5, sigmaColor=75, sigmaSpace=75)
|
|
except Exception:
|
|
pass
|
|
|
|
if blur and blur > 0 and cv2 is not None:
|
|
k = int(blur) if int(blur) % 2 == 1 else int(blur) + 1
|
|
try:
|
|
gray = cv2.GaussianBlur(gray, (k, k), 0)
|
|
except Exception:
|
|
pass
|
|
|
|
if use_threshold and cv2 is not None:
|
|
try:
|
|
gray = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
|
|
cv2.THRESH_BINARY, 25, 15)
|
|
except Exception:
|
|
pass
|
|
|
|
if invert:
|
|
gray = 255 - gray
|
|
|
|
return gray
|
|
|
|
|
|
def _to_gray(img_rgb):
|
|
if img_rgb.ndim == 3 and img_rgb.shape[2] == 3:
|
|
return np.dot(img_rgb[..., :3], [0.2989, 0.5870, 0.1140]).astype('uint8')
|
|
return img_rgb.astype('uint8')
|
|
|
|
|
|
def crop(img_rgb, bbox):
|
|
x1, y1, x2, y2 = bbox
|
|
h, w = img_rgb.shape[:2]
|
|
x1 = max(0, min(w - 1, int(x1)))
|
|
y1 = max(0, min(h - 1, int(y1)))
|
|
x2 = max(0, min(w, int(x2)))
|
|
y2 = max(0, min(h, int(y2)))
|
|
return img_rgb[y1:y2, x1:x2]
|