autopsy-flatpak/Experimental/src/org/sleuthkit/autopsy/experimental/objectdetection/ObjectDetectectionFileIngestModule.java

/*
 * Autopsy Forensic Browser
 *
 * Copyright 2018 Basis Technology Corp.
 * Contact: carrier <at> sleuthkit <dot> org
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.sleuthkit.autopsy.experimental.objectdetection;

import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.logging.Level;
import org.apache.commons.io.FilenameUtils;
import org.apache.commons.io.IOUtils;
import org.opencv.core.CvException;
import org.opencv.core.Mat;
import org.opencv.core.MatOfByte;
import org.opencv.core.MatOfRect;
import org.opencv.highgui.Highgui;
import org.opencv.objdetect.CascadeClassifier;
import org.openide.util.NbBundle.Messages;
import org.sleuthkit.autopsy.casemodule.Case;
import org.sleuthkit.autopsy.casemodule.NoCurrentCaseException;
import org.sleuthkit.autopsy.casemodule.services.Blackboard;
import org.sleuthkit.autopsy.corelibs.OpenCvLoader;
import org.sleuthkit.autopsy.coreutils.ImageUtils;
import org.sleuthkit.autopsy.coreutils.Logger;
import org.sleuthkit.autopsy.coreutils.PlatformUtil;
import org.sleuthkit.autopsy.ingest.FileIngestModuleAdapter;
import org.sleuthkit.autopsy.ingest.IngestJobContext;
import org.sleuthkit.autopsy.ingest.IngestMessage;
import org.sleuthkit.autopsy.ingest.IngestModule;
import org.sleuthkit.autopsy.ingest.IngestModuleReferenceCounter;
import org.sleuthkit.autopsy.ingest.IngestServices;
import org.sleuthkit.autopsy.ingest.ModuleDataEvent;
import org.sleuthkit.datamodel.AbstractFile;
import org.sleuthkit.datamodel.BlackboardArtifact;
import static org.sleuthkit.datamodel.BlackboardArtifact.ARTIFACT_TYPE.TSK_OBJECT_DETECTED;
import org.sleuthkit.datamodel.BlackboardAttribute;
import org.sleuthkit.datamodel.ReadContentInputStream;
import org.sleuthkit.datamodel.TskCoreException;

/**
 * Data source module to detect objects in images.
 */
public class ObjectDetectectionFileIngestModule extends FileIngestModuleAdapter {

    private final static Logger logger = Logger.getLogger(ObjectDetectectionFileIngestModule.class.getName());
    private static final IngestModuleReferenceCounter refCounter = new IngestModuleReferenceCounter();
    private long jobId;
    private Map<String, CascadeClassifier> classifiers;
    private final IngestServices services = IngestServices.getInstance();
    private Blackboard blackboard;

    @Messages({"ObjectDetectionFileIngestModule.noClassifiersFound.subject=No classifiers found.",
        "# {0} - classifierDir", "ObjectDetectionFileIngestModule.noClassifiersFound.message=No classifiers were found in {0}, object detection will not be executed."})
    @Override
    public void startUp(IngestJobContext context) throws IngestModule.IngestModuleException {
        jobId = context.getJobId();
        File classifierDir = new File(PlatformUtil.getObjectDetectionClassifierPath());
        classifiers = new HashMap<>();
        //Load all classifiers found in PlatformUtil.getObjectDetectionClassifierPath()
        if (OpenCvLoader.isOpenCvLoaded() && classifierDir.exists() && classifierDir.isDirectory()) {
            for (File classifier : classifierDir.listFiles()) {
                if (classifier.isFile() && FilenameUtils.getExtension(classifier.getName()).equalsIgnoreCase("xml")) {
                    classifiers.put(classifier.getName(), new CascadeClassifier(classifier.getAbsolutePath()));
                }
            }
        } else {
            throw new IngestModule.IngestModuleException("Unable to load classifiers for object detection module.");
        }
        if (refCounter.incrementAndGet(jobId) == 1 && classifiers.isEmpty()) {
            services.postMessage(IngestMessage.createWarningMessage(ObjectDetectionModuleFactory.getModuleName(),
                    Bundle.ObjectDetectionFileIngestModule_noClassifiersFound_subject(),
                    Bundle.ObjectDetectionFileIngestModule_noClassifiersFound_message(PlatformUtil.getObjectDetectionClassifierPath())));
        }
        try {
            blackboard = Case.getCurrentCaseThrows().getServices().getBlackboard();
        } catch (NoCurrentCaseException ex) {
            throw new IngestModule.IngestModuleException("Exception while getting open case.", ex);
        }
    }

    @Messages({"# {0} - detectionCount", "ObjectDetectionFileIngestModule.classifierDetection.text=Classifier detected {0} object(s)"})
    @Override
    public ProcessResult process(AbstractFile file) {
        if (!classifiers.isEmpty() && ImageUtils.isImageThumbnailSupported(file)) {
            //Any image we can create a thumbnail for is one we should apply the classifiers to
            InputStream inputStream = new ReadContentInputStream(file);
            byte[] imageInMemory;
            try {
                imageInMemory = IOUtils.toByteArray(inputStream);
            } catch (IOException ex) {
                logger.log(Level.WARNING, "Unable to perform object detection on " + file.getName(), ex);
                return IngestModule.ProcessResult.ERROR;
            }
            Mat originalImage = Highgui.imdecode(new MatOfByte(imageInMemory), Highgui.IMREAD_GRAYSCALE);
            MatOfRect detectionRectangles = new MatOfRect(); //the rectangles which reprent the coordinates on the image for where objects were detected
            for (String classifierKey : classifiers.keySet()) {
                //apply each classifier to the file
                try {
                    classifiers.get(classifierKey).detectMultiScale(originalImage, detectionRectangles);
                } catch (CvException ignored) {
                    //The image was likely an image which we are unable to generate a thumbnail for, and the classifier was likely one where that is not acceptable
                    logger.log(Level.INFO, String.format("Classifier '%s' could not be applied to file '%s'.", classifierKey, file.getParentPath() + file.getName())); //NON-NLS
                    continue;
                }

                if (!detectionRectangles.empty()) {
                    //if any detections occurred create an artifact for this classifier and file combination
                    try {
                        BlackboardArtifact artifact = file.newArtifact(TSK_OBJECT_DETECTED);
                        artifact.addAttribute(new BlackboardAttribute(BlackboardAttribute.ATTRIBUTE_TYPE.TSK_DESCRIPTION,
                                ObjectDetectionModuleFactory.getModuleName(),
                                classifierKey));
                        artifact.addAttribute(new BlackboardAttribute(BlackboardAttribute.ATTRIBUTE_TYPE.TSK_COMMENT,
                                ObjectDetectionModuleFactory.getModuleName(),
                                Bundle.ObjectDetectionFileIngestModule_classifierDetection_text((int) detectionRectangles.size().height)));

                        try {
                            /*
                             * Index the artifact for keyword search.
                             */
                            blackboard.indexArtifact(artifact);
                        } catch (Blackboard.BlackboardException ex) {
                            logger.log(Level.SEVERE, "Unable to index blackboard artifact " + artifact.getArtifactID(), ex); //NON-NLS
                        }

                        /*
                         * Send an event to update the view with the new result.
                         */
                        services.fireModuleDataEvent(new ModuleDataEvent(ObjectDetectionModuleFactory.getModuleName(), TSK_OBJECT_DETECTED, Collections.singletonList(artifact)));

                    } catch (TskCoreException ex) {
                        logger.log(Level.SEVERE, String.format("Failed to create blackboard artifact for '%s'.", file.getParentPath() + file.getName()), ex); //NON-NLS
                        return IngestModule.ProcessResult.ERROR;
                    }
                }
            }
        }

        return IngestModule.ProcessResult.OK;
    }

    @Override
    public void shutDown() {
        refCounter.decrementAndGet(jobId);
    }
}