cuizhibin
/
znpt-backend
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

1.施工信息问题修复

This commit is contained in:
cuizhibin 2025-07-03 16:34:43 +08:00
parent 3dd13474b6
commit 8788c6f7e4
1 changed files with 0 additions and 45 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

45
core/src/main/java/com/dite/znpt/config/YoloModelConfig.java
View File

@ -1,45 +0,0 @@
package com.dite.znpt.config;
import ai.onnxruntime.OrtEnvironment;
import ai.onnxruntime.OrtException;
import ai.onnxruntime.OrtSession;
import lombok.Getter;
import org.springframework.context.annotation.Configuration;
import javax.annotation.PostConstruct;
import java.util.ArrayList;
import java.util.List;
import java.util.Random;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
@Getter
@Configuration
public class YoloModelConfig {
private OrtEnvironment env;
private OrtSession session;
private String[] labels;
private List<double[]> colors;
@PostConstruct
public void init() throws OrtException {
env = OrtEnvironment.getEnvironment();
OrtSession.SessionOptions options = new OrtSession.SessionOptions();
session = env.createSession("d:\\tmp\\best.onnx", options);
// 解析 labels
String names = session.getMetadata().getCustomMetadata().get("names");
Pattern pattern = Pattern.compile("'([^']*)'");
Matcher matcher = pattern.matcher(names);
List<String> labelList = new ArrayList<>();
colors = new ArrayList<>();
Random random = new Random();
while (matcher.find()) {
labelList.add(matcher.group(1));
colors.add(new double[]{random.nextDouble()*256, random.nextDouble()*256, random.nextDouble()*256});
}
labels = labelList.toArray(new String[0]);
}
}