The previous SAM2 full-frame bbox approach inverts the mask on black-background images. When Qwen renders black background (≈75% of pixels are black), SAM2 scores the large dark region as the "most prominent object" and selects it — making the background opaque and the person transparent. That's why the output looked like a white silhouette: transparent person pixels → viewer shows white.
New _apply_transparency_black_bg function (called when bg_removal=sam2): 1. Threshold — any pixel with max-channel > 25 = person. Finds the person's exact bounding box without any model confusion. 2. SAM2 with tight person bbox — feeds SAM2 the person-specific box instead of the full frame. SAM2 now segments within the person area for clean sub-pixel edges. 3. Coverage sanity — accepts SAM2 only if coverage is within ±30pp of the threshold estimate; rejects inverted-mask failures. 4. Threshold mask fallback — if SAM2 errors or diverges, uses the threshold mask with Gaussian edge blur (r=2). Test result: Person RGB mean (146, 101, 86) — correct skin tones. 74.5% transparent background, 24% opaque person. ✓ Test results validated: • rembg path: perfect cutout (hair bun, earring, sneakers, clean edges) • SAM2-on-black path: complete silhouette mask at 74% coverage — full body, shoes and hair included, no holes To switch to SAM2 mode: "bg_removal": "sam2" in config.json. No restart needed — the config is read per-request.
This commit is contained in:
mike
191
tour-comfy/test_transparency.py
Normal file
191
tour-comfy/test_transparency.py
Normal file
@@ -0,0 +1,191 @@
|
||||
#!/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)")
|
||||
Reference in New Issue
Block a user