IMG_Worker/modules/test/test_external_inpaint.py

"""외부 인페인팅 서버 테스트 스크립트"""
import sys
import os
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))

import cv2
import numpy as np
import requests
import base64

# 테스트 서버 URL
SERVER_URL = "http://192.168.0.146:8008"

def test_health_check():
    """서버 상태 확인"""
    print("=" * 50)
    print("1. 서버 상태 확인 (Health Check)")
    print("=" * 50)
    try:
        response = requests.get(f"{SERVER_URL}/health", timeout=5)
        print(f"Status: {response.status_code}")
        print(f"Response: {response.json()}")
        return response.status_code == 200
    except Exception as e:
        print(f"Error: {e}")
        return False

def test_inpaint_api(model_name="migan"):
    """인페인팅 API 테스트"""
    print("=" * 50)
    print(f"2. 인페인팅 API 테스트 (model: {model_name})")
    print("=" * 50)
    
    # 테스트 이미지 로드
    test_image_path = os.path.join(os.path.dirname(__file__), "1.jpg")
    if not os.path.exists(test_image_path):
        print(f"테스트 이미지 없음: {test_image_path}")
        return False
    
    image = cv2.imread(test_image_path)
    print(f"이미지 크기: {image.shape}")
    
    # 간단한 테스트 마스크 생성 (중앙에 사각형)
    h, w = image.shape[:2]
    mask = np.zeros((h, w), dtype=np.uint8)
    # 중앙에 100x100 마스크 영역
    cx, cy = w // 2, h // 2
    mask[cy-50:cy+50, cx-50:cx+50] = 255
    print(f"마스크 크기: {mask.shape}, 마스크 영역: 100x100 중앙")
    
    # Base64 인코딩
    _, img_encoded = cv2.imencode('.png', image)
    _, mask_encoded = cv2.imencode('.png', mask)
    img_b64 = base64.b64encode(img_encoded).decode('utf-8')
    mask_b64 = base64.b64encode(mask_encoded).decode('utf-8')
    
    # API 요청
    api_url = f"{SERVER_URL}/api/v1/inpaint"
    params = {
        "response_format": "binary",
        "image_format": "webp",
    }
    payload = {
        "image": img_b64,
        "mask": mask_b64,
        "model_name": model_name
    }
    
    print(f"요청 URL: {api_url}")
    print(f"파라미터: {params}")
    print(f"모델: {model_name}")
    
    try:
        response = requests.post(api_url, params=params, json=payload, timeout=(5, 60))
        print(f"Status: {response.status_code}")
        print(f"Content-Type: {response.headers.get('content-type', 'N/A')}")
        print(f"응답 크기: {len(response.content)} bytes")
        
        if response.status_code == 200:
            # 이미지 디코딩 테스트
            nparr = np.frombuffer(response.content, np.uint8)
            result = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
            
            if result is not None:
                print(f"✅ 성공! 결과 이미지 크기: {result.shape}")
                
                # 결과 저장
                output_path = os.path.join(os.path.dirname(__file__), f"outputs_test/inpaint_result_{model_name}.webp")
                os.makedirs(os.path.dirname(output_path), exist_ok=True)
                cv2.imwrite(output_path, result)
                print(f"결과 저장: {output_path}")
                return True
            else:
                print("❌ 이미지 디코딩 실패")
                return False
        else:
            print(f"❌ 서버 에러: {response.text}")
            return False
            
    except Exception as e:
        print(f"❌ 요청 실패: {e}")
        return False

def test_with_request_ai_server():
    """Request_AI_Server 클래스를 통한 테스트"""
    print("=" * 50)
    print("3. Request_AI_Server 클래스 테스트")
    print("=" * 50)
    
    from loggerModule import Logger
    from modules.request_inpaint import Request_AI_Server
    
    logger = Logger()
    server = Request_AI_Server(logger)
    
    # 테스트 이미지 로드
    test_image_path = os.path.join(os.path.dirname(__file__), "1.jpg")
    image = cv2.imread(test_image_path)
    
    # 마스크 생성
    h, w = image.shape[:2]
    mask = np.zeros((h, w), dtype=np.uint8)
    cx, cy = w // 2, h // 2
    mask[cy-50:cy+50, cx-50:cx+50] = 255
    
    # 외부 인페인팅 요청
    print(f"서버 URL: {SERVER_URL}")
    result = server.request_external_inpaint(
        image=image,
        mask=mask,
        server_url=SERVER_URL,
        model_name="migan"
    )
    
    if result is not None:
        print(f"✅ 성공! 결과 이미지 크기: {result.shape}")
        output_path = os.path.join(os.path.dirname(__file__), "outputs_test/inpaint_result_class_test.webp")
        cv2.imwrite(output_path, result)
        print(f"결과 저장: {output_path}")
        return True
    else:
        print("❌ 인페인팅 실패")
        return False

if __name__ == "__main__":
    print("\n🔧 외부 인페인팅 서버 테스트 시작\n")
    
    # 1. 서버 상태 확인
    if not test_health_check():
        print("\n❌ 서버에 연결할 수 없습니다. 테스트 중단.")
        sys.exit(1)
    
    print()
    
    # 2. migan 모델 테스트
    test_inpaint_api("migan")
    
    print()
    
    # 3. simple-lama 모델 테스트
    test_inpaint_api("simple-lama")
    
    print()
    
    # 4. Request_AI_Server 클래스 테스트
    test_with_request_ai_server()
    
    print("\n🎉 테스트 완료!")