ImageQueryParser.java

package org.elasticsearch.index.query.image;


import net.semanticmetadata.lire.imageanalysis.LireFeature;
import net.semanticmetadata.lire.indexing.hashing.BitSampling;
import net.semanticmetadata.lire.indexing.hashing.LocalitySensitiveHashing;
import net.semanticmetadata.lire.utils.ImageUtils;
import org.apache.lucene.index.Term;
import org.apache.lucene.search.BooleanClause;
import org.apache.lucene.search.BooleanQuery;
import org.apache.lucene.search.Query;
import org.elasticsearch.ElasticsearchImageProcessException;
import org.elasticsearch.action.get.GetRequest;
import org.elasticsearch.action.get.GetResponse;
import org.elasticsearch.client.Client;
import org.elasticsearch.common.bytes.BytesReference;
import org.elasticsearch.common.inject.Inject;
import org.elasticsearch.common.io.stream.BytesStreamInput;
import org.elasticsearch.common.xcontent.XContentParser;
import org.elasticsearch.index.get.GetField;
import org.elasticsearch.index.mapper.image.FeatureEnum;
import org.elasticsearch.index.mapper.image.HashEnum;
import org.elasticsearch.index.mapper.image.ImageMapper;
import org.elasticsearch.index.query.QueryParseContext;
import org.elasticsearch.index.query.QueryParser;
import org.elasticsearch.index.query.QueryParsingException;

import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.IOException;

public class ImageQueryParser implements QueryParser {

    public static final String NAME = "image";

    private Client client;

    @Inject
    public ImageQueryParser(Client client) {
        this.client = client;
    }

    @Override
    public String[] names() {
        return new String[] {NAME};
    }

    @Override
    public Query parse(QueryParseContext parseContext) throws IOException, QueryParsingException {
        XContentParser parser = parseContext.parser();

        XContentParser.Token token = parser.nextToken();
        if (token != XContentParser.Token.FIELD_NAME) {
            throw new QueryParsingException(parseContext.index(), "[image] query malformed, no field");
        }


        String fieldName = parser.currentName();
        FeatureEnum featureEnum = null;
        byte[] image = null;
        HashEnum hashEnum = null;
        float boost = 1.0f;
        int limit = -1;

        String lookupIndex = parseContext.index().name();
        String lookupType = null;
        String lookupId = null;
        String lookupPath = null;
        String lookupRouting = null;


        token = parser.nextToken();
        if (token == XContentParser.Token.START_OBJECT) {
            String currentFieldName = null;
            while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) {
                if (token == XContentParser.Token.FIELD_NAME) {
                    currentFieldName = parser.currentName();
                } else {
                    if ("feature".equals(currentFieldName)) {
                        featureEnum = FeatureEnum.getByName(parser.text());
                    } else if ("image".equals(currentFieldName)) {
                        image = parser.binaryValue();
                    } else if ("hash".equals(currentFieldName)) {
                        hashEnum = HashEnum.getByName(parser.text());
                    } else if ("boost".equals(currentFieldName)) {
                        boost = parser.floatValue();
                    } else if ("limit".equals(currentFieldName)) {
                        limit = parser.intValue();
                    }else if ("index".equals(currentFieldName)) {
                        lookupIndex = parser.text();
                    } else if ("type".equals(currentFieldName)) {
                        lookupType = parser.text();
                    } else if ("id".equals(currentFieldName)) {
                        lookupId = parser.text();
                    } else if ("path".equals(currentFieldName)) {
                        lookupPath = parser.text();
                    } else if ("routing".equals(currentFieldName)) {
                        lookupRouting = parser.textOrNull();
                    } else {
                        throw new QueryParsingException(parseContext.index(), "[image] query does not support [" + currentFieldName + "]");
                    }
                }
            }
            parser.nextToken();
        }

        if (featureEnum == null) {
            throw new QueryParsingException(parseContext.index(), "No feature specified for image query");
        }

        String luceneFieldName = fieldName + "." + featureEnum.name();
        LireFeature feature = null;

        if (image != null) {
            try {
                feature = featureEnum.getFeatureClass().newInstance();
                BufferedImage img = ImageIO.read(new BytesStreamInput(image, false));
                if (Math.max(img.getHeight(), img.getWidth()) > ImageMapper.MAX_IMAGE_DIMENSION) {
                    img = ImageUtils.scaleImage(img, ImageMapper.MAX_IMAGE_DIMENSION);
                }
                feature.extract(img);
            } catch (Exception e) {
                throw new ElasticsearchImageProcessException("Failed to parse image", e);
            }
        } else if (lookupIndex != null && lookupType != null && lookupId != null && lookupPath != null) {
            String lookupFieldName = lookupPath + "." + featureEnum.name();
            GetResponse getResponse = client.get(new GetRequest(lookupIndex, lookupType, lookupId).preference("_local").routing(lookupRouting).fields(lookupFieldName).realtime(false)).actionGet();
            if (getResponse.isExists()) {
                GetField getField = getResponse.getField(lookupFieldName);
                if (getField != null) {
                    BytesReference bytesReference = (BytesReference) getField.getValue();
                    try {
                        feature = featureEnum.getFeatureClass().newInstance();
                        feature.setByteArrayRepresentation(bytesReference.array(), bytesReference.arrayOffset(), bytesReference.length());
                    } catch (Exception e) {
                        throw new ElasticsearchImageProcessException("Failed to parse image", e);
                    }
                }
            }
        }
        if (feature == null) {
            throw new QueryParsingException(parseContext.index(), "No image specified for image query");
        }


        if (hashEnum == null) {  // no hash, need to scan all documents
            return new ImageQuery(luceneFieldName, feature, boost);
        } else {  // query by hash first
            int[] hash = null;
            if (hashEnum.equals(HashEnum.BIT_SAMPLING)) {
                hash = BitSampling.generateHashes(feature.getDoubleHistogram());
            } else if (hashEnum.equals(HashEnum.LSH)) {
                hash = LocalitySensitiveHashing.generateHashes(feature.getDoubleHistogram());
            }
            String hashFieldName = luceneFieldName + "." + ImageMapper.HASH + "." + hashEnum.name();

            if (limit > 0) {  // has max result limit, use ImageHashLimitQuery
                return new ImageHashLimitQuery(hashFieldName, hash, limit, luceneFieldName, feature, boost);
            } else {  // no max result limit, use ImageHashQuery
                BooleanQuery query = new BooleanQuery(true);
                ImageScoreCache imageScoreCache = new ImageScoreCache();

                for (int h : hash) {
                    query.add(new BooleanClause(new ImageHashQuery(new Term(hashFieldName, Integer.toString(h)), luceneFieldName, feature, imageScoreCache, boost), BooleanClause.Occur.SHOULD));
                }
                return query;
            }

        }
    }
}