shangguanyanyan
/

flux1-fill-dev-custom

Diffusers

Safetensors

FluxFillPipeline

Model card Files Files and versions

xet

Community

toby007 commited on Apr 26

Commit

92fc27b

1 Parent(s): 3c71e14

update gene mask script

Browse files

Files changed (1) hide show

generated_mask.py +163 -0

generated_mask.py ADDED Viewed

	@@ -0,0 +1,163 @@

+# 一键生成mask工具：上传原图 -> 自动做人像分割 -> 生成mask图
+import os
+import uuid
+import numpy as np
+from PIL import Image
+import mediapipe as mp
+from flask import Flask, request, jsonify, send_file, render_template_string
+app = Flask(__name__)
+# 初始化 mediapipe 模型
+mp_selfie_segmentation = mp.solutions.selfie_segmentation.SelfieSegmentation(model_selection=1)
+UPLOAD_FOLDER = "uploads"
+MASK_FOLDER = "masks"
+os.makedirs(UPLOAD_FOLDER, exist_ok=True)
+os.makedirs(MASK_FOLDER, exist_ok=True)
+# 简单的HTML上传表单
+HTML_TEMPLATE = """
+<!DOCTYPE html>
+<html>
+<head>
+    <title>生成图像分割Mask</title>
+    <style>
+        body { font-family: Arial, sans-serif; max-width: 800px; margin: 0 auto; padding: 20px; }
+        h1 { color: #333; }
+        .container { margin-top: 20px; }
+        .result { margin-top: 20px; }
+        img { max-width: 100%; margin-top: 10px; }
+    </style>
+</head>
+<body>
+    <h1>上传图像生成Mask</h1>
+    <div class="container">
+        <form action="/upload" method="post" enctype="multipart/form-data">
+            <input type="file" name="file" accept="image/*" required>
+            <button type="submit">生成Mask</button>
+        </form>
+    </div>
+    {% if original_image %}
+    <div class="result">
+        <h2>处理结果</h2>
+        <div>
+            <h3>原始图像</h3>
+            <img src="{{ original_image }}" alt="原始图像">
+        </div>
+        <div>
+            <h3>生成的Mask</h3>
+            <img src="{{ mask_image }}" alt="生成的Mask">
+        </div>
+    </div>
+    {% endif %}
+</body>
+</html>
+"""
+@app.route("/", methods=["GET"])
+def index():
+    # 显示一个简单的上传表单
+    return render_template_string(HTML_TEMPLATE)
+@app.route("/upload", methods=["POST"])
+def upload_image():
+    try:
+        if 'file' not in request.files:
+            return jsonify({"error": "No file uploaded."}), 400
+        file = request.files['file']
+        if file.filename == '':
+            return jsonify({"error": "Empty filename."}), 400
+        # 检查是否为图像文件
+        try:
+            img = Image.open(file.stream)
+            img.verify()  # 验证图像
+            file.stream.seek(0)  # 重置文件流位置
+        except:
+            return jsonify({"error": "Invalid image file."}), 400
+        # 保存上传的文件
+        filename = f"{uuid.uuid4().hex}.png"
+        file_path = os.path.join(UPLOAD_FOLDER, filename)
+        file.save(file_path)
+        print(f"Saved uploaded image to {file_path}")
+        # 处理图像生成mask
+        try:
+            mask_filename = generate_mask(file_path)
+            mask_path = os.path.join(MASK_FOLDER, mask_filename)
+            print(f"Generated mask saved to {mask_path}")
+            # 判断是API请求还是网页请求
+            if request.headers.get('Accept') == 'application/json':
+                return jsonify({
+                    "status": "success",
+                    "original_image": f"/uploads/{filename}",
+                    "mask_image": f"/masks/{mask_filename}"
+                })
+            else:
+                # 返回HTML页面展示结果
+                return render_template_string(
+                    HTML_TEMPLATE,
+                    original_image=f"/uploads/{filename}",
+                    mask_image=f"/masks/{mask_filename}"
+                )
+        except Exception as e:
+            print(f"Error generating mask: {str(e)}")
+            return jsonify({"error": f"Failed to generate mask: {str(e)}"}), 500
+    except Exception as e:
+        print(f"Unexpected error: {str(e)}")
+        return jsonify({"error": f"Server error: {str(e)}"}), 500
+@app.route("/uploads/<filename>")
+def serve_upload(filename):
+    return send_file(os.path.join(UPLOAD_FOLDER, filename))
+@app.route("/masks/<filename>")
+def serve_mask(filename):
+    return send_file(os.path.join(MASK_FOLDER, filename))
+def generate_mask(image_path):
+    print(f"Generating mask for {image_path}")
+    # 读取图片
+    image = Image.open(image_path).convert("RGB")
+    image_np = np.array(image)
+    height, width = image_np.shape[:2]
+    print(f"Image dimensions: {width}x{height}")
+    # 生成分割mask
+    results = mp_selfie_segmentation.process(image_np)
+    if results.segmentation_mask is None:
+        raise ValueError("Segmentation failed. No mask generated.")
+    mask = results.segmentation_mask
+    print(f"Mask shape: {mask.shape}")
+    # 根据分割结果生成二值mask（背景白，前景黑）
+    binary_mask = (mask < 0.5).astype(np.uint8) * 255
+    # 确保mask是RGB模式
+    mask_img = Image.fromarray(binary_mask).convert("RGB")
+    # 保存mask
+    mask_filename = f"mask_{os.path.basename(image_path)}"
+    mask_path = os.path.join(MASK_FOLDER, mask_filename)
+    mask_img.save(mask_path)
+    print(f"Mask saved to {mask_path}")
+    return mask_filename
+if __name__ == "__main__":
+    print("Starting mask generation server on http://127.0.0.1:5555")
+    app.run(host="0.0.0.0", port=5555, debug=True)