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

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.Path;
import java.nio.file.StandardCopyOption;
import java.util.HashSet;
import java.util.Set;

public class YoloSdk implements AutoCloseable {

    private long nativeHandle; // ָ<><D6B8> C++ YoloDetector ʵ<><CAB5><EFBFBD><EFBFBD>ָ<EFBFBD><D6B8>
    private static final Set<String> loadedLibraries = new HashSet<>();
    private static Path tempDir; // <20><><EFBFBD>ڴ<EFBFBD><DAB4><EFBFBD><EFBFBD><EFBFBD><EFBFBD>б<EFBFBD><D0B1>ؿ<EFBFBD><D8BF><EFBFBD>Ψһ<CEA8><D2BB>ʱĿ¼

    /**
     * <EFBFBD><EFBFBD>̬<EFBFBD><EFBFBD>ʼ<EFBFBD><EFBFBD><EFBFBD>飬<EFBFBD><EFBFBD><EFBFBD>ڼ<EFBFBD><EFBFBD>ر<EFBFBD><EFBFBD>ؿ⡣
     */
    static {
        try {
            // 1. <20><><EFBFBD><EFBFBD>һ<EFBFBD><D2BB>Ψһ<CEA8><D2BB><EFBFBD><EFBFBD>ʱĿ¼<C4BF><C2BC><EFBFBD><EFBFBD><EFBFBD>ڴ<EFBFBD><DAB4><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>DLL/SO<53>ļ<EFBFBD>
            tempDir = Files.createTempDirectory("yolo_sdk_native_libs_");
            // 2. ȷ<><C8B7><EFBFBD><EFBFBD>JVM<56>˳<EFBFBD>ʱɾ<CAB1><C9BE><EFBFBD><EFBFBD>Ŀ¼
            tempDir.toFile().deleteOnExit();
            
            // 3. <20><><EFBFBD><EFBFBD><EFBFBD>п<EFBFBD><D0BF><EFBFBD>ѹ<EFBFBD><D1B9><EFBFBD><EFBFBD><EFBFBD>ص<EFBFBD><D8B5><EFBFBD>Ŀ¼
            loadSdkLibrary(tempDir);

        } catch (IOException e) {
            throw new RuntimeException("CRITICAL: Failed to create temp directory for native libraries", e);
        }
    }

    /**
     * <EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>ѹ<EFBFBD>ͼ<EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>б<EFBFBD><EFBFBD>ؿ⡣
     * @param tempDir 
     */
    private static void loadSdkLibrary(Path tempDir) 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")) {
            libPathInJar = "/lib/win-x64/";
            // <20><><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>밴<EFBFBD><EBB0B4><EFBFBD><EFBFBD>˳<EFBFBD><CBB3><EFBFBD>г<EFBFBD>
            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/";
            // ȷ<><C8B7><EFBFBD><EFBFBD>Щ<EFBFBD><D0A9><EFBFBD><EFBFBD>Linux<75><78><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD> .so <20><><EFBFBD><EFBFBD><EFBFBD><EFBFBD>
            dependencyLibs = new String[]{
                    "libonnxruntime.so.1.23.2", // ʾ<><CABE><EFBFBD><EFBFBD><EFBFBD><EFBFBD>ʹ<EFBFBD><CAB9><EFBFBD><EFBFBD>ȷ<EFBFBD>е<EFBFBD><D0B5>ļ<EFBFBD><C4BC><EFBFBD>
                    "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);
        }

        // --- <20>ؼ<EFBFBD><D8BC><EFBFBD><EFBFBD>޸<EFBFBD><DEB8>߼<EFBFBD> ---

        // <20><><EFBFBD><EFBFBD> 1: <20><><EFBFBD>ȣ<EFBFBD><C8A3><EFBFBD><EFBFBD><EFBFBD><EFBFBD>п<EFBFBD><D0BF>ļ<EFBFBD><C4BC><EFBFBD>JAR<41><52>ѹ<EFBFBD><D1B9>ͬһ<CDAC><D2BB><EFBFBD><EFBFBD>ʱĿ¼<C4BF><C2BC><EFBFBD><EFBFBD><EFBFBD><EFBFBD>ԭʼ<D4AD>ļ<EFBFBD><C4BC><EFBFBD>
        for (String lib : dependencyLibs) {
            extractLibraryFromJar(libPathInJar + lib, tempDir);
        }
        // <20><>ѹ<EFBFBD><D1B9>SDK<44><4B>
        extractLibraryFromJar(libPathInJar + sdkLibName, tempDir);

        // <20><><EFBFBD><EFBFBD> 2: Ȼ<>󣬰<EFBFBD>˳<EFBFBD><CBB3><EFBFBD>Ӹ<EFBFBD><D3B8><EFBFBD>ʱĿ¼<C4BF><C2BC><EFBFBD>ؿ<EFBFBD>
        // <20><><EFBFBD><EFBFBD><EFBFBD><EFBFBD>Windows/Linux<75><78><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>ͬһĿ¼<C4BF><C2BC><EFBFBD>ҵ<EFBFBD><D2B5><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>
        try {
            for (String lib : dependencyLibs) {
                if (!loadedLibraries.contains(lib)) {
                    System.load(tempDir.resolve(lib).toAbsolutePath().toString());
                    loadedLibraries.add(lib);
                }
            }
            if (!loadedLibraries.contains(sdkLibName)) {
                System.load(tempDir.resolve(sdkLibName).toAbsolutePath().toString());
                loadedLibraries.add(sdkLibName);
            }
        } catch (UnsatisfiedLinkError e) {
            // <20>ṩ<EFBFBD><E1B9A9><EFBFBD><EFBFBD>ϸ<EFBFBD>Ĵ<EFBFBD><C4B4><EFBFBD><EFBFBD><EFBFBD>Ϣ
            System.err.println("--- NATIVE LIBRARY LOAD FAILED ---");
            System.err.println("Failed to load native library. This often means a dependency is missing on the host system.");
            System.err.println("For Windows: Ensure 'Visual C++ Redistributable for Visual Studio (x64)' is installed.");
            System.err.println("Libraries were extracted to: " + tempDir.toAbsolutePath());
            System.err.println("-----------------------------------");
            throw e; // <20><><EFBFBD><EFBFBD><EFBFBD>׳<EFBFBD>ԭʼ<D4AD><CABC><EFBFBD><EFBFBD>
        }
    }

    /**
     * * @param pathInJar 
     * @param tempDir 
     */
    private static void extractLibraryFromJar(String pathInJar, Path tempDir) throws IOException {
        String libName = new File(pathInJar).getName();
        
        try (InputStream in = YoloSdk.class.getResourceAsStream(pathInJar)) {
            if (in == null) {
                throw new FileNotFoundException("Library " + pathInJar + " not found in JAR.");
            }

            // Ŀ<><C4BF><EFBFBD>ļ<EFBFBD>·<EFBFBD><C2B7><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>ԭʼ<D4AD>ļ<EFBFBD><C4BC><EFBFBD>
            Path targetFile = tempDir.resolve(libName);
            
            // <20><><EFBFBD>ļ<EFBFBD><C4BC><EFBFBD>JAR<41><52><EFBFBD>Ƶ<EFBFBD><C6B5><EFBFBD>ʱĿ¼
            Files.copy(in, targetFile, StandardCopyOption.REPLACE_EXISTING);
            
            // Ϊ<><CEAA>ѹ<EFBFBD><D1B9>Ŀ¼/<2F>ļ<EFBFBD><C4BC><EFBFBD><EFBFBD><EFBFBD><EFBFBD>˳<EFBFBD>ʱɾ<CAB1><C9BE><EFBFBD><EFBFBD>˫<EFBFBD>ر<EFBFBD><D8B1>գ<EFBFBD>
            targetFile.toFile().deleteOnExit();
        }
    }

    // --- Native JNI <20><><EFBFBD><EFBFBD> ---
    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
    );

    // --- <20><><EFBFBD><EFBFBD> Java API ---

    /**
     * <EFBFBD><EFBFBD>ʼ<EFBFBD><EFBFBD>YOLO SDK<EFBFBD><EFBFBD>
     * @param modelPath 
     * @param inputWidth 
     * @param inputHeight 
     */
    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 for details.");
        }
    }

    /**
     * @param image
     * @param confThreshold
     * @param iouThreshold
     * @return 
     */
    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
        );
    }

    /**
     * <EFBFBD>ͷű<EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>Դ<EFBFBD><EFBFBD>
     */
    @Override
    public void close() {
        if (this.nativeHandle != 0) {
            nativeRelease(this.nativeHandle);
            this.nativeHandle = 0;
        }
    }

    /**
     * * @param image 
     * @return 
     */
    private byte[] getBgrBytes(BufferedImage image) {
        // <20><><EFBFBD><EFBFBD>ͼ<EFBFBD><CDBC><EFBFBD>Ѿ<EFBFBD><D1BE><EFBFBD> 3-byte BGR<47><52>ֱ<EFBFBD>ӷ<EFBFBD><D3B7><EFBFBD>
        if (image.getType() == BufferedImage.TYPE_3BYTE_BGR) {
            return ((DataBufferByte) image.getRaster().getDataBuffer()).getData();
        }
        
        // <20><><EFBFBD>򣬴<EFBFBD><F2A3ACB4><EFBFBD>һ<EFBFBD><D2BB> BGR <20><><EFBFBD><EFBFBD>
        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();
    }
}