import cv2, numpy as np, base64, requests
import time, os

img_path = 'img/6.jpg'
img = cv2.imread(img_path)

if img is None:
    raise FileNotFoundError(f'이미지를 찾을 수 없습니다: {img_path}')

_, buf = cv2.imencode('.png', img)
img_b64 = base64.b64encode(buf).decode()

url = "http://192.168.0.150:35756/api/v1/run_plugin_gen_image"
# url = "http://59.26.209.89:47396//api/v1/run_plugin_gen_image"
payload = {
      "name": "RemoveBG",
      "image": f"data:image/png;base64,{img_b64}",
      "scale": 1
}

start = time.perf_counter()
resp = requests.post(url, json=payload)
elapsed = time.perf_counter() - start

print("STATUS", resp.status_code)
print("ELAPSED_SEC", elapsed)

if resp.status_code == 200:
    # 결과 저장
    os.makedirs("output", exist_ok=True)
    out_path = "output/6_removed.png"
    with open(out_path, "wb") as f:
        f.write(resp.content)

    # 수신된 바이너리 PNG를 numpy 배열로 변환
    nparr = np.frombuffer(resp.content, np.uint8)
    result_img = cv2.imdecode(nparr, cv2.IMREAD_UNCHANGED)

    # 결과 이미지가 있을 때 대상 영역을 찾아 자르고 흰색으로 합성
    if result_img is not None and result_img.ndim == 3:
        # 1) 초기 마스크 계산 (알파 있으면 알파 > 200, 없으면 밝기 < 230)
        if result_img.shape[2] == 4:
            mask_init = (result_img[:, :, 3] > 200).astype(np.uint8)
            rgba_img = result_img
        else:
            gray = cv2.cvtColor(result_img[:, :, :3], cv2.COLOR_BGR2GRAY)
            mask_init = (gray < 230).astype(np.uint8)
            alpha_channel = (mask_init * 255).astype(np.uint8)
            rgba_img = np.dstack([result_img, alpha_channel])

        # 2) 모폴로지로 잡티 제거 (erode 1 → dilate 2)
        kernel = np.ones((3, 3), np.uint8)
        mask = cv2.erode(mask_init, kernel, iterations=1)
        mask = cv2.dilate(mask, kernel, iterations=2)

        # 3) 가장 큰 연결 요소만 선택
        num, labels, stats, _ = cv2.connectedComponentsWithStats(mask, connectivity=8)
        if num > 1:
            largest = 1 + np.argmax(stats[1:, cv2.CC_STAT_AREA])
            mask = (labels == largest).astype(np.uint8)

        ys, xs = np.where(mask > 0)
        if len(xs) > 0 and len(ys) > 0:
            top, left = ys.min(), xs.min()
            bottom, right = ys.max(), xs.max()

            crop_rgba = rgba_img[top:bottom + 1, left:right + 1]

            # 객체 크기에 비례한 테두리 마진 (10%)
            ch, cw = crop_rgba.shape[:2]
            margin = int(max(ch, cw) * 0.1)

            crop_rgba = cv2.copyMakeBorder(
                crop_rgba, margin, margin, margin, margin,
                borderType=cv2.BORDER_CONSTANT,
                value=[255, 255, 255, 0]
            )

            # RGBA -> 흰 배경 BGR
            bgr_crop = crop_rgba[:, :, :3].astype(np.float32)
            alpha_crop = crop_rgba[:, :, 3:4].astype(np.float32) / 255.0
            white_bg = np.full_like(bgr_crop, 255.0)
            result_img = (bgr_crop * alpha_crop + white_bg * (1 - alpha_crop)).astype(np.uint8)
        else:
            # mask 가 거의 없으면 객체 없음으로 판단
            white_bg = np.full_like(result_img[:, :, :3], 255)
            result_img = white_bg.astype(np.uint8)

    # 흰배경 결과를 파일로 저장
    white_path = "output/6_removed_white.png"
    cv2.imwrite(white_path, result_img)

    # 화면에 표시
    cv2.imshow('Removed BG (White BG)', result_img)
    cv2.waitKey(0)
    cv2.destroyAllWindows()
else:
    print("ERROR", resp.text)