com.vision4j.classification.GrpcClassifier Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of grpc-classifier Show documentation
GRPC classification model for the classification problem
The newest version!
//Autogenerated
package com.vision4j.classification;

import com.google.protobuf.ByteString;
import com.vision4j.classification.grpc.ClassificationGrpc;
import com.vision4j.classification.grpc.Image;
import com.vision4j.classification.grpc.Prediction;
import com.vision4j.utils.Categories;
import com.vision4j.utils.Category;
import com.vision4j.utils.SimpleImageInfo;
import com.vision4j.utils.VisionUtils;
import io.grpc.ManagedChannel;
import io.grpc.ManagedChannelBuilder;

import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.net.URL;

public class GrpcClassifier implements ImageClassifier {

    private final Categories categories;

    private final ClassificationGrpc.ClassificationBlockingStub classificationStub;

    public GrpcClassifier(Categories categories, ManagedChannel channel) {
        this.categories = categories;
        this.classificationStub = ClassificationGrpc.newBlockingStub(channel);
    }

    public GrpcClassifier(String[] categoriesArray, ManagedChannel channel) {
        this.categories = new Categories(categoriesArray);
        this.classificationStub = ClassificationGrpc.newBlockingStub(channel);
    }

    public GrpcClassifier(String[] categoriesArray, String host, int port) {
        this(categoriesArray, ManagedChannelBuilder.forAddress(host, port).usePlaintext().build());
    }

    public GrpcClassifier(String[] categoriesArray) {
        this(categoriesArray, "localhost", 50051);
    }

    @Override()
    public Categories getAcceptableCategories() {
        return this.categories;
    }

    @Override()
    public Category predict(InputStream inputStream) throws IOException {
        return this.predict(VisionUtils.toByteArray(inputStream));
    }

    @Override()
    public Category predict(File file) throws IOException {
        return this.predict(VisionUtils.toByteArray(file));
    }

    @Override()
    public Category predict(byte[] imageBytes) throws IOException {
        SimpleImageInfo simpleImageInfo = new SimpleImageInfo(imageBytes);
        ByteString imageData = ByteString.copyFrom(imageBytes);
        Image.Builder imageBuilder = Image.newBuilder();
        imageBuilder.setWidth(simpleImageInfo.getWidth());
        imageBuilder.setHeight(simpleImageInfo.getHeight());
        imageBuilder.setChannels(3);
        imageBuilder.setImageData(imageData);
        Image image = imageBuilder.build();
        Prediction prediction = classificationStub.predict(image);
        return this.convert(prediction);
    }

    @Override()
    public Category predict(URL imageURL) throws IOException {
        return this.predict(VisionUtils.toByteArray(imageURL));
    }

    private Category convert(Prediction prediction) {
        return getAcceptableCategories().get(prediction.getIndex());
    }
}