sishu-yolo-sdk/src/main/java/com/bonus/sdk/YoloSdk.java

140 lines
4.6 KiB
Java
Raw Normal View History

2025-11-14 20:21:04 +08:00
package com.bonus.sdk;
import java.awt.image.BufferedImage;
import java.awt.image.DataBufferByte;
import java.io.*;
import java.nio.file.Files;
import java.nio.file.StandardCopyOption;
import java.util.HashSet;
import java.util.Set;
public class YoloSdk implements AutoCloseable {
2025-11-14 20:21:04 +08:00
private long nativeHandle; //
private static final Set<String> loadedLibraries = new HashSet<>();
2025-11-14 20:21:04 +08:00
/**
* */
static {
try {
loadSdkLibrary();
} catch (IOException e) {
throw new RuntimeException("CRITICAL: Failed to load native YOLO SDK libraries", e);
}
}
private static void loadSdkLibrary() throws IOException {
String osName = System.getProperty("os.name").toLowerCase();
String osArch = System.getProperty("os.arch").toLowerCase();
String libPathInJar;
String sdkLibName;
String[] dependencyLibs = {};
if (osName.contains("win") && osArch.contains("64")) {
2025-11-14 20:21:04 +08:00
libPathInJar = "/lib/win-x64/";
dependencyLibs = new String[]{
"onnxruntime.dll",
"abseil_dll.dll",
"libprotobuf.dll",
"zlib1.dll",
"opencv_core4.dll",
"opencv_imgproc4.dll",
"opencv_dnn4.dll"
};
sdkLibName = "my_yolo_sdk.dll";
} else if ((osName.contains("nix") || osName.contains("nux")) && osArch.contains("64")) {
libPathInJar = "/lib/linux-x86_64/";
2025-11-14 20:21:04 +08:00
dependencyLibs = new String[]{
2025-11-14 20:21:04 +08:00
"libonnxruntime.so.1.23.2",
"libopencv_core.so.4.6.0",
"libopencv_imgproc.so.4.6.0",
"libopencv_dnn.so.4.6.0"
};
sdkLibName = "libmy_yolo_sdk.so";
} else {
throw new UnsupportedOperationException("Unsupported OS/Arch: " + osName + "/" + osArch);
}
2025-11-14 20:21:04 +08:00
// 1.
for (String lib : dependencyLibs) {
loadLibraryFromJar(libPathInJar + lib);
}
2025-11-14 20:21:04 +08:00
// 2.
loadLibraryFromJar(libPathInJar + sdkLibName);
}
2025-11-14 20:21:04 +08:00
/**
* */
private static void loadLibraryFromJar(String path) throws IOException {
String libName = new File(path).getName();
if (loadedLibraries.contains(libName)) {
2025-11-14 20:21:04 +08:00
return; //
}
try (InputStream in = YoloSdk.class.getResourceAsStream(path)) {
if (in == null) {
throw new FileNotFoundException("Library " + path + " not found in JAR.");
}
File tempFile = File.createTempFile(libName, ".tmp");
tempFile.deleteOnExit(); //
Files.copy(in, tempFile.toPath(), StandardCopyOption.REPLACE_EXISTING);
System.load(tempFile.getAbsolutePath());
loadedLibraries.add(libName);
}
}
2025-11-14 20:21:04 +08:00
// ---
private native long nativeInit(String modelPath, int inputWidth, int inputHeight);
private native void nativeRelease(long handle);
private native Detection[] nativePredict(
long handle, byte[] bgrBytes, int imageWidth, int imageHeight,
float confThreshold, float iouThreshold
);
// ---
public YoloSdk(String modelPath, int inputWidth, int inputHeight) {
this.nativeHandle = nativeInit(modelPath, inputWidth, inputHeight);
if (this.nativeHandle == 0) {
throw new RuntimeException("Failed to initialize native YOLO SDK. Check logs.");
}
}
public Detection[] predict(BufferedImage image, float confThreshold, float iouThreshold) {
if (this.nativeHandle == 0) {
throw new IllegalStateException("SDK already closed or failed to initialize.");
}
byte[] bgrBytes = getBgrBytes(image);
return nativePredict(
this.nativeHandle, bgrBytes, image.getWidth(), image.getHeight(),
confThreshold, iouThreshold
);
}
@Override
public void close() {
if (this.nativeHandle != 0) {
nativeRelease(this.nativeHandle);
this.nativeHandle = 0;
}
}
2025-11-14 20:21:04 +08:00
/**
* */
private byte[] getBgrBytes(BufferedImage image) {
if (image.getType() == BufferedImage.TYPE_3BYTE_BGR) {
return ((DataBufferByte) image.getRaster().getDataBuffer()).getData();
}
BufferedImage bgrImage = new BufferedImage(image.getWidth(), image.getHeight(), BufferedImage.TYPE_3BYTE_BGR);
bgrImage.getGraphics().drawImage(image, 0, 0, null);
return ((DataBufferByte) bgrImage.getRaster().getDataBuffer()).getData();
}
}