Add AI subject replace: file-picker compositing with LAB color matching
- Backend: POST /api/image/replace-subject Uses rembg to extract the subject from a source photo, scales it to fit the selection mask bounding box (or canvas centre when no selection is active), applies a partial LAB color transfer (blend=0.45) so the subject's lighting matches the background, then composites the result. Also adds POST /api/image/extract-subject for standalone subject extraction. - Frontend: "Replace with subject from file" button in the selection panel (selection_actions.js) — opens a native file picker so no clipboard API or HTTPS is required. Calls the new endpoint with the current SAM mask. - Frontend: Image > Replace Subject (AI)... menu entry backed by modules/image/replace_subject.js — a full dialog with file picker, thumbnail preview, and "match background lighting" toggle. Works with or without a prior Smart Select; if a selection exists it confines the subject to that region. https://claude.ai/code/session_01UtrvbisMp1yGu6PqeLmrFr
This commit is contained in:
@@ -745,6 +745,179 @@ async def enhance(req: EnhanceRequest):
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
# ─── Subject replace ─────────────────────────────────────────────────────────
|
||||
|
||||
class ExtractSubjectRequest(BaseModel):
|
||||
image: str # base64
|
||||
|
||||
|
||||
class ReplaceSubjectRequest(BaseModel):
|
||||
background_image: str # base64 — image whose background we keep
|
||||
subject_image: str # base64 — image whose subject we extract
|
||||
mask: Optional[str] = None # base64 — white = where the subject should land
|
||||
match_colors: bool = True # blend subject color stats toward background
|
||||
|
||||
|
||||
def _extract_subject_bytes(image_bytes: bytes) -> bytes:
|
||||
"""Remove background from image using rembg; return RGBA PNG bytes."""
|
||||
if rembg_available():
|
||||
return remove_background_rembg(image_bytes)
|
||||
raise RuntimeError(
|
||||
"rembg is not installed. Run: pip install rembg (or add it to requirements.txt)"
|
||||
)
|
||||
|
||||
|
||||
def _color_transfer_lab(subj_rgba: "Image", bg_rgb: "Image", blend: float = 0.45) -> "Image":
|
||||
"""
|
||||
Partial LAB color transfer: nudge subject color statistics 'blend' fraction
|
||||
toward the background's statistics so it looks like it belongs in the scene.
|
||||
"""
|
||||
import cv2
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
src_arr = np.array(subj_rgba.convert("RGB"), dtype=np.float32)
|
||||
tgt_arr = np.array(bg_rgb.convert("RGB"), dtype=np.float32)
|
||||
|
||||
alpha = np.array(subj_rgba.split()[3])
|
||||
subject_mask = alpha > 10
|
||||
|
||||
if not subject_mask.any():
|
||||
return subj_rgba
|
||||
|
||||
src_lab = cv2.cvtColor(src_arr.astype(np.uint8), cv2.COLOR_RGB2LAB).astype(np.float32)
|
||||
tgt_lab = cv2.cvtColor(tgt_arr.astype(np.uint8), cv2.COLOR_RGB2LAB).astype(np.float32)
|
||||
|
||||
for ch in range(3):
|
||||
src_ch = src_lab[:, :, ch]
|
||||
src_pixels = src_ch[subject_mask]
|
||||
tgt_pixels = tgt_lab[:, :, ch].flatten()
|
||||
|
||||
src_mean, src_std = float(src_pixels.mean()), float(src_pixels.std()) + 1e-6
|
||||
tgt_mean, tgt_std = float(tgt_pixels.mean()), float(tgt_pixels.std()) + 1e-6
|
||||
|
||||
adjusted_std = src_std + blend * (tgt_std - src_std)
|
||||
adjusted = (src_ch - src_mean) * (adjusted_std / src_std) + src_mean + blend * (tgt_mean - src_mean)
|
||||
src_lab[:, :, ch] = np.clip(adjusted, 0, 255)
|
||||
|
||||
result_rgb = cv2.cvtColor(src_lab.astype(np.uint8), cv2.COLOR_LAB2RGB)
|
||||
r, g, b = result_rgb[:, :, 0], result_rgb[:, :, 1], result_rgb[:, :, 2]
|
||||
return Image.merge("RGBA", [
|
||||
Image.fromarray(r), Image.fromarray(g),
|
||||
Image.fromarray(b), Image.fromarray(alpha),
|
||||
])
|
||||
|
||||
|
||||
def _do_replace_subject(
|
||||
bg_bytes: bytes,
|
||||
subj_bytes: bytes,
|
||||
mask_bytes: Optional[bytes],
|
||||
match_colors: bool,
|
||||
) -> bytes:
|
||||
"""Core compositing: extract subject → scale → color-match → paste onto background."""
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
bg_img = Image.open(BytesIO(bg_bytes)).convert("RGBA")
|
||||
|
||||
subj_rgba = Image.open(BytesIO(_extract_subject_bytes(subj_bytes))).convert("RGBA")
|
||||
|
||||
# Determine target placement bounding box from mask or full canvas
|
||||
if mask_bytes:
|
||||
mask_img = Image.open(BytesIO(mask_bytes)).convert("L")
|
||||
if mask_img.size != bg_img.size:
|
||||
mask_img = mask_img.resize(bg_img.size, Image.LANCZOS)
|
||||
mask_arr = np.array(mask_img)
|
||||
ys, xs = np.where(mask_arr > 128)
|
||||
else:
|
||||
mask_img = None
|
||||
mask_arr = None
|
||||
ys, xs = np.array([]), np.array([])
|
||||
|
||||
if len(xs) > 0:
|
||||
minx, maxx = int(xs.min()), int(xs.max())
|
||||
miny, maxy = int(ys.min()), int(ys.max())
|
||||
else:
|
||||
minx, miny = 0, 0
|
||||
maxx, maxy = bg_img.width - 1, bg_img.height - 1
|
||||
|
||||
target_w = maxx - minx + 1
|
||||
target_h = maxy - miny + 1
|
||||
|
||||
# Scale subject to fit target area, preserving aspect ratio
|
||||
sw, sh = subj_rgba.size
|
||||
scale = min(target_w / sw, target_h / sh)
|
||||
new_w = max(1, round(sw * scale))
|
||||
new_h = max(1, round(sh * scale))
|
||||
subj_scaled = subj_rgba.resize((new_w, new_h), Image.LANCZOS)
|
||||
|
||||
# Optional color transfer to blend lighting/tone
|
||||
if match_colors:
|
||||
subj_scaled = _color_transfer_lab(subj_scaled, bg_img.convert("RGB"))
|
||||
|
||||
# Center in target area
|
||||
px = minx + (target_w - new_w) // 2
|
||||
py = miny + (target_h - new_h) // 2
|
||||
|
||||
result = bg_img.copy()
|
||||
|
||||
if mask_img is not None and len(xs) > 0:
|
||||
# Build a full-canvas RGBA layer for the subject
|
||||
subj_canvas = Image.new("RGBA", bg_img.size, (0, 0, 0, 0))
|
||||
subj_canvas.paste(subj_scaled, (px, py), subj_scaled.split()[3])
|
||||
# Clip subject's alpha to the selection mask
|
||||
sc_arr = np.array(subj_canvas)
|
||||
sc_arr[:, :, 3] = np.minimum(sc_arr[:, :, 3], mask_arr).astype(np.uint8)
|
||||
subj_canvas = Image.fromarray(sc_arr)
|
||||
result.paste(subj_canvas, (0, 0), subj_canvas.split()[3])
|
||||
else:
|
||||
result.paste(subj_scaled, (px, py), subj_scaled.split()[3])
|
||||
|
||||
out = BytesIO()
|
||||
result.convert("RGB").save(out, format="PNG")
|
||||
return out.getvalue()
|
||||
|
||||
|
||||
@router.post("/image/extract-subject")
|
||||
async def extract_subject(req: ExtractSubjectRequest):
|
||||
"""
|
||||
Remove background from an image and return the subject with transparency (RGBA PNG).
|
||||
Uses rembg (AI-powered) when available.
|
||||
"""
|
||||
try:
|
||||
result = await asyncio.get_event_loop().run_in_executor(
|
||||
None, _extract_subject_bytes, _decode(req.image)
|
||||
)
|
||||
return {"result": _encode(result)}
|
||||
except Exception as e:
|
||||
import traceback; traceback.print_exc()
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.post("/image/replace-subject")
|
||||
async def replace_subject(req: ReplaceSubjectRequest):
|
||||
"""
|
||||
Extract the primary subject from `subject_image` (via rembg background removal),
|
||||
scale it to fit the `mask` selection on `background_image`, apply optional LAB
|
||||
color transfer for lighting consistency, and composite the result.
|
||||
|
||||
Returns the composited image as base64 PNG.
|
||||
"""
|
||||
try:
|
||||
result = await asyncio.get_event_loop().run_in_executor(
|
||||
None,
|
||||
_do_replace_subject,
|
||||
_decode(req.background_image),
|
||||
_decode(req.subject_image),
|
||||
_decode(req.mask) if req.mask else None,
|
||||
req.match_colors,
|
||||
)
|
||||
return {"result": _encode(result)}
|
||||
except Exception as e:
|
||||
import traceback; traceback.print_exc()
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
# ─── Extract colors ───────────────────────────────────────────────────────────
|
||||
|
||||
class ExtractColorsRequest(BaseModel):
|
||||
|
||||
Reference in New Issue
Block a user