package com.zhgd.xmgl.modules.basicdata.api.image_ai.common;

import ai.djl.inference.Predictor;
import ai.djl.modality.cv.Image;
import ai.djl.modality.cv.ImageFactory;
import ai.djl.modality.cv.output.BoundingBox;
import ai.djl.modality.cv.output.DetectedObjects;
import ai.djl.repository.zoo.Criteria;
import ai.djl.repository.zoo.ModelZoo;
import ai.djl.repository.zoo.ZooModel;
import com.zhgd.xmgl.modules.basicdata.api.image_ai.utils.SmallSafetyHelmetDetect;
import lombok.experimental.UtilityClass;
import lombok.extern.slf4j.Slf4j;

import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.List;

/**
 * 安全帽检测例子
 *
 * 目录：http://aias.top/
 *
 * @author Calvin
 */
@Slf4j
@UtilityClass
public final class SmallSafetyHelmetDetectExample {

  public String imageAi(String imagePath) throws Exception {
    Path imageFile = Paths.get(imagePath);
    Image image = ImageFactory.getInstance().fromFile(imageFile);

    Criteria<Image, DetectedObjects> criteria = new SmallSafetyHelmetDetect().criteria(image);

    try (ZooModel model = ModelZoo.loadModel(criteria);
        Predictor<Image, DetectedObjects> predictor = model.newPredictor()) {
      DetectedObjects detections = predictor.predict(image);
      List<DetectedObjects.DetectedObject> items = detections.items();

      List<String> names = new ArrayList<>();
      List<Double> prob = new ArrayList<>();
      List<BoundingBox> boxes = new ArrayList<>();
      for (int i = 0; i < items.size(); i++) {
        DetectedObjects.DetectedObject item = items.get(i);
        if (item.getProbability() < 0.3f) {
          continue;
        }
        names.add(item.getClassName() + " " + item.getProbability());
        prob.add(item.getProbability());
        boxes.add(item.getBoundingBox());
      }

      detections = new DetectedObjects(names, prob, boxes);
//      ImageUtils.saveBoundingBoxImage(image, detections, "safety_helmet_result_s.png", "build/output");
      log.info("{}", detections);
      return detections.toJson();
    }
  }
}
