Naaturel/board-mate
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Updates

Browse Source
This commit is contained in:
Laurent
2025-12-21 12:06:28 +01:00
parent c0374b3890
commit 94b83fb018
32 changed files with 281 additions and 212 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
rpi/board-detector/realtime_detect.py
View File

@@ -1,11 +1,14 @@
from ultralytics import YOLO
from paths import * # make sure model_path is defined here
import cv2
if __name__ == "__main__":
corner_model_path = "../assets/models/corner.pt"
pieces_model_path = "../assets/models/unified-nano-refined.pt"
print("Initializing model...")
model = YOLO(model_path)
corner_model = YOLO(corner_model_path)
pieces_model = YOLO(pieces_model_path)
print("Initializing camera...")
cap = cv2.VideoCapture(0)
@@ -29,11 +32,13 @@ if __name__ == "__main__":
# Optional: resize frame to improve YOLO performance
# frame = cv2.resize(frame, (416, 416))
results = model.predict(source=frame, conf=0.5)
corner_result = corner_model.predict(source=frame, conf=0.6)
pieces_result = pieces_model.predict(source=frame, conf=0.6)
annotated_frame = results[0].plot() # annotated frame as NumPy array
corner_annotated_frame = corner_result[0].plot()
pieces_annotated_frame = pieces_result[0].plot(img=corner_annotated_frame)
cv2.imshow("Predictions", annotated_frame)
cv2.imshow("Predictions", pieces_annotated_frame)
cv2.resizeWindow("Predictions", 640, 640)
if cv2.waitKey(1) & 0xFF == ord('q'):
break