com.github.chen0040.objdetect.utils.SavedModelUtils Maven / Gradle / Ivy

Go to download

Show more of this group Show more artifacts with this name
Show all versions of java-ssd-object-detection Show documentation

Java implementation of SSD MultiBox Object Detector

The newest version!

package com.github.chen0040.objdetect.utils;

import org.tensorflow.SavedModelBundle;
import org.tensorflow.framework.MetaGraphDef;
import org.tensorflow.framework.SignatureDef;
import org.tensorflow.framework.TensorInfo;

import java.util.Map;

public class SavedModelUtils {
    public static String getSignature(SavedModelBundle model) throws Exception {
        MetaGraphDef m = MetaGraphDef.parseFrom(model.metaGraphDef());
        SignatureDef sig = m.getSignatureDefOrThrow("serving_default");
        int numInputs = sig.getInputsCount();
        int i = 1;
        StringBuilder sb = new StringBuilder();
        sb.append("MODEL SIGNATURE\n");
        sb.append("Inputs:\n");
        for (Map.Entry entry : sig.getInputsMap().entrySet()) {
            TensorInfo t = entry.getValue();
            sb.append(String.format(
                    "%d of %d: %-20s (Node name in graph: %-20s, type: %s)\n",
                    i++, numInputs, entry.getKey(), t.getName(), t.getDtype()));
        }
        int numOutputs = sig.getOutputsCount();
        i = 1;
        System.out.println("Outputs:");
        for (Map.Entry entry : sig.getOutputsMap().entrySet()) {
            TensorInfo t = entry.getValue();
            sb.append(String.format(
                    "%d of %d: %-20s (Node name in graph: %-20s, type: %s)\n",
                    i++, numOutputs, entry.getKey(), t.getName(), t.getDtype()));
        }
        sb.append("-----------------------------------------------");
        return sb.toString();
    }
}