alternate_rgb_follower

"""Module makes a rover follow a line using a webcam for color detection."""

import asyncio

from viam.robot.client import RobotClient
from viam.rpc.dial import Credentials, DialOptions
from viam.components.base import Base, Vector3
from viam.components.camera import Camera
from viam.services.vision import VisionServiceClient
from viam.media.video import CameraMimeType
from viam.media.utils.pil import viam_to_pil_image, pil_to_viam_image
from enum import Enum, auto


# Copy and paste the following from the Connect tab's Python SDK section
async def connect():
    """
    Connects code to robot.
    """
    opts = RobotClient.Options.with_api_key(
      # Replace "<API-KEY>" (including brackets) with your robot's API key
      api_key='<API-KEY>',
      # Replace "<API-KEY-ID>" (including brackets) with your robot's API key ID
      api_key_id='<API-KEY-ID>'
    )
    return await RobotClient.at_address('ADDRESS FROM THE VIAM APP', opts)


class Direction(Enum):
    STOP = auto()
    FORWARD = auto()
    LEFT = auto()
    RIGHT = auto()


async def move_to_color(frame, vis, detector_name) -> Direction:
    x, y = frame.size[0], frame.size[1]

    # Crop the image to get only the middle fifth of the top third of the original image
    top_frame = frame.crop((x / 2.5, 0, x / 1.25, y / 3))
    top_det = await vis.get_detections(pil_to_viam_image(top_frame, CameraMimeType.JPEG), detector_name)
    if top_det:
        return Direction.FORWARD

    # Crop image to get only the left two fifths of the original image
    l_frame = frame.crop((0, 0, x / 2.5, y))
    l_det = await vis.get_detections(pil_to_viam_image(l_frame, CameraMimeType.JPEG), detector_name)

    # Crop image to get only the right two fifths of the original image
    r_frame = frame.crop((x / 1.25, 0, x, y))
    r_det = await vis.get_detections(pil_to_viam_image(r_frame, CameraMimeType.JPEG), detector_name)

    if not (r_det or l_det):
        return Direction.STOP

    # Color detected in either right, left, or both; choose the section with the larger green area
    _max = max([r_det, l_det], key=lambda xy: (xy[0].x_max - xy[0].x_min) * (xy[0].y_max - xy[0].y_min) if xy else 0)
    return Direction.LEFT if _max == l_det else Direction.RIGHT


async def main():
    """
    Main line follower function.
    """
    robot = await connect()
    camera = Camera.from_robot(robot, "my_camera")
    base = Base.from_robot(robot, "scuttlebase")
    vision = VisionServiceClient.from_robot(robot)

    linear_power = 0.4
    angular_power = 0.35
    try:
        while True:
            frame = await camera.get_image()
            direction = await move_to_color(viam_to_pil_image(frame), vision, "green_detector")
            print(f"move:{direction}")

            if Direction.FORWARD is direction:
                await base.set_power(Vector3(y=linear_power), Vector3())
            elif Direction.LEFT is direction:
                await base.set_power(Vector3(), Vector3(z=angular_power))
            elif Direction.RIGHT is direction:
                await base.set_power(Vector3(), Vector3(z=-angular_power))
            else:
                assert Direction.STOP is direction
                break
    finally:
        await base.stop()
        await robot.close()


if __name__ == "__main__":
    print("Starting up...")
    asyncio.run(main())
    print("Done.")