RemBG-api

Runtime error

App Files Files Community

soiz1 commited on May 29, 2025

Commit

075fc85

verified ·

1 Parent(s): 9ce2677

Update app.py

Browse files

Files changed (1) hide show

app.py +345 -86

app.py CHANGED Viewed

@@ -1,109 +1,368 @@
-import gradio as gr
 import os
 import cv2
 from rembg import new_session, remove
 from rembg.sessions import sessions_class
-def inference(file, mask, model, x, y):
-    im = cv2.imread(file, cv2.IMREAD_COLOR)
-    input_path = "input.png"
-    output_path = "output.png"
     cv2.imwrite(input_path, im)
     with open(input_path, 'rb') as i:
         with open(output_path, 'wb') as o:
-            input = i.read()
             session = new_session(model)
             output = remove(
-                input,
                 session=session,
-                **{ "sam_prompt": [{"type": "point", "data": [x, y], "label": 1}] },
                 only_mask=(mask == "Mask only")
             )
             o.write(output)
     return output_path
-title = "RemBG"
-description = "Gradio demo for **[RemBG](https://github.com/danielgatis/rembg)**. To use it, simply upload your image, select a model, click Process, and wait."
-badge = """
-    <div style="position: fixed; left: 50%; text-align: center;">
-        <a href="https://github.com/danielgatis/rembg" target="_blank" style="text-decoration: none;">
-            <img src="https://img.shields.io/badge/RemBG-Github-blue" alt="RemBG Github" />
-        </a>
-    </div>
-"""
-def get_coords(evt: gr.SelectData) -> tuple:
-    return evt.index[0], evt.index[1]
-def show_coords(model: str):
-    visible = model == "sam"
-    return gr.update(visible=visible), gr.update(visible=visible), gr.update(visible=visible)
-for session in sessions_class:
-    session.download_models()
-with gr.Blocks() as app:
-    gr.Markdown(f"# {title}")
-    gr.Markdown(description)
-    with gr.Row():
-        inputs = gr.Image(type="filepath", label="Input Image")
-        outputs = gr.Image(type="filepath", label="Output Image")
-    with gr.Row():
-        mask_option = gr.Radio(
-            ["Default", "Mask only"],
-            value="Default",
-            label="Output Type"
-        )
-        model_selector = gr.Dropdown(
-            [
-                "u2net",
-                "u2netp",
-                "u2net_human_seg",
-                "u2net_cloth_seg",
-                "silueta",
-                "isnet-general-use",
-                "isnet-anime",
-                "sam",
-                "birefnet-general",
-                "birefnet-general-lite",
-                "birefnet-portrait",
-                "birefnet-dis",
-                "birefnet-hrsod",
-                "birefnet-cod",
-                "birefnet-massive"
-            ],
-            value="isnet-general-use",
-            label="Model Selection"
-        )
-    extra = gr.Markdown("## Click on the image to capture coordinates (for SAM model)", visible=False)
-    x = gr.Number(label="Mouse X Coordinate", visible=False)
-    y = gr.Number(label="Mouse Y Coordinate", visible=False)
-    model_selector.change(show_coords, inputs=model_selector, outputs=[x, y, extra])
-    inputs.select(get_coords, None, [x, y])
-    gr.Button("Process Image").click(
-        inference,
-        inputs=[inputs, mask_option, model_selector, x, y],
-        outputs=outputs
-    )
-    gr.Examples(
-        examples=[
-            ["lion.png", "Default", "u2net", None, None],
-            ["girl.jpg", "Default", "u2net", None, None],
-            ["anime-girl.jpg", "Default", "isnet-anime", None, None]
-        ],
-        inputs=[inputs, mask_option, model_selector, x, y],
-        outputs=outputs
-    )
-    gr.HTML(badge)
-app.launch(share=True)

+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
+app = Flask(__name__)
+# セッションの初期化
+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)