main ¶

def main():
    """
    The main function.
    Here, the connection to the robot is established, the `read_faces` method is started and the faces are scanned in the
    cv module and set as faces for the 3d representation of the cube.
    The resulting cube is then used to calculate the solution using the solver module.
    The solution is then optimised and used by the SimulatingCuber to find the shortest non-colliding
    robot paths and actions needed to solve the cube.
    """
    sys.setrecursionlimit(5000)
    I_WANT_THAT_PICKLE = False
    SIMULATION_PROGRAM_PICKLE_FILEPATH = 'sim_program.pickle'
    FACE_IMAGES_FOLDER: str = 'images/lab_set_franz_const_exp/'
    USE_DUMMY_SOLUTION = False
    RUN_ON_ROBOT = False
    desired_run_mode = robolink.RUNMODE_RUN_ROBOT if RUN_ON_ROBOT else robolink.RUNMODE_SIMULATE

    logging.basicConfig(level=logging.INFO)
    logging.getLogger("core.manipulator").setLevel(logging.DEBUG)

    # move to idle
    RoboDKInterface.robot.MoveJ(RoboDKInterface.idle_target)
    RoboDKInterface.link.setRunMode(desired_run_mode)

    # initialize the cuber
    cuber = manipulator.AnimatingCuber()

    RoboDKInterface.setup_collisions()

    sim_program = None
    if I_WANT_THAT_PICKLE:
        with open(SIMULATION_PROGRAM_PICKLE_FILEPATH, 'rb') as file:
            initial_cube_cs, initial_cube_segments, sim_program = pickle.load(file)
            ui.UIUtilities.set_pose(
                RoboDKInterface.ui_cube.ui_cube_frame,
                cs=initial_cube_cs
            )
            RoboDKInterface.ui_cube.segments = initial_cube_segments
            RoboDKInterface.ui_cube.update_ui()
            print(f'Read simulation from "{SIMULATION_PROGRAM_PICKLE_FILEPATH}".')
    else:
        solution = "U2 L2 R2"
        initial_cube_cs = spatial.CoordinateSystem()
        initial_cube_segments = RoboDKInterface.ui_cube.segments

        if not USE_DUMMY_SOLUTION:

            # choose an image source
            with cv.PyUEyeImageSource(0) if RUN_ON_ROBOT else cv.DirImageSource(FACE_IMAGES_FOLDER) as image_source:
                # take photos
                face_images = read_faces(cuber, image_source)
                faces = cv.process_images(face_images)

                # update the pose of the cube and its segments
                base = RoboDKInterface.ui_cube.set_faces(faces)
                initial_cube_cs = spatial.CoordinateSystem(origin=None, base=base)
                initial_cube_segments = RoboDKInterface.ui_cube.segments
                ui.UIUtilities.set_pose(
                    RoboDKInterface.ui_cube.ui_cube_frame,
                    cs=initial_cube_cs
                )
                RoboDKInterface.ui_cube.update_ui()

                # this is necessary because there might be some color tiles add from stl file, which might have
                # an unwanted collision map applied by default -> thats why we override them here after modifying
                RoboDKInterface.setup_collisions()

                solution = solver.solve(solver.Cube.from_3d_representation(RoboDKInterface.ui_cube), "kociemba")

        print('SOLUTION: ', solution)
        start_time = time.time()

        cube_state = manipulator.CubeState(
            cube_is_gripped=False,
            cube_coordinate_system=initial_cube_cs,
            grip_position=None
        )
        sim_program = manipulator.SimulationProgram(
            name="perform_solution",
            initial_cube_state=cube_state,
            initial_robot_joints=RoboDKInterface.robot.Joints().list()
        )
        RoboDKInterface.link.setRunMode(robolink.RUNMODE_SIMULATE)
        RoboDKInterface.link.Render(False)
        sim_program = cuber.perform_rotations(
            rotations=spatial.Rotation.from_string_notation(solution),
            simulation_program=sim_program
        )
        sim_program = cuber.exit_mount(sim_program, detach_before_exiting=True)
        sim_program = cuber.move_to_idle(sim_program)
        end_time = time.time()
        print('SIMULATION COMPLETE ({:.1f}s).'.format(end_time - start_time))

        with open(SIMULATION_PROGRAM_PICKLE_FILEPATH, 'wb') as file:
            pickle.dump((initial_cube_cs, initial_cube_segments, sim_program), file)
            print(f'Wrote simulation program to "{SIMULATION_PROGRAM_PICKLE_FILEPATH}".')

    RoboDKInterface.link.setRunMode(desired_run_mode)
    RoboDKInterface.link.Render(True)
    sim_program.execute()

def read_faces(simulating_cuber: manipulator.SimulatingCuber,
               image_source: cv.ImageSource) -> List[np.ndarray]:
    """
    Reads all six faces of the physical or simulated cube.
    Args:
        simulating_cuber: The simulating cuber
        image_source: The image source to use

    Returns:
        Six images of the six faces of the cube, in order of appearance
    """
    assert isinstance(simulating_cuber, manipulator.SimulatingCuber)
    assert isinstance(image_source, cv.ImageSource)



    face_images = []

    def get_image(*args, **kwargs):
        face_images.append(image_source.get_image())

    get_image_instruction = manipulator.PythonInstruction('get_image', get_image)

    manipulator.GripInstruction(False).execute()
    initial_cube_state = manipulator.CubeState(
        cube_is_gripped=False,
        cube_coordinate_system=spatial.CoordinateSystem(),
        grip_position=None
    )

    init_and_show_left = manipulator.SimulationProgram('init_and_show_left', initial_cube_state, initial_robot_joints=RoboDKInterface.robot.Joints())
    init_and_show_left = simulating_cuber.move_to_idle(init_and_show_left)
    factory = grip_target_generation.GripTargetFactory(
        face_to_mount=grip_target_generation.GripperFace.MAIN,
        num_rotations=1
    )


    init_and_show_left = simulating_cuber.enter_mount(factory, init_and_show_left, attach_after_entering=True)
    init_and_show_left = simulating_cuber.exit_mount(init_and_show_left)

    # show first three sides
    simulating_cuber.move(RoboDKInterface.show_left_target, init_and_show_left)
    init_and_show_left.append_instruction(get_image_instruction)

    show_main = manipulator.SimulationProgram('show_main', parent=init_and_show_left, initial_robot_joints=init_and_show_left.final_robot_position())
    simulating_cuber.move(RoboDKInterface.show_main_target, show_main)
    show_main.append_instruction(get_image_instruction)

    show_right = manipulator.SimulationProgram('show_right', parent=show_main, initial_robot_joints=show_main.final_robot_position())
    simulating_cuber.move(RoboDKInterface.show_right_target, show_right)
    show_right.append_instruction(get_image_instruction)

    regrip_program = manipulator.SimulationProgram('regrip', parent=show_right, initial_robot_joints=show_right.final_robot_position())

    # do a regrip
    factory = grip_target_generation.GripTargetFactory(
        face_to_mount=grip_target_generation.GripperFace.Y_NEGATIVE,
        num_rotations=1
    )
    regrip_program = simulating_cuber.enter_mount(factory, regrip_program, prepare_linear=True)
    regrip_program = simulating_cuber.exit_mount(regrip_program, detach_before_exiting=True)

    regrip_program = simulating_cuber.move_to_idle(regrip_program)

    factory = grip_target_generation.GripTargetFactory(
        face_to_mount=grip_target_generation.GripperFace.GRIPPED_X_NEGATIVE,
        num_rotations=-1
    )
    regrip_program = simulating_cuber.enter_mount(factory, regrip_program, attach_after_entering=True)
    regrip_program = simulating_cuber.exit_mount(regrip_program)

    # show last three sides
    init_and_show_left = manipulator.SimulationProgram('init_and_show_left', parent=regrip_program, initial_robot_joints=regrip_program.final_robot_position())
    simulating_cuber.move(RoboDKInterface.show_left_target, init_and_show_left)
    init_and_show_left.append_instruction(get_image_instruction)

    show_main = manipulator.SimulationProgram('show_main', parent=init_and_show_left, initial_robot_joints=init_and_show_left.final_robot_position())
    simulating_cuber.move(RoboDKInterface.show_main_target, show_main)
    show_main.append_instruction(get_image_instruction)

    show_right = manipulator.SimulationProgram('show_right', parent=show_main, initial_robot_joints=show_main.final_robot_position())
    simulating_cuber.move(RoboDKInterface.show_right_target, show_right)
    show_right.append_instruction(get_image_instruction)

    put_cube_back = manipulator.SimulationProgram('put_cube_back', parent=show_right, initial_robot_joints=show_right.final_robot_position())

    # put the cube back to where it belongs :-)
    factory = grip_target_generation.GripTargetFactory(
        face_to_mount=grip_target_generation.GripperFace.MAIN,
        num_rotations=0
    )

    put_cube_back = simulating_cuber.enter_mount(factory, put_cube_back, prepare_linear=True)
    simulating_cuber.detach_cube(put_cube_back)
    put_cube_back = simulating_cuber.move_to_idle(put_cube_back)
    put_cube_back.execute()

    return face_images