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)