AutoPercenty3/test/server_inpaint_smoke.py

import os
import re
import sys
import ast
import json
import cv2
import numpy as np

#사용법

# ./Scripts/python.exe -X utf8 test/server_inpaint_smoke.py --image D:\py\AutoPercenty3_311\src\temp_images\detail_image_000_O1CN016JhqPc28B31bgShOk_!!3823937893_resized.jpg_Q75_resized.jpg


def _project_root():
    here = os.path.abspath(os.path.dirname(__file__))
    return os.path.abspath(os.path.join(here, '..'))


def parse_case_from_log(log_path: str, target_image: str | None = None):
    """Edit_PartTimer3.log 에서 마지막 케이스의
    - 로컬 이미지 경로
    - filter_ocr_results (polygon 포함)
    - 요청 인페인트 서버 URL
    을 추출한다.
    """
    image_path = None
    ocr_results = None
    inpaint_endpoint = None

    if not os.path.exists(log_path):
        raise FileNotFoundError(f'로그 파일을 찾을 수 없음: {log_path}')

    with open(log_path, 'r', encoding='utf-8', errors='ignore') as f:
        lines = f.readlines()

    def _parse_block(start_idx: int):
        nonlocal image_path, ocr_results, inpaint_endpoint
        # 앞/뒤 모두에서 찾을 수 있으니 주변에서 스캔
        for i in range(start_idx, min(start_idx + 800, len(lines))):
            line = lines[i]
            if ocr_results is None and 'filter_ocr_results:' in line:
                try:
                    txt = line.split('filter_ocr_results:', 1)[1].strip()
                    ocr_results = ast.literal_eval(txt)
                except Exception:
                    buf = [line.split('filter_ocr_results:', 1)[1]]
                    j = i + 1
                    while j < len(lines) and ']' not in ''.join(buf):
                        buf.append(lines[j])
                        j += 1
                    try:
                        ocr_results = ast.literal_eval(''.join(buf))
                    except Exception:
                        pass
            if inpaint_endpoint is None and '요청 인페인팅 서버 :' in line:
                ep = line.split('요청 인페인팅 서버 :', 1)[1].strip()
                inpaint_endpoint = ep.split()[0]
            if ocr_results is not None and inpaint_endpoint:
                break

    if target_image:
        # 정확히 해당 이미지 경로를 포함한 라인을 찾음
        target_idx = None
        for i, line in enumerate(lines):
            if target_image in line:
                target_idx = i
        if target_idx is None:
            raise RuntimeError(f'지정한 이미지 경로를 로그에서 찾을 수 없음: {target_image}')
        image_path = target_image
        _parse_block(target_idx)
    else:
        # 뒤에서부터 스캔하여 마지막 케이스 사용
        for i in range(len(lines) - 1, -1, -1):
            line = lines[i]
            if image_path is None:
                m = re.search(r'로컬 저장위치\(.*\):\s*(.+)$', line)
                if m:
                    image_path = m.group(1).strip()
                    continue
            if ocr_results is None and 'filter_ocr_results:' in line:
                try:
                    txt = line.split('filter_ocr_results:', 1)[1].strip()
                    ocr_results = ast.literal_eval(txt)
                except Exception:
                    buf = [line.split('filter_ocr_results:', 1)[1]]
                    j = i + 1
                    while j < len(lines) and ']' not in ''.join(buf):
                        buf.append(lines[j])
                        j += 1
                    try:
                        ocr_results = ast.literal_eval(''.join(buf))
                    except Exception:
                        pass
                continue
            if inpaint_endpoint is None and '요청 인페인팅 서버 :' in line:
                ep = line.split('요청 인페인팅 서버 :', 1)[1].strip()
                inpaint_endpoint = ep.split()[0]
                continue
            if image_path and ocr_results is not None and inpaint_endpoint:
                break

    if not image_path:
        raise RuntimeError('로그에서 로컬 이미지 경로를 찾지 못했습니다.')
    if ocr_results is None:
        raise RuntimeError('로그에서 filter_ocr_results 를 찾지 못했습니다.')

    # endpoint 가 전체 경로(/api/v1/inpaint 포함)라면 그대로 사용
    return image_path, ocr_results, inpaint_endpoint


def build_mask(image_path: str, ocr_results: list, expansion_size: int = 9, blur_size: int = 12):
    # 모듈 임포트 경로 설정
    sys.path.append(_project_root())
    from src.modules.mask_module_for_paddle import MaskModule

    class _DummyLogger:
        def log(self, msg, level=None, exc_info=False):
            print(msg)

    mm = MaskModule(logger=_DummyLogger(), base_dir=_project_root())
    mask = mm.create_masks(
        image_path=image_path,
        ocr_results=ocr_results,
        mask_option='basic',
        expansion_size=expansion_size,
        blur_size=blur_size,
    )
    return mask


def request_inpaint(endpoint: str, image_path: str, mask: np.ndarray, invert_mask: bool = False, save_path: str = None):
    sys.path.append(_project_root())
    from src.modules.request_inpaint import Request_AI_Server

    class _DummyLogger:
        def log(self, msg, level=None, exc_info=False):
            print(msg)

    # endpoint 가 /api/v1/inpaint 까지 포함되어 오면 base_url 을 잘라낸다
    base_url = endpoint
    if '/api/v1/inpaint' in endpoint:
        base_url = endpoint.split('/api/v1/inpaint', 1)[0]

    client = Request_AI_Server(logger=_DummyLogger(), inpaint_server_url=base_url)
    result = client.request_inpaint(image_path, mask, invert_mask=invert_mask)
    if result is not None and save_path:
        os.makedirs(os.path.dirname(save_path), exist_ok=True)
        cv2.imwrite(save_path, result)
    return result


def main():
    import argparse
    ap = argparse.ArgumentParser()
    ap.add_argument('--image', type=str, default=None, help='테스트할 원본 이미지 경로(로그에 존재해야 함)')
    ap.add_argument('--invert', action='store_true', help='반전 only 단일 테스트')
    args = ap.parse_args()
    root = _project_root()
    log_path = os.path.join(root, 'src', 'logs', 'Edit_PartTimer3.log')
    debug_dir = os.path.join(root, 'src', 'debug_images')
    os.makedirs(debug_dir, exist_ok=True)

    image_path, ocr_results, endpoint = parse_case_from_log(log_path, target_image=args.image)
    print('[SMOKE] image_path =', image_path)
    print('[SMOKE] endpoint   =', endpoint)
    print('[SMOKE] ocr count  =', len(ocr_results))

    mask = build_mask(image_path, ocr_results)
    mask_path = os.path.join(debug_dir, 'smoke_mask.png')
    cv2.imwrite(mask_path, (mask > 0).astype(np.uint8) * 255)
    print('[SMOKE] mask saved =', mask_path)

    if args.invert:
        out_path = os.path.join(debug_dir, 'smoke_inpaint_invert_1.png')
        res = request_inpaint(endpoint, image_path, mask, invert_mask=True, save_path=out_path)
        print(f'[SMOKE] invert=True →', 'OK' if res is not None else 'FAILED', '→', out_path)
    else:
        for invert in (False, True):
            out_path = os.path.join(debug_dir, f'smoke_inpaint_invert_{int(invert)}.png')
            res = request_inpaint(endpoint, image_path, mask, invert_mask=invert, save_path=out_path)
            print(f'[SMOKE] invert={invert} →', 'OK' if res is not None else 'FAILED', '→', out_path)


if __name__ == '__main__':
    main()