AutoPercenty3/mask_cli.py

555 lines
21 KiB
Python

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
마스크 CLI 도구 - 마스크 모듈의 모든 기능을 커맨드라인에서 사용
"""
import argparse
import sys
import os
import json
from pathlib import Path
# 모듈 경로 추가
sys.path.append(str(Path(__file__).parent / "test" / "img_test2" / "modules"))
from ocr_module import OCRModule
from mask_module import MaskModule
def create_parser():
"""커맨드라인 파서 생성"""
parser = argparse.ArgumentParser(
description="마스크 생성 및 처리 도구",
formatter_class=argparse.RawDescriptionHelpFormatter,
epilog="""
사용 예시:
# 기본 마스크 생성
python mask_cli.py basic -i image.jpg -o output/
# 강화된 블러 마스크
python mask_cli.py enhanced -i image.jpg -o output/ --expansion 15 --blur 25
# 초강력 블러 마스크
python mask_cli.py ultra -i image.jpg -o output/ --expansion 30 --blur 40
# 고급 마스크 (모폴로지)
python mask_cli.py advanced -i image.jpg -o output/ --method morphology
# 페더링 마스크
python mask_cli.py feathered -i image.jpg -o output/ --feather 20
# 마스크 통계 확인
python mask_cli.py stats -i image.jpg
# 여러 방법으로 일괄 처리
python mask_cli.py batch -i image.jpg -o output/ --methods basic,enhanced,ultra
# OCR 결과 저장
python mask_cli.py ocr-only -i image.jpg -o results.json
"""
)
subparsers = parser.add_subparsers(dest='command', help='사용할 마스크 방법')
# 공통 인수
def add_common_args(parser):
parser.add_argument('-i', '--input', required=True, help='입력 이미지 파일')
parser.add_argument('-o', '--output', help='출력 디렉토리 또는 파일')
parser.add_argument('--ocr-results', help='OCR 결과 JSON 파일 (선택사항)')
parser.add_argument('--lang', default='korean', help='OCR 언어 (기본: korean)')
parser.add_argument('--verbose', '-v', action='store_true', help='상세 출력')
# 1. 기본 마스크
basic_parser = subparsers.add_parser('basic', help='기본 마스크 생성')
add_common_args(basic_parser)
basic_parser.add_argument('--expansion', type=int, default=10, help='확장 크기 (기본: 10)')
basic_parser.add_argument('--blur', type=int, default=15, help='블러 크기 (기본: 15)')
# 2. 강화된 블러 마스크
enhanced_parser = subparsers.add_parser('enhanced', help='강화된 블러 마스크 생성')
add_common_args(enhanced_parser)
enhanced_parser.add_argument('--expansion', type=int, default=12, help='확장 크기 (기본: 12)')
enhanced_parser.add_argument('--blur', type=int, default=20, help='블러 크기 (기본: 20)')
# 3. 초강력 블러 마스크
ultra_parser = subparsers.add_parser('ultra', help='초강력 블러 마스크 생성')
add_common_args(ultra_parser)
ultra_parser.add_argument('--expansion', type=int, default=25, help='확장 크기 (기본: 25)')
ultra_parser.add_argument('--blur', type=int, default=35, help='블러 크기 (기본: 35)')
# 4. 고급 마스크
advanced_parser = subparsers.add_parser('advanced', help='고급 마스크 생성')
add_common_args(advanced_parser)
advanced_parser.add_argument('--method', choices=['morphology', 'adaptive', 'contour'],
default='morphology', help='고급 마스크 방법')
# 5. 페더링 마스크
feathered_parser = subparsers.add_parser('feathered', help='페더링 마스크 생성')
add_common_args(feathered_parser)
feathered_parser.add_argument('--feather', type=int, default=10, help='페더링 크기 (기본: 10)')
feathered_parser.add_argument('--expansion', type=int, default=10, help='기본 확장 크기 (기본: 10)')
feathered_parser.add_argument('--blur', type=int, default=15, help='기본 블러 크기 (기본: 15)')
# 6. 마스크 통계
stats_parser = subparsers.add_parser('stats', help='마스크 통계 정보')
add_common_args(stats_parser)
stats_parser.add_argument('--expansion', type=int, default=10, help='확장 크기 (기본: 10)')
stats_parser.add_argument('--blur', type=int, default=15, help='블러 크기 (기본: 15)')
# 7. 일괄 처리
batch_parser = subparsers.add_parser('batch', help='여러 방법으로 일괄 처리')
add_common_args(batch_parser)
batch_parser.add_argument('--methods', default='basic,enhanced,ultra',
help='사용할 방법들 (쉼표로 구분, 기본: basic,enhanced,ultra)')
batch_parser.add_argument('--expansion', type=int, default=15, help='확장 크기 (기본: 15)')
batch_parser.add_argument('--blur', type=int, default=20, help='블러 크기 (기본: 20)')
# 8. OCR만 실행
ocr_parser = subparsers.add_parser('ocr-only', help='OCR만 실행하고 결과 저장')
add_common_args(ocr_parser)
# 9. 시각화
viz_parser = subparsers.add_parser('visualize', help='마스크 시각화')
add_common_args(viz_parser)
viz_parser.add_argument('--expansion', type=int, default=10, help='확장 크기 (기본: 10)')
viz_parser.add_argument('--blur', type=int, default=15, help='블러 크기 (기본: 15)')
viz_parser.add_argument('--alpha', type=float, default=0.5, help='투명도 (기본: 0.5)')
# 10. 테스트
test_parser = subparsers.add_parser('test', help='마스크 모듈 테스트')
test_parser.add_argument('-o', '--output', default='test_output', help='테스트 출력 디렉토리')
return parser
def load_ocr_results(ocr_results_path):
"""OCR 결과 JSON 파일 로드"""
try:
with open(ocr_results_path, 'r', encoding='utf-8') as f:
return json.load(f)
except Exception as e:
print(f"❌ OCR 결과 파일 로드 실패: {e}")
return None
def save_ocr_results(ocr_results, output_path):
"""OCR 결과를 JSON 파일로 저장"""
try:
with open(output_path, 'w', encoding='utf-8') as f:
json.dump(ocr_results, f, ensure_ascii=False, indent=2)
print(f"💾 OCR 결과 저장: {output_path}")
except Exception as e:
print(f"❌ OCR 결과 저장 실패: {e}")
def run_ocr(image_path, lang='korean', ocr_results_path=None):
"""OCR 실행"""
if ocr_results_path and os.path.exists(ocr_results_path):
print(f"📁 기존 OCR 결과 사용: {ocr_results_path}")
return load_ocr_results(ocr_results_path)
print(f"🔍 OCR 실행 중: {image_path}")
ocr_module = OCRModule(lang=lang)
ocr_results = ocr_module.detect_text(image_path, method='polygon')
if not ocr_results:
print("❌ OCR 결과가 없습니다.")
return None
print(f"✅ OCR 완료: {len(ocr_results)}개 텍스트 영역 발견")
return ocr_results
def print_mask_stats(mask_module, mask):
"""마스크 통계 출력"""
stats = mask_module.get_mask_statistics(mask)
print(f"\n📊 마스크 통계:")
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']:,}")
def basic_mask(args):
"""기본 마스크 생성"""
print(f"🎭 기본 마스크 생성 시작...")
ocr_results = run_ocr(args.input, args.lang, args.ocr_results)
if not ocr_results:
return
mask_module = MaskModule()
mask = mask_module.create_masks(args.input, ocr_results, args.expansion, args.blur)
if mask is None:
print("❌ 마스크 생성 실패")
return
if args.verbose:
print_mask_stats(mask_module, mask)
if args.output:
os.makedirs(args.output, exist_ok=True)
mask, mask_path = mask_module.create_masks_with_save(
args.input, ocr_results, args.output, args.expansion, args.blur
)
print(f"✅ 기본 마스크 완료: {args.output}")
def enhanced_mask(args):
"""강화된 블러 마스크 생성"""
print(f"🌟 강화된 블러 마스크 생성 시작...")
ocr_results = run_ocr(args.input, args.lang, args.ocr_results)
if not ocr_results:
return
mask_module = MaskModule()
# 기본 마스크 생성 후 강화된 블러 적용
import cv2
import numpy as np
image = cv2.imread(args.input)
height, width = image.shape[:2]
base_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(base_mask, [polygon_np], 255)
enhanced_mask = mask_module.process_mask_with_enhanced_blur(
base_mask, args.expansion, args.blur
)
if args.verbose:
print_mask_stats(mask_module, enhanced_mask)
if args.output:
os.makedirs(args.output, exist_ok=True)
base_name = os.path.splitext(os.path.basename(args.input))[0]
# 마스크 저장
mask_path = os.path.join(args.output, f"{base_name}_enhanced_mask.jpg")
cv2.imwrite(mask_path, enhanced_mask)
print(f"💾 강화된 마스크 저장: {mask_path}")
# 시각화 저장
viz_path = os.path.join(args.output, f"{base_name}_enhanced_visualization.jpg")
mask_module.visualize_mask(args.input, enhanced_mask, viz_path)
print(f"✅ 강화된 블러 마스크 완료: {args.output}")
def ultra_mask(args):
"""초강력 블러 마스크 생성"""
print(f"🚀 초강력 블러 마스크 생성 시작...")
ocr_results = run_ocr(args.input, args.lang, args.ocr_results)
if not ocr_results:
return
mask_module = MaskModule()
ultra_mask = mask_module.create_ultra_blur_mask(
args.input, ocr_results, args.expansion, args.blur
)
if ultra_mask is None:
print("❌ 초강력 블러 마스크 생성 실패")
return
if args.verbose:
print_mask_stats(mask_module, ultra_mask)
if args.output:
os.makedirs(args.output, exist_ok=True)
base_name = os.path.splitext(os.path.basename(args.input))[0]
# 마스크 저장
import cv2
mask_path = os.path.join(args.output, f"{base_name}_ultra_mask.jpg")
cv2.imwrite(mask_path, ultra_mask)
print(f"💾 초강력 마스크 저장: {mask_path}")
# 시각화 저장
viz_path = os.path.join(args.output, f"{base_name}_ultra_visualization.jpg")
mask_module.visualize_mask(args.input, ultra_mask, viz_path)
print(f"✅ 초강력 블러 마스크 완료: {args.output}")
def advanced_mask(args):
"""고급 마스크 생성"""
print(f"🔧 고급 마스크 생성 시작 ({args.method} 방법)...")
ocr_results = run_ocr(args.input, args.lang, args.ocr_results)
if not ocr_results:
return
mask_module = MaskModule()
advanced_mask = mask_module.create_advanced_mask(
args.input, ocr_results, args.method
)
if advanced_mask is None:
print("❌ 고급 마스크 생성 실패")
return
if args.verbose:
print_mask_stats(mask_module, advanced_mask)
if args.output:
os.makedirs(args.output, exist_ok=True)
base_name = os.path.splitext(os.path.basename(args.input))[0]
# 마스크 저장
import cv2
mask_path = os.path.join(args.output, f"{base_name}_{args.method}_mask.jpg")
cv2.imwrite(mask_path, advanced_mask)
print(f"💾 {args.method} 마스크 저장: {mask_path}")
# 시각화 저장
viz_path = os.path.join(args.output, f"{base_name}_{args.method}_visualization.jpg")
mask_module.visualize_mask(args.input, advanced_mask, viz_path)
print(f"✅ 고급 마스크 완료 ({args.method}): {args.output}")
def feathered_mask(args):
"""페더링 마스크 생성"""
print(f"✨ 페더링 마스크 생성 시작...")
ocr_results = run_ocr(args.input, args.lang, args.ocr_results)
if not ocr_results:
return
mask_module = MaskModule()
# 기본 마스크 생성
base_mask = mask_module.create_masks(args.input, ocr_results, args.expansion, args.blur)
if base_mask is None:
print("❌ 기본 마스크 생성 실패")
return
# 페더링 적용
feathered_mask = mask_module.create_feathered_mask(base_mask, args.feather)
if args.verbose:
print_mask_stats(mask_module, feathered_mask)
if args.output:
os.makedirs(args.output, exist_ok=True)
base_name = os.path.splitext(os.path.basename(args.input))[0]
# 마스크 저장
import cv2
mask_path = os.path.join(args.output, f"{base_name}_feathered_mask.jpg")
cv2.imwrite(mask_path, feathered_mask)
print(f"💾 페더링 마스크 저장: {mask_path}")
# 시각화 저장
viz_path = os.path.join(args.output, f"{base_name}_feathered_visualization.jpg")
mask_module.visualize_mask(args.input, feathered_mask, viz_path)
print(f"✅ 페더링 마스크 완료: {args.output}")
def stats_mask(args):
"""마스크 통계 정보"""
print(f"📊 마스크 통계 분석 시작...")
ocr_results = run_ocr(args.input, args.lang, args.ocr_results)
if not ocr_results:
return
mask_module = MaskModule()
mask = mask_module.create_masks(args.input, ocr_results, args.expansion, args.blur)
if mask is None:
print("❌ 마스크 생성 실패")
return
print_mask_stats(mask_module, mask)
# OCR 결과 정보도 출력
print(f"\n🔍 OCR 결과 정보:")
for i, result in enumerate(ocr_results, 1):
print(f" 영역 #{i}: '{result['text']}' (신뢰도: {result['confidence']:.1%})")
def batch_mask(args):
"""일괄 처리"""
print(f"🔄 일괄 처리 시작...")
ocr_results = run_ocr(args.input, args.lang, args.ocr_results)
if not ocr_results:
return
methods = args.methods.split(',')
mask_module = MaskModule()
if args.output:
os.makedirs(args.output, exist_ok=True)
base_name = os.path.splitext(os.path.basename(args.input))[0]
for method in methods:
method = method.strip()
print(f"\n🔹 {method} 방법 처리 중...")
if method == 'basic':
mask = mask_module.create_masks(args.input, ocr_results, args.expansion, args.blur)
elif method == 'enhanced':
import cv2
import numpy as np
image = cv2.imread(args.input)
height, width = image.shape[:2]
base_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(base_mask, [polygon_np], 255)
mask = mask_module.process_mask_with_enhanced_blur(base_mask, args.expansion, args.blur)
elif method == 'ultra':
mask = mask_module.create_ultra_blur_mask(args.input, ocr_results, args.expansion, args.blur)
elif method in ['morphology', 'adaptive', 'contour']:
mask = mask_module.create_advanced_mask(args.input, ocr_results, method)
elif method == 'feathered':
base_mask = mask_module.create_masks(args.input, ocr_results, args.expansion, args.blur)
mask = mask_module.create_feathered_mask(base_mask, 10)
else:
print(f" ❌ 알 수 없는 방법: {method}")
continue
if mask is not None and args.output:
import cv2
mask_path = os.path.join(args.output, f"{base_name}_{method}_mask.jpg")
cv2.imwrite(mask_path, mask)
viz_path = os.path.join(args.output, f"{base_name}_{method}_visualization.jpg")
mask_module.visualize_mask(args.input, mask, viz_path)
print(f"{method} 완료")
print(f"\n🎉 일괄 처리 완료: {args.output}")
def ocr_only(args):
"""OCR만 실행"""
print(f"🔍 OCR만 실행...")
ocr_results = run_ocr(args.input, args.lang)
if not ocr_results:
return
# 결과 출력
print(f"\n📋 OCR 결과 ({len(ocr_results)}개 영역):")
for i, result in enumerate(ocr_results, 1):
print(f" {i}. '{result['text']}' (신뢰도: {result['confidence']:.1%})")
print(f" 위치: {result['bbox']}")
print(f" 폴리곤: {result['polygon']}")
# JSON 저장
if args.output:
save_ocr_results(ocr_results, args.output)
def visualize_mask(args):
"""마스크 시각화"""
print(f"🎨 마스크 시각화 시작...")
ocr_results = run_ocr(args.input, args.lang, args.ocr_results)
if not ocr_results:
return
mask_module = MaskModule()
mask = mask_module.create_masks(args.input, ocr_results, args.expansion, args.blur)
if mask is None:
print("❌ 마스크 생성 실패")
return
if args.output:
os.makedirs(os.path.dirname(args.output) if os.path.dirname(args.output) else '.', exist_ok=True)
mask_module.visualize_mask(args.input, mask, args.output, args.alpha)
print(f"✅ 시각화 완료: {args.output}")
else:
# 기본 출력 경로
base_name = os.path.splitext(os.path.basename(args.input))[0]
output_path = f"{base_name}_visualization.jpg"
mask_module.visualize_mask(args.input, mask, output_path, args.alpha)
print(f"✅ 시각화 완료: {output_path}")
def test_mask(args):
"""마스크 모듈 테스트"""
print(f"🧪 마스크 모듈 테스트 시작...")
mask_module = MaskModule()
# 출력 디렉토리 설정
original_test_dir = "test_output"
if args.output != original_test_dir:
# test_module 메서드에서 사용하는 경로를 변경하기 위해 임시로 변경
import tempfile
import shutil
# 임시 디렉토리에서 테스트 실행
with tempfile.TemporaryDirectory() as temp_dir:
temp_test_dir = os.path.join(temp_dir, "test_output")
os.makedirs(temp_test_dir, exist_ok=True)
# 현재 디렉토리 변경
original_cwd = os.getcwd()
os.chdir(temp_dir)
try:
mask_module.test_module()
# 결과를 원하는 출력 디렉토리로 복사
os.makedirs(args.output, exist_ok=True)
for file in os.listdir(temp_test_dir):
shutil.copy2(os.path.join(temp_test_dir, file), args.output)
print(f"✅ 테스트 완료: {args.output}")
finally:
os.chdir(original_cwd)
else:
mask_module.test_module()
print(f"✅ 테스트 완료: {args.output}")
def main():
"""메인 함수"""
parser = create_parser()
args = parser.parse_args()
if not args.command:
parser.print_help()
return
# 입력 파일 확인 (test 명령어 제외)
if args.command != 'test' and hasattr(args, 'input') and not os.path.exists(args.input):
print(f"❌ 입력 파일이 존재하지 않습니다: {args.input}")
return
# 명령어 실행
try:
if args.command == 'basic':
basic_mask(args)
elif args.command == 'enhanced':
enhanced_mask(args)
elif args.command == 'ultra':
ultra_mask(args)
elif args.command == 'advanced':
advanced_mask(args)
elif args.command == 'feathered':
feathered_mask(args)
elif args.command == 'stats':
stats_mask(args)
elif args.command == 'batch':
batch_mask(args)
elif args.command == 'ocr-only':
ocr_only(args)
elif args.command == 'visualize':
visualize_mask(args)
elif args.command == 'test':
test_mask(args)
else:
print(f"❌ 알 수 없는 명령어: {args.command}")
parser.print_help()
except KeyboardInterrupt:
print("\n⚠️ 사용자에 의해 중단되었습니다.")
except Exception as e:
print(f"❌ 오류 발생: {e}")
if args.verbose if hasattr(args, 'verbose') else False:
import traceback
traceback.print_exc()
if __name__ == "__main__":
main()