基于YOLO10预测视频流（Java代码调用python脚本实现）

01 Java代码部分

public class YoloInference {

    public static void main(String[] args) {
        String modelPath = "D:/yolov10/yolov10n.pt";
        String streamPath = "rtmp://192.168.200.201:1935/live/hahaha";
        runPythonScript(modelPath, streamPath);
    }

    public static void runPythonScript(String modelPath, String streamPath) {
        try {
            // Command to execute the Python script
            String[] command = new String[] {
                "C:/Users/Administrator/anaconda3/envs/yolov10/python.exe",
                "D:/yolov10/tests/yolo_inference.py", // Replace with the actual path to your Python script
                modelPath,
                streamPath
            };

            // Start the process
            ProcessBuilder processBuilder = new ProcessBuilder(command);
            processBuilder.redirectErrorStream(true);
            Process process = processBuilder.start();

            // Read the output of the Python script
            BufferedReader reader = new BufferedReader(new InputStreamReader(process.getInputStream()));
            String line;
            while ((line = reader.readLine()) != null) {
                System.out.println(line);
            }

            // Wait for the process to finish
            int exitCode = process.waitFor();
            System.out.println("Exited with code: " + exitCode);
        } catch (Exception e) {
            e.printStackTrace();
        }
    }
}

02 python脚本

import sys
import cv2
from ultralytics import YOLOv10
from collections import Counter

def main(model_path, stream_path):
    # Load the YOLOv10 model
    model = YOLOv10(model_path)

    # Open the video file
    cap = cv2.VideoCapture(stream_path)

    # Get class names from the model
    class_names = model.names

    # Initialize a counter for the total counts
    total_class_counts = Counter()

    # Loop through the video frames
    while cap.isOpened():
        # Read a frame from the video
        success, frame = cap.read()

        if success:
            # Run YOLOv10 inference on the frame
            results = model(frame)

            # Visualize the results on the frame
            annotated_frame = results[0].plot()

            # Extract class IDs from the results
            class_ids = results[0].boxes.cls.numpy().astype(int)

            # Convert class IDs to class names
            class_names_detected = [class_names[class_id] for class_id in class_ids]

            # Count occurrences of each class name for the current frame
            class_counts = Counter(class_names_detected)

            # Update total counts
            total_class_counts.update(class_counts)

            # Print the class names and their counts for the current frame
            # for class_name, count in class_counts.items():
            #     print(f"{class_name}: {count}")

            # Display the annotated frame
            cv2.imshow("YOLOv10_Prediction", annotated_frame)

            # Break the loop if 'q' is pressed
            if cv2.waitKey(1) & 0xFF == ord("q"):
                break
        else:
            # Break the loop if the end of the video is reached
            break

    # Release the video capture object and close the display window
    cap.release()
    cv2.destroyAllWindows()

    # Print the total counts
    print("Total counts:===============")
    for class_name, count in total_class_counts.items():
        print(f"{class_name}: {count}")

if __name__ == "__main__":
    model_path = sys.argv[1]
    stream_path = sys.argv[2]
    main(model_path, stream_path)