Add conversion to FEN

This commit is contained in:
2025-12-22 16:30:07 +01:00
parent 86dea774e4
commit 0aaea36586
12 changed files with 246 additions and 105 deletions

View File

View File

@@ -0,0 +1,31 @@
#!/usr/bin/env python3
import cv2
import numpy as np
from ultralytics import YOLO
from paths import *
def main():
model = YOLO(model_path)
# Load image
image = cv2.imread(img_path)
if image is None:
print(f"Failed to read {img_path}")
return
height, width = image.shape[:2]
warped = image # For now assume top-down view
# Run YOLO detection
results = model(warped)
res = results[0]
debug_img = res.plot() # draws boxes around detected objects
cv2.imshow("Detections", debug_img)
cv2.waitKey(0)
cv2.destroyAllWindows()
if __name__ == "__main__":
main()

View File

@@ -0,0 +1,46 @@
from ultralytics import YOLO
import cv2
if __name__ == "__main__":
corner_model_path = "../../assets/models/edges.pt"
pieces_model_path = "../../assets/models/unified-nano-refined.pt"
print("Initializing model...")
corner_model = YOLO(corner_model_path)
pieces_model = YOLO(pieces_model_path)
print("Initializing camera...")
cap = cv2.VideoCapture(0)
cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
cap.set(cv2.CAP_PROP_FPS, 30)
print("Initialized")
if not cap.isOpened():
print("Error: Could not open camera")
exit()
cv2.namedWindow("Predictions", cv2.WINDOW_NORMAL)
while True:
ret, frame = cap.read()
if not ret:
print("Error: Failed to grab frame")
break
# Optional: resize frame to improve YOLO performance
# frame = cv2.resize(frame, (416, 416))
corner_result = corner_model.predict(source=frame, conf=0.6)
pieces_result = pieces_model.predict(source=frame, conf=0.6)
corner_annotated_frame = corner_result[0].plot()
pieces_annotated_frame = pieces_result[0].plot(img=corner_annotated_frame)
cv2.imshow("Predictions", pieces_annotated_frame)
cv2.resizeWindow("Predictions", 640, 640)
if cv2.waitKey(1) & 0xFF == ord('q'):
break
cap.release()
cv2.destroyAllWindows()