PredictorFactory.java

package com.neuronrobotics.bowlerkernel.djl;

import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;

import ai.djl.MalformedModelException;
import ai.djl.ModelException;
import ai.djl.engine.Engine;
import ai.djl.inference.Predictor;
import ai.djl.modality.cv.Image;
import ai.djl.modality.cv.output.DetectedObjects;
import ai.djl.modality.cv.translator.YoloV5TranslatorFactory;
import ai.djl.pytorch.jni.JniUtils;
import ai.djl.repository.zoo.Criteria;
import ai.djl.repository.zoo.ModelNotFoundException;
import ai.djl.repository.zoo.ZooModel;
import ai.djl.training.util.ProgressBar;

public class PredictorFactory {
	static {
		Engine.getEngine("PyTorch"); // Make sure PyTorch engine is loaded
		Engine.getEngine("OnnxRuntime"); // Make sure PyTorch engine is loaded
	}
	private static HashMap<ImagePredictorType, Predictor<Image, DetectedObjects>> preloaded = new HashMap<>();
	private static Predictor<Image, float[]> features = null;

	public static Predictor<Image, DetectedObjects> imageContentsFactory(ImagePredictorType type)
			throws ModelNotFoundException, MalformedModelException, IOException {
		JniUtils.setGraphExecutorOptimize(false);

		if (preloaded.get(type) == null) {

			switch (type) {
			case retinaface:
				double confThreshretinaface = 0.85f;
				double nmsThreshretinaface = 0.45f;
				double[] varianceretinaface = { 0.1f, 0.2f };
				int topKretinaface = 5000;
				int[][] scalesretinaface = { { 16, 32 }, { 64, 128 }, { 256, 512 } };
				int[] stepsretinaface = { 8, 16, 32 };
				FaceDetectionTranslator translatorretinaface = new FaceDetectionTranslator(confThreshretinaface,
						nmsThreshretinaface, varianceretinaface, topKretinaface, scalesretinaface, stepsretinaface);

				Criteria<Image, DetectedObjects> criteriaretinaface = Criteria.builder()
						.setTypes(Image.class, DetectedObjects.class)
						.optModelUrls("https://resources.djl.ai/test-models/pytorch/retinaface.zip")
						// Load model from local file, e.g:
						.optModelName("retinaface") // specify model file prefix
						.optTranslator(translatorretinaface).optProgress(new ProgressBar()).optEngine("PyTorch") // Use
																													// PyTorch
																													// engine
						.build();

				preloaded.put(type, criteriaretinaface.loadModel().newPredictor());
				break;
			case ultranet:
				double confThresh = 0.85f;
				double nmsThresh = 0.45f;
				double[] variance = { 0.1f, 0.2f };
				int topK = 5000;
				int[][] scales = { { 10, 16, 24 }, { 32, 48 }, { 64, 96 }, { 128, 192, 256 } };
				int[] steps = { 8, 16, 32, 64 };
				FaceDetectionTranslator translator = new FaceDetectionTranslator(confThresh, nmsThresh, variance, topK,
						scales, steps);

				Criteria<Image, DetectedObjects> criteria = Criteria.builder()
						.setTypes(Image.class, DetectedObjects.class)
						.optModelUrls("https://resources.djl.ai/test-models/pytorch/ultranet.zip")
						.optTranslator(translator).optProgress(new ProgressBar()).optEngine("PyTorch") // Use PyTorch
																										// engine
						.build();

				preloaded.put(type, criteria.loadModel().newPredictor());
				break;
			case yolov5:
				String MODEL_URL = "https://mlrepo.djl.ai/model/cv/object_detection/ai/djl/onnxruntime/yolo5s/0.0.1/yolov5s.zip";

				Criteria<Image, DetectedObjects> criteria2 = Criteria.builder()
						.setTypes(Image.class, DetectedObjects.class).optModelUrls(MODEL_URL).optEngine("OnnxRuntime")
						.optTranslatorFactory(new YoloV5TranslatorFactory()).build();

				try { 
					YoloManager ym = new YoloManager(criteria2);
					preloaded.put(type,ym.predictor());
					
				} catch (IOException e) {
					// TODO Auto-generated catch block
					e.printStackTrace();
				} catch (ModelException e) {
					// TODO Auto-generated catch block
					e.printStackTrace();
				}
				break;
			default:
				throw new RuntimeException("No Model available of type " + type);

			}
		}
		return preloaded.get(type);
	}

	public static Predictor<Image, float[]> faceFeatureFactory()
			throws ModelNotFoundException, MalformedModelException, IOException {
		JniUtils.setGraphExecutorOptimize(false);
		if (features == null) {
			Criteria<Image, float[]> criteria = Criteria.builder().setTypes(Image.class, float[].class)
					.optModelUrls("https://resources.djl.ai/test-models/pytorch/face_feature.zip")
					.optModelName("face_feature") // specify model file prefix
					.optTranslator(new FaceFeatureTranslator()).optProgress(new ProgressBar()).optEngine("PyTorch") // Use
																													// PyTorch
																													// engine
					.build();
			ZooModel<Image, float[]> model = criteria.loadModel();
			features = model.newPredictor();
		}
		return features;
	}

	public static float calculSimilarFaceFeature(float[] feature1, ArrayList<float[]> people) {
		float ret = 0.0f;
		float mod1 = 0.0f;
		float mod2 = 0.0f;
		int length = feature1.length;
		for (int j = 0; j < people.size(); j++) {
			float[] feature2 = people.get(j);
			for (int i = 0; i < length; ++i) {
				ret += feature1[i] * feature2[i];
				mod1 += feature1[i] * feature1[i];
				mod2 += feature2[i] * feature2[i];
			}
		}
		return (float) ((ret / Math.sqrt(mod1) / Math.sqrt(mod2) + 1) / 2.0f);
	}

	public static float calculSimilarFaceFeature(float[] feature1, float[] feature2) {
		float ret = 0.0f;
		float mod1 = 0.0f;
		float mod2 = 0.0f;
		int length = feature1.length;
		for (int i = 0; i < length; ++i) {
			ret += feature1[i] * feature2[i];
			mod1 += feature1[i] * feature1[i];
			mod2 += feature2[i] * feature2[i];
		}
		return (float) ((ret / Math.sqrt(mod1) / Math.sqrt(mod2) + 1) / 2.0f);
	}

}