diff --git a/backend/app/routers/tools.py b/backend/app/routers/tools.py index d5f664f..03a6fc7 100644 --- a/backend/app/routers/tools.py +++ b/backend/app/routers/tools.py @@ -204,18 +204,19 @@ _u2net_model = None async def _download_u2net_model(models_dir): - """Auto-download full U2Net model (~176MB) for best quality background removal""" + """Auto-download U2Net PyTorch model (~176MB) for background removal""" import urllib.request from pathlib import Path models_dir = Path(models_dir) models_dir.mkdir(parents=True, exist_ok=True) - # Download full U2Net model (176MB) for best quality + # Download U2Net PyTorch model (avoids ONNX executable stack issues in Docker) + # Using the PyTorch state dict format url = "https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2net.onnx" dest_path = models_dir / "u2net.onnx" - print(f"Downloading full U2Net model from {url} (~176MB)...") + print(f"Downloading U2Net model from {url} (~176MB)...") print("This may take a few minutes...") def download_progress(count, block_size, total_size): @@ -227,40 +228,38 @@ async def _download_u2net_model(models_dir): print(f" Download progress: {percent}% ({downloaded_mb:.1f}/{total_mb:.1f} MB)") urllib.request.urlretrieve(url, str(dest_path), download_progress) - print(f"Full U2Net model downloaded to {dest_path}") + print(f"U2Net model downloaded to {dest_path}") return dest_path async def _remove_background_u2net(img: Image.Image) -> bytes: """ - Remove background using U2Net model directly. - This avoids rembg dependency issues while providing good quality. + Remove background using U2Net model via OpenCV DNN. + Uses OpenCV's DNN module which doesn't have executable stack issues. """ global _u2net_model - import torch from pathlib import Path # Check for U2Net model models_dir = Path('/app/data/models') - u2net_path = models_dir / 'u2net.pth' + u2net_path = None - # Also check alternative names - if not u2net_path.exists(): - for alt_name in ['u2net.onnx', 'u2netp.pth', 'u2net_human_seg.pth']: - alt_path = models_dir / alt_name - if alt_path.exists(): - u2net_path = alt_path - break + # Check for ONNX model (preferred for OpenCV DNN) + for alt_name in ['u2net.onnx', 'u2netp.onnx']: + alt_path = models_dir / alt_name + if alt_path.exists(): + u2net_path = alt_path + break - if not u2net_path.exists(): + if u2net_path is None: # Try to auto-download the model print("U2Net model not found, attempting to download...") try: await _download_u2net_model(models_dir) # Check again - for alt_name in ['u2net.onnx', 'u2netp.onnx', 'u2net.pth']: + for alt_name in ['u2net.onnx', 'u2netp.onnx']: alt_path = models_dir / alt_name if alt_path.exists(): u2net_path = alt_path @@ -268,57 +267,49 @@ async def _remove_background_u2net(img: Image.Image) -> bytes: except Exception as download_error: print(f"Auto-download failed: {download_error}") - if not u2net_path.exists(): + if u2net_path is None: raise FileNotFoundError( "U2Net model not found. To fix this, run:\n" " docker exec -it ai-photo-edit-backend python /scripts/download_u2net_model.py\n" "Or manually download from: https://github.com/danielgatis/rembg/releases" ) - # Load model if not cached + # Load model if not cached (using OpenCV DNN - no executable stack issues) if _u2net_model is None: - print(f"Loading U2Net model from {u2net_path}") - - if str(u2net_path).endswith('.onnx'): - # Use ONNX runtime - import onnxruntime as ort - _u2net_model = ort.InferenceSession(str(u2net_path)) - else: - # Use PyTorch - from app.services.u2net_model import U2NET - _u2net_model = U2NET(3, 1) - _u2net_model.load_state_dict(torch.load(str(u2net_path), map_location='cpu')) - _u2net_model.eval() - - print("U2Net model loaded") + print(f"Loading U2Net model from {u2net_path} using OpenCV DNN") + try: + _u2net_model = cv2.dnn.readNetFromONNX(str(u2net_path)) + print("U2Net model loaded successfully with OpenCV DNN") + except Exception as e: + print(f"Failed to load with OpenCV DNN: {e}") + raise # Preprocess image - img_np = np.array(img) original_size = img.size - - # Resize to model input size input_size = 320 + + # Resize and convert to blob img_resized = img.resize((input_size, input_size), Image.Resampling.BILINEAR) img_np = np.array(img_resized).astype(np.float32) - # Normalize + # Normalize (ImageNet normalization) img_np = img_np / 255.0 img_np = (img_np - [0.485, 0.456, 0.406]) / [0.229, 0.224, 0.225] - img_np = img_np.transpose(2, 0, 1) # HWC to CHW - img_np = np.expand_dims(img_np, 0) # Add batch dimension + + # Create blob (NCHW format) + blob = cv2.dnn.blobFromImage( + img_np.astype(np.float32), + scalefactor=1.0, + size=(input_size, input_size), + swapRB=False + ) # Run inference - if hasattr(_u2net_model, 'run'): - # ONNX runtime - input_name = _u2net_model.get_inputs()[0].name - outputs = _u2net_model.run(None, {input_name: img_np}) - mask = outputs[0][0, 0] - else: - # PyTorch - with torch.no_grad(): - input_tensor = torch.from_numpy(img_np).float() - d1, d2, d3, d4, d5, d6, d7 = _u2net_model(input_tensor) - mask = d1[0, 0].numpy() + _u2net_model.setInput(blob) + outputs = _u2net_model.forward() + + # Get mask from first output + mask = outputs[0, 0] # Post-process mask mask = (mask - mask.min()) / (mask.max() - mask.min() + 1e-8) diff --git a/backend/requirements.txt b/backend/requirements.txt index 413fbfb..9649ea5 100644 --- a/backend/requirements.txt +++ b/backend/requirements.txt @@ -18,8 +18,8 @@ torch==2.1.2 torchvision==0.16.2 segment-anything @ git+https://github.com/facebookresearch/segment-anything.git -# ONNX Runtime for U2Net background removal (works with numpy<2.0) -onnxruntime==1.16.3 +# Note: Using OpenCV DNN instead of onnxruntime for U2Net +# (onnxruntime has executable stack issues in some Docker environments) # NOTE: rembg (background removal) disabled due to dependency conflicts # rembg>=2.0.70 requires: diff --git a/frontend/src/css/layout.css b/frontend/src/css/layout.css index 2274e2b..8c6980a 100644 --- a/frontend/src/css/layout.css +++ b/frontend/src/css/layout.css @@ -867,6 +867,63 @@ canvas{ background-color: #ffc107; color: #000; } + +/* Alertify dialog styling - fix text color issues */ +.alertify .ajs-dialog { + background-color: #3a3f44; + color: #f4f3f3; + border-radius: 6px; + box-shadow: 0 4px 20px rgba(0,0,0,0.4); +} +.alertify .ajs-header { + background-color: #4a5058; + color: #f4f3f3; + border-bottom: 1px solid #555; + padding: 10px 15px; + font-weight: bold; +} +.alertify .ajs-body { + color: #f4f3f3; + padding: 15px; +} +.alertify .ajs-body .ajs-content { + color: #f4f3f3; +} +.alertify .ajs-footer { + background-color: #3a3f44; + border-top: 1px solid #555; + padding: 10px 15px; +} +.alertify .ajs-footer .ajs-buttons .ajs-button { + background-color: #4a5058; + color: #f4f3f3; + border: 1px solid #666; + border-radius: 4px; + padding: 6px 16px; + margin: 0 4px; + cursor: pointer; +} +.alertify .ajs-footer .ajs-buttons .ajs-button:hover { + background-color: #5a6068; +} +.alertify .ajs-footer .ajs-buttons .ajs-button.ajs-ok { + background-color: #28a745; + border-color: #28a745; +} +.alertify .ajs-footer .ajs-buttons .ajs-button.ajs-ok:hover { + background-color: #218838; +} +.alertify .ajs-input { + background-color: #2a2f34; + color: #f4f3f3; + border: 1px solid #555; + padding: 8px; + border-radius: 4px; + width: 100%; +} +.alertify .ajs-input::placeholder { + color: #999; +} .effectsPreview{ cursor: pointer; background-color: #ddd; diff --git a/frontend/src/js/tools/brush_select.js b/frontend/src/js/tools/brush_select.js index 00b4ad5..6d6d4b9 100644 --- a/frontend/src/js/tools/brush_select.js +++ b/frontend/src/js/tools/brush_select.js @@ -262,7 +262,8 @@ class Brush_select_class extends Base_tools_class { if (isAdditive && this.currentMask) { alertify.success('Added to selection! Shift+brush to add more.'); } else { - alertify.success('Selection complete! Ctrl+C to copy, Ctrl+X to cut.'); + // Offer to float the selection for immediate manipulation (Canva-like) + this.offerFloatSelection(); } } catch (error) { @@ -287,7 +288,8 @@ class Brush_select_class extends Base_tools_class { var maskCanvas = await this.decodeMask(result.mask); this.applyMaskCanvas(maskCanvas, false); - alertify.success('Selection complete! Brush over more to add.'); + // Offer to float the selection for immediate manipulation + this.offerFloatSelection(); } catch (error) { console.error('Select error:', error); @@ -297,6 +299,34 @@ class Brush_select_class extends Base_tools_class { } } + /** + * Offer to float the selection to a new layer for manipulation (Canva-like workflow) + */ + offerFloatSelection() { + var _this = this; + + alertify.confirm( + 'Selection Complete', + 'Would you like to move/scale this selection? This will copy it to a new layer.', + function() { + // Yes - copy to layer and switch to Select tool + _this.copyToLayer(); + + // Switch to Select tool + setTimeout(function() { + var selectTool = document.querySelector('.sidebar_left .item[data-tool="select"]'); + if (selectTool) { + selectTool.click(); + } + }, 100); + }, + function() { + // No - just keep the selection + alertify.message('Tip: Use Ctrl+C to copy or Ctrl+X to cut the selection.'); + } + ).set('labels', {ok: 'Yes, Move/Scale', cancel: 'Keep Selection'}); + } + /** * Sample key points from brush stroke for SAM */