Minor fixes
This commit is contained in:
@@ -1,46 +0,0 @@
|
||||
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()
|
||||
Reference in New Issue
Block a user