Tool_Agent/tests/test_rembg/test_rembg.py

"""测试 rembg 背景移除 API"""

import sys
import io
sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding="utf-8")

import httpx
from pathlib import Path

REMBG_URL = "http://47.84.182.56:8001"
TEST_DIR = Path(__file__).parent
INPUT_IMAGE = TEST_DIR / "20260421-175044.jpg"
OUTPUT_IMAGE = TEST_DIR / "output.png"


def test_rembg():
    print(f"输入图片: {INPUT_IMAGE} ({INPUT_IMAGE.stat().st_size / 1024:.1f} KB)")

    with open(INPUT_IMAGE, "rb") as f:
        response = httpx.post(
            f"{REMBG_URL}/api/remove",
            files={"file": ("input.jpg", f, "image/jpeg")},
            timeout=120.0,
        )

    print(f"HTTP Status: {response.status_code}")

    if response.status_code == 200:
        OUTPUT_IMAGE.write_bytes(response.content)
        print(f"输出图片: {OUTPUT_IMAGE} ({len(response.content) / 1024:.1f} KB)")
        print("背景移除成功!")
    else:
        print(f"失败: {response.text[:200]}")


if __name__ == "__main__":
    test_rembg()