root/PaintPlus
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge pull request #43 from outis1one/claude/fix-ai-paint-tool-yz6q8

Browse Source 
Fix U2Net, alertify dialogs, and improve AI Paint workflow
This commit is contained in:
Outis
2026-01-28 08:38:06 -05:00
committed by GitHub
parent 9a305126ed 871fc696f5
commit d5898dd054
4 changed files with 132 additions and 54 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/app/routers/tools.py
+41 -50
View File
@@ -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)
backend/requirements.txt
+2 -2
View File
@@ -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:
frontend/src/css/layout.css
+57
View File
@@ -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;
frontend/src/js/tools/brush_select.js
+32 -2
View File
@@ -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
*/