RMBG-2.0 企业级集成：API 封装、Flask 后端与前端拖拽上传方案

RMBG-2.0 企业级集成：API 封装、Flask 后端与前端拖拽上传方案 | 极客日志

# 创建虚拟环境
python -m venv rmbg_env
source rmbg_env/bin/activate # Linux/macOS
# 或
rmbg_env\Scripts\activate # Windows

# 安装核心依赖
pip install torch torchvision --extra-index-url https://download.pytorch.org/whl/cu113
pip install flask flask-cors pillow rembg
pip install opencv-python numpy requests

import os
from rembg import download_model

# 创建模型存储目录
model_dir = os.path.expanduser("~/.u2net")
os.makedirs(model_dir, exist_ok=True)

# 下载 RMBG-2.0 模型
model_url = "https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2net.onnx"
model_path = os.path.join(model_dir, "u2net.onnx")

if not os.path.exists(model_path):
    print("正在下载 RMBG-2.0 模型...")
    download_model(model_url, model_path)
    print("模型下载完成！")
else:
    print("模型已存在，跳过下载")

from flask import Flask, request, jsonify, send_file
from rembg import remove
from PIL import Image
import io
import os

app = Flask(__name__)

@app.route('/api/remove-bg', methods=['POST'])
def remove_background():
    try:
        # 检查是否有文件上传
        if 'image' not in request.files:
            return jsonify({'error': '没有上传图片'}), 400
        file = request.files['image']
        if file.filename == '':
            return jsonify({'error': '未选择文件'}), 400

        # 读取图片并处理
        input_image = Image.open(file.stream)
        output_image = remove(input_image)

        # 将结果保存到内存中
        img_io = io.BytesIO()
        output_image.save(img_io, 'PNG')
        img_io.seek(0)
        return send_file(img_io, mimetype='image/png', as_attachment=True, download_name='no-bg.png')
    except Exception as e:
        return jsonify({'error': f'处理失败：{str(e)}'}), 500

if __name__ == '__main__':
    app.run(debug=True, host='0.0.0.0', port=5000)

import uuid
import time
from datetime import datetime
from flask import Flask, request, jsonify, send_file, send_from_directory
from werkzeug.utils import secure_filename
import os
from PIL import Image
from rembg import remove

app = Flask(__name__)
app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024 # 16MB 限制
app.config['UPLOAD_FOLDER'] = 'uploads'
app.config['RESULT_FOLDER'] = 'results'

os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)
os.makedirs(app.config['RESULT_FOLDER'], exist_ok=True)

def allowed_file(filename):
    return '.' in filename and \
           filename.rsplit('.', 1)[1].lower() in {'png', 'jpg', 'jpeg', 'webp'}

@app.route('/api/advanced/remove-bg', methods=['POST'])
def advanced_remove_bg():
    start_time = time.time()
    process_id = str(uuid.uuid4())[:8]
    try:
        # 验证文件
        if 'image' not in request.files:
            return jsonify({'error': 'No image provided', 'process_id': process_id}), 400
        file = request.files['image']
        if file.filename == '':
            return jsonify({'error': 'No file selected', 'process_id': process_id}), 400
        if not allowed_file(file.filename):
            return jsonify({'error': 'Invalid file type', 'process_id': process_id}), 400

        # 保存原始文件
        original_filename = secure_filename(file.filename)
        original_path = os.path.join(app.config['UPLOAD_FOLDER'], f"{process_id}_{original_filename}")
        file.save(original_path)

        # 处理图片
        input_image = Image.open(original_path)

        # 获取处理参数
        size_limit = request.form.get('size_limit', '1024')
        try:
            size_limit = int(size_limit)
        except:
            size_limit = 1024

        # 调整大小（如果需要）
        if max(input_image.size) > size_limit:
            ratio = size_limit / max(input_image.size)
            new_size = (int(input_image.size[0] * ratio), int(input_image.size[1] * ratio))
            input_image = input_image.resize(new_size, Image.Resampling.LANCZOS)

        # 去除背景
        output_image = remove(input_image)

        # 保存结果
        result_filename = f"{process_id}_result.png"
        result_path = os.path.join(app.config['RESULT_FOLDER'], result_filename)
        output_image.save(result_path, 'PNG')

        processing_time = round(time.time() - start_time, 2)

        # 返回结果信息
        return jsonify({
            'success': True,
            'process_id': process_id,
            'original_size': input_image.size,
            'processing_time': processing_time,
            'result_url': f'/api/download/{result_filename}',
            'message': f'背景去除完成，耗时 {processing_time} 秒'
        })
    except Exception as e:
        return jsonify({
            'error': str(e),
            'process_id': process_id
        }), 500

@app.route('/api/download/<filename>')
def download_file(filename):
    return send_from_directory(app.config['RESULT_FOLDER'], filename, as_attachment=True)

@app.route('/api/status')
def api_status():
    return jsonify({
        'status': 'running',
        'version': '1.0.0',
        'model': 'RMBG-2.0'
    })

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000, threaded=True)

<!DOCTYPE html>
<html lang="zh-CN">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>RMBG-2.0 背景去除工具</title>
    <style>
        * {
            margin: 0;
            padding: 0;
            box-sizing: border-box;
        }
        body {
            font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif;
            background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
            min-height: 100vh;
            padding: 20px;
        }
        .container {
            max-width: 1000px;
            margin: 0 auto;
            background: white;
            border-radius: 15px;
            box-shadow: 0 20px 40px rgba(0,0,0,0.1);
            overflow: hidden;
        }
        .header {
            background: #4a5568;
            color: white;
            padding: 30px;
            text-align: center;
        }
        .header h1 {
            font-size: 2.5em;
            margin-bottom: 10px;
        }
        .header p {
            opacity: 0.9;
            font-size: 1.1em;
        }
        .upload-area {
            padding: 40px;
            text-align: center;
        }
        .drop-zone {
            border: 3px dashed #cbd5e0;
            border-radius: 12px;
            padding: 60px 40px;
            cursor: pointer;
            transition: all 0.3s ease;
            background: #f7fafc;
        }
        .drop-zone:hover, .drop-zone.dragover {
            border-color: #667eea;
            background: #ebf4ff;
        }
        .drop-zone i {
            font-size: 3em;
            color: #a0aec0;
            margin-bottom: 15px;
            display: block;
        }
        .file-input {
            display: none;
        }
        .btn {
            background: #667eea;
            color: white;
            padding: 12px 30px;
            border: none;
            border-radius: 8px;
            cursor: pointer;
            font-size: 1.1em;
            transition: background 0.3s ease;
        }
        .btn:hover {
            background: #5a67d8;
        }
        .btn:disabled {
            background: #cbd5e0;
            cursor: not-allowed;
        }
        .preview-area {
            display: none;
            padding: 30px;
            background: #f7fafc;
            border-top: 1px solid #e2e8f0;
        }
        .preview-container {
            display: flex;
            gap: 20px;
            flex-wrap: wrap;
            justify-content: center;
        }
        .preview-box {
            flex: 1;
            min-width: 300px;
            text-align: center;
        }
        .preview-box h3 {
            margin-bottom: 15px;
            color: #4a5568;
        }
        .preview-img {
            max-width: 100%;
            max-height: 400px;
            border-radius: 8px;
            box-shadow: 0 4px 12px rgba(0,0,0,0.1);
        }
        .progress-bar {
            height: 6px;
            background: #e2e8f0;
            border-radius: 3px;
            margin: 20px 0;
            overflow: hidden;
        }
        .progress {
            height: 100%;
            background: linear-gradient(90deg, #667eea, #764ba2);
            width: 0%;
            transition: width 0.3s ease;
        }
        .status {
            margin: 15px 0;
            font-weight: 500;
        }
        .success { color: #38a169; }
        .error { color: #e53e3e; }
        .info { color: #3182ce; }
        @media (max-width: 768px) {
            .preview-container { flex-direction: column; }
            .header h1 { font-size: 2em; }
        }
    </style>
</head>
<body>
    <div class="container">
        <div class="header">
            <h1>RMBG-2.0 智能背景去除</h1>
            <p>轻量高效 · 精准抠图 · 一键处理</p>
        </div>
        <div class="upload-area">
            <div id="dropZone" class="drop-zone">
                <i></i>
                <h3>拖拽图片到此处</h3>
                <p>或</p>
                <button onclick="document.getElementById('fileInput').click()">选择图片文件</button>
                <input type="file" accept="image/*" class="file-input" id="fileInput">
            </div>
        </div>
        <div class="preview-area" id="previewArea">
            <div class="preview-container">
                <div class="preview-box">
                    <h3>原始图片</h3>
                    <img id="originalPreview" class="preview-img">
                </div>
                <div class="preview-box">
                    <h3>处理结果</h3>
                    <img id="resultPreview" class="preview-img">
                </div>
            </div>
            <div class="progress-bar">
                <div class="progress" id="progressBar"></div>
            </div>
            <div class="status" id="statusMessage">准备就绪</div>
            <button id="downloadBtn" class="btn" disabled>下载结果图片</button>
        </div>
    </div>
    <script>
        // JavaScript 逻辑将在下一节展示
    </script>
</body>
</html>

// 在之前的 HTML 文件中的<script>标签内添加以下代码
const dropZone = document.getElementById('dropZone');
const fileInput = document.getElementById('fileInput');
const previewArea = document.getElementById('previewArea');
const originalPreview = document.getElementById('originalPreview');
const resultPreview = document.getElementById('resultPreview');
const progressBar = document.getElementById('progressBar');
const statusMessage = document.getElementById('statusMessage');
const downloadBtn = document.getElementById('downloadBtn');
let currentProcessId = null;

// 拖拽功能
dropZone.addEventListener('dragover', (e) => {
    e.preventDefault();
    dropZone.classList.add('dragover');
});
dropZone.addEventListener('dragleave', () => {
    dropZone.classList.remove('dragover');
});
dropZone.addEventListener('drop', (e) => {
    e.preventDefault();
    dropZone.classList.remove('dragover');
    const files = e.dataTransfer.files;
    if (files.length > 0) {
        handleFile(files[0]);
    }
});

// 文件选择处理
fileInput.addEventListener('change', (e) => {
    if (e.target.files.length > 0) {
        handleFile(e.target.files[0]);
    }
});

function handleFile(file) {
    if (!file.type.startsWith('image/')) {
        showStatus('请选择图片文件', 'error');
        return;
    }
    // 显示预览
    const reader = new FileReader();
    reader.onload = function(e) {
        originalPreview.src = e.target.result;
        previewArea.style.display = 'block';
        processImage(file);
    };
    reader.readAsDataURL(file);
}

function processImage(file) {
    showStatus('正在处理中...', 'info');
    updateProgress(30);
    const formData = new FormData();
    formData.append('image', file);
    formData.append('size_limit', '2048');

    fetch('http://localhost:5000/api/advanced/remove-bg', {
        method: 'POST',
        body: formData
    })
    .then(response => response.json())
    .then(data => {
        if (data.success) {
            updateProgress(100);
            showStatus(`处理完成！耗时 ${data.processing_time} 秒`, 'success');
            currentProcessId = data.process_id;
            // 显示处理结果
            resultPreview.src = data.result_url;
            // 启用下载按钮
            downloadBtn.disabled = false;
            downloadBtn.onclick = () => {
                window.location.href = data.result_url;
            };
        } else {
            throw new Error(data.error);
        }
    })
    .catch(error => {
        showStatus('处理失败：' + error.message, 'error');
        updateProgress(0);
    });
}

function updateProgress(percent) {
    progressBar.style.width = percent + '%';
}

function showStatus(message, type) {
    statusMessage.textContent = message;
    statusMessage.className = 'status ' + type;
}

// 初始化
showStatus('准备就绪，请上传图片', 'info');

# production.py
import os
from gevent import pywsgi
from werkzeug.middleware.proxy_fix import ProxyFix

app.wsgi_app = ProxyFix(app.wsgi_app, x_proto=1, x_host=1)

# 生产环境配置
class ProductionConfig:
    DEBUG = False
    TESTING = False
    MAX_CONTENT_LENGTH = 16 * 1024 * 1024 # 16MB
    UPLOAD_FOLDER = '/var/www/uploads'
    RESULT_FOLDER = '/var/www/results'

    os.makedirs(UPLOAD_FOLDER, exist_ok=True)
    os.makedirs(RESULT_FOLDER, exist_ok=True)

app.config.from_object(ProductionConfig())

if __name__ == '__main__':
    # 使用 gevent 作为生产服务器
    server = pywsgi.WSGIServer(('0.0.0.0', 5000), app)
    print("生产服务器启动在端口 5000")
    server.serve_forever()

# Dockerfile
FROM python:3.9-slim
WORKDIR /app

# 安装系统依赖
RUN apt-get update && apt-get install -y \
    libgl1 \
    libglib2.0-0 \
    && rm -rf /var/lib/apt/lists/*

# 复制依赖文件
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

# 复制应用代码
COPY . .

# 创建必要的目录
RUN mkdir -p uploads results

# 下载模型（构建时下载，避免每次启动下载）
RUN python -c "
from rembg import download_model
import os
model_dir = os.path.expanduser('~/.u2net')
os.makedirs(model_dir, exist_ok=True)
download_model('https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2net.onnx', os.path.join(model_dir, 'u2net.onnx'))
"

EXPOSE 5000
CMD ["gunicorn", "-w", "4", "-b", "0.0.0.0:5000", "app:app"]

# docker-compose.yml
version: '3.8'
services:
  rmbg-api:
    build: .
    ports:
      - "5000:5000"
    volumes:
      - uploads_data:/app/uploads
      - results_data:/app/results
    environment:
      - PYTHONUNBUFFERED=1
    restart: unless-stopped
  nginx:
    image: nginx:alpine
    ports:
      - "80:80"
    volumes:
      - ./nginx.conf:/etc/nginx/conf.d/default.conf
      - uploads_data:/app/uploads
      - results_data:/app/results
    depends_on:
      - rmbg-api
    restart: unless-stopped
volumes:
  uploads_data:
  results_data:

# optimizations.py
from concurrent.futures import ThreadPoolExecutor
from functools import lru_cache
import hashlib

# 使用线程池处理并发请求
executor = ThreadPoolExecutor(max_workers=4)

# 缓存处理结果（根据实际需求调整）
@lru_cache(maxsize=100)
def get_image_hash(image_path):
    """计算图片哈希值用于缓存"""
    with open(image_path, 'rb') as f:
        return hashlib.md5(f.read()).hexdigest()

# 批量处理功能
@app.route('/api/batch/remove-bg', methods=['POST'])
def batch_remove_bg():
    if 'images' not in request.files:
        return jsonify({'error': '没有上传图片'}), 400
    files = request.files.getlist('images')
    if len(files) > 20: # 限制批量处理数量
        return jsonify({'error': '一次最多处理 20 张图片'}), 400
    results = []
    for file in files:
        if file and allowed_file(file.filename):
            result = process_single_image(file)
            results.append(result)
    return jsonify({'results': results})

def process_single_image(file):
    """处理单张图片的辅助函数"""
    # 这里实现单张图片的处理逻辑
    pass

RMBG-2.0 企业级集成：API 封装、Flask 后端与前端拖拽上传方案

RMBG-2.0 企业级集成：API 封装、Flask 后端与前端拖拽上传方案

1. 引言：为什么选择 RMBG-2.0？

2. 环境准备与快速部署

2.1 系统要求与依赖安装

2.2 RMBG-2.0 模型下载与配置

3. 核心 API 封装与 Flask 后端实现

3.1 基础 API 服务搭建

3.2 增强型 API 服务

4. 前端拖拽上传界面实现

4.1 HTML 页面结构

4.2 JavaScript 交互逻辑

5. 企业级部署与优化建议

5.1 生产环境配置

5.2 Docker 容器化部署

5.3 性能优化建议

6. 总结与下一步建议

更多推荐文章

相关免费在线工具

RMBG-2.0 企业级集成：API 封装、Flask 后端与前端拖拽上传方案

RMBG-2.0 企业级集成：API 封装、Flask 后端与前端拖拽上传方案

1. 引言：为什么选择 RMBG-2.0？

2. 环境准备与快速部署

2.1 系统要求与依赖安装

2.2 RMBG-2.0 模型下载与配置

3. 核心 API 封装与 Flask 后端实现

3.1 基础 API 服务搭建

3.2 增强型 API 服务

4. 前端拖拽上传界面实现

4.1 HTML 页面结构

4.2 JavaScript 交互逻辑

5. 企业级部署与优化建议

5.1 生产环境配置

5.2 Docker 容器化部署

5.3 性能优化建议

6. 总结与下一步建议

微信扫一扫，关注极客日志

更多推荐文章

相关免费在线工具