56 lines
1.7 KiB
Python
56 lines
1.7 KiB
Python
#!/usr/bin/env python3
|
|
import os
|
|
import random
|
|
|
|
import numpy as np
|
|
from ultralytics import YOLO
|
|
import cv2
|
|
from ultralytics.engine.results import Results
|
|
|
|
|
|
class Detector :
|
|
|
|
def __init__(self, model_path):
|
|
self.model = YOLO(model_path)
|
|
self.used_height = 640
|
|
self.used_width = 640
|
|
|
|
def make_prediction(self, image : str | np.ndarray) -> list[Results]:
|
|
return self.model.predict(source=image, conf=0.6)
|
|
|
|
if __name__ == "__main__":
|
|
corner_model_path = "../../assets/models/edges.pt"
|
|
pieces_model_path = "../../assets/models/unified-nano-refined.pt"
|
|
|
|
corner_model = YOLO(corner_model_path)
|
|
pieces_model = YOLO(pieces_model_path)
|
|
|
|
img_folder = "../training/datasets/pieces/unified/test/images/"
|
|
save_folder = "./results"
|
|
os.makedirs(save_folder, exist_ok=True)
|
|
|
|
test_images = os.listdir(img_folder)
|
|
|
|
for i in range(0, 10):
|
|
rnd = random.randint(0, len(test_images) - 1)
|
|
img_path = os.path.join(img_folder, test_images[rnd])
|
|
save_path = os.path.join(save_folder, test_images[rnd])
|
|
|
|
image = cv2.imread(img_path)
|
|
height, width = image.shape[:2]
|
|
|
|
#fen = prediction_to_fen(results, height, width)
|
|
#print("Predicted FEN:", fen)
|
|
|
|
corner_result = corner_model.predict(source=image, conf=0.6)
|
|
pieces_result = pieces_model.predict(source=image, conf=0.6)
|
|
|
|
corner_annotated_image = corner_result[0].plot()
|
|
pieces_annotated_image = pieces_result[0].plot(img=corner_annotated_image)
|
|
|
|
cv2.imwrite(save_path, pieces_annotated_image)
|
|
#cv2.namedWindow("YOLO Predictions", cv2.WINDOW_NORMAL)
|
|
#cv2.imshow("YOLO Predictions", annotated_image)
|
|
|
|
cv2.waitKey(0)
|
|
cv2.destroyAllWindows() |