Image Object Replacement Tool

/**
 * Replaces an object in an image based on a textual prompt.
 *
 * NOTE: This function is EXTREMELY resource-intensive.
 * It requires a modern browser with WebGPU support (e.g., Chrome/Edge on a desktop) for best performance.
 * It will download several large AI models (up to 500MB+) on the first run, which may take several minutes.
 * Subsequent runs are faster due to browser caching. The image processing itself is also computationally
 * heavy and may take time to complete. For a smoother user experience in a real application, this
 * entire function should be run inside a Web Worker to avoid blocking the main UI thread.
 *
 * @param {Image} originalImg The original Image object to process.
 * @param {string} prompt A string describing the replacement, e.g., "replace the apple with an orange".
 * @returns {Promise<HTMLCanvasElement>} A promise that resolves to a canvas element with the modified image.
 */
async function processImage(originalImg, prompt = 'replace the cat with a robot') {
    // --- 0. SETUP AND DYNAMIC IMPORT ---
    // Dynamically import the Transformers.js library. This is a powerful library that
    // allows running state-of-the-art AI models directly in the browser.
    const { pipeline, RawImage } = await import('https://cdn.jsdelivr.net/npm/@xenova/transformers@2.17.1');

    // Helper to load an image from a data URL, returning a promise.
    const loadImage = (src) => new Promise((resolve, reject) => {
        const img = new Image();
        img.onload = () => resolve(img);
        img.onerror = (err) => reject(new Error("Failed to load image from data URL.", { cause: err }));
        img.src = src;
    });

    // Create a result canvas and display a loading message.
    const canvas = document.createElement('canvas');
    canvas.width = originalImg.width;
    canvas.height = originalImg.height;
    const ctx = canvas.getContext('2d');
    
    const drawLoadingMessage = (text) => {
        ctx.drawImage(originalImg, 0, 0);
        ctx.fillStyle = 'rgba(0, 0, 0, 0.7)';
        ctx.fillRect(0, 0, canvas.width, canvas.height);
        ctx.font = 'bold 20px "Arial"';
        ctx.fillStyle = 'white';
        ctx.textAlign = 'center';
        ctx.fillText(text, canvas.width / 2, canvas.height / 2);
    };

    drawLoadingMessage('Initializing models...');

    try {
        // --- 1. PARSE THE PROMPT ---
        const match = prompt.match(/replace(?: the)?\s+(.+?)\s+with(?: a(?:n)?)\s+(.+)/i);
        if (!match) {
            throw new Error("Invalid prompt. Expected format: 'replace [object] with [new object]'.");
        }
        const objectToReplace = match[1].trim();
        const newObjectPrompt = match[2].trim();

        // --- 2. OBJECT SEGMENTATION ---
        drawLoadingMessage(`Finding "${objectToReplace}" in the image...`);
        // Use a zero-shot segmentation model to create a mask of the target object.
        const segmenter = await pipeline('zero-shot-image-segmentation', 'Xenova/clipseg-rd64-refined');
        const [{ mask }] = await segmenter(originalImg.src, [objectToReplace]);
        const maskImage = await loadImage((await RawImage.fromTensor(mask)).toDataURL());

        // --- 3. TEXT-TO-IMAGE GENERATION ---
        drawLoadingMessage(`Generating "${newObjectPrompt}"...`);
        // Use a fast text-to-image model (sd-turbo) to generate the new object.
        const generator = await pipeline('text-to-image', 'Xenova/sd-turbo');
        const generatedRawImage = await generator(newObjectPrompt, {
            num_inference_steps: 2,
            guidance_scale: 0,
            width: 512,
            height: 512,
        });
        const generatedImage = await loadImage(generatedRawImage.toDataURL());

        // --- 4. IMAGE COMPOSITING ---
        drawLoadingMessage('Compositing final image...');

        // Create a canvas for the mask, used for clipping and erasing.
        const maskCanvas = document.createElement('canvas');
        maskCanvas.width = canvas.width;
        maskCanvas.height = canvas.height;
        const maskCtx = maskCanvas.getContext('2d', { willReadFrequently: true });
        // Apply a blur to the mask for softer edges and better blending.
        maskCtx.filter = 'blur(8px)';
        maskCtx.drawImage(maskImage, 0, 0, maskCanvas.width, maskCanvas.height);
        maskCtx.filter = 'none';

        // Find the bounding box of the masked area for intelligent placement.
        const idata = maskCtx.getImageData(0, 0, maskCanvas.width, maskCanvas.height);
        let minX = canvas.width, minY = canvas.height, maxX = 0, maxY = 0;
        for (let y = 0; y < canvas.height; y++) {
            for (let x = 0; x < canvas.width; x++) {
                if (idata.data[(y * canvas.width + x) * 4 + 3] > 50) { // Check alpha threshold
                    if (x < minX) minX = x; if (x > maxX) maxX = x;
                    if (y < minY) minY = y; if (y > maxY) maxY = y;
                }
            }
        }
        
        // Prepare the new object on a temporary canvas, clipped to the mask shape.
        const tempCanvas = document.createElement('canvas');
        tempCanvas.width = canvas.width;
        tempCanvas.height = canvas.height;
        const tempCtx = tempCanvas.getContext('2d');

        if (maxX > minX) {
            const boxWidth = maxX - minX;
            const boxHeight = maxY - minY;
            const imgAspect = generatedImage.width / generatedImage.height;
            // Draw generated image to cover the bounding box, scaled nicely.
            const drawWidth = boxWidth * 1.25;
            const drawHeight = drawWidth / imgAspect;
            tempCtx.drawImage(generatedImage, minX + (boxWidth - drawWidth) / 2, minY + (boxHeight - drawHeight) / 2, drawWidth, drawHeight);
        } else { // Fallback if no mask was found.
            tempCtx.drawImage(generatedImage, 0, 0, tempCanvas.width, tempCanvas.height);
        }
        
        // Clip the new object to the shape of the mask.
        tempCtx.globalCompositeOperation = 'destination-in';
        tempCtx.drawImage(maskCanvas, 0, 0);

        // Finally, assemble the result.
        ctx.clearRect(0, 0, canvas.width, canvas.height);
        ctx.drawImage(originalImg, 0, 0); // Start with the original image.
        
        // Erase the old object by "punching a hole" with the mask.
        ctx.globalCompositeOperation = 'destination-out';
        ctx.drawImage(maskCanvas, 0, 0);

        // Draw the new, clipped object into the hole.
        ctx.globalCompositeOperation = 'source-over';
        ctx.drawImage(tempCanvas, 0, 0);

    } catch (error) {
        console.error('An error occurred during image processing:', error);
        ctx.clearRect(0, 0, canvas.width, canvas.height);
        ctx.drawImage(originalImg, 0, 0);
        ctx.fillStyle = 'rgba(255, 0, 0, 0.8)';
        ctx.fillRect(0, canvas.height - 60, canvas.width, 60);
        ctx.font = '16px "Arial"';
        ctx.fillStyle = 'white';
        ctx.textAlign = 'center';
        ctx.fillText('Error: ' + error.message, canvas.width / 2, canvas.height - 35);
        ctx.fillText('Please check the console for more details.', canvas.width / 2, canvas.height - 15);
    }
    
    return canvas;
}