java调用python yolo onnx模型AI视频识别支持yolov5yolov8yolov7，包含处理 和 后处理

package cn.ck; import ai.onnxruntime.OnnxTensor; import ai.onnxruntime.OrtEnvironment; import ai.onnxruntime.OrtException; import ai.onnxruntime.OrtSession; import cn.ck.config.ODConfig; import cn.ck.domain.Detection; import cn.ck.utils.Letterbox; import org.opencv.core.Mat; import org.opencv.core.Point; import org.opencv.core.Scalar; import org.opencv.highgui.HighGui; import org.opencv.imgcodecs.Imgcodecs; import org.opencv.imgproc.Imgproc; import java.io.File; import java.nio.FloatBuffer; import java.util.*; import java.util.stream.Collectors; /** * yolov8 * * 作者：常康 */ public class ObjectDetection_1_n_8400 { static { // 加载opencv动态库， //System.load(ClassLoader.getSystemResource("lib/opencv_java470-无用.dll").getPath()); nu.pattern.OpenCV.loadLocally(); } public static void main(String[] args) throws OrtException { String model_path = "src\\main\\resources\\model\\yolov8s.onnx"; float confThreshold = 0.35F; float nmsThreshold = 0.55F; String[] labels = { "person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light", "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow", "elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee", "skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", "skateboard", "surfboard", "tennis racket", "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple", "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch", "potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote", "keyboard", "cell phone", "microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors", "teddy bear", "hair drier", "toothbrush"}; // 加载ONNX模型 OrtEnvironment environment = OrtEnvironment.getEnvironment(); OrtSession.SessionOptions sessionOptions = new OrtSession.SessionOptions(); // 使用gpu,需要本机按钻过cuda，并修改pom.xml，不安装也能运行本程序 // sessionOptions.addCUDA(0); OrtSession session = environment.createSession(model_path, sessionOptions); // 输出基本信息 session.getInputInfo().keySet().forEach(x-> { try { System.out.println("input name = " + x); System.out.println(session.getInputInfo().get(x).getInfo().toString()); } catch (OrtException e) { throw new RuntimeException(e); } }); // 要检测的图片所在目录 String imagePath = "images"; // 加载标签及颜色 ODConfig odConfig = new ODConfig(); Map<String, String> map = getImagePathMap(imagePath); for(String fileName : map.keySet()){ String imageFilePath = map.get(fileName); System.out.println(imageFilePath); // 读取 image Mat img = Imgcodecs.imread(imageFilePath); Mat image = img.clone(); Imgproc.cvtColor(image, image, Imgproc.COLOR_BGR2RGB); // 在这里先定义下框的粗细、字的大小、字的类型、字的颜色(按比例设置大小粗细比较好一些) int minDwDh = Math.min(img.width(), img.height()); int thickness = minDwDh/ODConfig.lineThicknessRatio; long start_time = System.currentTimeMillis(); // 更改 image 尺寸 Letterbox letterbox = new Letterbox(); image = letterbox.letterbox(image); double ratio = letterbox.getRatio(); double dw = letterbox.getDw(); double dh = letterbox.getDh(); int rows = letterbox.getHeight(); int cols = letterbox.getWidth(); int channels = image.channels(); // 将Mat对象的像素值赋值给Float[]对象 float[] pixels = new float[channels * rows * cols]; for (int i = 0; i < rows; i++) { for (int j = 0; j < cols; j++) { double[] pixel = image.get(j,i); for (int k = 0; k < channels; k++) { // 这样设置相当于同时做了image.transpose((2, 0, 1))操作 pixels[rows*cols*k+j*cols+i] = (float) pixel[k]/255.0f; } } } // 创建OnnxTensor对象 long[] shape = { 1L, (long)channels, (long)rows, (long)cols }; OnnxTensor tensor = OnnxTensor.createTensor(environment, FloatBuffer.wrap(pixels), shape); HashMap<String, OnnxTensor> stringOnnxTensorHashMap = new HashMap<>(); stringOnnxTensorHashMap.put(session.getInputInfo().keySet().iterator().next(), tensor); // 运行推理 OrtSession.Result output = session.run(stringOnnxTensorHashMap); float[][] outputData = ((float[][][])output.get(0).getValue())[0]; outputData = transposeMatrix(outputData); Map<Integer, List<float[]>> class2Bbox = new HashMap<>(); for (float[] bbox : outputData) { float[] conditionalProbabilities = Arrays.copyOfRange(bbox, 4, outputData.length); int label = argmax(conditionalProbabilities); float conf = conditionalProbabilities[label]; if (conf < confThreshold) continue; bbox[4] = conf; // xywh to (x1, y1, x2, y2) xywh2xyxy(bbox); // skip invalid predictions if (bbox[0] >= bbox[2] || bbox[1] >= bbox[3]) continue; class2Bbox.putIfAbsent(label, new ArrayList<>()); class2Bbox.get(label).add(bbox); } List<Detection> detections = new ArrayList<>(); for (Map.Entry<Integer, List<float[]>> entry : class2Bbox.entrySet()) { int label = entry.getKey(); List<float[]> bboxes = entry.getValue(); bboxes = nonMaxSuppression(bboxes, nmsThreshold); for (float[] bbox : bboxes) { String labelString = labels[label]; detections.add(new Detection(labelString,entry.getKey(), Arrays.copyOfRange(bbox, 0, 4), bbox[4])); } } for (Detection detection : detections) { float[] bbox = detection.getBbox(); System.out.println(detection.toString()); // 画框 Point topLeft = new Point((bbox[0]-dw)/ratio, (bbox[1]-dh)/ratio); Point bottomRight = new Point((bbox[2]-dw)/ratio, (bbox[3]-dh)/ratio); Scalar color = new Scalar(odConfig.getOtherColor(detection.getClsId())); Imgproc.rectangle(img, topLeft, bottomRight, color, thickness); // 框上写文字 Point boxNameLoc = new Point((bbox[0]-dw)/ratio, (bbox[1]-dh)/ratio-3); Imgproc.putText(img, detection.getLabel(), boxNameLoc, Imgproc.FONT_HERSHEY_SIMPLEX, 0.7, color, thickness); } System.out.printf("time：%d ms.", (System.currentTimeMillis() - start_time)); System.out.println(); //服务器部署：由于服务器没有桌面，所以无法弹出画面预览，主要注释一下代码 // 保存图像到同级目录 // Imgcodecs.imwrite(ODConfig.savePicPath, img); // 弹窗展示图像 HighGui.imshow("Display Image", img); // 按任意按键关闭弹窗画面，结束程序