src/pipeline_hf.py

import cv2
import numpy as np
import pandas as pd

from pathlib import Path
from src.config_loader import load_config
from src.ocr_utils_demo import ocr_title, ocr_sensors

def process_single_image(img: np.ndarray, color_ranges: dict):
    """
    Обрабатывает одно изображение:
    - вытаскивает титул;
    - находит и оцифровывает сенсоры;
    - возвращает структуру с результатом.
    """

    # ---------- 1. Оцифровка титула ----------
    h, w = img.shape[:2]
    title_roi = img[:45, :int(w / 2.4)]
    title_text = ocr_title(title_roi)

    # ---------- 2. Оцифровка сенсоров ----------
    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)

    mask = None

    for (lo, hi) in color_ranges.values():
        lo_np = np.array(lo, dtype=np.uint8)
        hi_np = np.array(hi, dtype=np.uint8)
        cur = cv2.inRange(hsv, lo_np, hi_np)
        mask = cur if mask is None else cv2.bitwise_or(mask, cur)

    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    rois, positions = [], []
    for cnt in contours:
        x, y, ww, hh = cv2.boundingRect(cnt)
        if ww < 90 or hh < 17:
            continue
        if hh / ww > 1.5:
            continue

        hh_clamped = min(hh, 17)
        roi = img[y:y + hh_clamped, x:x + ww]
        rois.append(roi)
        positions.append((x, y, ww, hh))

    # ---------- 3. OCR сенсоров ----------
    sensors = []
    if rois:
        ocr_results = ocr_sensors(rois)

        for (x, y, ww, hh), result in zip(positions, ocr_results):
            sensors.append({
                "text": result["text"],
                "score": result["score"],
                "x": x,
                "y": y,
                "w": ww,
                "h": hh
            })

    return title_text, sensors

# ---------------------------------------------------------
# Основной pipeline → Excel
# ---------------------------------------------------------
def process_image_to_excel(cfg: dict):
    input_dir = Path(cfg["paths"]["input"])
    output_dir = Path(cfg["paths"]["output"])
    excel_name = cfg["export"]["excel_filename"]

    output_dir.mkdir(parents=True, exist_ok=True)
    excel_path = output_dir / excel_name

    color_ranges = cfg["colors"]

    results = []

    for img_path in input_dir.glob("*.png"):
        name = img_path.name
        img = cv2.imread(str(img_path))

        if img is None:
            print(f"Не могу прочитать файл {name}")
            continue

        title_text, sensors = process_single_image(img, color_ranges)

        for sen in sensors:
            results.append({
                "filename": name,
                "title": title_text,
                "sensor_name": sen["text"],
                "score": sen["score"]
            })

    if not results:
        print("⚠ Нет данных для записи.")
        return

    df = pd.DataFrame(results)
    df.to_excel(excel_path, index=False, engine="openpyxl")

    print(f"✅ Готово! Excel сохранён: {excel_path.resolve()}")


# ---------------------------------------------------------
# Запуск
# ---------------------------------------------------------
if __name__ == "__main__":
    cfg = load_config("configs/config.yaml")
    process_image_to_excel(cfg)