Image Background Remover

// This promise variable should be defined outside the processImage function,
// ideally in a scope that persists across calls to processImage if it can be called multiple times.
let mediaPipeSelfieSegmentationLoadPromise = null;

async function processImage(originalImg, modelSelection = 1) {
    // default modelSelection: 1 (landscape model, often more detailed), 0 for general model.

    // 1. Dynamically import MediaPipe SelfieSegmentation script
    // This ensures the script is loaded only once using a shared promise.
    if (typeof window.SelfieSegmentation === 'undefined') {
        if (!mediaPipeSelfieSegmentationLoadPromise) {
            mediaPipeSelfieSegmentationLoadPromise = new Promise((resolve, reject) => {
                const script = document.createElement('script');
                script.src = 'https://cdn.jsdelivr.net/npm/@mediapipe/selfie_segmentation/selfie_segmentation.js';
                script.crossOrigin = 'anonymous';
                script.onload = () => {
                    resolve();
                };
                script.onerror = (errorEvent) => {
                    // Keep the promise rejected to prevent retries on this failed load.
                    reject(new Error("Failed to load MediaPipe SelfieSegmentation script."));
                };
                document.head.appendChild(script);
            });
        }
        try {
            await mediaPipeSelfieSegmentationLoadPromise;
        } catch (scriptLoadError) {
            console.error(scriptLoadError.message);
            // Create a fallback canvas to display an error message
            const errorCanvas = document.createElement('canvas');
            const w = (originalImg.naturalWidth && originalImg.naturalWidth > 0) ? originalImg.naturalWidth : (originalImg.width || 300);
            const h = (originalImg.naturalHeight && originalImg.naturalHeight > 0) ? originalImg.naturalHeight : (originalImg.height || 150);
            errorCanvas.width = w;
            errorCanvas.height = h;
            const errCtx = errorCanvas.getContext('2d');
            try { // Attempt to draw original image as background for the error message
                if (originalImg.complete && originalImg.naturalWidth > 0) {
                    errCtx.drawImage(originalImg, 0, 0, w, h);
                }
            } catch(e) { /* ignore if drawing originalImg fails, canvas might be blank */ }
            
            const errMsg = "Error: Background removal library failed to load.";
            errCtx.fillStyle = "rgba(0, 0, 0, 0.7)"; // Semi-transparent black banner
            errCtx.fillRect(0, h / 2 - 20, w, 40); 
            errCtx.fillStyle = "white";
            errCtx.font = "bold 14px Arial";
            errCtx.textAlign = "center";
            errCtx.textBaseline = "middle";
            errCtx.fillText(errMsg, w / 2, h / 2);
            return errorCanvas; // Return canvas with error message
        }
    }

    // 2. Prepare canvas for output
    const canvas = document.createElement('canvas');
    // Use natural dimensions if available and valid, otherwise fall back to element's width/height or defaults
    const canvasWidth = (originalImg.naturalWidth && originalImg.naturalWidth > 0) ? originalImg.naturalWidth : (originalImg.width || 300);
    const canvasHeight = (originalImg.naturalHeight && originalImg.naturalHeight > 0) ? originalImg.naturalHeight : (originalImg.height || 150);
    canvas.width = canvasWidth;
    canvas.height = canvasHeight;
    const ctx = canvas.getContext('2d');

    // Initial check if image seems problematic (e.g., not loaded, zero dimensions)
    if (!(originalImg.complete && originalImg.naturalWidth > 0 && originalImg.naturalHeight > 0)) {
        console.warn("Input image may not be fully loaded or is invalid. Attempting processing, but results may be unexpected or errors may occur.");
        // Draw whatever we have of the image. This might be nothing if severely broken.
        try {
            ctx.drawImage(originalImg, 0, 0, canvasWidth, canvasHeight);
        } catch(e) {
            console.error("Error drawing potentially unloaded/invalid image initially:", e);
        }
        // Display a temporary warning message on the canvas. This will be cleared if processing succeeds.
        const warnMsg = "Warning: Problem with input image.";
        ctx.save();
        ctx.fillStyle = "rgba(255, 165, 0, 0.7)"; // Orange, semi-transparent banner
        ctx.fillRect(0, 0, canvasWidth, 30);
        ctx.fillStyle = "black";
        ctx.font = "bold 14px Arial";
        ctx.textAlign = "center";
        ctx.textBaseline = "middle";
        ctx.fillText(warnMsg, canvasWidth / 2, 15);
        ctx.restore();
        // Continue to MediaPipe; it might handle it or fail.
    }

    // 3. Perform segmentation using MediaPipe SelfieSegmentation
    return new Promise((resolveOuter, rejectOuter) => {
        let selfieSegmentationInstance = null; 
        try {
            selfieSegmentationInstance = new window.SelfieSegmentation({
                locateFile: (file) => `https://cdn.jsdelivr.net/npm/@mediapipe/selfie_segmentation/${file}`
            });

            selfieSegmentationInstance.setOptions({
                modelSelection: parseInt(modelSelection, 10) === 0 ? 0 : 1, // Ensure 0 or 1
            });

            selfieSegmentationInstance.onResults((results) => {
                const currentInstanceToClose = selfieSegmentationInstance; // Capture instance for safe closure
                selfieSegmentationInstance = null; // Prevent re-closing if errors occur later

                try {
                    ctx.save();
                    ctx.clearRect(0, 0, canvasWidth, canvasHeight); // Clear canvas (e.g. removes initial warning)

                    try {
                         ctx.drawImage(originalImg, 0, 0, canvasWidth, canvasHeight);
                    } catch (drawError) {
                        console.error("Error drawing original image to canvas in onResults:", drawError);
                        // If drawing fails here, the canvas might be blank or incomplete.
                        // This implies a significant issue with originalImg.
                    }

                    if (results.segmentationMask) {
                        ctx.globalCompositeOperation = 'destination-in'; // Keep where new (mask) overlaps existing (original image)
                        ctx.drawImage(results.segmentationMask, 0, 0, canvasWidth, canvasHeight);
                    } else {
                        console.warn("Segmentation mask not present in results. Background removal may not have occurred.");
                        // If no mask, original image (if drawn) remains.
                    }
                    
                    ctx.globalCompositeOperation = 'source-over'; // Reset composite operation
                    ctx.restore();
                    resolveOuter(canvas);
                } catch (e) {
                    console.error("Error during MediaPipe onResults processing:", e);
                    // In case of error here, canvas may be in an intermediate state.
                    // Depending on requirements, one might clear/reset canvas or resolve/reject.
                    rejectOuter(new Error(`Processing segmentation results failed: ${e.message}`));
                } finally {
                    if (currentInstanceToClose) {
                        currentInstanceToClose.close().catch(closeError => console.error("Error closing SelfieSegmentation in onResults:", closeError));
                    }
                }
            });
            
            // Send the image for processing. This is an async operation.
            selfieSegmentationInstance.send({ image: originalImg })
                .catch(sendError => {
                    console.error("Error sending image to SelfieSegmentation:", sendError);
                    const currentInstanceToClose = selfieSegmentationInstance;
                    selfieSegmentationInstance = null;

                    // Fallback: resolve with canvas (may contain original image or warning)
                    // Clear canvas and attempt to draw original image as a last resort.
                    ctx.clearRect(0, 0, canvasWidth, canvasHeight);
                    try {
                        ctx.drawImage(originalImg, 0, 0, canvasWidth, canvasHeight);
                    } catch(e) { /* ignore if this also fails */ }
                    const errMsg = "Error: Failed to process image.";
                    ctx.fillStyle = "rgba(255, 0, 0, 0.7)";
                    ctx.fillRect(0, canvasHeight / 2 - 20, canvasWidth, 40);
                    ctx.fillStyle = "white";
                    ctx.font = "bold 14px Arial";
                    ctx.textAlign = "center";
                    ctx.textBaseline = "middle";
                    ctx.fillText(errMsg, canvasWidth / 2, canvasHeight / 2);
                    resolveOuter(canvas); // Resolve with the canvas showing the error.
                                        // or rejectOuter(new Error(`Failed to send image: ${sendError.message}`));
                    if (currentInstanceToClose) {
                         currentInstanceToClose.close().catch(closeError => console.error("Error closing SelfieSegmentation after send error:", closeError));
                    }
                });

        } catch (initError) {
            console.error("Error initializing SelfieSegmentation:", initError);
            const currentInstanceToClose = selfieSegmentationInstance; // Might be null or partially init
            selfieSegmentationInstance = null;

            // Fallback: draw original image on canvas with an error message
            ctx.clearRect(0, 0, canvasWidth, canvasHeight);
             try {
                ctx.drawImage(originalImg, 0, 0, canvasWidth, canvasHeight);
            } catch(e) { /* ignore */ }
            
            const errMsg = "Error: Segmentation tool initialization failed.";
            ctx.fillStyle = "rgba(255, 0, 0, 0.7)";
            ctx.fillRect(0, canvasHeight / 2 - 20, canvasWidth, 40);
            ctx.fillStyle = "white";
            ctx.font = "bold 14px Arial";
            ctx.textAlign = "center";
            ctx.textBaseline = "middle";
            ctx.fillText(errMsg, canvasWidth / 2, canvasHeight / 2);
            resolveOuter(canvas); // Resolve with this error-displaying canvas
            
            if (currentInstanceToClose) { 
                currentInstanceToClose.close().catch(e => console.error("Error closing SelfieSegmentation after init error:", e));
            }
        }
    });
}