src/models/modelLoader.js

/**
 * Model Loading - Load and initialize ONNX models
 */

import * as ort from 'https://cdn.jsdelivr.net/npm/onnxruntime-web@1.17.0/dist/esm/ort.min.js';
import { MODEL_PATHS, ONNX_CONFIG, LABEL_MAPPINGS } from '../config.js';
import { setModel, appState } from '../state.js';
import { updateLoadingStatus } from '../ui/loading.js';

/**
 * Initialize ONNX Runtime
 */
export async function initONNXRuntime() {
    ort.env.wasm.wasmPaths = ONNX_CONFIG.wasmPaths;
    ort.env.wasm.numThreads = ONNX_CONFIG.numThreads;
}

/**
 * Load label mappings from config files
 */
export async function loadLabelMappings() {
    try {
        const graderConfig = await fetch('./models/grader_model_compressed/config.json');
        const graderData = await graderConfig.json();
        LABEL_MAPPINGS.grader = graderData.id2label || {};
        
        console.log('Label mappings loaded:', LABEL_MAPPINGS);
    } catch (error) {
        console.warn('Could not load label mappings, using defaults:', error);
    }
}

/**
 * Load all ONNX models
 */
export async function loadAllModels() {
    try {
        updateLoadingStatus('Loading Classifier Model...', 25);
        const classifier = await ort.InferenceSession.create(MODEL_PATHS.classifier);
        setModel('classifier', classifier);
        
        updateLoadingStatus('Loading Quality Assessment Model...', 50);
        const poorGood = await ort.InferenceSession.create(MODEL_PATHS.poorGood);
        setModel('poorGood', poorGood);
        
        updateLoadingStatus('Loading Grader Model...', 75);
        const grader = await ort.InferenceSession.create(MODEL_PATHS.grader);
        setModel('grader', grader);
        
        updateLoadingStatus('Loading YOLO Detection Model...', 90);
        try {
            const yolo = await ort.InferenceSession.create(MODEL_PATHS.yolo);
            setModel('yolo', yolo);
        } catch (error) {
            console.warn('YOLO model not available:', error);
        }
        
        updateLoadingStatus('Models loaded successfully!', 100);
    } catch (error) {
        throw new Error(`Model loading failed: ${error.message}`);
    }
}

/**
 * Get model status for UI display
 */
export function getModelStatus() {
    return {
        classifier: !!appState.models.classifier,
        poorGood: !!appState.models.poorGood,
        grader: !!appState.models.grader,
        yolo: !!appState.models.yolo
    };
}