road_follower.py

import cv2
import numpy as np
import time
from OminiBot_HV import ominibothv

def region_of_interest(img, vertices):
    mask = np.zeros_like(img)
    cv2.fillPoly(mask, vertices, (255, 255, 255))
    return cv2.bitwise_and(img, mask)

# hsv color ranges (both lane lines are YELLOW)
low_yellow = np.array([26, 77, 100])
high_yellow = np.array([34, 255, 255])
# red (stop line) wraps around hue 0/180
low_red1 = np.array([0, 60, 90]);    high_red1 = np.array([12, 255, 255])
low_red2 = np.array([165, 60, 90]);  high_red2 = np.array([180, 255, 255])

# ---- open the camera (auto-find a working index) ----
cap = None
for i in range(6):
    c = cv2.VideoCapture(i)
    if c.isOpened():
        ok, _ = c.read()
        if ok:
            cap = c
            print("camera opened at index", i)
            break
    c.release()
if cap is None:
    raise SystemExit("no camera found (checked /dev/video0..5)")

# capture at high resolution (full/wide FOV), then downscale to 320x240 for
# processing so all the pixel coordinates below stay unchanged
CAP_W, CAP_H = 640, 480
cap.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter_fourcc('M', 'J', 'P', 'G'))
cap.set(cv2.CAP_PROP_FRAME_WIDTH, CAP_W)
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, CAP_H)

robot_control = ominibothv('/dev/ominibot', 115200)
time.sleep(3)

CENTER_X = 148            # calibrated from "calib" (was 160)
BASE_SPEED = 0.15
TURN_GAIN = BASE_SPEED / 80.0
SMOOTH = 0.5
HALF_LANE = 70            # offset from one line when only one is seen
MARGIN = 20              # dead zone around center for left/right classing
Y0, Y1 = 145, 205        # lookahead scan band
CORR_MAX = 0.16
MAX_STEP = 35            # max target move per frame (anti-jump)
LOST_STOP = 40
MIN_AREA = 20
RED_STOP_AREA = 800

interest_vertices = [np.array([
    [0, 240], [0, 200], [80, 120], [240, 120], [320, 200], [320, 240]
])]

def blob_centroids(band):
    num, _, stats, cents = cv2.connectedComponentsWithStats(band, connectivity=8)
    return [cents[i][0] for i in range(1, num) if stats[i, cv2.CC_STAT_AREA] >= MIN_AREA]

stopped = False
show_ok = True
prev_correction = 0.0
prev_target = CENTER_X
lost = 0
calib_sum = 0.0
calib_n = 0

while True:
    ret, frame = cap.read()
    if not ret or frame is None:
        robot_control.motor_speed(0.0, 0.0, 0.0, 0.0)
        continue

    if frame.shape[1] != 320 or frame.shape[0] != 240:
        frame = cv2.resize(frame, (320, 240))

    cropped = region_of_interest(frame.copy(), interest_vertices)
    hsv = cv2.cvtColor(cropped, cv2.COLOR_BGR2HSV)
    yellow_mask = cv2.inRange(hsv, low_yellow, high_yellow)
    red_mask = cv2.inRange(hsv, low_red1, high_red1) | cv2.inRange(hsv, low_red2, high_red2)

    # red stop line: latch once enough red is seen ahead
    red_count = int((red_mask[Y0:, :] > 0).sum() / 255)
    if red_count > RED_STOP_AREA:
        stopped = True

    # two yellow lines: group into blobs, classify by side, take inner one each side
    yblobs = blob_centroids(yellow_mask[Y0:Y1, :])
    left_side = [x for x in yblobs if x < CENTER_X - MARGIN]
    right_side = [x for x in yblobs if x > CENTER_X + MARGIN]
    left_b = max(left_side) if left_side else None
    right_b = min(right_side) if right_side else None

    if left_b is not None and right_b is not None:
        target_x = int((left_b + right_b) / 2)
        src = "both"
        calib_sum += (left_b + right_b) / 2
        calib_n += 1
    elif left_b is not None:
        target_x = int(left_b + HALF_LANE)
        src = "left"
    elif right_b is not None:
        target_x = int(right_b - HALF_LANE)
        src = "right"
    else:
        target_x = prev_target
        src = "hold"

    have_signal = (left_b is not None) or (right_b is not None)
    lost = 0 if have_signal else lost + 1

    target_x = max(prev_target - MAX_STEP, min(prev_target + MAX_STEP, target_x))
    target_x = max(0, min(319, target_x))
    prev_target = target_x

    error = target_x - CENTER_X
    correction = error * TURN_GAIN
    correction = SMOOTH * correction + (1 - SMOOTH) * prev_correction
    # both lines gone = deep in a sharp corner -> commit HARDER to finish the turn
    if src == "hold":
        correction *= 1.5
    correction = max(min(correction, CORR_MAX), -CORR_MAX)
    prev_correction = correction

    forward = max(0.07, BASE_SPEED - 0.6 * abs(correction))
    if lost > LOST_STOP:
        forward = 0.0
        correction = 0.0
        src = "LOST"

    if stopped:
        robot_speed_l = 0.0
        robot_speed_r = 0.0
        src = "RED-STOP"
    else:
        robot_speed_l = max(min(forward + correction, 0.30), -0.12)
        robot_speed_r = max(min(forward - correction, 0.30), -0.12)

    calib = int(calib_sum / calib_n) if calib_n else CENTER_X
    print("{:>8} L:{} R:{} tgt:{} calib:{} red:{} l:{:.2f} r:{:.2f}".format(
        src, None if left_b is None else int(left_b),
        None if right_b is None else int(right_b),
        target_x, calib, red_count, robot_speed_l, robot_speed_r))

    robot_control.motor_speed(robot_speed_l * -1, robot_speed_r, 0.0, 0.0)

    # ---- visualization ----
    ty = (Y0 + Y1) // 2
    cv2.rectangle(frame, (0, Y0), (319, Y1), (80, 80, 80), 1)
    if left_b is not None:
        cv2.circle(frame, (int(left_b), ty), 6, (255, 0, 0), -1)
    if right_b is not None:
        cv2.circle(frame, (int(right_b), ty), 6, (255, 0, 0), -1)
    cv2.line(frame, (CENTER_X, 0), (CENTER_X, 240), (0, 0, 255), 1)
    cv2.arrowedLine(frame, (CENTER_X, 240), (target_x, ty), (0, 255, 255), 3, tipLength=0.3)
    cv2.circle(frame, (target_x, ty), 6, (255, 0, 255), -1)
    if stopped:
        cv2.putText(frame, "RED STOP", (90, 120), cv2.FONT_HERSHEY_SIMPLEX, 1.0, (0, 0, 255), 3)
    cv2.imwrite('debug_live.jpg', frame)
    if show_ok:
        try:
            cv2.imshow('frame', frame)
            if (cv2.waitKey(1) & 0xFF) == ord('q'):
                break
        except cv2.error:
            show_ok = False

robot_control.motor_speed(0.0, 0.0, 0.0, 0.0)
cap.release()
cv2.destroyAllWindows()