@salama/image-finder/TemplateMatcher.js

Advanced template matching tool with OpenCV.js featuring color sensitivity, batch processing, and performance optimization

const fs = require('fs').promises;
const path = require('path');
const sharp = require('sharp');

let cv = null;

class TemplateMatcher {
    constructor(options = {}) {
        this.threshold = options.threshold || 0.85;
        this.multiscale = options.multiscale || false;
        this.scales = options.scales || [0.8, 0.85, 0.9, 0.95, 1.0, 1.05, 1.1, 1.15, 1.2,2,3];
        this.grayscale = options.grayscale !== false;
        this.useEdges = options.useEdges || false;
        this.useBlur = options.useBlur || false;
        this.useInvert = options.useInvert || true;
        this.useChannel = options.useChannel || null;
        this.annotationColor = options.annotationColor || { r: 33, g: 255, b: 0, a: 255 };
        this.lineThickness = options.lineThickness || 4;
        this.method = options.method || null;
        this.maxMatches = options.maxMatches || 10;
        
        // New checkerboard pattern options
        this.useCheckerboard = options.useCheckerboard || false;
        this.checkerboardSize = options.checkerboardSize || 20; // Minimum area size to apply pattern
        this.checkerboardCellSize = options.checkerboardCellSize || 4; // Size of each checker cell
        this.solidColorThreshold = options.solidColorThreshold || 10; // Color variance threshold for "solid" detection
        
        this.cvReady = false;
        this.cv = null;
    }

    async ensureOpenCVReady() {
        if (!cv) {
            const cvLoader = require('@techstark/opencv-js');
            if (typeof cvLoader === 'function' || (typeof cvLoader === 'object' && cvLoader.then)) {
                cv = await cvLoader;
            } else {
                cv = cvLoader;
                if (cv.onRuntimeInitialized && !cv.Mat) {
                    await new Promise((resolve) => {
                        cv.onRuntimeInitialized = () => resolve();
                    });
                }
            }
            if (!cv || !cv.Mat) throw new Error('Failed to initialize OpenCV.js');
        }

        this.cv = cv;
        this.cvReady = true;
    }

    setMethodByName(methodName) {
        if (!this.cvReady || !this.cv) {
            throw new Error('OpenCV must be initialized before setting method');
        }
        
        const methodMap = {
            'CCOEFF': this.cv.TM_CCOEFF,
            'CCOEFF_NORMED': this.cv.TM_CCOEFF_NORMED,
            'CCORR': this.cv.TM_CCORR,
            'CCORR_NORMED': this.cv.TM_CCORR_NORMED,
            'SQDIFF': this.cv.TM_SQDIFF,
            'SQDIFF_NORMED': this.cv.TM_SQDIFF_NORMED
        };
        
        if (methodMap[methodName]) {
            this.method = methodMap[methodName];
            return true;
        } else {
            console.error(`Invalid method: ${methodName}. Available methods: ${Object.keys(methodMap).join(', ')}`);
            return false;
        }
    }

    async loadImageFromPath(imagePath) {
        const image = sharp(imagePath);
        const { data, info } = await image.ensureAlpha().raw().toBuffer({ resolveWithObject: true });
        const mat = new cv.Mat(info.height, info.width, cv.CV_8UC4);
        mat.data.set(data);
        return mat;
    }

    // New method: Detect solid color regions
    detectSolidRegions(img) {
        const gray = new this.cv.Mat();
        if (img.channels() > 1) {
            this.cv.cvtColor(img, gray, this.cv.COLOR_RGBA2GRAY);
        } else {
            gray = img.clone();
        }

        // Apply Gaussian blur to reduce noise
        const blurred = new this.cv.Mat();
        this.cv.GaussianBlur(gray, blurred, new this.cv.Size(5, 5), 0);

        // Calculate local standard deviation to find uniform regions
        const mean = new this.cv.Mat();
        const sqMean = new this.cv.Mat();
        const variance = new this.cv.Mat();
        const stdDev = new this.cv.Mat();
        
        // Calculate mean and variance in local neighborhoods
        const kernelSize = new this.cv.Size(this.checkerboardSize, this.checkerboardSize);
        this.cv.blur(blurred, mean, kernelSize);
        
        const blurredSq = new this.cv.Mat();
        this.cv.multiply(blurred, blurred, blurredSq);
        this.cv.blur(blurredSq, sqMean, kernelSize);
        
        const meanSq = new this.cv.Mat();
        this.cv.multiply(mean, mean, meanSq);
        this.cv.subtract(sqMean, meanSq, variance);
        this.cv.sqrt(variance, stdDev);
        
        meanSq.delete();

        // Create mask for solid regions (low standard deviation)
        const solidMask = new this.cv.Mat();
        this.cv.threshold(stdDev, solidMask, this.solidColorThreshold, 255, this.cv.THRESH_BINARY_INV);

        // Clean up
        gray.delete();
        blurred.delete();
        blurredSq.delete();
        mean.delete();
        sqMean.delete();
        variance.delete();
        stdDev.delete();

        return solidMask;
    }

    // New method: Apply checkerboard pattern to solid regions
    applyCheckerboardPattern(img) {
        const result = img.clone();
        const solidMask = this.detectSolidRegions(img);

        // Create checkerboard pattern using rectangle drawing (safer approach)
        const checkerboard = new this.cv.Mat.zeros(img.rows, img.cols, this.cv.CV_8UC1);
        
        const cellSize = this.checkerboardCellSize;
        for (let y = 0; y < img.rows; y += cellSize) {
            for (let x = 0; x < img.cols; x += cellSize) {
                const cellX = Math.floor(x / cellSize);
                const cellY = Math.floor(y / cellSize);
                const isWhite = (cellX + cellY) % 2 === 0;
                
                if (isWhite) {
                    const pt1 = new this.cv.Point(x, y);
                    const pt2 = new this.cv.Point(
                        Math.min(x + cellSize, img.cols),
                        Math.min(y + cellSize, img.rows)
                    );
                    this.cv.rectangle(checkerboard, pt1, pt2, new this.cv.Scalar(255), -1);
                }
            }
        }

        // Convert to floating point for calculations
        const resultFloat = new this.cv.Mat();
        result.convertTo(resultFloat, this.cv.CV_32F);
        
        const checkerFloat = new this.cv.Mat();
        checkerboard.convertTo(checkerFloat, this.cv.CV_32F, 1.0/255.0);
        
        // Create intensity multiplier (1.1 for white, 0.9 for black)
        const intensityMult = new this.cv.Mat();
        this.cv.addWeighted(checkerFloat, 0.2, checkerFloat, 0.0, 0.9, intensityMult);
        
        // Apply pattern based on solid mask
        const solidFloat = new this.cv.Mat();
        solidMask.convertTo(solidFloat, this.cv.CV_32F, 1.0/255.0);
        
        const channels = result.channels();
        
        if (channels === 1) {
            // Grayscale: multiply by intensity where solid
            const modified = new this.cv.Mat();
            this.cv.multiply(resultFloat, intensityMult, modified);
            
            // Blend original and modified based on solid mask
            const finalResult = new this.cv.Mat();
            this.cv.addWeighted(resultFloat, 1.0, modified, 0.0, 0.0, finalResult);
            
            // Use solidFloat as alpha for blending
            const invSolid = new this.cv.Mat();
            this.cv.subtract(new this.cv.Scalar(1.0), solidFloat, invSolid);
            
            const part1 = new this.cv.Mat();
            const part2 = new this.cv.Mat();
            this.cv.multiply(resultFloat, invSolid, part1);
            this.cv.multiply(modified, solidFloat, part2);
            
            this.cv.add(part1, part2, finalResult);
            finalResult.convertTo(result, this.cv.CV_8U);
            
            modified.delete();
            finalResult.delete();
            invSolid.delete();
            part1.delete();
            part2.delete();
        } else {
            // Multi-channel: apply to each RGB channel
            const channels_vec = new this.cv.MatVector();
            this.cv.split(resultFloat, channels_vec);
            
            const numColorChannels = Math.min(3, channels);
            for (let c = 0; c < numColorChannels; c++) {
                const channel = channels_vec.get(c);
                const modified = new this.cv.Mat();
                this.cv.multiply(channel, intensityMult, modified);
                
                // Blend original and modified based on solid mask
                const invSolid = new this.cv.Mat();
                this.cv.subtract(new this.cv.Scalar(1.0), solidFloat, invSolid);
                
                const part1 = new this.cv.Mat();
                const part2 = new this.cv.Mat();
                this.cv.multiply(channel, invSolid, part1);
                this.cv.multiply(modified, solidFloat, part2);
                
                const blended = new this.cv.Mat();
                this.cv.add(part1, part2, blended);
                
                // Replace the channel
                blended.copyTo(channel);
                
                modified.delete();
                invSolid.delete();
                part1.delete();
                part2.delete();
                blended.delete();
            }
            
            // Merge back
            const merged = new this.cv.Mat();
            this.cv.merge(channels_vec, merged);
            merged.convertTo(result, this.cv.CV_8U);
            
            // Cleanup
            for (let c = 0; c < channels_vec.size(); c++) {
                channels_vec.get(c).delete();
            }
            channels_vec.delete();
            merged.delete();
        }

        // Cleanup
        checkerboard.delete();
        resultFloat.delete();
        checkerFloat.delete();
        intensityMult.delete();
        solidFloat.delete();
        solidMask.delete();
        
        return result;
    }

    // Enhanced preprocessing with checkerboard pattern
    preprocessImage(img) {
        let processed = img.clone();

        // Apply checkerboard pattern to solid regions first
        if (this.useCheckerboard) {
            const withPattern = this.applyCheckerboardPattern(processed);
            processed.delete();
            processed = withPattern;
        }

        // Convert to grayscale if enabled
        if (this.grayscale && processed.channels() > 1) {
            const gray = new this.cv.Mat();
            this.cv.cvtColor(processed, gray, this.cv.COLOR_RGBA2GRAY);
            processed.delete();
            processed = gray;
        }

        // Apply Gaussian blur if enabled
        if (this.useBlur) {
            const blurred = new this.cv.Mat();
            this.cv.GaussianBlur(processed, blurred, new this.cv.Size(3, 3), 0);
            processed.delete();
            processed = blurred;
        }

        // Apply Canny edge detection if enabled
        if (this.useEdges) {
            const edges = new this.cv.Mat();
            this.cv.Canny(processed, edges, 50, 150);
            processed.delete();
            processed = edges;
        }

        return processed;
    }

    extractBestMatch(result, templateWidth, templateHeight, method) {
        const isMinMethod = method === this.cv.TM_SQDIFF || method === this.cv.TM_SQDIFF_NORMED;
        const matches = [];
        
        const resultCopy = result.clone();
        
        for (let i = 0; i < this.maxMatches; i++) {
            const minMax = this.cv.minMaxLoc(resultCopy);
            const confidence = isMinMethod ? minMax.minVal : minMax.maxVal;
            const loc = isMinMethod ? minMax.minLoc : minMax.maxLoc;

            const matchPass = (!isMinMethod && confidence >= this.threshold) || (isMinMethod && confidence <= this.threshold);

            if (matchPass) {
                matches.push({
                    x: loc.x,
                    y: loc.y,
                    width: templateWidth,
                    height: templateHeight,
                    confidence: this.normalizeConfidence(confidence, method),
                    rawConfidence: confidence,
                    scale: 1.0,
                    method: method
                });

                const suppressValue = isMinMethod ? 1.0 : 0.0;
                const suppressSize = Math.max(templateWidth, templateHeight) / 3;
                
                const pt1 = new this.cv.Point(
                    Math.max(0, loc.x - suppressSize), 
                    Math.max(0, loc.y - suppressSize)
                );
                const pt2 = new this.cv.Point(
                    Math.min(resultCopy.cols, loc.x + templateWidth + suppressSize), 
                    Math.min(resultCopy.rows, loc.y + templateHeight + suppressSize)
                );
                
                const suppressColor = new this.cv.Scalar(suppressValue);
                this.cv.rectangle(resultCopy, pt1, pt2, suppressColor, -1);
            } else {
                break;
            }
        }

        resultCopy.delete();
        return matches;
    }

    normalizeConfidence(rawConfidence, method) {
        const isMinMethod = method === this.cv.TM_SQDIFF || method === this.cv.TM_SQDIFF_NORMED;
        const isNormalized = method === this.cv.TM_CCOEFF_NORMED || 
                           method === this.cv.TM_SQDIFF_NORMED || 
                           method === this.cv.TM_CCORR_NORMED;
        
        if (isNormalized) {
            return isMinMethod ? (1.0 - rawConfidence) : rawConfidence;
        } else {
            if (isMinMethod) {
                return Math.max(0, 1.0 - Math.min(1.0, rawConfidence / 100000));
            } else {
                return Math.min(1.0, Math.max(0, rawConfidence / 100000));
            }
        }
    }

    async singleScaleMatch(screenshot, template) {
        const processedScreenshot = this.preprocessImage(screenshot);
        const processedTemplate = this.preprocessImage(template);

        const result = new this.cv.Mat();
        const mask = new this.cv.Mat();

        const methods = this.method ? [this.method] : [this.cv.TM_CCOEFF_NORMED, this.cv.TM_SQDIFF_NORMED];
        let allMatches = [];

        for (const method of methods) {
            this.cv.matchTemplate(processedScreenshot, processedTemplate, result, method, mask);
            const matches = this.extractBestMatch(result, template.cols, template.rows, method);
            allMatches = allMatches.concat(matches);
        }

        allMatches.sort((a, b) => b.confidence - a.confidence);
        const nmsMatches = this.applyNonMaxSuppression(allMatches);

        processedScreenshot.delete();
        processedTemplate.delete();
        result.delete();
        mask.delete();

        return nmsMatches.slice(0, this.maxMatches);
    }

    async multiScaleMatch(screenshot, template) {
        let allMatches = [];

        for (const scale of this.scales) {
            const scaledTemplate = new this.cv.Mat();
            const size = new this.cv.Size(Math.round(template.cols * scale), Math.round(template.rows * scale));
            this.cv.resize(template, scaledTemplate, size, 0, 0, this.cv.INTER_LINEAR);

            const matches = await this.singleScaleMatch(screenshot, scaledTemplate);
            matches.forEach(match => {
                match.scale = scale;
                match.x = Math.round(match.x);
                match.y = Math.round(match.y);
                match.width = Math.round(template.cols * scale);
                match.height = Math.round(template.rows * scale);
            });
            
            allMatches = allMatches.concat(matches);
            scaledTemplate.delete();
        }

        const nmsMatches = this.applyNonMaxSuppression(allMatches, 0.3);
        return nmsMatches.slice(0, this.maxMatches);
    }

    applyNonMaxSuppression(matches, overlapThreshold = 0.3) {
        if (matches.length === 0) return matches;

        matches.sort((a, b) => b.confidence - a.confidence);

        const kept = [];
        const suppressed = new Set();

        for (let i = 0; i < matches.length; i++) {
            if (suppressed.has(i)) continue;

            const current = matches[i];
            kept.push(current);

            for (let j = i + 1; j < matches.length; j++) {
                if (suppressed.has(j)) continue;

                const other = matches[j];
                const overlap = this.calculateIoU(current, other);
                
                if (overlap > overlapThreshold) {
                    suppressed.add(j);
                }
            }
        }

        return kept;
    }

    calculateIoU(match1, match2) {
        const x1 = Math.max(match1.x, match2.x);
        const y1 = Math.max(match1.y, match2.y);
        const x2 = Math.min(match1.x + match1.width, match2.x + match2.width);
        const y2 = Math.min(match1.y + match1.height, match2.y + match2.height);

        if (x2 <= x1 || y2 <= y1) return 0;

        const intersectionArea = (x2 - x1) * (y2 - y1);
        const area1 = match1.width * match1.height;
        const area2 = match2.width * match2.height;
        const unionArea = area1 + area2 - intersectionArea;

        return intersectionArea / unionArea;
    }

    async refineMatches(screenshot, template, initialMatches) {
        const refinedMatches = [];
        
        for (const match of initialMatches) {
            const padding = 10;
            const x = Math.max(0, match.x - padding);
            const y = Math.max(0, match.y - padding);
            const width = Math.min(screenshot.cols - x, match.width + 2 * padding);
            const height = Math.min(screenshot.rows - y, match.height + 2 * padding);
            
            const roi = screenshot.roi(new this.cv.Rect(x, y, width, height));
            
            const result = new this.cv.Mat();
            this.cv.matchTemplate(roi, template, result, match.method);
            
            const minMax = this.cv.minMaxLoc(result);
            const isMinMethod = match.method === this.cv.TM_SQDIFF || match.method === this.cv.TM_SQDIFF_NORMED;
            const newConfidence = isMinMethod ? minMax.minVal : minMax.maxVal;
            const newLoc = isMinMethod ? minMax.minLoc : minMax.maxLoc;
            
            refinedMatches.push({
                x: x + newLoc.x,
                y: y + newLoc.y,
                width: match.width,
                height: match.height,
                confidence: this.normalizeConfidence(newConfidence, match.method),
                rawConfidence: newConfidence,
                scale: match.scale,
                method: match.method
            });
            
            roi.delete();
            result.delete();
        }
        
        return refinedMatches;
    }

    async findMatches(screenshotPath, templatePath) {
        await this.ensureOpenCVReady();

        const screenshot = await this.loadImageFromPath(screenshotPath);
        const template = await this.loadImageFromPath(templatePath);

        let matches;
        if (this.multiscale) {
            matches = await this.multiScaleMatch(screenshot, template);
        } else {
            matches = await this.singleScaleMatch(screenshot, template);
        }

        const annotatedImage = this.annotateMatches(screenshot, matches);

        screenshot.delete();
        template.delete();

        return {
            matches: matches,
            annotatedImage: annotatedImage
        };
    }

    annotateMatches(screenshot, matches) {
        const annotated = screenshot.clone();
        if (annotated.channels() === 1) {
            this.cv.cvtColor(annotated, annotated, this.cv.COLOR_GRAY2RGBA);
        }

        matches.forEach((match, idx) => {
            const pt1 = new this.cv.Point(Math.round(match.x), Math.round(match.y));
            const pt2 = new this.cv.Point(Math.round(match.x + match.width), Math.round(match.y + match.height));
            
            const colors = [
                { r: 33, g: 255, b: 0, a: 255 },   // Green
                { r: 255, g: 33, b: 0, a: 255 },   // Red
                { r: 0, g: 33, b: 255, a: 255 },   // Blue
                { r: 255, g: 255, b: 0, a: 255 },  // Yellow
                { r: 255, g: 0, b: 255, a: 255 },  // Magenta
                { r: 0, g: 255, b: 255, a: 255 },  // Cyan
                { r: 255, g: 128, b: 0, a: 255 },  // Orange
                { r: 128, g: 0, b: 255, a: 255 },  // Purple
                { r: 255, g: 192, b: 203, a: 255 },// Pink
                { r: 128, g: 128, b: 128, a: 255 } // Gray
            ];
            
            const colorIndex = idx % colors.length;
            const color = new this.cv.Scalar(colors[colorIndex].r, colors[colorIndex].g, colors[colorIndex].b, colors[colorIndex].a);

            this.cv.rectangle(annotated, pt1, pt2, color, this.lineThickness, this.cv.LINE_8, 0);

            const label = `#${idx} ${(match.confidence * 100).toFixed(1)}% [${match.method}]`;
            const textPt = new this.cv.Point(Math.round(match.x), Math.max(0, Math.round(match.y) - 5));

            this.cv.putText(annotated, label, textPt, this.cv.FONT_HERSHEY_SIMPLEX, 0.5, color, 1, this.cv.LINE_AA);
        });

        return annotated;
    }

    async matToBuffer(mat) {
        const width = mat.cols;
        const height = mat.rows;
        const channels = mat.channels();
        let buffer;

        if (channels === 4) {
            buffer = Buffer.from(mat.data);
        } else {
            buffer = Buffer.allocUnsafe(width * height * 4);
            for (let i = 0, j = 0; i < mat.data.length; i++, j += 4) {
                if (channels === 1) {
                    const val = mat.data[i];
                    buffer[j] = val;
                    buffer[j + 1] = val;
                    buffer[j + 2] = val;
                } else {
                    buffer[j] = mat.data[i * 3];
                    buffer[j + 1] = mat.data[i * 3 + 1];
                    buffer[j + 2] = mat.data[i * 3 + 2];
                }
                buffer[j + 3] = 255;
            }
        }

        return sharp(buffer, { raw: { width, height, channels: 4 } }).png().toBuffer();
    }

    async saveAnnotatedImage(annotatedMat, outputPath) {
        const buffer = await this.matToBuffer(annotatedMat);
        await fs.writeFile(outputPath, buffer);
    }
}

// Enhanced CLI
async function main() {
    const args = process.argv.slice(2);
    if (args.length < 2 || args.includes('--help')) {
        console.log(`Usage: node enhanced-cv.js [options] <screenshot> <template>
Options:
  --threshold <value>        Matching threshold (default 0.85)
  --multiscale              Enable multi-scale matching
  --edges                   Enable edge-based preprocessing
  --blur                    Enable Gaussian blur preprocessing
  --checkerboard            Enable checkerboard pattern for solid regions
  --checker-size <value>    Minimum area size for checkerboard pattern (default 20)
  --checker-cell <value>    Size of each checker cell (default 4)
  --solid-threshold <value> Color variance threshold for solid detection (default 10)
  --method <name>           Force specific method (CCOEFF, CCOEFF_NORMED, CCORR, CCORR_NORMED, SQDIFF, SQDIFF_NORMED)
  --output <path>           Save annotated image to path
Example:
  node enhanced-cv.js screen.png template.png --threshold 0.8 --multiscale --checkerboard --output result.png
`);
        process.exit(0);
    }

    const options = { 
        threshold: 0.85, 
        multiscale: false, 
        grayscale: true, 
        useEdges: false, 
        useBlur: false,
        useCheckerboard: false,
        checkerboardSize: 20,
        checkerboardCellSize: 4,
        solidColorThreshold: 10
    };
    const positional = [];

    for (let i = 0; i < args.length; i++) {
        const arg = args[i];
        if (arg === '--threshold') options.threshold = parseFloat(args[++i]);
        else if (arg === '--multiscale') options.multiscale = true;
        else if (arg === '--edges') options.useEdges = true;
        else if (arg === '--blur') options.useBlur = true;
        else if (arg === '--checkerboard') options.useCheckerboard = true;
        else if (arg === '--checker-size') options.checkerboardSize = parseInt(args[++i]);
        else if (arg === '--checker-cell') options.checkerboardCellSize = parseInt(args[++i]);
        else if (arg === '--solid-threshold') options.solidColorThreshold = parseInt(args[++i]);
        else if (arg === '--method') {
            const methodName = args[++i];
            const methodMap = {
                'CCOEFF': 'CCOEFF',
                'CCOEFF_NORMED': 'CCOEFF_NORMED',
                'CCORR': 'CCORR',
                'CCORR_NORMED': 'CCORR_NORMED',
                'SQDIFF': 'SQDIFF',
                'SQDIFF_NORMED': 'SQDIFF_NORMED'
            };
            if (methodMap[methodName]) {
                options.methodName = methodName;
            } else {
                console.error(`Invalid method: ${methodName}. Available methods: ${Object.keys(methodMap).join(', ')}`);
                process.exit(1);
            }
        }
        else if (arg === '--output') options.output = args[++i];
        else positional.push(arg);
    }

    const [screenshotPath, templatePath] = positional;

    if (!screenshotPath || !templatePath) {
        console.error('Error: Both screenshot and template paths are required');
        process.exit(1);
    }

    const matcher = new TemplateMatcher(options);
    
    // Set method if specified
    if (options.methodName) {
        await matcher.ensureOpenCVReady();
        matcher.setMethodByName(options.methodName);
    }

    try {
        console.log('Running enhanced template matching with checkerboard pattern...');
        const result = await matcher.findMatches(screenshotPath, templatePath);

        console.log(`Found ${result.matches.length} match(es):`);
        result.matches.forEach((m, idx) => {
            console.log(`#${idx}: [${m.x}, ${m.y}] Confidence: ${m.confidence.toFixed(4)} (Raw: ${m.rawConfidence.toFixed(4)}) Scale: ${m.scale}`);
        });

        if (options.output) {
            await matcher.saveAnnotatedImage(result.annotatedImage, options.output);
            console.log(`Annotated image saved to: ${options.output}`);
        }

        result.annotatedImage.delete();
    } catch (err) {
        console.error('Error:', err);
        process.exit(1);
    }
}

module.exports = TemplateMatcher;

if (require.main === module) {
    main().catch(console.error);
}