AutoPercenty3/mask_cli.py

#!/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()