106 lines
5.2 KiB
Python
106 lines
5.2 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""
|
|
마스크 모듈 - OCR 결과를 기반으로 마스크 생성 및 처리 (라이브러리화)
|
|
확장, 블러 등 basic 마스크 처리만 지원
|
|
"""
|
|
|
|
import cv2
|
|
import numpy as np
|
|
from typing import List, Dict, Any
|
|
from shapely.geometry import Polygon
|
|
import os
|
|
import logging
|
|
|
|
class MaskModule:
|
|
def __init__(self, logger, base_dir):
|
|
self.logger = logger
|
|
self.base_dir = base_dir
|
|
self.logger.log("마스크 모듈 초기화 완료", level=logging.DEBUG)
|
|
|
|
def create_masks(self, image_path: str, ocr_results: List[Dict], expansion_size: int = 15, blur_size: int = 15, mask_option: str = "basic") -> np.ndarray:
|
|
image = cv2.imread(image_path)
|
|
if image is None:
|
|
self.logger.log(f"이미지를 읽을 수 없습니다: {image_path}", level=logging.ERROR)
|
|
return None
|
|
height, width = image.shape[:2]
|
|
mask = np.zeros((height, width), dtype=np.uint8)
|
|
|
|
# self.logger.log(f"[마스크] OCR 결과 개수: {len(ocr_results)}, 이미지 크기: {width}x{height}", level=logging.DEBUG)
|
|
|
|
for i, result in enumerate(ocr_results, 1):
|
|
polygon = result['polygon']
|
|
text = result.get('text', '')
|
|
# self.logger.log(f"[마스크] {i}번째 텍스트: '{text}', 원본 polygon: {polygon}", level=logging.DEBUG)
|
|
|
|
# 확장 없이 원본 폴리곤 그대로 사용
|
|
original_poly = np.array(polygon, dtype=np.int32)
|
|
# self.logger.log(f"[마스크] {i}번째 원본 polygon shape: {original_poly.shape}, 값: {original_poly}", level=logging.DEBUG)
|
|
|
|
# polygon이 (N,2) np.int32인지 체크
|
|
if original_poly.ndim != 2 or original_poly.shape[1] != 2 or original_poly.shape[0] < 3:
|
|
# self.logger.log(f"[마스크] 잘못된 폴리곤 shape: {original_poly.shape}, 값: {original_poly}", level=logging.WARNING)
|
|
continue
|
|
|
|
cv2.fillPoly(mask, [original_poly], 255)
|
|
# self.logger.log(f"[마스크] {i}번째 폴리곤 마스크 적용 완료", level=logging.DEBUG)
|
|
|
|
# 마스크에 실제로 255 값이 있는지 확인
|
|
mask_pixels_before = np.sum(mask > 0)
|
|
# self.logger.log(f"[마스크] Erosion 전 마스크 픽셀 개수: {mask_pixels_before} / {height * width}", level=logging.DEBUG)
|
|
|
|
# Erosion 적용 (마스크를 글자 주변으로 축소)
|
|
erosion_kernel_size = 6 # 3x3 커널로 적당히 축소
|
|
erosion_kernel = np.ones((erosion_kernel_size, erosion_kernel_size), np.uint8)
|
|
eroded_mask = cv2.erode(mask, erosion_kernel, iterations=1)
|
|
|
|
# Erosion 후 마스크 픽셀 개수 확인
|
|
mask_pixels_after = np.sum(eroded_mask > 0)
|
|
# self.logger.log(f"[마스크] Erosion 후 마스크 픽셀 개수: {mask_pixels_after} / {height * width} (감소: {mask_pixels_before - mask_pixels_after}픽셀)", level=logging.DEBUG)
|
|
|
|
# 기본 마스크만 사용 (erosion 적용된 마스크 반환)
|
|
if mask_option == "basic":
|
|
return eroded_mask
|
|
else:
|
|
# 다른 옵션일 경우에만 추가 후처리 적용
|
|
processed_mask = self.process_mask(eroded_mask, expansion_size, blur_size)
|
|
return processed_mask
|
|
|
|
def expand_polygon(self, polygon, offset=15):
|
|
try:
|
|
poly = Polygon(polygon)
|
|
expanded = poly.buffer(offset)
|
|
if expanded.is_empty:
|
|
return np.array(polygon, dtype=np.int32)
|
|
|
|
# exterior.coords는 마지막 점이 첫 번째 점과 동일하므로 제거
|
|
coords = list(expanded.exterior.coords)[:-1]
|
|
|
|
# 정수 좌표로 변환
|
|
coords = [[int(x), int(y)] for x, y in coords]
|
|
|
|
# 최소 3점 이상인지 확인
|
|
if len(coords) < 3:
|
|
return np.array(polygon, dtype=np.int32)
|
|
|
|
return np.array(coords, dtype=np.int32)
|
|
except Exception as e:
|
|
self.logger.log(f"[마스크] expand_polygon 오류: {e}, 원본 polygon 사용", level=logging.WARNING)
|
|
return np.array(polygon, dtype=np.int32)
|
|
|
|
def process_mask(self, mask: np.ndarray, expansion_size: int = 5, blur_size: int = 3) -> np.ndarray:
|
|
processed_mask = mask.copy()
|
|
if expansion_size > 0:
|
|
kernel = np.ones((expansion_size, expansion_size), np.uint8)
|
|
processed_mask = cv2.dilate(processed_mask, kernel, iterations=1)
|
|
if blur_size > 0:
|
|
blur_size = blur_size if blur_size % 2 == 1 else blur_size + 1
|
|
processed_mask = cv2.GaussianBlur(processed_mask, (blur_size, blur_size), 0)
|
|
return processed_mask
|
|
|
|
def save_mask(self, mask: np.ndarray, out_path: str) -> str:
|
|
"""마스크(np.uint8 0/255)를 파일로 저장하고 경로 반환"""
|
|
os.makedirs(os.path.dirname(out_path), exist_ok=True)
|
|
cv2.imwrite(out_path, mask)
|
|
self.logger.log(f"마스크를 {out_path} 에 저장했습니다.", level=logging.INFO)
|
|
return out_path
|