app.py

# app.py
import cv2
import insightface
import numpy as np
from flask import Flask, request, jsonify
from flask_cors import CORS
import base64
import os
import tempfile
import uuid
from werkzeug.utils import secure_filename
from insightface.app import FaceAnalysis

app = Flask(__name__)
CORS(app)

# Configuration
app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024  # 16MB max file size
app.config['UPLOAD_FOLDER'] = tempfile.gettempdir()
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'webp'}

# Initialize face analysis globally
face_app = FaceAnalysis(providers=['CPUExecutionProvider'])
face_app.prepare(ctx_id=0, det_size=(640, 640))

# Initialize face swapper
swapper = insightface.model_zoo.get_model('inswapper_128.onnx')

def allowed_file(filename):
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS

def decode_base64_image(base64_string):
    """Decode base64 string to numpy image"""
    if ',' in base64_string:
        base64_string = base64_string.split(',')[1]
    img_data = base64.b64decode(base64_string)
    nparr = np.frombuffer(img_data, np.uint8)
    return cv2.imdecode(nparr, cv2.IMREAD_COLOR)

def encode_image_to_base64(image):
    """Encode numpy image to base64 string"""
    _, buffer = cv2.imencode('.jpg', image)
    return base64.b64encode(buffer).decode('utf-8')

@app.route('/health', methods=['GET'])
def health_check():
    return jsonify({'status': 'healthy', 'message': 'Face Swap API is running'})

@app.route('/swap', methods=['POST'])
def swap_faces():
    """
    Swap faces between source and target images
    Expected JSON payload:
    {
        "source_image": "base64_string or file_path",
        "target_image": "base64_string or file_path",
        "source_is_base64": true,  # optional, defaults to false
        "target_is_base64": true   # optional, defaults to false
    }
    Or use multipart/form-data with files:
    - source_file (image file)
    - target_file (image file)
    """
    
    try:
        source_img = None
        target_img = None
        
        # Handle JSON request
        if request.is_json:
            data = request.get_json()
            
            # Get source image
            source_is_base64 = data.get('source_is_base64', False)
            if source_is_base64:
                source_img = decode_base64_image(data['source_image'])
            else:
                source_img = cv2.imread(data['source_image'])
            
            # Get target image
            target_is_base64 = data.get('target_is_base64', False)
            if target_is_base64:
                target_img = decode_base64_image(data['target_image'])
            else:
                target_img = cv2.imread(data['target_image'])
        
        # Handle multipart form request
        elif 'source_file' in request.files and 'target_file' in request.files:
            source_file = request.files['source_file']
            target_file = request.files['target_file']
            
            if source_file and allowed_file(source_file.filename):
                source_filename = secure_filename(source_file.filename)
                source_path = os.path.join(app.config['UPLOAD_FOLDER'], f"source_{uuid.uuid4()}_{source_filename}")
                source_file.save(source_path)
                source_img = cv2.imread(source_path)
                os.remove(source_path)  # Clean up
            else:
                return jsonify({'error': 'Invalid source file type'}), 400
            
            if target_file and allowed_file(target_file.filename):
                target_filename = secure_filename(target_file.filename)
                target_path = os.path.join(app.config['UPLOAD_FOLDER'], f"target_{uuid.uuid4()}_{target_filename}")
                target_file.save(target_path)
                target_img = cv2.imread(target_path)
                os.remove(target_path)  # Clean up
            else:
                return jsonify({'error': 'Invalid target file type'}), 400
        
        else:
            return jsonify({'error': 'Invalid request. Provide source_image/target_image or source_file/target_file'}), 400
        
        # Validate images
        if source_img is None:
            return jsonify({'error': 'Could not read source image'}), 400
        if target_img is None:
            return jsonify({'error': 'Could not read target image'}), 400
        
        # Detect faces
        source_faces = face_app.get(source_img)
        target_faces = face_app.get(target_img)
        
        if len(source_faces) == 0:
            return jsonify({'error': 'No face found in source image'}), 400
        if len(target_faces) == 0:
            return jsonify({'error': 'No face found in target image'}), 400
        
        # Perform face swap
        swapped_image = swapper.get(target_img, target_faces[0], source_faces[0], paste_back=True)
        
        # Prepare response
        return_type = request.args.get('return_type', 'base64')
        
        if return_type == 'file':
            # Save and return file path
            output_filename = f"swapped_{uuid.uuid4()}.jpg"
            output_path = os.path.join(app.config['UPLOAD_FOLDER'], output_filename)
            cv2.imwrite(output_path, swapped_image)
            
            return jsonify({
                'success': True,
                'message': 'Face swap completed successfully',
                'output_path': output_path,
                'filename': output_filename
            })
        else:
            # Return base64 encoded image
            encoded_image = encode_image_to_base64(swapped_image)
            return jsonify({
                'success': True,
                'message': 'Face swap completed successfully',
                'swapped_image': encoded_image
            })
    
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/swap/batch', methods=['POST'])
def batch_swap_faces():
    """
    Batch face swap with multiple targets
    Expected JSON payload:
    {
        "source_image": "base64_string or path",
        "target_images": ["base64_string1", "base64_string2", ...],
        "source_is_base64": true,
        "targets_are_base64": true
    }
    """
    try:
        data = request.get_json()
        
        # Get source image
        source_is_base64 = data.get('source_is_base64', False)
        if source_is_base64:
            source_img = decode_base64_image(data['source_image'])
        else:
            source_img = cv2.imread(data['source_image'])
        
        if source_img is None:
            return jsonify({'error': 'Could not read source image'}), 400
        
        # Detect source face once
        source_faces = face_app.get(source_img)
        if len(source_faces) == 0:
            return jsonify({'error': 'No face found in source image'}), 400
        
        source_face = source_faces[0]
        
        # Process all target images
        results = []
        target_images = data.get('target_images', [])
        targets_are_base64 = data.get('targets_are_base64', False)
        
        for idx, target_img_data in enumerate(target_images):
            try:
                if targets_are_base64:
                    target_img = decode_base64_image(target_img_data)
                else:
                    target_img = cv2.imread(target_img_data)
                
                if target_img is None:
                    results.append({'index': idx, 'error': 'Could not read target image'})
                    continue
                
                target_faces = face_app.get(target_img)
                if len(target_faces) == 0:
                    results.append({'index': idx, 'error': 'No face found in target image'})
                    continue
                
                swapped_image = swapper.get(target_img, target_faces[0], source_face, paste_back=True)
                encoded_image = encode_image_to_base64(swapped_image)
                
                results.append({
                    'index': idx,
                    'success': True,
                    'swapped_image': encoded_image
                })
            except Exception as e:
                results.append({'index': idx, 'error': str(e)})
        
        return jsonify({
            'success': True,
            'message': f'Processed {len(results)} images',
            'results': results
        })
    
    except Exception as e:
        return jsonify({'error': str(e)}), 500

if __name__ == '__main__':
    print("Starting Face Swap API Server...")
    print("Make sure 'inswapper_128.onnx' is in the current directory")
    print("API endpoints:")
    print("  GET  /health - Health check")
    print("  POST /swap - Single face swap")
    print("  POST /swap/batch - Batch face swap")
    print("\nStarting server on http://localhost:5000")
    app.run(host='0.0.0.0', port=5000, debug=True)