Files
ocr_gateway_test/utils/preprocessor.py
2025-10-27 09:18:24 +09:00

61 lines
1.9 KiB
Python

import cv2
import numpy as np
from PIL import Image
import logging
logger = logging.getLogger(__name__)
def to_rgb_uint8(img_np: np.ndarray) -> np.ndarray:
"""
입력 이미지를 3채널 RGB, uint8 [0,255] 로 표준화
허용 입력: HxW, HxWx1, HxWx3, HxWx4, float[0..1]/[0..255], int 등
"""
if img_np is None:
raise ValueError("Input image is None")
# dtype/범위 표준화
if img_np.dtype != np.uint8:
arr = img_np.astype(np.float32)
if arr.max() <= 1.0: # [0,1]로 보이면 스케일업
arr *= 255.0
arr = np.clip(arr, 0, 255).astype(np.uint8)
img_np = arr
# 채널 표준화
if img_np.ndim == 2: # HxW
img_np = cv2.cvtColor(img_np, cv2.COLOR_GRAY2RGB)
elif img_np.ndim == 3:
h, w, c = img_np.shape
if c == 1:
img_np = cv2.cvtColor(img_np, cv2.COLOR_GRAY2RGB)
elif c == 4:
img_np = cv2.cvtColor(img_np, cv2.COLOR_RGBA2RGB)
elif c == 3:
pass # 그대로 사용
else:
raise ValueError(f"Unsupported channel count: {c}")
else:
raise ValueError(f"Unsupported ndim: {img_np.ndim}")
return img_np
# tesseract 전처리 함수
def tess_prep_cv2(pil_img):
logger.info("[UTILS-OCR] 이미지 전처리 시작")
img = np.array(pil_img.convert("RGB")) # PIL → OpenCV 변환
img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY) # 그레이스케일 변환
img = cv2.bilateralFilter(img, 9, 75, 75) # 노이즈 제거
img = cv2.adaptiveThreshold(
img,
255,
cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
cv2.THRESH_BINARY,
31,
10, # 대비 향상
)
img = cv2.resize(
img, None, fx=2, fy=2, interpolation=cv2.INTER_LINEAR
) # 해상도 확대
return Image.fromarray(img)