cube drawing finalized

test.py → cubeplot.py

@@ -15,11 +15,12 @@ def draw_face(ax, origin, normal, colors):
 
     for i in range(rows):
         for j in range(cols):
-            x_offset = j * step
-            y_offset = i * step
             color = colors[i][j]
-
+            x_offset = (cols - 1 - j) * step
+            y_offset = (rows - 1 - i) * step
             if normal == [0, 0, 1]:  # XY plane (Top)
+                x_offset = (cols - 1 - j) * step  # Standard column order
+                y_offset = i * step 
                 vertices = [
                     [origin[0] + x_offset, origin[1] + y_offset, origin[2]],
                     [origin[0] + x_offset + step, origin[1] + y_offset, origin[2]],
@@ -34,6 +35,8 @@ def draw_face(ax, origin, normal, colors):
                     [origin[0] + x_offset, origin[1] + y_offset + step, origin[2]]
                 ]
             elif normal == [0, 1, 0]:  # XZ plane (Front)
+                x_offset = (cols - 1 - j) * step  # Reverse column order
+                y_offset = (rows - 1 - i) * step  # Reverse row order
                 vertices = [
                     [origin[0] + x_offset, origin[1], origin[2] + y_offset],
                     [origin[0] + x_offset + step, origin[1], origin[2] + y_offset],
@@ -41,6 +44,7 @@ def draw_face(ax, origin, normal, colors):
                     [origin[0] + x_offset, origin[1], origin[2] + y_offset + step]
                 ]
             elif normal == [0, -1, 0]:  # XZ plane (Back)
+                x_offset = j * step 
                 vertices = [
                     [origin[0] + x_offset, origin[1], origin[2] + y_offset],
                     [origin[0] + x_offset + step, origin[1], origin[2] + y_offset],
@@ -48,6 +52,7 @@ def draw_face(ax, origin, normal, colors):
                     [origin[0] + x_offset, origin[1], origin[2] + y_offset + step]
                 ]
             elif normal == [1, 0, 0]:  # YZ plane (Right)
+                x_offset = j * step 
                 vertices = [
                     [origin[0], origin[1] + x_offset, origin[2] + y_offset],
                     [origin[0], origin[1] + x_offset + step, origin[2] + y_offset],
@@ -86,6 +91,20 @@ def draw_cube(state):
     for face, (origin, normal) in face_definitions.items():
         draw_face(ax, origin, normal, state[face])
 
+        # Add text at the center of each face
+        if normal == [0, 0, 1]:  # Top
+            ax.text(0, 0, 0.7, face, color='black', fontsize=12, ha='center', va='center')
+        elif normal == [0, 0, -1]:  # Bottom
+            ax.text(0, 0, -0.7, face, color='black', fontsize=12, ha='center', va='center')
+        elif normal == [0, 1, 0]:  # Front
+            ax.text(0, 0.7, 0, face, color='black', fontsize=12, ha='center', va='center')
+        elif normal == [0, -1, 0]:  # Back
+            ax.text(0, -0.7, 0, face, color='black', fontsize=12, ha='center', va='center')
+        elif normal == [1, 0, 0]:  # Right
+            ax.text(0.7, 0, 0, face, color='black', fontsize=12, ha='center', va='center')
+        elif normal == [-1, 0, 0]:  # Left
+            ax.text(-0.7, 0, 0, face, color='black', fontsize=12, ha='center', va='center')
+
     # Set the aspect of the plot to be equal
     ax.set_box_aspect([1, 1, 1])  # Aspect ratio is 1:1:1
 
@@ -101,13 +120,11 @@ def draw_cube(state):
     # 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)
+
+# cube = {'Front': [['red', 'red', 'white'], ['white', 'white', 'red'], ['red', 'red', 'red']], 
+#  'Top': [['white', 'blue', 'white'], ['white', 'white', 'red'], ['white', 'white', 'red']], 
+#  'Bottom': [['white', 'blue', 'white'], ['red', 'blue', 'red'], ['red', 'white', 'white']], 
+#  'Left': [['white', 'white', 'red'], ['white', 'red', 'white'], ['white', 'white', 'red']], 
+#  'Right': [['white', 'red', 'white'], ['red', 'red', 'blue'], ['blue', 'blue', 'blue']], 
+#  'Back': [['white', 'yellow', 'red'], ['white', 'yellow', 'blue'], ['blue', 'white', 'blue']]}
+# draw_cube(cube)
\ No newline at end of file

main.py

-import matplotlib.pyplot as plt
-from mpl_toolkits.mplot3d.art3d import Poly3DCollection
-
+from cubeplot import draw_cube, draw_face
 import numpy as np
 from ultralytics import YOLO
 import cv2
@@ -28,105 +26,6 @@ color_map = {
     '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")
 
@@ -145,19 +44,23 @@ 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)
+# cube_state = {'Front': [['red', 'red', 'white'], ['white', 'white', 'red'], ['red', 'red', 'red']], 
+#  'Top': [['white', 'blue', 'white'], ['white', 'white', 'red'], ['white', 'white', 'red']], 
+#  'Bottom': [['white', 'blue', 'white'], ['red', 'blue', 'red'], ['red', 'white', 'white']], 
+#  'Left': [['white', 'white', 'red'], ['white', 'red', 'white'], ['white', 'white', 'red']], 
+#  'Right': [['white', 'red', 'white'], ['red', 'red', 'blue'], ['blue', 'blue', 'blue']], 
+#  'Back': [['white', 'yellow', 'red'], ['white', 'yellow', 'blue'], ['blue', 'white', 'blue']]}
+
+# draw_cube(cube_state)
 
-if cube_state ==None:
+if cube_state ==None or True:
     while True:
         ret, frame = cap.read()
         if not ret:
@@ -194,11 +97,15 @@ if cube_state ==None:
                     predicted_colors = raspberryPi.predict_colors(model_color, cropped_frame)
                     colors = raspberryPi.indices_to_colors(predicted_colors)
                     save_colors(sequence[current_step], colors)
+                    print(colors)
                     current_step += 1
 
                     if current_step >= len(sequence):
                         print("All sides scanned successfully!")
                         print("Final Cube Colors:", cube_state)
+                        fig = plt.figure()
+                        ax = fig.add_subplot(111, projection='3d')
+                        draw_cube(cube_state)  # Update visualization
                         break
                 elif key_confirm == ord('r'):
                     print("Retaking the picture...")