Image Capital Identifier

/**
 * Identifies the capital city in an image by recognizing famous landmarks.
 * This function uses the TensorFlow.js and MobileNet pre-trained model for image classification.
 * It loads the necessary scripts and model on the first run.
 * 
 * @param {Image} originalImg The original image object to be processed.
 * @returns {Promise<HTMLCanvasElement>} A promise that resolves with a canvas element 
 *   containing the original image and an overlay with the identified capital city's name.
 */
const processImage = (function() {
    // Use a closure to ensure the model is loaded only once.
    let modelPromise = null;

    // A map of recognizable landmarks (from MobileNet model) to their capital cities.
    const landmarkToCapitalMap = {
        // Paris, France
        'eiffel tower': 'Paris, France',
        'louvre': 'Paris, France',
        'arc de triomphe': 'Paris, France',
        // London, UK
        'tower bridge': 'London, UK',
        'big ben': 'London, UK',
        'palace of westminster': 'London, UK',
        'buckingham palace': 'London, UK',
        // Rome, Italy
        'colosseum': 'Rome, Italy',
        'roman forum': 'Rome, Italy',
        'trevi fountain': 'Rome, Italy',
        'pantheon': 'Rome, Italy',
        // Washington D.C., USA
        'white house': 'Washington D.C., USA',
        'washington monument': 'Washington D.C., USA',
        'united states capitol': 'Washington D.C., USA',
        // Moscow, Russia
        'kremlin': 'Moscow, Russia',
        'st. basil\'s cathedral': 'Moscow, Russia',
        // Tokyo, Japan
        'tokyo tower': 'Tokyo, Japan',
        // Beijing, China
        'forbidden city': 'Beijing, China',
        'tiananmen square': 'Beijing, China',
        // Cairo, Egypt
        'great pyramid of giza': 'Cairo, Egypt',
        'great sphinx of giza': 'Cairo, Egypt',
        // Berlin, Germany
        'brandenburg gate': 'Berlin, Germany',
        'reichstag building': 'Berlin, Germany',
        // Athens, Greece
        'parthenon': 'Athens, Greece',
        'acropolis': 'Athens, Greece',
        // Special cases for famous landmarks in non-capital cities
        'statue of liberty': 'New York City, USA',
        'sydney opera house': 'Canberra, Australia (Sydney shown)',
        'taj mahal': 'New Delhi, India (Agra shown)',
        'leaning tower of pisa': 'Rome, Italy (Pisa shown)',
        'christ the redeemer': 'Brasília, Brazil (Rio de Janeiro shown)'
    };


    // Helper function to dynamically load a script and return a promise.
    function loadScript(src) {
        return new Promise((resolve, reject) => {
            if (document.querySelector(`script[src="${src}"]`)) {
                resolve();
                return;
            }
            const script = document.createElement('script');
            script.src = src;
            script.onload = () => resolve();
            script.onerror = () => reject(new Error(`Script load error for ${src}`));
            document.head.appendChild(script);
        });
    }

    // Initializes and loads the machine learning model.
    async function getModel() {
        if (!modelPromise) {
            modelPromise = new Promise(async (resolve, reject) => {
                try {
                    await loadScript('https://cdn.jsdelivr.net/npm/@tensorflow/tfjs@latest/dist/tf.min.js');
                    await loadScript('https://cdn.jsdelivr.net/npm/@tensorflow-models/mobilenet@latest/dist/mobilenet.min.js');
                    // mobilenet is now available on the window object.
                    const model = await window.mobilenet.load();
                    resolve(model);
                } catch (error) {
                    console.error("Failed to load ML model.", error);
                    reject(error);
                }
            });
        }
        return modelPromise;
    }

    // The main function exposed as processImage.
    return async function(originalImg) {
        const canvas = document.createElement('canvas');
        canvas.width = originalImg.naturalWidth;
        canvas.height = originalImg.naturalHeight;
        const ctx = canvas.getContext('2d');
        ctx.drawImage(originalImg, 0, 0);

        // Helper function to draw text overlay on the canvas.
        const drawTextOverlay = (text) => {
            const fontSize = Math.max(20, Math.min(50, Math.floor(canvas.width / 20)));
            const boxHeight = fontSize * 1.8;
            ctx.fillStyle = 'rgba(0, 0, 0, 0.7)';
            ctx.fillRect(0, canvas.height - boxHeight, canvas.width, boxHeight);
            ctx.fillStyle = 'white';
            ctx.font = `bold ${fontSize}px Arial`;
            ctx.textAlign = 'center';
            ctx.textBaseline = 'middle';
            ctx.fillText(text, canvas.width / 2, canvas.height - (boxHeight / 2));
        };

        drawTextOverlay('Loading recognition model...');

        try {
            const model = await getModel();
            drawTextOverlay('Analyzing image...');
            const predictions = await model.classify(originalImg);

            let foundCapital = 'Capital not identified';
            let bestMatchConfidence = 0.0;

            for (const prediction of predictions) {
                // Predictions can have multiple labels, e.g., "tower, watchtower"
                const labels = prediction.className.toLowerCase().split(/, | or /);
                for (const label of labels) {
                    const cleanLabel = label.trim();
                    if (landmarkToCapitalMap[cleanLabel] && prediction.probability > bestMatchConfidence) {
                        foundCapital = landmarkToCapitalMap[cleanLabel];
                        bestMatchConfidence = prediction.probability;
                    }
                }
            }

            // Redraw image to clear previous message, then draw the final result.
            ctx.drawImage(originalImg, 0, 0);
            drawTextOverlay(foundCapital);

        } catch (error) {
            console.error('Failed to identify capital in image:', error);
            ctx.drawImage(originalImg, 0, 0);
            drawTextOverlay('Error during analysis');
        }

        return canvas;
    };
})();