| from flask import Flask, request, jsonify, send_file, render_template_string |
| import os |
| import cv2 |
| from rembg import new_session, remove |
| from rembg.sessions import sessions_class |
| import base64 |
| import uuid |
| from flask_cors import CORS |
|
|
| app = Flask(__name__) |
| CORS(app) |
|
|
| |
| for session in sessions_class: |
| session.download_models() |
|
|
| def process_image(file_path, mask, model, x, y): |
| im = cv2.imread(file_path, cv2.IMREAD_COLOR) |
| input_path = f"temp_input_{uuid.uuid4().hex}.png" |
| output_path = f"temp_output_{uuid.uuid4().hex}.png" |
| cv2.imwrite(input_path, im) |
|
|
| with open(input_path, 'rb') as i: |
| with open(output_path, 'wb') as o: |
| input_data = i.read() |
| session = new_session(model) |
|
|
| output = remove( |
| input_data, |
| session=session, |
| **{"sam_prompt": [{"type": "point", "data": [x, y], "label": 1}]}, |
| only_mask=(mask == "Mask only") |
| ) |
| o.write(output) |
|
|
| |
| if os.path.exists(input_path): |
| os.remove(input_path) |
| |
| return output_path |
|
|
| @app.route('/api/process', methods=['POST']) |
| def api_process(): |
| if 'file' not in request.files: |
| return jsonify({'error': 'No file uploaded'}), 400 |
| |
| file = request.files['file'] |
| mask = request.form.get('mask', 'Default') |
| model = request.form.get('model', 'isnet-general-use') |
| x = request.form.get('x', None) |
| y = request.form.get('y', None) |
| |
| try: |
| x = float(x) if x is not None else None |
| y = float(y) if y is not None else None |
| except (TypeError, ValueError): |
| x = None |
| y = None |
| |
| |
| temp_input = f"temp_{uuid.uuid4().hex}.png" |
| file.save(temp_input) |
| |
| try: |
| output_path = process_image(temp_input, mask, model, x, y) |
| return send_file(output_path, mimetype='image/png') |
| except Exception as e: |
| return jsonify({'error': str(e)}), 500 |
| finally: |
| |
| if os.path.exists(temp_input): |
| os.remove(temp_input) |
|
|
| HTML_TEMPLATE = """ |
| <!DOCTYPE html> |
| <html> |
| <head> |
| <title>RemBG API</title> |
| <style> |
| body { |
| font-family: Arial, sans-serif; |
| max-width: 800px; |
| margin: 0 auto; |
| padding: 20px; |
| } |
| .container { |
| display: flex; |
| flex-direction: column; |
| gap: 20px; |
| } |
| .row { |
| display: flex; |
| gap: 20px; |
| } |
| .column { |
| flex: 1; |
| } |
| img { |
| max-width: 100%; |
| height: auto; |
| border: 1px solid #ddd; |
| } |
| .form-group { |
| margin-bottom: 15px; |
| } |
| label { |
| display: block; |
| margin-bottom: 5px; |
| font-weight: bold; |
| } |
| select, input, button { |
| width: 100%; |
| padding: 8px; |
| box-sizing: border-box; |
| } |
| button { |
| background-color: #4CAF50; |
| color: white; |
| border: none; |
| cursor: pointer; |
| padding: 10px; |
| } |
| button:hover { |
| background-color: #45a049; |
| } |
| #fetch-code { |
| width: 100%; |
| height: 150px; |
| font-family: monospace; |
| padding: 10px; |
| box-sizing: border-box; |
| background-color: #f5f5f5; |
| border: 1px solid #ddd; |
| } |
| .coords-input { |
| display: none; |
| } |
| </style> |
| </head> |
| <body> |
| <h1>RemBG API</h1> |
| <p>Upload an image to process with RemBG. Select options and click "Process Image".</p> |
| |
| <div class="container"> |
| <div class="row"> |
| <div class="column"> |
| <div class="form-group"> |
| <label for="file">Input Image:</label> |
| <input type="file" id="file" accept="image/*"> |
| </div> |
| <img id="input-image" src="" alt="Input image will appear here"> |
| </div> |
| <div class="column"> |
| <div class="form-group"> |
| <label>Output Image:</label> |
| <img id="output-image" src="" alt="Output image will appear here"> |
| </div> |
| </div> |
| </div> |
| |
| <div class="row"> |
| <div class="column"> |
| <div class="form-group"> |
| <label for="mask">Output Type:</label> |
| <select id="mask"> |
| <option value="Default">Default</option> |
| <option value="Mask only">Mask only</option> |
| </select> |
| </div> |
| </div> |
| <div class="column"> |
| <div class="form-group"> |
| <label for="model">Model Selection:</label> |
| <select id="model"> |
| <option value="u2net">u2net</option> |
| <option value="u2netp">u2netp</option> |
| <option value="u2net_human_seg">u2net_human_seg</option> |
| <option value="u2net_cloth_seg">u2net_cloth_seg</option> |
| <option value="silueta">silueta</option> |
| <option value="isnet-general-use" selected>isnet-general-use</option> |
| <option value="isnet-anime">isnet-anime</option> |
| <option value="sam">sam</option> |
| <option value="birefnet-general">birefnet-general</option> |
| <option value="birefnet-general-lite">birefnet-general-lite</option> |
| <option value="birefnet-portrait">birefnet-portrait</option> |
| <option value="birefnet-dis">birefnet-dis</option> |
| <option value="birefnet-hrsod">birefnet-hrsod</option> |
| <option value="birefnet-cod">birefnet-cod</option> |
| <option value="birefnet-massive">birefnet-massive</option> |
| </select> |
| </div> |
| </div> |
| </div> |
| |
| <div id="coords-section" style="display: none;"> |
| <h3>SAM Model Coordinates</h3> |
| <p>Click on the image to set coordinates (for SAM model only)</p> |
| <div class="row"> |
| <div class="column"> |
| <div class="form-group"> |
| <label for="x">X Coordinate:</label> |
| <input type="number" id="x" class="coords-input"> |
| </div> |
| </div> |
| <div class="column"> |
| <div class="form-group"> |
| <label for="y">Y Coordinate:</label> |
| <input type="number" id="y" class="coords-input"> |
| </div> |
| </div> |
| </div> |
| </div> |
| |
| <button id="process-btn">Process Image</button> |
| |
| <div class="form-group"> |
| <label for="fetch-code">Fetch Code:</label> |
| <textarea id="fetch-code" readonly></textarea> |
| </div> |
| </div> |
| |
| <script> |
| const fileInput = document.getElementById('file'); |
| const inputImage = document.getElementById('input-image'); |
| const outputImage = document.getElementById('output-image'); |
| const maskSelect = document.getElementById('mask'); |
| const modelSelect = document.getElementById('model'); |
| const xInput = document.getElementById('x'); |
| const yInput = document.getElementById('y'); |
| const coordsSection = document.getElementById('coords-section'); |
| const processBtn = document.getElementById('process-btn'); |
| const fetchCodeTextarea = document.getElementById('fetch-code'); |
| |
| // 画像プレビュー |
| fileInput.addEventListener('change', function(e) { |
| const file = e.target.files[0]; |
| if (file) { |
| const reader = new FileReader(); |
| reader.onload = function(event) { |
| inputImage.src = event.target.result; |
| updateFetchCode(); |
| }; |
| reader.readAsDataURL(file); |
| } |
| }); |
| |
| // モデル選択でSAMの場合は座標入力表示 |
| modelSelect.addEventListener('change', function() { |
| const isSam = modelSelect.value === 'sam'; |
| coordsSection.style.display = isSam ? 'block' : 'none'; |
| document.querySelectorAll('.coords-input').forEach(el => { |
| el.style.display = isSam ? 'block' : 'none'; |
| }); |
| updateFetchCode(); |
| }); |
| |
| // 画像クリックで座標取得 (SAMモデルのみ) |
| inputImage.addEventListener('click', function(e) { |
| if (modelSelect.value === 'sam') { |
| const rect = e.target.getBoundingClientRect(); |
| const x = e.clientX - rect.left; |
| const y = e.clientY - rect.top; |
| |
| xInput.value = Math.round(x); |
| yInput.value = Math.round(y); |
| updateFetchCode(); |
| } |
| }); |
| |
| // その他の入力変更時 |
| [maskSelect, xInput, yInput].forEach(el => { |
| el.addEventListener('change', updateFetchCode); |
| }); |
| |
| // 画像処理 |
| processBtn.addEventListener('click', async function() { |
| if (!fileInput.files || fileInput.files.length === 0) { |
| alert('Please select an image file'); |
| return; |
| } |
| |
| const formData = new FormData(); |
| formData.append('file', fileInput.files[0]); |
| formData.append('mask', maskSelect.value); |
| formData.append('model', modelSelect.value); |
| |
| if (modelSelect.value === 'sam' && xInput.value && yInput.value) { |
| formData.append('x', xInput.value); |
| formData.append('y', yInput.value); |
| } |
| |
| try { |
| const response = await fetch('/api/process', { |
| method: 'POST', |
| body: formData |
| }); |
| |
| if (!response.ok) { |
| const error = await response.json(); |
| throw new Error(error.error || 'Failed to process image'); |
| } |
| |
| const blob = await response.blob(); |
| outputImage.src = URL.createObjectURL(blob); |
| } catch (error) { |
| alert('Error: ' + error.message); |
| console.error(error); |
| } |
| }); |
| |
| // Fetchコード生成 |
| function updateFetchCode() { |
| const file = fileInput.files && fileInput.files[0]; |
| if (!file) { |
| fetchCodeTextarea.value = '// Select an image first'; |
| return; |
| } |
| |
| const mask = maskSelect.value; |
| const model = modelSelect.value; |
| const x = xInput.value; |
| const y = yInput.value; |
| |
| let code = `const formData = new FormData();\n`; |
| code += `formData.append('file', fileInput.files[0]);\n`; |
| code += `formData.append('mask', '${mask}');\n`; |
| code += `formData.append('model', '${model}');\n`; |
| |
| if (model === 'sam' && x && y) { |
| code += `formData.append('x', '${x}');\n`; |
| code += `formData.append('y', '${y}');\n`; |
| } |
| |
| code += `\n`; |
| code += `fetch('http://${window.location.host}/api/process', {\n`; |
| code += ` method: 'POST',\n`; |
| code += ` body: formData\n`; |
| code += `})\n`; |
| code += `.then(response => {\n`; |
| code += ` if (!response.ok) {\n`; |
| code += ` return response.json().then(err => { throw new Error(err.error); });\n`; |
| code += ` }\n`; |
| code += ` return response.blob();\n`; |
| code += `})\n`; |
| code += `.then(blob => {\n`; |
| code += ` // Handle the processed image blob\n`; |
| code += ` const imgUrl = URL.createObjectURL(blob);\n`; |
| code += ` document.getElementById('output-image').src = imgUrl;\n`; |
| code += `})\n`; |
| code += `.catch(error => {\n`; |
| code += ` console.error('Error:', error);\n`; |
| code += ` alert('Error: ' + error.message);\n`; |
| code += `});`; |
| |
| fetchCodeTextarea.value = code; |
| } |
| |
| // 初期化 |
| updateFetchCode(); |
| </script> |
| </body> |
| </html> |
| """ |
|
|
| @app.route('/') |
| def index(): |
| return render_template_string(HTML_TEMPLATE) |
|
|
| if __name__ == '__main__': |
| app.run(host='0.0.0.0', port=7860) |
