53 lines
2.6 KiB
Python
53 lines
2.6 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
|
|
|
|
class MaskModule:
|
|
def __init__(self):
|
|
print("마스크 모듈 초기화 완료")
|
|
def create_masks(self, image_path: str, ocr_results: List[Dict], expansion_size: int = 10, blur_size: int = 15, mask_option: str = "basic") -> np.ndarray:
|
|
image = cv2.imread(image_path)
|
|
if image is None:
|
|
print(f"이미지를 읽을 수 없습니다: {image_path}")
|
|
return None
|
|
height, width = image.shape[:2]
|
|
mask = np.zeros((height, width), dtype=np.uint8)
|
|
for i, result in enumerate(ocr_results, 1):
|
|
polygon = result['polygon']
|
|
expanded_poly = self.expand_polygon(polygon, offset=5)
|
|
cv2.fillPoly(mask, [expanded_poly], 255)
|
|
processed_mask = self.process_mask(mask, expansion_size, blur_size)
|
|
return processed_mask
|
|
def expand_polygon(self, polygon, offset=15):
|
|
poly = Polygon(polygon)
|
|
expanded = poly.buffer(offset)
|
|
if expanded.is_empty:
|
|
return np.array(polygon, dtype=np.int32)
|
|
return np.array(expanded.exterior.coords, 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 create_masks_with_save(self, image_path: str, ocr_results: List[Dict], output_dir: str = None, expansion_size: int = 10, blur_size: int = 15, mask_option: str = "basic") -> tuple:
|
|
mask = self.create_masks(image_path, ocr_results, expansion_size, blur_size, mask_option)
|
|
mask_image_path = None
|
|
if output_dir and mask is not None:
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
base_name = os.path.splitext(os.path.basename(image_path))[0]
|
|
mask_filename = f"{base_name}_mask_combined.jpg"
|
|
mask_image_path = os.path.join(output_dir, mask_filename)
|
|
cv2.imwrite(mask_image_path, mask)
|
|
return mask, mask_image_path |