567 lines
23 KiB
Python
567 lines
23 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""
|
|
마스크 모듈 - OCR 결과를 기반으로 마스크 생성 및 처리
|
|
확장, 블러, 모폴로지 연산 등 다양한 마스크 처리 기능을 제공합니다.
|
|
"""
|
|
|
|
import cv2
|
|
import numpy as np
|
|
from typing import List, Dict, Any, Tuple
|
|
import os
|
|
from shapely.geometry import Polygon
|
|
|
|
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:
|
|
"""
|
|
OCR 결과를 기반으로 마스크 생성 (블러 확장 강화)
|
|
|
|
Args:
|
|
image_path (str): 원본 이미지 경로
|
|
ocr_results (List[Dict]): OCR 결과
|
|
expansion_size (int): 마스크 확장 크기 (기본값 증가)
|
|
blur_size (int): 블러 크기 (기본값 증가)
|
|
mask_option (str): 마스크 옵션 (기본값 "basic")
|
|
Returns:
|
|
np.ndarray: 생성된 마스크 (0-255)
|
|
"""
|
|
# 원본 이미지 로드
|
|
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)
|
|
|
|
print(f"\n🎭 마스크 생성 시작 - 이미지 크기: {width}x{height}")
|
|
print(f"📏 확장 크기: {expansion_size}px, 블러 크기: {blur_size}px")
|
|
|
|
# OCR 결과로부터 마스크 생성
|
|
for i, result in enumerate(ocr_results, 1):
|
|
polygon = result['polygon']
|
|
expanded_poly = self.expand_polygon(polygon, offset=5)
|
|
text = result['text']
|
|
|
|
print(f" 🔷 영역 #{i}: '{text}' 마스크 생성 중...")
|
|
|
|
# 폴리곤을 마스크에 그리기
|
|
cv2.fillPoly(mask, [expanded_poly], 255)
|
|
|
|
# 마스크 후처리 (블러 확장 강화)
|
|
if mask_option == "basic":
|
|
processed_mask = self.process_mask(mask, expansion_size, blur_size)
|
|
# elif mask_option == "enhanced":
|
|
# processed_mask = self.process_mask_with_enhanced_blur(processed_mask)
|
|
elif mask_option == "morphology":
|
|
processed_mask = self.create_advanced_mask(image_path, ocr_results, method="morphology")
|
|
elif mask_option == "adaptive":
|
|
processed_mask = self.create_advanced_mask(image_path, ocr_results, method="adaptive")
|
|
elif mask_option == "contour":
|
|
processed_mask = self.create_advanced_mask(image_path, ocr_results, method="contour")
|
|
elif mask_option == "ultra_blur":
|
|
processed_mask = self.create_ultra_blur_mask(image_path, ocr_results)
|
|
# elif mask_option == "feathered":
|
|
# processed_mask = self.create_feathered_mask(processed_mask)
|
|
|
|
print(f"✅ 마스크 생성 완료: {len(ocr_results)}개 영역")
|
|
return processed_mask
|
|
|
|
def expand_polygon(self, polygon, offset=15):
|
|
poly = Polygon(polygon)
|
|
expanded = poly.buffer(offset)
|
|
if expanded.is_empty:
|
|
return polygon
|
|
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:
|
|
"""
|
|
기본 마스크 후처리 (확장, 블러 등)
|
|
|
|
Args:
|
|
mask (np.ndarray): 원본 마스크
|
|
expansion_size (int): 확장 크기
|
|
blur_size (int): 블러 크기
|
|
|
|
Returns:
|
|
np.ndarray: 처리된 마스크
|
|
"""
|
|
processed_mask = mask.copy()
|
|
|
|
# 1. 마스크 확장 (Dilation)
|
|
if expansion_size > 0:
|
|
kernel = np.ones((expansion_size, expansion_size), np.uint8)
|
|
processed_mask = cv2.dilate(processed_mask, kernel, iterations=1)
|
|
|
|
# 2. 가우시안 블러 적용
|
|
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 process_mask_with_enhanced_blur(self, mask: np.ndarray, expansion_size: int = 10,
|
|
blur_size: int = 15) -> np.ndarray:
|
|
"""
|
|
마스크 후처리 (강화된 블러 확장)
|
|
|
|
Args:
|
|
mask (np.ndarray): 원본 마스크
|
|
expansion_size (int): 확장 크기
|
|
blur_size (int): 블러 크기
|
|
|
|
Returns:
|
|
np.ndarray: 처리된 마스크
|
|
"""
|
|
print(f"🔄 강화된 마스크 후처리 시작...")
|
|
processed_mask = mask.copy()
|
|
|
|
# 1. 다단계 마스크 확장 (Dilation)
|
|
if expansion_size > 0:
|
|
print(f" 📈 1단계: 기본 확장 ({expansion_size}px)")
|
|
kernel = np.ones((expansion_size, expansion_size), np.uint8)
|
|
processed_mask = cv2.dilate(processed_mask, kernel, iterations=1)
|
|
|
|
# 추가 확장 (더 부드러운 경계를 위해)
|
|
print(f" 📈 2단계: 추가 확장 ({expansion_size//2}px)")
|
|
kernel2 = np.ones((expansion_size//2, expansion_size//2), np.uint8)
|
|
processed_mask = cv2.dilate(processed_mask, kernel2, iterations=2)
|
|
|
|
# 2. 강화된 가우시안 블러 적용
|
|
if blur_size > 0:
|
|
blur_size = blur_size if blur_size % 2 == 1 else blur_size + 1 # 홀수로 만들기
|
|
print(f" 🌫️ 1단계 블러: 가우시안 블러 ({blur_size}px)")
|
|
processed_mask = cv2.GaussianBlur(processed_mask, (blur_size, blur_size), 0)
|
|
|
|
# 추가 블러 (더 부드러운 효과)
|
|
additional_blur = max(5, blur_size // 2)
|
|
additional_blur = additional_blur if additional_blur % 2 == 1 else additional_blur + 1
|
|
print(f" 🌫️ 2단계 블러: 추가 블러 ({additional_blur}px)")
|
|
processed_mask = cv2.GaussianBlur(processed_mask, (additional_blur, additional_blur), 0)
|
|
|
|
# 3. 모폴로지 연산으로 마무리
|
|
print(f" 🔧 3단계: 모폴로지 연산")
|
|
kernel_morph = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
|
|
processed_mask = cv2.morphologyEx(processed_mask, cv2.MORPH_CLOSE, kernel_morph)
|
|
|
|
print(f"✅ 강화된 마스크 후처리 완료")
|
|
return processed_mask
|
|
|
|
def create_ultra_blur_mask(self, image_path: str, ocr_results: List[Dict],
|
|
expansion_size: int = 20, blur_size: int = 25) -> np.ndarray:
|
|
"""
|
|
초강력 블러 확장 마스크 생성
|
|
|
|
Args:
|
|
image_path (str): 원본 이미지 경로
|
|
ocr_results (List[Dict]): OCR 결과
|
|
expansion_size (int): 대형 확장 크기
|
|
blur_size (int): 대형 블러 크기
|
|
|
|
Returns:
|
|
np.ndarray: 초블러 마스크
|
|
"""
|
|
image = cv2.imread(image_path)
|
|
if image is None:
|
|
return None
|
|
|
|
height, width = image.shape[:2]
|
|
mask = np.zeros((height, width), dtype=np.uint8)
|
|
|
|
print(f"\n🌟 초강력 블러 마스크 생성 시작")
|
|
print(f"📏 초대형 확장: {expansion_size}px, 초대형 블러: {blur_size}px")
|
|
|
|
# 기본 마스크 생성
|
|
for result in ocr_results:
|
|
polygon = result['polygon']
|
|
polygon_np = np.array(polygon, dtype=np.int32)
|
|
cv2.fillPoly(mask, [polygon_np], 255)
|
|
|
|
# 다단계 확장
|
|
for i in range(3):
|
|
kernel_size = expansion_size - (i * 3)
|
|
if kernel_size > 0:
|
|
kernel = np.ones((kernel_size, kernel_size), np.uint8)
|
|
mask = cv2.dilate(mask, kernel, iterations=1)
|
|
print(f" 📈 확장 단계 {i+1}: {kernel_size}px")
|
|
|
|
# 다단계 블러
|
|
blur_steps = [blur_size, blur_size//2, blur_size//4]
|
|
for i, blur in enumerate(blur_steps):
|
|
if blur > 0:
|
|
blur = blur if blur % 2 == 1 else blur + 1
|
|
mask = cv2.GaussianBlur(mask, (blur, blur), 0)
|
|
print(f" 🌫️ 블러 단계 {i+1}: {blur}px")
|
|
|
|
print(f"✅ 초강력 블러 마스크 완료")
|
|
return mask
|
|
|
|
def create_advanced_mask(self, image_path: str, ocr_results: List[Dict],
|
|
method: str = "morphology") -> np.ndarray:
|
|
"""
|
|
고급 마스크 생성 방법들
|
|
|
|
Args:
|
|
image_path (str): 원본 이미지 경로
|
|
ocr_results (List[Dict]): OCR 결과
|
|
method (str): 마스크 생성 방법
|
|
- "morphology": 모폴로지 연산
|
|
- "adaptive": 적응적 확장
|
|
- "contour": 컨투어 기반
|
|
|
|
Returns:
|
|
np.ndarray: 생성된 마스크
|
|
"""
|
|
image = cv2.imread(image_path)
|
|
if image is None:
|
|
return None
|
|
|
|
height, width = image.shape[:2]
|
|
mask = np.zeros((height, width), dtype=np.uint8)
|
|
|
|
if method == "morphology":
|
|
return self._create_morphology_mask(image, ocr_results)
|
|
elif method == "adaptive":
|
|
return self._create_adaptive_mask(image, ocr_results)
|
|
elif method == "contour":
|
|
return self._create_contour_mask(image, ocr_results)
|
|
else:
|
|
return self.create_masks(image_path, ocr_results)
|
|
|
|
def _create_morphology_mask(self, image: np.ndarray,
|
|
ocr_results: List[Dict]) -> np.ndarray:
|
|
"""모폴로지 연산을 사용한 마스크 생성"""
|
|
height, width = image.shape[:2]
|
|
mask = np.zeros((height, width), dtype=np.uint8)
|
|
|
|
for result in ocr_results:
|
|
polygon = result['polygon']
|
|
polygon_np = np.array(polygon, dtype=np.int32)
|
|
cv2.fillPoly(mask, [polygon_np], 255)
|
|
|
|
# 모폴로지 연산
|
|
kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (7, 7))
|
|
mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)
|
|
mask = cv2.dilate(mask, kernel, iterations=1)
|
|
|
|
return mask
|
|
|
|
def _create_adaptive_mask(self, image: np.ndarray,
|
|
ocr_results: List[Dict]) -> np.ndarray:
|
|
"""적응적 확장을 사용한 마스크 생성"""
|
|
height, width = image.shape[:2]
|
|
mask = np.zeros((height, width), dtype=np.uint8)
|
|
|
|
for result in ocr_results:
|
|
polygon = result['polygon']
|
|
polygon_np = np.array(polygon, dtype=np.int32)
|
|
|
|
# 텍스트 크기에 따라 확장 크기 조정
|
|
bbox = result['bbox']
|
|
text_area = bbox[2] * bbox[3]
|
|
expansion_size = max(3, int(np.sqrt(text_area) / 20))
|
|
|
|
# 임시 마스크 생성
|
|
temp_mask = np.zeros((height, width), dtype=np.uint8)
|
|
cv2.fillPoly(temp_mask, [polygon_np], 255)
|
|
|
|
# 적응적 확장
|
|
kernel = np.ones((expansion_size, expansion_size), np.uint8)
|
|
temp_mask = cv2.dilate(temp_mask, kernel, iterations=1)
|
|
|
|
# 메인 마스크에 추가
|
|
mask = cv2.bitwise_or(mask, temp_mask)
|
|
|
|
return mask
|
|
|
|
def _create_contour_mask(self, image: np.ndarray,
|
|
ocr_results: List[Dict]) -> np.ndarray:
|
|
"""컨투어 기반 마스크 생성"""
|
|
height, width = image.shape[:2]
|
|
mask = np.zeros((height, width), dtype=np.uint8)
|
|
|
|
for result in ocr_results:
|
|
polygon = result['polygon']
|
|
polygon_np = np.array(polygon, dtype=np.int32)
|
|
|
|
# 컨투어로 변환
|
|
contour = polygon_np.reshape(-1, 1, 2)
|
|
|
|
# 컨투어 확장
|
|
epsilon = 0.02 * cv2.arcLength(contour, True)
|
|
approx = cv2.approxPolyDP(contour, epsilon, True)
|
|
|
|
# 확장된 컨투어로 마스크 생성
|
|
cv2.fillPoly(mask, [approx], 255)
|
|
|
|
return mask
|
|
|
|
def create_feathered_mask(self, mask: np.ndarray, feather_size: int = 10) -> np.ndarray:
|
|
"""
|
|
페더링된 마스크 생성 (부드러운 경계)
|
|
|
|
Args:
|
|
mask (np.ndarray): 원본 마스크
|
|
feather_size (int): 페더링 크기
|
|
|
|
Returns:
|
|
np.ndarray: 페더링된 마스크
|
|
"""
|
|
# 거리 변환
|
|
dist_transform = cv2.distanceTransform(mask, cv2.DIST_L2, 5)
|
|
|
|
# 페더링 적용
|
|
feathered = np.clip(dist_transform / feather_size * 255, 0, 255).astype(np.uint8)
|
|
|
|
return feathered
|
|
|
|
def visualize_mask(self, image_path: str, mask: np.ndarray,
|
|
output_path: str = None, alpha: float = 0.5) -> np.ndarray:
|
|
"""
|
|
마스크를 원본 이미지에 오버레이하여 시각화
|
|
|
|
Args:
|
|
image_path (str): 원본 이미지 경로
|
|
mask (np.ndarray): 마스크
|
|
output_path (str): 저장할 경로
|
|
alpha (float): 투명도
|
|
|
|
Returns:
|
|
np.ndarray: 시각화된 이미지
|
|
"""
|
|
image = cv2.imread(image_path)
|
|
if image is None:
|
|
return None
|
|
|
|
# 마스크를 3채널로 변환
|
|
if len(mask.shape) == 2:
|
|
mask_colored = cv2.cvtColor(mask, cv2.COLOR_GRAY2BGR)
|
|
else:
|
|
mask_colored = mask
|
|
|
|
# 마스크에 색상 적용 (빨간색)
|
|
mask_colored[:, :, 0] = 0 # Blue
|
|
mask_colored[:, :, 1] = 0 # Green
|
|
# Red 채널은 그대로 유지
|
|
|
|
# 오버레이
|
|
overlay = cv2.addWeighted(image, 1-alpha, mask_colored, alpha, 0)
|
|
|
|
if output_path:
|
|
cv2.imwrite(output_path, overlay)
|
|
print(f"마스크 시각화 결과 저장: {output_path}")
|
|
|
|
return overlay
|
|
|
|
def get_mask_statistics(self, mask: np.ndarray) -> Dict[str, Any]:
|
|
"""
|
|
마스크 통계 정보 반환
|
|
|
|
Args:
|
|
mask (np.ndarray): 마스크
|
|
|
|
Returns:
|
|
Dict: 통계 정보
|
|
"""
|
|
total_pixels = mask.shape[0] * mask.shape[1]
|
|
mask_pixels = np.sum(mask > 0)
|
|
coverage = mask_pixels / total_pixels * 100
|
|
|
|
# 컨투어 찾기
|
|
contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
|
num_regions = len(contours)
|
|
|
|
# 가장 큰 영역 크기
|
|
if contours:
|
|
largest_area = max(cv2.contourArea(contour) for contour in contours)
|
|
else:
|
|
largest_area = 0
|
|
|
|
stats = {
|
|
'total_pixels': total_pixels,
|
|
'mask_pixels': mask_pixels,
|
|
'coverage_percent': coverage,
|
|
'num_regions': num_regions,
|
|
'largest_region_area': largest_area
|
|
}
|
|
|
|
return stats
|
|
|
|
def test_module(self):
|
|
"""마스크 모듈 테스트 (강화된 블러 확장 포함)"""
|
|
print("🎭 마스크 모듈 테스트 시작...")
|
|
|
|
# 테스트용 가상 OCR 결과 생성
|
|
fake_ocr_results = [
|
|
{
|
|
'text': '测试文本1',
|
|
'confidence': 0.95,
|
|
'polygon': [[100, 50], [200, 50], [200, 80], [100, 80]],
|
|
'bbox': (100, 50, 100, 30)
|
|
},
|
|
{
|
|
'text': '测试文本2',
|
|
'confidence': 0.88,
|
|
'polygon': [[150, 120], [300, 120], [300, 160], [150, 160]],
|
|
'bbox': (150, 120, 150, 40)
|
|
}
|
|
]
|
|
|
|
# 테스트 이미지 생성 (없을 경우)
|
|
test_image_path = "test_output/test_mask_image.jpg"
|
|
os.makedirs("test_output", exist_ok=True)
|
|
|
|
# 간단한 테스트 이미지 생성
|
|
test_image = np.ones((200, 400, 3), dtype=np.uint8) * 255
|
|
cv2.putText(test_image, "Test Image for Mask", (50, 100),
|
|
cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 0), 2)
|
|
cv2.imwrite(test_image_path, test_image)
|
|
|
|
# 1. 기본 마스크 생성 테스트
|
|
print("\n📋 1. 기본 마스크 테스트")
|
|
mask = self.create_masks(test_image_path, fake_ocr_results)
|
|
if mask is not None:
|
|
print("✅ 기본 마스크 생성 성공")
|
|
|
|
# 마스크 통계
|
|
stats = self.get_mask_statistics(mask)
|
|
print(f"📊 마스크 통계:")
|
|
print(f" - 총 픽셀: {stats['total_pixels']:,}")
|
|
print(f" - 마스크 픽셀: {stats['mask_pixels']:,}")
|
|
print(f" - 커버리지: {stats['coverage_percent']:.2f}%")
|
|
print(f" - 영역 수: {stats['num_regions']}")
|
|
print(f" - 최대 영역 크기: {stats['largest_region_area']:,}")
|
|
|
|
# 시각화
|
|
self.visualize_mask(test_image_path, mask, "test_output/mask_basic.jpg")
|
|
|
|
# 2. 강화된 블러 확장 마스크 테스트
|
|
print("\n🌟 2. 강화된 블러 확장 마스크 테스트")
|
|
image = cv2.imread(test_image_path)
|
|
height, width = image.shape[:2]
|
|
base_mask = np.zeros((height, width), dtype=np.uint8)
|
|
|
|
# 기본 마스크 생성
|
|
for result in fake_ocr_results:
|
|
polygon = result['polygon']
|
|
polygon_np = np.array(polygon, dtype=np.int32)
|
|
cv2.fillPoly(base_mask, [polygon_np], 255)
|
|
|
|
# 강화된 블러 적용
|
|
enhanced_mask = self.process_mask_with_enhanced_blur(base_mask, 12, 20)
|
|
self.visualize_mask(test_image_path, enhanced_mask, "test_output/mask_enhanced_blur.jpg")
|
|
print("✅ 강화된 블러 마스크 완료")
|
|
|
|
# 3. 초강력 블러 마스크 테스트
|
|
print("\n🚀 3. 초강력 블러 마스크 테스트")
|
|
ultra_mask = self.create_ultra_blur_mask(test_image_path, fake_ocr_results, 25, 35)
|
|
if ultra_mask is not None:
|
|
self.visualize_mask(test_image_path, ultra_mask, "test_output/mask_ultra_blur.jpg")
|
|
print("✅ 초강력 블러 마스크 완료")
|
|
|
|
# 4. 고급 마스크 테스트
|
|
print("\n🔧 4. 고급 마스크 방법 테스트")
|
|
methods = ["morphology", "adaptive", "contour"]
|
|
for method in methods:
|
|
print(f" 🔹 {method} 방법 테스트 중...")
|
|
advanced_mask = self.create_advanced_mask(
|
|
test_image_path, fake_ocr_results, method
|
|
)
|
|
if advanced_mask is not None:
|
|
output_path = f"test_output/mask_{method}.jpg"
|
|
self.visualize_mask(test_image_path, advanced_mask, output_path)
|
|
print(f" ✅ {method} 마스크 완료")
|
|
|
|
# 5. 페더링 마스크 테스트
|
|
print("\n✨ 5. 페더링 마스크 테스트")
|
|
if mask is not None:
|
|
feathered_mask = self.create_feathered_mask(mask, 15)
|
|
self.visualize_mask(test_image_path, feathered_mask,
|
|
"test_output/mask_feathered.jpg")
|
|
print("✅ 페더링 마스크 완료")
|
|
|
|
print("\n🎉 마스크 모듈 테스트 완료!")
|
|
print("📁 결과 파일들이 test_output/ 폴더에 저장되었습니다.")
|
|
print(" - mask_basic.jpg: 기본 마스크")
|
|
print(" - mask_enhanced_blur.jpg: 강화된 블러 마스크")
|
|
print(" - mask_ultra_blur.jpg: 초강력 블러 마스크")
|
|
print(" - mask_morphology.jpg: 모폴로지 마스크")
|
|
print(" - mask_adaptive.jpg: 적응적 마스크")
|
|
print(" - mask_contour.jpg: 컨투어 마스크")
|
|
print(" - mask_feathered.jpg: 페더링 마스크")
|
|
|
|
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:
|
|
"""
|
|
OCR 결과를 기반으로 마스크 생성 및 저장
|
|
|
|
Args:
|
|
image_path (str): 원본 이미지 경로
|
|
ocr_results (List[Dict]): OCR 결과
|
|
output_dir (str): 출력 디렉토리 (마스크 이미지 저장용)
|
|
expansion_size (int): 마스크 확장 크기
|
|
blur_size (int): 블러 크기
|
|
mask_option (str): 마스크 옵션 (기본값 "basic")
|
|
Returns:
|
|
tuple: (mask, mask_image_path)
|
|
"""
|
|
# 기본 마스크 생성
|
|
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)
|
|
print(f"💾 마스크 이미지 저장: {mask_image_path}")
|
|
|
|
# 개별 마스크들도 저장
|
|
self._save_individual_masks(image_path, ocr_results, output_dir)
|
|
|
|
# 마스크 시각화 이미지도 저장
|
|
visualization_path = os.path.join(output_dir, f"{base_name}_mask_visualization.jpg")
|
|
self.visualize_mask(image_path, mask, visualization_path)
|
|
print(f"💾 마스크 시각화 이미지 저장: {visualization_path}")
|
|
|
|
return mask, mask_image_path
|
|
|
|
def _save_individual_masks(self, image_path: str, ocr_results: List[Dict], output_dir: str):
|
|
"""개별 텍스트 영역의 마스크를 각각 저장"""
|
|
image = cv2.imread(image_path)
|
|
if image is None:
|
|
return
|
|
|
|
height, width = image.shape[:2]
|
|
base_name = os.path.splitext(os.path.basename(image_path))[0]
|
|
|
|
for i, result in enumerate(ocr_results, 1):
|
|
# 개별 마스크 생성
|
|
individual_mask = np.zeros((height, width), dtype=np.uint8)
|
|
polygon = result['polygon']
|
|
polygon_np = np.array(polygon, dtype=np.int32)
|
|
cv2.fillPoly(individual_mask, [polygon_np], 255)
|
|
|
|
# 개별 마스크 저장
|
|
individual_filename = f"{base_name}_mask_area_{i:02d}.jpg"
|
|
individual_path = os.path.join(output_dir, individual_filename)
|
|
cv2.imwrite(individual_path, individual_mask)
|
|
|
|
print(f" 💾 개별 마스크 #{i} 저장: {individual_filename}")
|
|
|
|
if __name__ == "__main__":
|
|
mask_module = MaskModule()
|
|
mask_module.test_module() |