com.flash3388.flashlib.vision.cv.template.MultipleTemplateMatcher Maven / Gradle / Ivy

Go to download
package com.flash3388.flashlib.vision.cv.template;

import org.opencv.core.Mat;

import java.util.ArrayList;
import java.util.Collection;
import java.util.concurrent.Callable;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Future;
import java.util.function.Function;

public class MultipleTemplateMatcher implements TemplateMatcher {
    // TODO: CONSIDER USING A FORK-JOIN POOL TO SEPARATE WORK FROM matchWithScaling TASKS

    private final Collection mTemplates;
    private final TemplateMatchingMethod mTemplateMatchingMethod;
    private final ExecutorService mExecutorService;

    public MultipleTemplateMatcher(Collection templates, TemplateMatchingMethod templateMatchingMethod, ExecutorService executorService) {
        mTemplates = templates;
        mTemplateMatchingMethod = templateMatchingMethod;
        mExecutorService = executorService;
    }

    @Override
    public TemplateMatchingResult match(Mat scene) throws TemplateMatchingException {
        try {
            return runMatchOnTemplates((template) ->
                    new TemplateMatchingTask(
                            new SingleTemplateMatcher(template, mTemplateMatchingMethod),
                            scene));
        } catch (InterruptedException e) {
            throw new TemplateMatchingException(e);
        }
    }

    @Override
    public ScaledTemplateMatchingResult matchWithScaling(Mat scene, double initialScaleFactor) throws TemplateMatchingException {
        try {
            return runMatchOnTemplates((template) ->
                    new ScaledTemplateMatchingTask(
                            new SingleTemplateMatcher(template, mTemplateMatchingMethod),
                            scene,
                            initialScaleFactor));
        } catch (InterruptedException e) {
            throw new TemplateMatchingException(e);
        }
    }

    private  T runMatchOnTemplates(Function> taskFromTemplate) throws InterruptedException, TemplateMatchingException {
        Collection> futures = new ArrayList<>();
        try {
            for (Mat template : mTemplates) {
                Callable task = taskFromTemplate.apply(template);
                Future future = mExecutorService.submit(task);
                futures.add(future);
            }

            return getBestMatch(futures);
        } finally {
            cancelRunningFutures(futures);
        }
    }

    private  T getBestMatch(Collection> futures) throws InterruptedException, TemplateMatchingException {
        T bestMatch = null;

        for (Future future : futures) {
            try {
                T result = future.get();

                if (bestMatch == null || result.compareTo(bestMatch) > 0) {
                    bestMatch = result;
                }
            } catch (ExecutionException e) {
                throw new TemplateMatchingException(e);
            }
        }

        if (bestMatch == null) {
            throw new NoTemplateMatchException();
        }

        return bestMatch;
    }

    private  void cancelRunningFutures(Collection> futures) {
        for (Future future : futures) {
            if (!future.isDone()) {
                future.cancel(true);
            }
        }
    }

    private static class TemplateMatchingTask implements Callable {

        private final TemplateMatcher mTemplateMatcher;
        private final Mat mScene;

        private TemplateMatchingTask(TemplateMatcher templateMatcher, Mat scene) {
            mTemplateMatcher = templateMatcher;
            mScene = scene;
        }

        @Override
        public TemplateMatchingResult call() throws Exception {
            return mTemplateMatcher.match(mScene);
        }
    }

    private static class ScaledTemplateMatchingTask implements Callable {

        private final TemplateMatcher mTemplateMatcher;
        private final Mat mScene;
        private final double mInitialScaleFactor;

        private ScaledTemplateMatchingTask(TemplateMatcher templateMatcher, Mat scene, double initialScaleFactor) {
            mTemplateMatcher = templateMatcher;
            mScene = scene;
            mInitialScaleFactor = initialScaleFactor;
        }

        @Override
        public ScaledTemplateMatchingResult call() throws Exception {
            return mTemplateMatcher.matchWithScaling(mScene, mInitialScaleFactor);
        }
    }
}