前言
自从上次用onnx实现了java加载YOLOV10模型后一直在想搞点什么但是又不知道搞什么,后来再网上闲逛的时候发现有人用这个加载了u2net的模型,就很想搞个java版本经过几次尝试后终于实现了一个demo。
先看下效果
和目标检测和抠图之间的区别
- 模型信息包含的不一样检测模型有分类字典的,去背影是没有的。
- 目标检测返回的是可信度和坐标。抠图直接返回的是对应图片大小的多维数组(也就是每个像素点都对应的有值这里取值为1的就是我们要的去背景结果)。
实现demo的思路
- onnx的环境在yolov10的那篇文章做过简单说明(引入jar包,下载新版本的dll就可以了)。
- 模型的一些信息打印。
env = OrtEnvironment.getEnvironment();
session = env.createSession(weight, new OrtSession.SessionOptions());
OnnxModelMetadata metadata = session.getMetadata();
Map<String, NodeInfo> infoMap = session.getInputInfo();
TensorInfo nodeInfo = (TensorInfo)infoMap.get("input").getInfo();
log.info("-------打印模型信息开始--------");
log.info("getProducerName="+metadata.getProducerName());
log.info("getGraphName="+metadata.getGraphName());
log.info("getDescription="+metadata.getDescription());
log.info("getDomain="+metadata.getDomain());
log.info("getVersion="+metadata.getVersion());
log.info("getCustomMetadata="+metadata.getCustomMetadata());
log.info("getInputInfo="+infoMap);
log.info("nodeInfo="+nodeInfo);
count = nodeInfo.getShape()[0];//1 模型每次处理一张图片
channels = nodeInfo.getShape()[1];//3 模型通道数
netHeight = nodeInfo.getShape()[2];//640 模型高
netWidth = nodeInfo.getShape()[3];//640 模型宽
log.info("模型每次处理图片数:"+count);
log.info("模型通道数:"+channels);
log.info("模型输入宽高:"+netWidth+" "+netHeight);
log.info("-------打印模型信息结束--------");
- 关于模型加载数据的说明
Mat转数组把数组打平
数组格式转换whc2cwh
输入格式格式类型要和打平的一致count开头的是个数,看下面返回的四维数组就明白了
float[] whc = new float[channels*netWidth*netHeight];
dst.get(0, 0, whc);
float[] chw = whc2cwh(whc);
OnnxTensor tensor = OnnxTensor.createTensor(env, FloatBuffer.wrap(chw), new long[]{count,channels,netWidth,netHeight});
- 关于opencv的几个重要函数
2.1 图片的缩放问题:宽高模型是有要求的。我在没达到效果的时候犯了模型对图片有输入要求的错误。
2.2 归一化:我是这次才发现opencv也有归一化的函数的(dst.convertTo(dst, CvType.CV_32FC1, 1/255.0);)
2.3 均值和方差的处理。opencv也有直接的处理函数(Core.meanStdDev(dst,mean,std);)
- 模型进行推理,我是从python的demo那来用java改的用java处理数据是有点费劲,先看下数组内容的问题。
4.1 推理结果java中返回的是个四维数组:第一维度图片的个数,第二维度是通道数,第三第四是图片的宽高
4.2 第二维度通道数是1,效果中的第二张图片是单纯的把推理结果变为二值化图片的。推理结果返回的数据也是归一化的0-1之间,乘以255就可以了就是一张黑白图,上面是取反后再处理的。
4.3 由于0是黑色所以返回结果的数组坐标是1的就是推理出来的没有背影的原图,遍历坐标点判断当前推理处理出来的坐标是1的就是我们要的。
OnnxTensor tensor = transferTensor(Mat src);
OrtSession.Result result = session.run(Collections.singletonMap("input", tensor));
OnnxValue res = (OnnxValue)result.get(0);
float[][][][] outputArr = (float[][][][]) res.getValue();
float[][] output2= outputArr[0][0];
for(int i=0;i<output2.length;i++){
for(int j=0;j<output2[i].length;j++){
if(output2[i][j]==1f){
//从输入的图片中取像素点就可以了
}
}
}