diff --git a/backend/app/routers/ai_tools.py b/backend/app/routers/ai_tools.py index 7bf0cc1..bec517c 100644 --- a/backend/app/routers/ai_tools.py +++ b/backend/app/routers/ai_tools.py @@ -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): diff --git a/frontend/src/js/config-menu.js b/frontend/src/js/config-menu.js index 3efb940..48e42dd 100644 --- a/frontend/src/js/config-menu.js +++ b/frontend/src/js/config-menu.js @@ -344,6 +344,11 @@ const menuDefinition = [ ellipsis: true, target: 'image/remove_background.remove_background' }, + { + name: 'Replace Subject (AI)...', + ellipsis: true, + target: 'image/replace_subject.replace_subject' + }, { name: 'Prepare for Print...', ellipsis: true, diff --git a/frontend/src/js/modules/image/replace_subject.js b/frontend/src/js/modules/image/replace_subject.js new file mode 100644 index 0000000..796aadd --- /dev/null +++ b/frontend/src/js/modules/image/replace_subject.js @@ -0,0 +1,252 @@ +/** + * Replace Subject — extract the primary subject from a source photo and + * composite it onto the current layer's background. + * + * Workflow: + * 1. User selects the subject area via Smart Select (optional but recommended). + * 2. Opens this module → picks a source photo. + * 3. Backend removes the background from the source photo (rembg / AI), + * scales the extracted subject to fit the selection (or the canvas center), + * applies LAB color transfer so the lighting matches the background, and + * returns the composited image. + */ + +import app from './../../app.js'; +import config from './../../config.js'; +import Base_layers_class from './../../core/base-layers.js'; +import alertify from './../../../../node_modules/alertifyjs/build/alertify.min.js'; +import { showProgress, hideProgress } from './../../libs/progress_overlay.js'; + +const BASE = window.API_BASE_URL || ''; + +var instance = null; + +class Image_replace_subject_class { + + constructor() { + if (instance) return instance; + instance = this; + this.Base_layers = new Base_layers_class(); + this.isProcessing = false; + } + + replace_subject() { + if (this.isProcessing) { + alertify.warning('Already processing… please wait'); + return; + } + if (config.layer.type !== 'image') { + alertify.error('Current layer must be an image.'); + return; + } + this._showDialog(); + } + + // ── Private ─────────────────────────────────────────────────────────────── + + _showDialog() { + var hasSel = !!(window.smartSelectMask && window.smartSelectMask.canvas); + + // Build a dialog manually so we can embed a file input + var overlay = document.createElement('div'); + overlay.style.cssText = [ + 'position:fixed', 'inset:0', 'background:rgba(0,0,0,0.6)', + 'z-index:20000', 'display:flex', 'align-items:center', 'justify-content:center', + ].join(';'); + + var box = document.createElement('div'); + box.style.cssText = [ + 'background:#1a1a2e', 'border:1px solid #3a3a6a', 'border-radius:12px', + 'padding:24px', 'min-width:380px', 'max-width:460px', + 'font-family:sans-serif', 'color:#d0d0e0', 'font-size:13px', + ].join(';'); + + // Title + var title = document.createElement('div'); + title.textContent = 'Replace Subject'; + title.style.cssText = 'font-size:16px;font-weight:bold;color:#aaaaff;margin-bottom:6px'; + box.appendChild(title); + + var sub = document.createElement('div'); + sub.textContent = hasSel + ? 'Subject will be placed inside your current selection.' + : 'No selection active — subject will be centred on the canvas. Use Smart Select first for precise placement.'; + sub.style.cssText = 'font-size:11px;color:#7777aa;margin-bottom:16px;line-height:1.4'; + box.appendChild(sub); + + // File picker label + preview row + var fileRow = document.createElement('div'); + fileRow.style.cssText = 'display:flex;align-items:center;gap:10px;margin-bottom:12px'; + var fileLabel = document.createElement('label'); + fileLabel.textContent = 'Source photo:'; + fileLabel.style.cssText = 'color:#aaa;width:100px;flex-shrink:0'; + var fileInput = document.createElement('input'); + fileInput.type = 'file'; + fileInput.accept = 'image/*'; + fileInput.style.cssText = 'flex:1;background:#0f0f1a;color:#ccc;border:1px solid #4a4a8a;border-radius:5px;padding:4px 8px;font-size:12px;cursor:pointer'; + fileRow.appendChild(fileLabel); + fileRow.appendChild(fileInput); + box.appendChild(fileRow); + + // Thumbnail preview + var preview = document.createElement('img'); + preview.style.cssText = 'display:none;max-width:100%;max-height:160px;border-radius:6px;margin-bottom:12px;border:1px solid #3a3a6a'; + box.appendChild(preview); + fileInput.addEventListener('change', () => { + var f = fileInput.files[0]; + if (!f) return; + var url = URL.createObjectURL(f); + preview.src = url; + preview.style.display = 'block'; + preview.onload = () => URL.revokeObjectURL(url); + }); + + // Match colors checkbox + var colorRow = document.createElement('div'); + colorRow.style.cssText = 'display:flex;align-items:center;gap:8px;margin-bottom:16px'; + var colorCheck = document.createElement('input'); + colorCheck.type = 'checkbox'; + colorCheck.checked = true; + colorCheck.id = 'rs-match-colors'; + var colorLabel = document.createElement('label'); + colorLabel.htmlFor = 'rs-match-colors'; + colorLabel.textContent = 'Match background lighting & color tone'; + colorLabel.style.cssText = 'color:#bbb;cursor:pointer'; + colorRow.appendChild(colorCheck); + colorRow.appendChild(colorLabel); + box.appendChild(colorRow); + + // Buttons + var btnRow = document.createElement('div'); + btnRow.style.cssText = 'display:flex;gap:8px;justify-content:flex-end'; + + var cancelBtn = _btn('Cancel', '#2a2a4a', '#8888aa'); + cancelBtn.onclick = () => { document.body.removeChild(overlay); }; + + var goBtn = _btn('Replace Subject', '#1a3a5a', '#88ccff'); + goBtn.style.fontWeight = 'bold'; + goBtn.onclick = async () => { + var file = fileInput.files[0]; + if (!file) { + alertify.warning('Please pick a source photo first.'); + return; + } + document.body.removeChild(overlay); + await this._run(file, colorCheck.checked); + }; + + btnRow.appendChild(cancelBtn); + btnRow.appendChild(goBtn); + box.appendChild(btnRow); + + overlay.appendChild(box); + overlay.addEventListener('click', (e) => { if (e.target === overlay) document.body.removeChild(overlay); }); + document.body.appendChild(overlay); + } + + async _run(file, matchColors) { + this.isProcessing = true; + showProgress('Extracting subject and compositing…', 20); + + try { + var subjectBase64 = await _fileToBase64(file); + var bgBase64 = _getLayerBase64(); + var maskBase64 = _getMaskBase64(); + + var res = await _post('/api/image/replace-subject', { + background_image: bgBase64, + subject_image: subjectBase64, + mask: maskBase64 || undefined, + match_colors: matchColors, + }); + + var img = new Image(); + img.onload = () => { + var canvas = document.createElement('canvas'); + canvas.width = img.width; + canvas.height = img.height; + canvas.getContext('2d').drawImage(img, 0, 0); + + app.State.do_action( + new app.Actions.Bundle_action('replace_subject', 'Replace Subject', [ + new app.Actions.Update_layer_image_action(canvas, config.layer.id) + ]) + ); + + // Clear selection if one was used + if (window.smartSelectMask) { + window.smartSelectMask = null; + config.need_render = true; + } + + hideProgress(); + alertify.success('Subject replaced!'); + this.isProcessing = false; + }; + img.onerror = () => { + hideProgress(); + alertify.error('Failed to load result image.'); + this.isProcessing = false; + }; + img.src = 'data:image/png;base64,' + res.result; + + } catch (e) { + hideProgress(); + alertify.error('Replace subject failed: ' + (e.message || e)); + this.isProcessing = false; + } + } +} + +// ── Helpers ─────────────────────────────────────────────────────────────────── + +function _getLayerBase64() { + var canvas = document.createElement('canvas'); + canvas.width = config.layer.width_original; + canvas.height = config.layer.height_original; + canvas.getContext('2d').drawImage(config.layer.link, 0, 0); + return canvas.toDataURL('image/png').split(',')[1]; +} + +function _getMaskBase64() { + var m = window.smartSelectMask; + if (!m || !m.canvas) return null; + var w = config.layer.width_original; + var h = config.layer.height_original; + var canvas = document.createElement('canvas'); + canvas.width = w; + canvas.height = h; + canvas.getContext('2d').drawImage(m.canvas, 0, 0, w, h); + return canvas.toDataURL('image/png').split(',')[1]; +} + +function _fileToBase64(file) { + return new Promise((resolve, reject) => { + var reader = new FileReader(); + reader.onload = (e) => resolve(e.target.result.split(',')[1]); + reader.onerror = reject; + reader.readAsDataURL(file); + }); +} + +function _btn(text, bg, color) { + var b = document.createElement('button'); + b.textContent = text; + b.style.cssText = 'background:' + bg + ';color:' + color + ';border:1px solid #3a3a6a;padding:6px 14px;border-radius:6px;cursor:pointer;font-size:12px'; + return b; +} + +async function _post(path, body) { + var r = await fetch(BASE + path, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify(body), + }); + if (!r.ok) { + var err = await r.json().catch(() => ({ detail: r.statusText })); + throw new Error(err.detail || 'Request failed'); + } + return r.json(); +} + +export default Image_replace_subject_class; diff --git a/frontend/src/js/tools/selection_actions.js b/frontend/src/js/tools/selection_actions.js index b7ec705..104cc8c 100644 --- a/frontend/src/js/tools/selection_actions.js +++ b/frontend/src/js/tools/selection_actions.js @@ -144,6 +144,14 @@ export class SelectionActions { () => this._pasteFromClipboard() )); + // Replace subject from file + panel.appendChild(_actionCard( + 'Replace subject from file', + '#1a2a2a', '#66ddcc', + 'Pick any photo — AI extracts its subject and places it here, matching background lighting', + () => this._replaceSubjectFromFile() + )); + // Custom AI edit prompt var aiWrap = document.createElement('div'); aiWrap.style.cssText = 'background:#16213e;border-radius:7px;padding:7px 10px'; @@ -323,6 +331,42 @@ export class SelectionActions { } } + async _replaceSubjectFromFile() { + if (!this._check()) return; + + // Open a file picker — no clipboard API required + var fileInput = document.createElement('input'); + fileInput.type = 'file'; + fileInput.accept = 'image/*'; + + fileInput.onchange = async () => { + var file = fileInput.files[0]; + if (!file) return; + + var subjectBase64 = await _fileToBase64(file); + this.hide(); + showProgress('Extracting subject and compositing…', 15); + + try { + var res = await _post('/api/image/replace-subject', { + background_image: this._imageData, + subject_image: subjectBase64, + mask: this._maskData, + match_colors: true, + }); + this.tool.updateLayerWithResult(res.result); + this.tool.clearSelection(); + hideProgress(); + alertify.success('Subject replaced with background color matching!'); + } catch (e) { + hideProgress(); + alertify.error('Replace subject failed: ' + e.message); + } + }; + + fileInput.click(); + } + _check() { if (!this._imageData || !this._maskData) { alertify.error('No selection data. Make a new selection first.'); @@ -411,3 +455,12 @@ function _blobToBase64(blob) { reader.readAsDataURL(blob); }); } + +function _fileToBase64(file) { + return new Promise((resolve, reject) => { + var reader = new FileReader(); + reader.onload = (e) => resolve(e.target.result.split(',')[1]); + reader.onerror = reject; + reader.readAsDataURL(file); + }); +}