AutoKachalka/vision.py

import cv2
import numpy as np

import config


def blue_mask(bgr):
    """Маска пикселей синего семейства оттенков - этим цветом окрашены и фон
    шкалы, и маркер курсора (они могут быть визуально одного цвета)."""
    hsv = cv2.cvtColor(bgr, cv2.COLOR_BGR2HSV)
    lower = np.array([config.BLUE_HUE_MIN, config.BLUE_SAT_MIN, config.BLUE_VAL_MIN])
    upper = np.array([config.BLUE_HUE_MAX, 255, 255])
    return cv2.inRange(hsv, lower, upper)


def _blue_hue_mask(bgr):
    """Маска по одному оттенку, без учёта насыщенности/яркости - зоны не бывают
    синеватого оттенка даже в приглушённом виде, поэтому она шире blue_mask
    и не даёт тусклым синим пикселям шкалы попасть в зоны."""
    hsv = cv2.cvtColor(bgr, cv2.COLOR_BGR2HSV)
    return cv2.inRange(hsv, (config.BLUE_HUE_MIN, 0, 0), (config.BLUE_HUE_MAX, 255, 255))


def _contiguous_runs(bool_row):
    """Возвращает список (start, end) индексов подряд идущих True."""
    runs = []
    in_run = False
    start = 0
    for i, v in enumerate(bool_row):
        if v and not in_run:
            start = i
            in_run = True
        elif not v and in_run:
            runs.append((start, i - 1))
            in_run = False
    if in_run:
        runs.append((start, len(bool_row) - 1))
    return runs


def _merge_runs(runs, max_gap):
    """Объединяет соседние (start, end) интервалы, если разрыв между ними <= max_gap."""
    if not runs:
        return []
    runs = sorted(runs)
    merged = [list(runs[0])]
    for s, e in runs[1:]:
        if s - merged[-1][1] <= max_gap:
            merged[-1][1] = max(merged[-1][1], e)
        else:
            merged.append([s, e])
    return [(s, e) for s, e in merged]


def locate_bar(frame_bgr):
    """Ищет на кадре вытянутую синюю шкалу - широкую горизонтальную полосу
    сплошного синего цвета (плотность пикселей в строке, а не просто диапазон,
    чтобы не путать со разреженным синеватым фоном сцены/панели).
    Возвращает (x, y, w, h) в координатах кадра или None."""
    h_frame, w_frame = frame_bgr.shape[:2]
    mask = blue_mask(frame_bgr) > 0

    row_count = mask.sum(axis=1)
    width_threshold = w_frame * config.BAR_MIN_WIDTH_RATIO
    max_height = h_frame * config.BAR_MAX_HEIGHT_RATIO
    max_gap = w_frame * config.BAR_COLUMN_MERGE_GAP_RATIO

    wide_rows = row_count > width_threshold

    candidates = []
    for y0, y1 in _contiguous_runs(wide_rows):
        bh = y1 - y0 + 1
        if bh > max_height:
            continue

        cols_any = mask[y0:y1 + 1].any(axis=0)
        col_runs = _contiguous_runs(cols_any)
        merged = _merge_runs(col_runs, max_gap)
        if not merged:
            continue
        cs, ce = max(merged, key=lambda r: r[1] - r[0])
        bx, bw = cs, ce - cs + 1

        if bw / bh < config.BAR_MIN_ASPECT_RATIO:
            continue
        if bw < width_threshold:
            continue

        # Шкала - почти сплошной синий прямоугольник, в отличие от похожих
        # по цвету фоновых объектов (земля, небо, листва), где синие пиксели
        # покрывают лишь часть площади.
        region = mask[y0:y1 + 1, bx:bx + bw]
        fill_ratio = region.sum() / (bw * bh)
        if fill_ratio < config.BAR_MIN_FILL_RATIO:
            continue

        candidates.append((bx, int(y0), bw, bh))

    if not candidates:
        return None

    # Дополнительно: внутри шкалы должны лежать два цветных пятна-зоны.
    # Если ни у одного кандидата таких зон нет, значит шкала сейчас не видна.
    with_zones = [c for c in candidates if find_target_zone(frame_bgr, c) is not None]
    if not with_zones:
        return None
    return max(with_zones, key=lambda c: c[2])


def locate_cursor(frame_bgr, bar_bbox):
    """Курсор - "хвостик" того же синего цвета, выступающий из-под шкалы снизу.
    Под шкалой могут быть и другие синеватые элементы интерфейса (иконки,
    текст), поэтому ищем именно пятно, примыкающее к нижней границе шкалы,
    и берём центр наибольшего из таких пятен.
    Возвращает x-координату центра курсора в координатах кадра или None."""
    bx, by, bw, bh = bar_bbox
    mask = blue_mask(frame_bgr)

    search_h = int(bh * config.CURSOR_SEARCH_HEIGHT_RATIO)
    y0 = by + bh
    y1 = min(frame_bgr.shape[0], y0 + search_h)
    x0, x1 = bx, min(frame_bgr.shape[1], bx + bw)
    if y1 <= y0 or x1 <= x0:
        return None

    region = mask[y0:y1, x0:x1]
    touch_tolerance = max(2, int(bh * 0.1))
    contours, _ = cv2.findContours(region, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    best = None
    best_area = 0
    for c in contours:
        area = cv2.contourArea(c)
        if area < config.ZONE_MIN_AREA:
            continue
        _, y, _, _ = cv2.boundingRect(c)
        if y > touch_tolerance:
            continue
        if area > best_area:
            best_area = area
            best = c
    if best is None:
        return None
    x, _, w, _ = cv2.boundingRect(best)
    return x0 + x + w / 2.0


def find_target_zone(frame_bgr, bar_bbox):
    """Находит самую узкую цветную зону внутри шкалы (целевую) - это яркое
    насыщенное пятно, отличное от синего фона шкалы.
    Возвращает (x_min, x_max) в координатах кадра или None."""
    bx, by, bw, bh = bar_bbox
    bar = frame_bgr[by:by + bh, bx:bx + bw]
    if bar.size == 0:
        return None

    hsv = cv2.cvtColor(bar, cv2.COLOR_BGR2HSV)
    not_blue = cv2.bitwise_not(_blue_hue_mask(bar))
    vivid = cv2.inRange(hsv, (0, config.ZONE_SAT_MIN, config.ZONE_VAL_MIN), (179, 255, 255))
    mask_zones = cv2.bitwise_and(not_blue, vivid)

    contours, _ = cv2.findContours(mask_zones, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    min_zone_height = bh * config.ZONE_MIN_HEIGHT_RATIO
    zones = []
    for c in contours:
        if cv2.contourArea(c) < config.ZONE_MIN_AREA:
            continue
        x, y, w, h = cv2.boundingRect(c)
        if h < min_zone_height:
            continue
        zones.append((x, x + w))

    # Должно быть минимум две зоны (узкая+широкая). Если найдена только одна
    # или ни одной - вероятно, шум или сбой детекции, лучше пропустить кадр.
    if len(zones) < 2:
        return None

    narrow = min(zones, key=lambda z: z[1] - z[0])
    return bx + narrow[0], bx + narrow[1]