192 lines
8.5 KiB
Python
192 lines
8.5 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Validate background removal strategies.
|
|
|
|
Usage:
|
|
python test_transparency.py [image.png ...]
|
|
|
|
Writes comparison files next to each input:
|
|
*_rembg.png — pure rembg (bg_removal=rembg path)
|
|
*_blackbg.png — simulated black-bg composite (what Qwen renders in sam2 mode)
|
|
*_thresh.png — threshold mask only (non-black pixels → person)
|
|
*_thresh_sam2.png — threshold bbox → SAM2 edge refinement (new sam2 mode path)
|
|
"""
|
|
import io, sys, os
|
|
import numpy as np
|
|
from PIL import Image, ImageFilter
|
|
|
|
OUTPUT_DIR = "/mnt/zim/tour-comfy/output"
|
|
VENV_SITE = "/home/mike/comfyui/venv/lib/python3.13/site-packages"
|
|
SAM2_CKPT = os.path.expanduser("~/.sam/sam2.1_hiera_base_plus.pt")
|
|
SAM2_CFG = "configs/sam2.1/sam2.1_hiera_b+.yaml"
|
|
|
|
|
|
# ── rembg ──────────────────────────────────────────────────────────────────────
|
|
def apply_rembg(png_bytes: bytes) -> bytes:
|
|
from rembg import remove
|
|
return remove(png_bytes)
|
|
|
|
|
|
# ── SAM2 loader ────────────────────────────────────────────────────────────────
|
|
_predictor = None
|
|
def load_sam2():
|
|
global _predictor
|
|
if _predictor is not None:
|
|
return _predictor
|
|
try:
|
|
import torch
|
|
from sam2.build_sam import build_sam2
|
|
from sam2.sam2_image_predictor import SAM2ImagePredictor
|
|
model = build_sam2(SAM2_CFG, SAM2_CKPT, device="cuda")
|
|
_predictor = SAM2ImagePredictor(model)
|
|
print("[sam2] loaded")
|
|
except Exception as e:
|
|
print(f"[sam2] FAILED: {e}")
|
|
_predictor = False
|
|
return _predictor
|
|
|
|
|
|
# ── Simulate black-bg Qwen output ─────────────────────────────────────────────
|
|
def make_black_bg(png_bytes: bytes) -> bytes:
|
|
"""Composite a rembg cutout onto pure black — simulates Qwen 'black background' output."""
|
|
rgba = Image.open(io.BytesIO(apply_rembg(png_bytes))).convert("RGBA")
|
|
bg = Image.new("RGBA", rgba.size, (0, 0, 0, 255))
|
|
bg.paste(rgba, mask=rgba.split()[3])
|
|
out = bg.convert("RGB")
|
|
buf = io.BytesIO(); out.save(buf, "PNG"); return buf.getvalue()
|
|
|
|
|
|
# ── Threshold-only mask ────────────────────────────────────────────────────────
|
|
def apply_threshold_mask(png_bytes: bytes, threshold: int = 25) -> bytes:
|
|
"""Find non-black pixels → person mask. No SAM2 needed."""
|
|
img = Image.open(io.BytesIO(png_bytes)).convert("RGB")
|
|
arr = np.array(img)
|
|
h, w = arr.shape[:2]
|
|
|
|
is_person = np.max(arr, axis=2) > threshold
|
|
coverage = is_person.sum() / (h * w)
|
|
print(f" [threshold] person coverage: {coverage:.1%}")
|
|
|
|
if not is_person.any():
|
|
print(" [threshold] all-black image — no person found")
|
|
return png_bytes
|
|
|
|
mask_np = is_person.astype(np.uint8) * 255
|
|
alpha_img = Image.fromarray(mask_np, "L").filter(ImageFilter.GaussianBlur(radius=2))
|
|
rgba = img.convert("RGBA"); r, g, b, _ = rgba.split()
|
|
out = Image.merge("RGBA", (r, g, b, alpha_img))
|
|
buf = io.BytesIO(); out.save(buf, "PNG"); return buf.getvalue()
|
|
|
|
|
|
# ── NEW: Threshold bbox → SAM2 refinement (sam2 mode path) ────────────────────
|
|
def apply_thresh_sam2(png_bytes: bytes, threshold: int = 25) -> bytes:
|
|
"""
|
|
For black-background Qwen output:
|
|
1. Threshold to find person bbox (non-black pixels)
|
|
2. Run SAM2 with that tight bbox for clean edge refinement
|
|
3. Fallback to threshold mask if SAM2 unavailable or mask looks wrong
|
|
"""
|
|
import torch
|
|
|
|
img = Image.open(io.BytesIO(png_bytes)).convert("RGB")
|
|
arr = np.array(img)
|
|
h, w = arr.shape[:2]
|
|
|
|
# Step 1 — threshold
|
|
is_person = np.max(arr, axis=2) > threshold
|
|
thresh_cov = is_person.sum() / (h * w)
|
|
print(f" [thresh_sam2] threshold person coverage: {thresh_cov:.1%}")
|
|
|
|
if not is_person.any():
|
|
print(" [thresh_sam2] all-black — fallback to rembg")
|
|
return apply_rembg(png_bytes)
|
|
|
|
rows = np.any(is_person, axis=1)
|
|
cols = np.any(is_person, axis=0)
|
|
rmin = int(np.where(rows)[0][0]); rmax = int(np.where(rows)[0][-1])
|
|
cmin = int(np.where(cols)[0][0]); cmax = int(np.where(cols)[0][-1])
|
|
|
|
margin = int(min(h, w) * 0.02)
|
|
y1 = max(0, rmin - margin); y2 = min(h, rmax + margin)
|
|
x1 = max(0, cmin - margin); x2 = min(w, cmax + margin)
|
|
print(f" [thresh_sam2] person bbox (+margin): ({x1},{y1})-({x2},{y2})")
|
|
|
|
# Step 2 — SAM2 with person-specific bbox
|
|
predictor = load_sam2()
|
|
if predictor is not False:
|
|
box = np.array([[x1, y1, x2, y2]], dtype=np.float32)
|
|
try:
|
|
with torch.inference_mode():
|
|
predictor.set_image(arr)
|
|
masks, scores, _ = predictor.predict(box=box, multimask_output=True)
|
|
|
|
if masks is not None and len(masks) > 0:
|
|
best = masks[int(np.argmax(scores))]
|
|
sam_cov = float(best.sum()) / (h * w)
|
|
print(f" [thresh_sam2] SAM2 coverage: {sam_cov:.1%} (threshold was {thresh_cov:.1%})")
|
|
|
|
# Accept SAM2 result if coverage is within reasonable range of threshold
|
|
if 0.03 < sam_cov < 0.95 and abs(sam_cov - thresh_cov) < 0.30:
|
|
mask_np = best.astype(np.uint8) * 255
|
|
alpha_img = Image.fromarray(mask_np, "L").filter(ImageFilter.GaussianBlur(radius=1))
|
|
rgba = img.convert("RGBA"); r, g, b, _ = rgba.split()
|
|
out = Image.merge("RGBA", (r, g, b, alpha_img))
|
|
buf = io.BytesIO(); out.save(buf, "PNG")
|
|
print(" [thresh_sam2] SAM2 result accepted ✓")
|
|
return buf.getvalue()
|
|
else:
|
|
print(f" [thresh_sam2] SAM2 coverage diverged from threshold — using threshold mask")
|
|
except Exception as e:
|
|
print(f" [thresh_sam2] SAM2 error: {e} — using threshold mask")
|
|
else:
|
|
print(" [thresh_sam2] SAM2 not available — using threshold mask")
|
|
|
|
# Step 3 — fallback: threshold mask with soft edges
|
|
mask_np = is_person.astype(np.uint8) * 255
|
|
alpha_img = Image.fromarray(mask_np, "L").filter(ImageFilter.GaussianBlur(radius=2))
|
|
rgba = img.convert("RGBA"); r, g, b, _ = rgba.split()
|
|
out = Image.merge("RGBA", (r, g, b, alpha_img))
|
|
buf = io.BytesIO(); out.save(buf, "PNG")
|
|
print(" [thresh_sam2] threshold mask used as fallback")
|
|
return buf.getvalue()
|
|
|
|
|
|
# ── main ───────────────────────────────────────────────────────────────────────
|
|
if __name__ == "__main__":
|
|
paths = sys.argv[1:] if len(sys.argv) > 1 else [
|
|
os.path.join(OUTPUT_DIR, "20260622_181910_0_20260619_124038_image.png"),
|
|
]
|
|
|
|
for path in paths:
|
|
if not os.path.exists(path):
|
|
print(f"SKIP (not found): {path}"); continue
|
|
stem = os.path.splitext(path)[0]
|
|
print(f"\n══ {os.path.basename(path)} ══")
|
|
with open(path, "rb") as f:
|
|
raw = f.read()
|
|
|
|
print("1. rembg (bg_removal=rembg path)...")
|
|
rb = apply_rembg(raw)
|
|
with open(stem + "_rembg.png", "wb") as f: f.write(rb)
|
|
print(f" → {os.path.basename(stem)}_rembg.png")
|
|
|
|
print("2. Simulate black-bg Qwen output...")
|
|
bb = make_black_bg(raw)
|
|
with open(stem + "_blackbg.png", "wb") as f: f.write(bb)
|
|
print(f" → {os.path.basename(stem)}_blackbg.png")
|
|
|
|
print("3. Threshold-only mask on black-bg image...")
|
|
tm = apply_threshold_mask(bb)
|
|
with open(stem + "_thresh.png", "wb") as f: f.write(tm)
|
|
print(f" → {os.path.basename(stem)}_thresh.png")
|
|
|
|
print("4. Threshold bbox → SAM2 refinement on black-bg image (NEW sam2 mode path)...")
|
|
ts = apply_thresh_sam2(bb)
|
|
with open(stem + "_thresh_sam2.png", "wb") as f: f.write(ts)
|
|
print(f" → {os.path.basename(stem)}_thresh_sam2.png")
|
|
|
|
print("\n── Done ──")
|
|
print(" *_rembg.png rembg on real background (bg_removal=rembg path)")
|
|
print(" *_thresh.png threshold-only on black bg")
|
|
print(" *_thresh_sam2.png threshold-bbox → SAM2 on black bg (NEW sam2 mode path)")
|