app.py

from flask import Flask, request, jsonify
from flask_cors import CORS
import base64
import io
from PIL import Image
import requests
import logging

# Import the withoutBG library
from withoutbg.core import WithoutBGOpenSource
from huggingface_hub import hf_hub_download
from pathlib import Path
import shutil

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = Flask(__name__)
CORS(app)

# Model directory
MODEL_DIR = Path("./models")
MODEL_DIR.mkdir(parents=True, exist_ok=True)

def _ensure_model_file(filename: str) -> Path:
    """Download model file from HuggingFace if not exists"""
    target = MODEL_DIR / filename
    if target.exists():
        return target
    logger.info(f"📥 Downloading model file: {filename}")
    downloaded = Path(hf_hub_download(repo_id="withoutbg/focus", filename=filename))
    shutil.copy2(downloaded, target)
    logger.info(f"✅ Model file downloaded: {filename}")
    return target

def _create_model() -> WithoutBGOpenSource:
    """Create WithoutBG model instance"""
    logger.info("🚀 Creating WithoutBG model...")
    return WithoutBGOpenSource(
        depth_model_path=_ensure_model_file("depth_anything_v2_vits_slim.onnx"),
        isnet_model_path=_ensure_model_file("isnet.onnx"),
        matting_model_path=_ensure_model_file("focus_matting_1.0.0.onnx"),
        refiner_model_path=_ensure_model_file("focus_refiner_1.0.0.onnx"),
    )

# Initialize the model once at startup
try:
    logger.info("🚀 Loading withoutBG model...")
    model = _create_model()
    logger.info("✅ Model loaded successfully!")
except Exception as e:
    logger.error(f"❌ Failed to load model: {e}")
    model = None


# ─────────────────────────────────────────────
# 共通処理: PIL Image → 背景削除 → (image_data, mask_data)
# ─────────────────────────────────────────────
def process_image(img: Image.Image, return_mask: bool = False) -> dict:
    """
    背景削除を実行し、結果を辞書で返す。

    Returns:
        {
            'success': True,
            'image_data': 'data:image/png;base64,...',   # 背景削除済み画像（RGBA）
            'mask_data':  'data:image/png;base64,...',   # マスク画像（return_mask=Trueのとき）
        }
    """
    if not model:
        raise RuntimeError("Model not initialized")

    logger.info(f"🖼️ Image loaded: {img.size}, mode: {img.mode}")

    # 背景削除
    logger.info("🔄 Removing background with WithoutBGOpenSource...")
    result = model.remove_background(img)
    logger.info(f"✅ Background removed! Result mode: {result.mode}, Size: {result.size}")

    # RGBA に変換（透過PNG）
    if result.mode != 'RGBA':
        result = result.convert('RGBA')

    # ── 背景削除済み画像を base64 エンコード ──
    out_buf = io.BytesIO()
    result.save(out_buf, format='PNG')
    out_buf.seek(0)
    image_b64 = base64.b64encode(out_buf.read()).decode('utf-8')

    response = {
        'success': True,
        'image_data': f'data:image/png;base64,{image_b64}'
    }

    # ── マスク画像（アルファチャンネル）を base64 エンコード ──
    if return_mask:
        # RGBAのアルファチャンネルをそのままマスクとして使用
        # 前景=白(255)、背景=黒(0) のグレースケール画像
        alpha = result.split()[3]          # Aチャンネル取得
        mask_img = alpha.convert('L')      # グレースケール化（念のため）

        mask_buf = io.BytesIO()
        mask_img.save(mask_buf, format='PNG')
        mask_buf.seek(0)
        mask_b64 = base64.b64encode(mask_buf.read()).decode('utf-8')

        response['mask_data'] = f'data:image/png;base64,{mask_b64}'
        logger.info(f"🎭 Mask generated: {mask_img.size}, mode: {mask_img.mode}")

    return response


# ─────────────────────────────────────────────
# GET / ヘルスチェック
# ─────────────────────────────────────────────
@app.route('/', methods=['GET'])
def health_check():
    return jsonify({
        'service': 'withoutBG API Server',
        'status': 'healthy' if model else 'unhealthy',
        'model': 'withoutBG Focus v1.0.0',
        'version': '1.0.0',
        'platform': 'Hugging Face Spaces',
        'endpoints': [
            'POST /api/remove-bg          – image_url or image_base64, return_mask(opt)',
            'POST /api/remove-bg-from-url – image_url, return_mask(opt)',
            'POST /api/remove-bg-base64   – image_base64, return_mask(opt)',
        ]
    })


# ─────────────────────────────────────────────
# POST /api/remove-bg  （既存エンドポイント・後方互換）
# Body: { image_url or image_base64, return_mask(opt, default false) }
# ─────────────────────────────────────────────
@app.route('/api/remove-bg', methods=['POST'])
def remove_background():
    try:
        data = request.get_json()
        if not data:
            return jsonify({'success': False, 'error': 'No JSON data provided'}), 400
        if not model:
            return jsonify({'success': False, 'error': 'Model not initialized'}), 500

        return_mask = bool(data.get('return_mask', False))

        if 'image_url' in data:
            logger.info(f"📥 Downloading image from URL: {data['image_url']}")
            r = requests.get(data['image_url'], timeout=30)
            r.raise_for_status()
            img = Image.open(io.BytesIO(r.content))

        elif 'image_base64' in data:
            logger.info("📥 Decoding base64 image")
            b64 = data['image_base64']
            if ',' in b64:
                b64 = b64.split(',')[1]
            img = Image.open(io.BytesIO(base64.b64decode(b64)))

        else:
            return jsonify({'success': False, 'error': 'Either image_url or image_base64 is required'}), 400

        return jsonify(process_image(img, return_mask=return_mask))

    except Exception as e:
        logger.error(f"❌ Error: {e}")
        import traceback; traceback.print_exc()
        return jsonify({'success': False, 'error': str(e)}), 500


# ─────────────────────────────────────────────
# POST /api/remove-bg-from-url  ← Web側がURLを渡すときに呼ぶ
# Body: { "image_url": "https://...", "return_mask": true }
# ─────────────────────────────────────────────
@app.route('/api/remove-bg-from-url', methods=['POST'])
def remove_background_from_url():
    try:
        data = request.get_json()
        if not data:
            return jsonify({'success': False, 'error': 'No JSON data provided'}), 400
        if not model:
            return jsonify({'success': False, 'error': 'Model not initialized'}), 500
        if 'image_url' not in data:
            return jsonify({'success': False, 'error': 'image_url is required'}), 400

        return_mask = bool(data.get('return_mask', False))

        logger.info(f"📥 [from-url] Downloading image: {data['image_url']}")
        r = requests.get(data['image_url'], timeout=30)
        r.raise_for_status()
        img = Image.open(io.BytesIO(r.content))

        return jsonify(process_image(img, return_mask=return_mask))

    except Exception as e:
        logger.error(f"❌ Error [from-url]: {e}")
        import traceback; traceback.print_exc()
        return jsonify({'success': False, 'error': str(e)}), 500


# ─────────────────────────────────────────────
# POST /api/remove-bg-base64  ← Web側がbase64を渡すときに呼ぶ
# Body: { "image_base64": "data:image/...;base64,...", "return_mask": true }
# ─────────────────────────────────────────────
@app.route('/api/remove-bg-base64', methods=['POST'])
def remove_background_base64():
    try:
        data = request.get_json()
        if not data:
            return jsonify({'success': False, 'error': 'No JSON data provided'}), 400
        if not model:
            return jsonify({'success': False, 'error': 'Model not initialized'}), 500
        if 'image_base64' not in data:
            return jsonify({'success': False, 'error': 'image_base64 is required'}), 400

        return_mask = bool(data.get('return_mask', False))

        logger.info("📥 [base64] Decoding base64 image")
        b64 = data['image_base64']
        if ',' in b64:
            b64 = b64.split(',')[1]
        img = Image.open(io.BytesIO(base64.b64decode(b64)))

        return jsonify(process_image(img, return_mask=return_mask))

    except Exception as e:
        logger.error(f"❌ Error [base64]: {e}")
        import traceback; traceback.print_exc()
        return jsonify({'success': False, 'error': str(e)}), 500


if __name__ == '__main__':
    import os
    port = int(os.environ.get('PORT', 7860))
    logger.info(f"🚀 Starting withoutBG API Server on port {port}...")
    app.run(host='0.0.0.0', port=port, debug=False)