Image Black And White To High-Resolution Color Converter

async function processImage(originalImg, upscaleFactor = 2) {
    /**
     * Converts a black and white image to color and increases its resolution.
     *
     * This function performs the following steps:
     * 1. Dynamically loads the TensorFlow.js library.
     * 2. Loads a pre-trained image colorization model.
     * 3. Resizes the input image to the model's required input size (256x256).
     * 4. Pre-processes the image data by converting it to the Lab color space and extracting the L (lightness) channel.
     * 5. Feeds the L channel to the model to predict the a and b (color) channels.
     * 6. Recombines the original L channel with the predicted a and b channels.
     * 7. Converts the resulting Lab image back to RGB to create a small, colorized version.
     * 8. Uses a clever compositing technique to apply the generated colors to the original full-resolution grayscale image.
     * 9. Upscales the final colorized image using the browser's high-quality image smoothing.
     *
     * @param {Image} originalImg - The input javascript Image object (should be black and white).
     * @param {number} [upscaleFactor=2] - The factor by which to increase the final image resolution (e.g., 2 means 2x the original dimensions).
     * @returns {Promise<HTMLCanvasElement>} A canvas element containing the final high-resolution, colorized image.
     */

    // --- Helper Functions for Color Space Conversion (RGB <-> Lab) ---
    // These functions are necessary because the model operates in the Lab color space.
    const labToRgb = (l, a, b) => {
        let y = (l + 16) / 116,
            x = a / 500 + y,
            z = y - b / 200,
            r, g, b_out;

        const pivot = 0.206893034; // (6/29)^3
        x = x > pivot ? x * x * x : (x - 16 / 116) / 7.787;
        y = y > pivot ? y * y * y : (y - 16 / 116) / 7.787;
        z = z > pivot ? z * z * z : (z - 16 / 116) / 7.787;
        
        // D65 reference white
        x = x * 0.95047;
        y = y * 1.00000;
        z = z * 1.08883;

        r = x * 3.2406 + y * -1.5372 + z * -0.4986;
        g = x * -0.9689 + y * 1.8758 + z * 0.0415;
        b_out = x * 0.0557 + y * -0.2040 + z * 1.0570;

        const srgb = (c) => (c > 0.0031308 ? 1.055 * Math.pow(c, 1 / 2.4) - 0.055 : 12.92 * c);
        r = srgb(r);
        g = srgb(g);
        b_out = srgb(b_out);

        return [Math.max(0, Math.min(255, r * 255)), Math.max(0, Math.min(255, g * 255)), Math.max(0, Math.min(255, b_out * 255))];
    };

    const rgbToLab = (r, g, b) => {
        r /= 255; g /= 255; b /= 255;

        const linear = (c) => (c > 0.04045 ? Math.pow((c + 0.055) / 1.055, 2.4) : c / 12.92);
        r = linear(r); g = linear(g); b = linear(b);

        let x = r * 0.4124 + g * 0.3576 + b * 0.1805;
        let y = r * 0.2126 + g * 0.7152 + b * 0.0722;
        let z = r * 0.0193 + g * 0.1192 + b * 0.9505;
        
        // D65 reference white
        x /= 0.95047;
        y /= 1.00000;
        z /= 1.08883;

        const pivot = 0.008856; // (6/29)^3
        const f = (t) => (t > pivot ? Math.pow(t, 1 / 3) : 7.787 * t + 16 / 116);
        x = f(x); y = f(y); z = f(z);

        return [116 * y - 16, 500 * (x - y), 200 * (y - z)];
    };


    // --- Main Function Logic ---

    // 1. Load TensorFlow.js library dynamically to avoid adding it to every page.
    if (typeof window.tf === 'undefined') {
        const script = document.createElement('script');
        script.src = 'https://cdn.jsdelivr.net/npm/@tensorflow/tfjs@latest/dist/tf.min.js';
        script.async = true;
        document.head.appendChild(script);
        await new Promise((resolve) => {
            script.onload = resolve;
        });
    }

    // 2. Load the pre-trained colorization model.
    // This model is hosted publicly as part of the ml5.js library examples.
    const MODEL_URL = 'https://raw.githubusercontent.com/ml5js/ml5-data-and-models/main/models/colorization/model.json';
    const model = await tf.loadGraphModel(MODEL_URL);

    const MODEL_INPUT_SIZE = 256;

    // 3. Prepare input for the model
    const inputCanvas = document.createElement('canvas');
    inputCanvas.width = MODEL_INPUT_SIZE;
    inputCanvas.height = MODEL_INPUT_SIZE;
    const inputCtx = inputCanvas.getContext('2d', { willReadFrequently: true });
    inputCtx.drawImage(originalImg, 0, 0, MODEL_INPUT_SIZE, MODEL_INPUT_SIZE);
    
    const imageData = inputCtx.getImageData(0, 0, MODEL_INPUT_SIZE, MODEL_INPUT_SIZE);
    const labPixels = [];
    for (let i = 0; i < imageData.data.length; i += 4) {
        const r = imageData.data[i];
        const g = imageData.data[i + 1];
        const b = imageData.data[i + 2];
        const [l, a, b_lab] = rgbToLab(r, g, b);
        labPixels.push(l);
    }
    
    // Normalize L channel for the model (0-100 -> 0-1) and create a tensor
    const l_channel = tf.tensor1d(labPixels).map(x => x.div(100));
    const modelInput = l_channel.reshape([1, MODEL_INPUT_SIZE, MODEL_INPUT_SIZE, 1]);

    // 4. Run prediction to get the 'a' and 'b' color channels.
    const ab_channels_t = model.predict(modelInput);
    const ab_channels = await ab_channels_t.data();
    tf.dispose([l_channel, modelInput, ab_channels_t]); // Clean up GPU memory

    // 5. Create a small colorized image by combining original L and predicted a,b
    const colorCanvas = document.createElement('canvas');
    colorCanvas.width = MODEL_INPUT_SIZE;
    colorCanvas.height = MODEL_INPUT_SIZE;
    const colorCtx = colorCanvas.getContext('2d');
    const colorizedImageData = colorCtx.createImageData(MODEL_INPUT_SIZE, MODEL_INPUT_SIZE);

    for (let i = 0; i < labPixels.length; i++) {
        const l = labPixels[i];
        // The model output for 'a' and 'b' is often scaled, denormalize it.
        // This model seems to output a range that works well when multiplied by 128.
        const a = ab_channels[i * 2] * 128;
        const b = ab_channels[i * 2 + 1] * 128;
        
        const [r, g, b_rgb] = labToRgb(l, a, b);

        colorizedImageData.data[i * 4] = r;
        colorizedImageData.data[i * 4 + 1] = g;
        colorizedImageData.data[i * 4 + 2] = b_rgb;
        colorizedImageData.data[i * 4 + 3] = 255;
    }
    colorCtx.putImageData(colorizedImageData, 0, 0);

    // 6. Composite the color onto the original full-resolution grayscale image
    const compositeCanvas = document.createElement('canvas');
    const w = originalImg.naturalWidth;
    const h = originalImg.naturalHeight;
    compositeCanvas.width = w;
    compositeCanvas.height = h;
    const compositeCtx = compositeCanvas.getContext('2d');
    
    compositeCtx.drawImage(originalImg, 0, 0, w, h);
    // 'color' blend mode applies the hue and saturation of the source (colorCanvas)
    // to the luminance of the destination (originalImg).
    compositeCtx.globalCompositeOperation = 'color';
    compositeCtx.drawImage(colorCanvas, 0, 0, w, h);
    compositeCtx.globalCompositeOperation = 'source-over'; // Reset blend mode

    // 7. Upscale the final result to the desired resolution
    // Note: This uses standard browser resampling. True AI super-resolution would
    // require another large model for the best quality.
    const finalCanvas = document.createElement('canvas');
    finalCanvas.width = w * upscaleFactor;
    finalCanvas.height = h * upscaleFactor;
    const finalCtx = finalCanvas.getContext('2d');
    
    // Use high-quality scaling
    finalCtx.imageSmoothingEnabled = true;
    finalCtx.imageSmoothingQuality = 'high';
    
    finalCtx.drawImage(compositeCanvas, 0, 0, finalCanvas.width, finalCanvas.height);

    return finalCanvas;
}