IMG_Worker/modules/test/old_test/test_backup_rembg.py

195 lines
7.8 KiB
Python

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Request_AI_Server의 _use_backup_rembg 함수 테스트 코드
"""
import os
import sys
import cv2
import logging
import numpy as np
from datetime import datetime
# 현재 파일의 디렉토리를 기준으로 src 폴더를 Python 경로에 추가
current_dir = os.path.dirname(os.path.abspath(__file__))
src_dir = os.path.dirname(current_dir) # src 폴더
project_dir = os.path.dirname(src_dir) # 프로젝트 루트
sys.path.insert(0, src_dir)
try:
from modules.request_inpaint import Request_AI_Server
from modules.gpu_utils import GPUManager
except ImportError as e:
print(f"모듈 import 실패: {e}")
print(f"현재 디렉토리: {current_dir}")
print(f"src 디렉토리: {src_dir}")
print(f"sys.path: {sys.path[:3]}...")
sys.exit(1)
class SimpleLogger:
"""간단한 로거 클래스"""
def __init__(self, name="TestLogger"):
self.name = name
def log(self, message, level=logging.INFO, exc_info=False):
timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
level_name = logging.getLevelName(level)
print(f"[{timestamp}] [{level_name}] {message}")
if exc_info:
import traceback
traceback.print_exc()
def test_backup_rembg():
"""_use_backup_rembg 함수 테스트"""
print("=== Request_AI_Server _use_backup_rembg 함수 테스트 ===")
# 1. 테스트 이미지 경로 확인
test_image_path = os.path.join(current_dir, "1.jpg")
if not os.path.exists(test_image_path):
print(f"❌ 테스트 이미지가 없습니다: {test_image_path}")
print("1.jpg 파일을 다음 경로에 복사해주세요:")
print(f" {current_dir}")
return False
print(f"✅ 테스트 이미지 발견: {test_image_path}")
# 2. 이미지 로드
try:
image_data = cv2.imread(test_image_path)
if image_data is None:
print(f"❌ 이미지 로드 실패: {test_image_path}")
return False
height, width = image_data.shape[:2]
print(f"✅ 이미지 로드 성공: {width}x{height}")
except Exception as e:
print(f"❌ 이미지 로드 중 오류: {e}")
return False
# 3. 로거 및 GPU 관리자 초기화
logger = SimpleLogger("BackupRemBGTest")
try:
gpu_manager = GPUManager(logger=logger)
# DirectML GPU 모드로 강제 설정 (테스트용)
gpu_manager.can_use_cuda = True
print(f"✅ GPU 관리자 초기화 성공: CUDA 사용 가능 = {gpu_manager.can_use_cuda} (DirectML 강제 활성화)")
except Exception as e:
print(f"⚠️ GPU 관리자 초기화 실패, None으로 설정: {e}")
gpu_manager = None
# 4. Request_AI_Server 초기화
try:
# 로컬 rembg 모델 경로 설정
local_rembg_model_path = os.path.join(src_dir, "modules", "rembg_models", "birefnet-general-lite.onnx")
if not os.path.exists(local_rembg_model_path):
print(f"⚠️ 로컬 rembg 모델이 없습니다: {local_rembg_model_path}")
local_rembg_model_path = None
else:
print(f"✅ 로컬 rembg 모델 발견: {local_rembg_model_path}")
request_ai_server = Request_AI_Server(
logger=logger,
inpaint_server_url="http://localhost:8008", # 테스트용 (실제로는 사용하지 않음)
rembg_server_url=None,
gpu_manager=gpu_manager,
local_rembg_model_path=local_rembg_model_path
)
print("✅ Request_AI_Server 초기화 성공")
except Exception as e:
print(f"❌ Request_AI_Server 초기화 실패: {e}")
return False
# 5. _use_backup_rembg 함수 테스트
print("\n--- 백업 rembg 모듈 테스트 시작 ---")
test_cases = [
{"model_name": "birefnet-general-lite", "object_ratio": 0.75, "force_cpu": False, "desc": "DirectML GPU 모드 (문제)"},
{"model_name": "birefnet-general-lite", "object_ratio": 0.75, "force_cpu": True, "desc": "CPU 모드 (정상)"},
{"model_name": "birefnet-general-lite", "object_ratio": 1, "force_cpu": True, "desc": "CPU 모드 + 다른 비율"},
]
results = []
for i, test_case in enumerate(test_cases):
print(f"\n테스트 케이스 {i+1}: {test_case}")
try:
start_time = datetime.now()
result_img = request_ai_server._use_backup_rembg(
image_data=image_data,
model_name=test_case["model_name"],
object_ratio=test_case["object_ratio"],
debug_save=True,
debug_prefix=f"test_{i+1}_{test_case['desc'].replace(' ', '_')}",
force_cpu=test_case["force_cpu"]
)
end_time = datetime.now()
processing_time = (end_time - start_time).total_seconds()
if result_img is not None:
# 결과 이미지 저장
# 안전한 파일명 생성 (한글 및 특수문자 제거)
safe_desc = test_case['desc'].replace('DirectML GPU 모드 (문제)', 'DirectML_GPU_problem').replace('CPU 모드 (정상)', 'CPU_normal').replace('CPU 모드 + 다른 비율', 'CPU_different_ratio')
output_filename = f"backup_rembg_result_{i+1}_{safe_desc}_ratio{int(test_case['object_ratio']*100)}.png"
output_path = os.path.join(current_dir, output_filename)
success = cv2.imwrite(output_path, result_img)
if success:
result_height, result_width = result_img.shape[:2]
print(f" ✅ 성공! 결과 이미지 크기: {result_width}x{result_height}")
print(f" 📁 저장 경로: {output_path}")
print(f" ⏱️ 처리 시간: {processing_time:.2f}")
results.append({
"test_case": i+1,
"success": True,
"output_path": output_path,
"size": (result_width, result_height),
"processing_time": processing_time
})
else:
print(f" ❌ 결과 이미지 저장 실패: {output_path}")
results.append({"test_case": i+1, "success": False, "error": "이미지 저장 실패"})
else:
print(f" ❌ 함수 실행 실패: 결과 이미지가 None")
results.append({"test_case": i+1, "success": False, "error": "결과 이미지 None"})
except Exception as e:
print(f" ❌ 테스트 케이스 {i+1} 실행 중 오류: {e}")
results.append({"test_case": i+1, "success": False, "error": str(e)})
# 6. 테스트 결과 요약
print("\n=== 테스트 결과 요약 ===")
success_count = sum(1 for r in results if r.get("success", False))
total_count = len(results)
print(f"총 테스트: {total_count}")
print(f"성공: {success_count}")
print(f"실패: {total_count - success_count}")
for result in results:
if result.get("success", False):
print(f" ✅ 케이스 {result['test_case']}: {result['size'][0]}x{result['size'][1]}, {result['processing_time']:.2f}")
else:
print(f" ❌ 케이스 {result['test_case']}: {result.get('error', '알 수 없는 오류')}")
return success_count > 0
if __name__ == "__main__":
success = test_backup_rembg()
if success:
print("\n🎉 테스트 완료! 결과 이미지를 확인해보세요.")
print(f"📁 결과 파일 위치: {current_dir}")
else:
print("\n💥 테스트 실패!")
sys.exit(0 if success else 1)