Skip to content
Snippets Groups Projects
Commit 34406627 authored by Rainer Lourens's avatar Rainer Lourens
Browse files

added cube drawing after scanning

parent 802e821c
Branches
No related tags found
No related merge requests found
Hide whitespace changes
Inline Side-by-side
main.py
+ 194
− 42
View file @ 34406627
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
import numpy as np
from ultralytics import YOLO
import cv2
import os
from tensorflow.keras.models import load_model
import numpy as np
import os
from Color_detection import raspberryPi
import os
# Initialize cube state
cube_state = {
'Front': [['white'] * 3 for _ in range(3)],
'Top': [['yellow'] * 3 for _ in range(3)],
'Bottom': [['blue'] * 3 for _ in range(3)],
'Left': [['green'] * 3 for _ in range(3)],
'Right': [['red'] * 3 for _ in range(3)],
'Back': [['orange'] * 3 for _ in range(3)],
}
# Define color mapping for matplotlib
color_map = {
'white': '#FFFFFF',
'yellow': '#FFFF00',
'blue': '#0000FF',
'green': '#00FF00',
'red': '#FF0000',
'orange': '#FFA500'
}
def draw_face(ax, origin, normal, colors):
"""
Draw a single face of the Rubik's cube with 3x3 sub-squares.
:param ax: Matplotlib 3D axis
:param origin: Bottom-left corner of the face
:param normal: Normal vector of the face (to determine orientation)
:param colors: 3x3 grid of colors for the face
"""
rows, cols = len(colors), len(colors[0])
step = 1 / rows
for i in range(rows):
for j in range(cols):
x_offset = j * step
y_offset = i * step
color = colors[i][j]
if normal == [0, 0, 1]: # XY plane (Top)
vertices = [
[origin[0] + x_offset, origin[1] + y_offset, origin[2]],
[origin[0] + x_offset + step, origin[1] + y_offset, origin[2]],
[origin[0] + x_offset + step, origin[1] + y_offset + step, origin[2]],
[origin[0] + x_offset, origin[1] + y_offset + step, origin[2]]
]
elif normal == [0, 0, -1]: # XY plane (Bottom)
vertices = [
[origin[0] + x_offset, origin[1] + y_offset, origin[2]],
[origin[0] + x_offset + step, origin[1] + y_offset, origin[2]],
[origin[0] + x_offset + step, origin[1] + y_offset + step, origin[2]],
[origin[0] + x_offset, origin[1] + y_offset + step, origin[2]]
]
elif normal == [0, 1, 0]: # XZ plane (Front)
vertices = [
[origin[0] + x_offset, origin[1], origin[2] + y_offset],
[origin[0] + x_offset + step, origin[1], origin[2] + y_offset],
[origin[0] + x_offset + step, origin[1], origin[2] + y_offset + step],
[origin[0] + x_offset, origin[1], origin[2] + y_offset + step]
]
elif normal == [0, -1, 0]: # XZ plane (Back)
vertices = [
[origin[0] + x_offset, origin[1], origin[2] + y_offset],
[origin[0] + x_offset + step, origin[1], origin[2] + y_offset],
[origin[0] + x_offset + step, origin[1], origin[2] + y_offset + step],
[origin[0] + x_offset, origin[1], origin[2] + y_offset + step]
]
elif normal == [1, 0, 0]: # YZ plane (Right)
vertices = [
[origin[0], origin[1] + x_offset, origin[2] + y_offset],
[origin[0], origin[1] + x_offset + step, origin[2] + y_offset],
[origin[0], origin[1] + x_offset + step, origin[2] + y_offset + step],
[origin[0], origin[1] + x_offset, origin[2] + y_offset + step]
]
elif normal == [-1, 0, 0]: # YZ plane (Left)
vertices = [
[origin[0], origin[1] + x_offset, origin[2] + y_offset],
[origin[0], origin[1] + x_offset + step, origin[2] + y_offset],
[origin[0], origin[1] + x_offset + step, origin[2] + y_offset + step],
[origin[0], origin[1] + x_offset, origin[2] + y_offset + step]
]
else:
continue
face = Poly3DCollection([vertices], color=color, edgecolor='black')
ax.add_collection3d(face)
def draw_cube(state):
"""Draw a 3D Rubik's cube based on the input state."""
fig = plt.figure(figsize=(8, 8))
ax = fig.add_subplot(111, projection='3d')
# Define face origins and normals
face_definitions = {
"Top": ([-0.5, -0.5, 0.5], [0, 0, 1]),
"Bottom": ([-0.5, -0.5, -0.5], [0, 0, -1]),
"Front": ([-0.5, 0.5, -0.5], [0, 1, 0]),
"Back": ([-0.5, -0.5, -0.5], [0, -1, 0]),
"Left": ([-0.5, -0.5, -0.5], [-1, 0, 0]),
"Right": ([0.5, -0.5, -0.5], [1, 0, 0]),
}
# Draw each face
for face, (origin, normal) in face_definitions.items():
draw_face(ax, origin, normal, state[face])
# Set the aspect of the plot to be equal
ax.set_box_aspect([1, 1, 1]) # Aspect ratio is 1:1:1
# Set limits and labels
ax.set_xlim([-0.75, 0.75])
ax.set_ylim([-0.75, 0.75])
ax.set_zlim([-0.75, 0.75])
ax.set_xlabel("X")
ax.set_ylabel("Y")
ax.set_zlabel("Z")
# Display the cube
plt.show()
# Load the YOLO model
model = YOLO(r"Cube_detection/runs/detect/best/weights/best.pt")
......@@ -16,48 +138,78 @@ model_path = r"Color_detection/rubiks_cube_model.h5"
if not os.path.exists(model_path):
print("Error: Model file not found.")
print("Loading model...")
print("Loading color detection model...")
model_color = load_model(model_path)
sequence = ['Front', 'Top', 'Bottom', 'Left', 'Right', 'Back']
current_step = 0
# Initialize matplotlib figure
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
draw_cube(cube_state)
def save_colors(face, colors):
"""Save the detected colors for the given face."""
cube_state[face] = np.array(colors).reshape(3, 3).tolist()
draw_cube(ax, cube_state) # Update visualization
print(f"Colors for {face} saved:", colors)
cube_state = {'Front': [['red', 'blue', 'blue'], ['white', 'white', 'white'], ['white', 'red', 'blue']], 'Top': [['red', 'blue', 'blue'], ['white', 'white', 'white'], ['white', 'red', 'blue']], 'Bottom': [['white', 'blue', 'white'], ['white', 'white', 'blue'], ['white', 'red', 'blue']], 'Left': [['red', 'blue', 'blue'], ['white', 'blue', 'white'], ['white', 'red', 'blue']], 'Right': [['red', 'blue', 'blue'], ['white', 'blue', 'white'], ['white', 'red', 'blue']], 'Back': [['red', 'blue', 'blue'], ['white', 'blue', 'white'], ['white', 'red', 'blue']]}
draw_cube(cube_state)
if cube_state ==None:
while True:
ret, frame = cap.read()
if not ret:
print("Failed to capture frame.")
break
# Perform detection
results = model(frame, save=False, conf=0.25)
annotated_frame = results[0].plot()
box = results[0].boxes.xyxy
# Ensure bounding box exists
if box.size(0) > 0:
x1, y1, x2, y2 = map(int, box[0])
cropped_frame = frame[y1:y2, x1:x2]
else:
cropped_frame = None
# Display annotated frame
cv2.imshow(f"Scanning: {sequence[current_step]} Face", annotated_frame)
key = cv2.waitKey(1) & 0xFF
if key == ord('c'):
if cropped_frame is not None:
# Display captured frame
cv2.imshow("Captured Image - Press 'y' to save or 'r' to retake", cropped_frame)
# Wait for user input
key_confirm = cv2.waitKey(0) & 0xFF
cv2.destroyWindow("Captured Image - Press 'y' to save or 'r' to retake") # Close this window after input
if key_confirm == ord('y'):
print(f"Processing {sequence[current_step]} face...")
predicted_colors = raspberryPi.predict_colors(model_color, cropped_frame)
colors = raspberryPi.indices_to_colors(predicted_colors)
save_colors(sequence[current_step], colors)
current_step += 1
if current_step >= len(sequence):
print("All sides scanned successfully!")
print("Final Cube Colors:", cube_state)
break
elif key_confirm == ord('r'):
print("Retaking the picture...")
else:
print("No valid region detected. Retake the capture.")
while True:
ret, frame = cap.read()
if not ret:
print("Failed to capture frame.")
break
# Perform detection
results = model(frame, save=False, conf=0.25) # Adjust confidence as needed
annotated_frame = results[0].plot() # Annotate frame for display
# Extract bounding boxes from results
box = results[0].boxes.xyxy
print(box)
if box.size(0) > 0:
# Convert box to integer for cropping
x1, y1, x2, y2 = map(int, box[0])
# Crop the detected object
cropped_frame = frame[y1:y2, x1:x2]
# Save the cropped frame
# Display frame
cv2.imshow("Rubix detection", annotated_frame)
key = cv2.waitKey(1) & 0xFF
if key == ord('c'):
captured_frame = cropped_frame
cv2.imshow("Captured Image - Press 'y' to proceed or 'r' to retake", captured_frame)
key_confirm = cv2.waitKey(0) & 0xFF
if key_confirm == ord('y'):
print("Processing captured image...")
predicted_colors = raspberryPi.predict_colors(model_color, captured_frame)
colors = raspberryPi.indices_to_colors(predicted_colors)
print(colors)
elif key_confirm == ord('r'):
print("Retaking the picture...")
break # Retake the picture
# Break the loop if 'q' is pressed
if cv2.waitKey(1) & 0xFF == ord('q'):
break
if key == ord('q'):
break
# Release resources
cap.release()
cv2.destroyAllWindows()
\ No newline at end of file
cv2.destroyAllWindows()
plt.show()
\ No newline at end of file
test.py
+ 113
− 7
View file @ 34406627
import cv2
while True:
key = cv2.waitKey(1) & 0xFF # Wait for a key press
print(f"Key pressed: {key}") # Debug to see the key code
if key == ord('c'):
print("Key 'c' pressed!")
\ No newline at end of file
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
import numpy as np
def draw_face(ax, origin, normal, colors):
"""
Draw a single face of the Rubik's cube with 3x3 sub-squares.
:param ax: Matplotlib 3D axis
:param origin: Bottom-left corner of the face
:param normal: Normal vector of the face (to determine orientation)
:param colors: 3x3 grid of colors for the face
"""
rows, cols = len(colors), len(colors[0])
step = 1 / rows
for i in range(rows):
for j in range(cols):
x_offset = j * step
y_offset = i * step
color = colors[i][j]
if normal == [0, 0, 1]: # XY plane (Top)
vertices = [
[origin[0] + x_offset, origin[1] + y_offset, origin[2]],
[origin[0] + x_offset + step, origin[1] + y_offset, origin[2]],
[origin[0] + x_offset + step, origin[1] + y_offset + step, origin[2]],
[origin[0] + x_offset, origin[1] + y_offset + step, origin[2]]
]
elif normal == [0, 0, -1]: # XY plane (Bottom)
vertices = [
[origin[0] + x_offset, origin[1] + y_offset, origin[2]],
[origin[0] + x_offset + step, origin[1] + y_offset, origin[2]],
[origin[0] + x_offset + step, origin[1] + y_offset + step, origin[2]],
[origin[0] + x_offset, origin[1] + y_offset + step, origin[2]]
]
elif normal == [0, 1, 0]: # XZ plane (Front)
vertices = [
[origin[0] + x_offset, origin[1], origin[2] + y_offset],
[origin[0] + x_offset + step, origin[1], origin[2] + y_offset],
[origin[0] + x_offset + step, origin[1], origin[2] + y_offset + step],
[origin[0] + x_offset, origin[1], origin[2] + y_offset + step]
]
elif normal == [0, -1, 0]: # XZ plane (Back)
vertices = [
[origin[0] + x_offset, origin[1], origin[2] + y_offset],
[origin[0] + x_offset + step, origin[1], origin[2] + y_offset],
[origin[0] + x_offset + step, origin[1], origin[2] + y_offset + step],
[origin[0] + x_offset, origin[1], origin[2] + y_offset + step]
]
elif normal == [1, 0, 0]: # YZ plane (Right)
vertices = [
[origin[0], origin[1] + x_offset, origin[2] + y_offset],
[origin[0], origin[1] + x_offset + step, origin[2] + y_offset],
[origin[0], origin[1] + x_offset + step, origin[2] + y_offset + step],
[origin[0], origin[1] + x_offset, origin[2] + y_offset + step]
]
elif normal == [-1, 0, 0]: # YZ plane (Left)
vertices = [
[origin[0], origin[1] + x_offset, origin[2] + y_offset],
[origin[0], origin[1] + x_offset + step, origin[2] + y_offset],
[origin[0], origin[1] + x_offset + step, origin[2] + y_offset + step],
[origin[0], origin[1] + x_offset, origin[2] + y_offset + step]
]
else:
continue
face = Poly3DCollection([vertices], color=color, edgecolor='black')
ax.add_collection3d(face)
def draw_cube(state):
"""Draw a 3D Rubik's cube based on the input state."""
fig = plt.figure(figsize=(8, 8))
ax = fig.add_subplot(111, projection='3d')
# Define face origins and normals
face_definitions = {
"Top": ([-0.5, -0.5, 0.5], [0, 0, 1]),
"Bottom": ([-0.5, -0.5, -0.5], [0, 0, -1]),
"Front": ([-0.5, 0.5, -0.5], [0, 1, 0]),
"Back": ([-0.5, -0.5, -0.5], [0, -1, 0]),
"Left": ([-0.5, -0.5, -0.5], [-1, 0, 0]),
"Right": ([0.5, -0.5, -0.5], [1, 0, 0]),
}
# Draw each face
for face, (origin, normal) in face_definitions.items():
draw_face(ax, origin, normal, state[face])
# Set the aspect of the plot to be equal
ax.set_box_aspect([1, 1, 1]) # Aspect ratio is 1:1:1
# Set limits and labels
ax.set_xlim([-0.75, 0.75])
ax.set_ylim([-0.75, 0.75])
ax.set_zlim([-0.75, 0.75])
ax.set_xlabel("X")
ax.set_ylabel("Y")
ax.set_zlabel("Z")
# Display the cube
plt.show()
if __name__ == "__main__":
cube_state = {
'Front': [['red', 'blue', 'blue'], ['white', 'white', 'white'], ['white', 'red', 'blue']],
'Top': [['red', 'blue', 'blue'], ['white', 'white', 'white'], ['white', 'red', 'blue']],
'Bottom': [['white', 'blue', 'white'], ['white', 'white', 'blue'], ['white', 'red', 'blue']],
'Left': [['red', 'blue', 'blue'], ['white', 'blue', 'white'], ['white', 'red', 'blue']],
'Right': [['red', 'blue', 'blue'], ['white', 'blue', 'white'], ['white', 'red', 'blue']],
'Back': [['red', 'blue', 'blue'], ['white', 'blue', 'white'], ['white', 'red', 'blue']]
}
draw_cube(cube_state)
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment