Merge pull request #60 from outis1one/claude/compassionate-albattani-s4lqb5
Add AI subject replace: file-picker compositing with LAB color matching
This commit is contained in:
Outis
@@ -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):
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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;
|
||||
@@ -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);
|
||||
});
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user