98 lines
3.5 KiB
Python
98 lines
3.5 KiB
Python
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) |