Photo Category Organizer

/**
 * Organizes a photo by identifying its content and assigning categories using a pre-trained AI model.
 * This function dynamically loads TensorFlow.js and the MobileNet model to perform image classification in the browser.
 *
 * @param {Image} originalImg The original javascript Image object to be categorized.
 * @param {number} topK The maximum number of top categories to display. Defaults to 3.
 * @param {number} minConfidence The minimum confidence score (from 0.0 to 1.0) required for a category to be shown. Defaults to 0.3 (30%).
 * @returns {Promise<HTMLElement>} A promise that resolves to an HTML div element containing the original image and a list of predicted categories.
 */
async function processImage(originalImg, topK = 3, minConfidence = 0.3) {
    // Helper function to dynamically load a script and return a promise
    const loadScript = (url) => {
        return new Promise((resolve, reject) => {
            // Check if the script is already on the page
            if (document.querySelector(`script[src="${url}"]`)) {
                return resolve();
            }
            const script = document.createElement('script');
            script.src = url;
            script.onload = () => resolve();
            script.onerror = (err) => reject(new Error(`Failed to load script: ${url}`));
            document.head.appendChild(script);
        });
    };

    // Create a container for the output
    const resultContainer = document.createElement('div');
    resultContainer.style.fontFamily = 'Arial, sans-serif';
    resultContainer.style.textAlign = 'center';
    resultContainer.style.padding = '20px';
    resultContainer.style.border = '1px solid #ddd';
    resultContainer.style.borderRadius = '8px';
    resultContainer.style.backgroundColor = '#f9f9f9';
    resultContainer.style.maxWidth = '500px';
    resultContainer.style.margin = '0 auto';

    // Display a loading message
    const loadingMessage = document.createElement('p');
    loadingMessage.textContent = 'Initializing AI model...';
    loadingMessage.style.color = '#555';
    resultContainer.appendChild(loadingMessage);

    try {
        // Dynamically load TensorFlow.js and the MobileNet model scripts
        await loadScript('https://cdn.jsdelivr.net/npm/@tensorflow/tfjs@3.11.0/dist/tf.min.js');
        await loadScript('https://cdn.jsdelivr.net/npm/@tensorflow-models/mobilenet@2.1.0/dist/mobilenet.min.js');

        // Load the model (cache it on the window object to avoid reloading on subsequent calls)
        loadingMessage.textContent = 'Loading AI model (this may take a moment on first run)...';
        if (!window.mobilenetModel) {
            window.mobilenetModel = await mobilenet.load();
        }
        const model = window.mobilenetModel;

        loadingMessage.textContent = 'Analyzing image...';

        // Classify the image
        const predictions = await model.classify(originalImg);

        // Clear the loading message and build the final output
        resultContainer.innerHTML = '';

        // Display the original image
        const imageElement = document.createElement('img');
        imageElement.src = originalImg.src;
        imageElement.style.maxWidth = '100%';
        imageElement.style.maxHeight = '300px';
        imageElement.style.borderRadius = '8px';
        imageElement.style.boxShadow = '0 4px 8px rgba(0,0,0,0.1)';
        imageElement.style.marginBottom = '20px';
        resultContainer.appendChild(imageElement);

        const title = document.createElement('h3');
        title.textContent = 'Predicted Categories';
        title.style.margin = '0 0 10px 0';
        title.style.color = '#333';
        resultContainer.appendChild(title);

        // Filter predictions based on confidence and take the top K
        const filteredPredictions = predictions
            .filter(p => p.probability >= minConfidence)
            .slice(0, Number(topK) || 3);

        if (filteredPredictions.length === 0) {
            const noResults = document.createElement('p');
            noResults.textContent = `No categories found with at least ${Math.round(minConfidence * 100)}% confidence.`;
            noResults.style.color = '#666';
            resultContainer.appendChild(noResults);
        } else {
            const list = document.createElement('ul');
            list.style.listStyleType = 'none';
            list.style.padding = '0';
            list.style.margin = '0';
            list.style.textAlign = 'left';
            list.style.display = 'inline-block';

            filteredPredictions.forEach(prediction => {
                const listItem = document.createElement('li');
                listItem.style.background = '#e9ecef';
                listItem.style.margin = '5px 0';
                listItem.style.padding = '10px 15px';
                listItem.style.borderRadius = '5px';

                const className = document.createElement('span');
                // Take the most specific class name before the comma
                className.textContent = prediction.className.split(',')[0];
                className.style.fontWeight = 'bold';
                className.style.color = '#0056b3';

                const probability = document.createElement('span');
                probability.textContent = ` (${Math.round(prediction.probability * 100)}%)`;
                probability.style.color = '#555';
                probability.style.fontSize = '0.9em';

                listItem.appendChild(className);
                listItem.appendChild(probability);
                list.appendChild(listItem);
            });
            resultContainer.appendChild(list);
        }

    } catch (error) {
        console.error('Error during image classification:', error);
        resultContainer.innerHTML = ''; // Clear previous content
        const errorElement = document.createElement('div');
        errorElement.style.color = 'red';
        errorElement.style.padding = '20px';
        errorElement.textContent = 'An error occurred while organizing the photo. Please check the browser console for more details.';
        resultContainer.appendChild(errorElement);
    }

    return resultContainer;
}