Fix U2Net, alertify dialogs, and improve AI Paint workflow
- Switch U2Net from onnxruntime to OpenCV DNN to avoid "cannot enable executable stack" error in Docker - Add alertify dialog styling to fix white text on white background issue in popups - Add offerFloatSelection() to AI Paint that prompts user after selection to move/scale it (Canva-like workflow) - Auto-switch to Select tool after floating selection https://claude.ai/code/session_01CLedz6CanT9t46KBvng3vz
This commit is contained in:
Claude
@@ -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)
|
||||
|
||||
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user