IcpTask.java

package cz.fidentis.analyst.batch;

import cz.fidentis.analyst.Logger;
import cz.fidentis.analyst.canvas.Canvas;
import cz.fidentis.analyst.core.ProgressDialog;
import cz.fidentis.analyst.face.HumanFace;
import cz.fidentis.analyst.face.HumanFaceFactory;
import cz.fidentis.analyst.face.HumanFaceUtils;
import cz.fidentis.analyst.mesh.core.MeshModel;
import cz.fidentis.analyst.mesh.io.MeshObjExporter;
import cz.fidentis.analyst.scene.DrawableFace;
import cz.fidentis.analyst.visitors.kdtree.AvgFaceConstructor;
import cz.fidentis.analyst.visitors.mesh.sampling.NoSampling;
import cz.fidentis.analyst.visitors.mesh.sampling.PointSampling;
import cz.fidentis.analyst.visitors.mesh.sampling.RandomSampling;
import java.io.File;
import java.io.IOException;
import java.nio.file.Path;
import java.util.List;
import javax.swing.SwingWorker;

/**
 * A task that registers multiple faces using ICP and, simultaneously,
 * computes average face. 
 * The exact computation parameters are taken from the {@code BatchPanel}.
 * 
 * @author Radek Oslejsek
 */
public class IcpTask extends SwingWorker<MeshModel, HumanFace> {
    
    private static final double ICP_ERROR = 0.3; //0.05;
    private static final int    ICP_ITERS = 100;

    private HumanFace avgFace;
    private final ProgressDialog progressDialog;
    private final BatchPanel controlPanel;
    private final Canvas canvas;
    
    private int faceSceneSlot = -1; // scene slot to show the transformed face
    
    /**
     * Index of initial face in the list of available faces
     */
    private final int initialFaceIndex; 
    
    private final Stopwatch totalTime = new Stopwatch("Total computation time:\t");
    private final Stopwatch avgFaceComputationTime = new Stopwatch("AVG face computation time:\t");
    private final Stopwatch icpComputationTime = new Stopwatch("ICP registration time:\t");
    private final Stopwatch loadTime = new Stopwatch("Disk access time:\t");
    private final Stopwatch kdTreeConstructionTime = new Stopwatch("KD trees construction time:\t");
    
    /**
     * Constructor.
     * 
     * @param initialFaceIndex Index to the {@code controlPanel.getFaces()} list of faces
     * @param progressDialog A window that show the progress of the computation. Must not be {@code null}
     * @param controlPanel A control panel with computation parameters. Must not be {@code null}
     * @param canvas A canvas with 3D scene. 
     */
    public IcpTask(int initialFaceIndex, ProgressDialog progressDialog, BatchPanel controlPanel, Canvas canvas) {
        if (initialFaceIndex < 0 || initialFaceIndex >= controlPanel.getFacePaths().size()) {
            throw new IllegalArgumentException("initFace");
        }
        
        this.initialFaceIndex = initialFaceIndex;
        
        this.progressDialog = progressDialog;
        this.controlPanel = controlPanel;
        this.canvas = canvas;
    }

    @Override
    protected MeshModel doInBackground() throws Exception {
        HumanFaceFactory factory = controlPanel.getHumanFaceFactory();
        List<Path> faces = controlPanel.getFacePaths();
        int undersampling = controlPanel.getIcpUndersampling();
        boolean computeICP = controlPanel.computeICP();
        boolean computeAvgFace = controlPanel.computeAvgFace();

        factory.setReuseDumpFile(false); // we can't use this optimization because changes are made in models
        factory.setStrategy(HumanFaceFactory.Strategy.MRU); // keep first X faces in the memory
        
        // We don't need to reaload the initFace periodically for two reasons:
        //   - It is never dumped from memory to disk because we use MRU
        //   - Even if dumped, the face keeps in the mempry until we hold the pointer to it
        loadTime.start();
        String initFaceId = factory.loadFace(faces.get(initialFaceIndex).toFile());
        HumanFace initFace = factory.getFace(initFaceId);
        loadTime.stop();
        
        totalTime.start();
        AvgFaceConstructor avgFaceConstructor = null;
        
        for (int i = 0; i < faces.size(); i++) {

            if (isCancelled()) {
                return null;
            }
            
            // Compute AVG template face. Use each tranfromed face only once. Skip the original face
            if (i != initialFaceIndex && (computeICP || computeAvgFace)) {

                loadTime.start();
                String faceId = factory.loadFace(faces.get(i).toFile());
                HumanFace face = factory.getFace(faceId);
                loadTime.stop();
                
                if (computeICP) { // ICP registration - face is transformed!
                    
                    /*
                    //////////////////////
                    SymmetryEstimator se = new SymmetryEstimator(new SymmetryConfig());
                    face.getMeshModel().compute(se);
                    face.setSymmetryPlane(se.getSymmetryPlane());
                    se = new SymmetryEstimator(new SymmetryConfig());
                    initFace.getMeshModel().compute(se);
                    initFace.setSymmetryPlane(se.getSymmetryPlane());
                    HumanFaceUtils.alignSymmetryPlanes(initFace, face, computeAvgFace, true);
                    //////////////////////
                    */
                    
                    PointSampling sampling = (undersampling == 0)? new NoSampling() : new RandomSampling(undersampling);
                    icpComputationTime.start();
                    HumanFaceUtils.alignMeshes(
                            initFace,
                            face, // is transformed
                            ICP_ITERS,  // max iterations
                            controlPanel.scaleIcp(),
                            ICP_ERROR,  // error
                            sampling,
                            false // drop k-d tree, if exists
                    );
                    icpComputationTime.stop();
                    //System.out.println(sampling);
                }
                
                if (computeAvgFace) { // AVG template face
                    kdTreeConstructionTime.start();
                    face.computeKdTree(computeICP); // if transformed by ICP, force k-d tree re-computation
                    kdTreeConstructionTime.stop();
                
                    avgFaceComputationTime.start();
                    if (avgFaceConstructor == null) {
                        avgFaceConstructor = new AvgFaceConstructor(initFace.getMeshModel());
                    }
                    face.getKdTree().accept(avgFaceConstructor);
                    avgFaceComputationTime.stop();
                }
                
                //face.removeKdTree(); // preserve k-d tree for the similarity computation phase
                publish(face);       // update progress bar and possibly render the transformed face
            }

            int progress = (int) Math.round(100.0 * (i + 1.0) / faces.size());
            progressDialog.setValue(progress);
            
            //Logger.print(factory.toString());
        }

        if (computeAvgFace) {
            File tempFile = File.createTempFile(this.getClass().getSimpleName(), ".obj");
            tempFile.deleteOnExit();
            avgFace = new HumanFace(avgFaceConstructor.getAveragedMeshModel(), tempFile.getCanonicalPath());
            try {
                new MeshObjExporter(avgFace.getMeshModel()).exportModelToObj(tempFile);
            } catch (IOException ex) {
                Logger.print(ex.toString());
            }
        }
        
        totalTime.stop();

        printTimeStats();

        return (avgFace == null) ? null : avgFace.getMeshModel();
    }

    @Override
    protected void done() {
        progressDialog.dispose(); // close progess bar
        if (isCancelled()) {
            avgFace = null;
        }
    }

    @Override
    protected void process(List<HumanFace> chunks) {
        chunks.stream().forEach(f -> {
            if (isCancelled()) {
                return;
            }
            if (controlPanel.showIcpPreview()) {
                if (faceSceneSlot == -1) {
                    faceSceneSlot = canvas.getScene().getFreeSlotForFace();
                }
                canvas.getScene().setDrawableFace(faceSceneSlot, f);
                canvas.getScene().getDrawableFace(faceSceneSlot).setTransparency(0.5f);
                canvas.getScene().getDrawableFace(faceSceneSlot).setColor(DrawableFace.SKIN_COLOR_SECONDARY);
                canvas.renderScene();
            }
        });
    }

    public HumanFace getAverageFace() {
        return this.avgFace;
    }

    protected void printTimeStats() {
        Logger.print(avgFaceComputationTime.toString());
        Logger.print(icpComputationTime.toString());
        Logger.print(loadTime.toString());
        Logger.print(kdTreeConstructionTime.toString());
        Logger.print(totalTime.toString());
    }
}