Update project structure, add pipeline and tests

configs/config.yaml

@@ -1,7 +1,17 @@
-input_dir: data/sample/input
-output_dir: outputs
-ocr:
-  engine: paddleocr
-  lang: en
+paths:
+  input: "input"
+  output: "outputs"
+  
+colors:
+  purple: [[120, 80, 50], [150, 255, 255]]
+  cyan: [[85, 150, 150], [100, 255, 255]]
+  white: [[0, 0, 250], [180, 15, 255]]
+  green_light: [[85, 118, 118], [95, 138, 138]]
+  green_dark: [[55, 245, 186], [65, 255, 206]]
+  gray_dark: [[0, 0, 187], [180, 5, 197]]
+  yellow: [[20, 100, 100], [40, 255, 255]]
+  red: [[0, 150, 150], [10, 255, 255]]
+
 export:
-  excel_filename: results.xlsx
+  excel_filename: "ocr_results.xlsx"
+>>>>>>> e710af8 (Update project structure, add pipeline and tests)

pytest.ini

@@ -0,0 +1,2 @@
+[pytest]
+pythonpath = src

src/config_loader.py

@@ -0,0 +1,16 @@
+"""
+Загрузка конфигурационного файла config.yaml
+"""
+
+import yaml
+from pathlib import Path
+
+
+def load_config(path: str | Path = "config.yaml") -> dict:
+    cfg_path = Path(path).resolve()
+
+    if not cfg_path.exists():
+        raise FileNotFoundError(f"Файл конфигурации не найден: {cfg_path}")
+
+    with open(cfg_path, "r", encoding="utf-8") as f:
+        return yaml.safe_load(f)

src/ocr_utils.py

@@ -0,0 +1,67 @@
+"""
+OCR-инструменты: pytesseract для титула и PaddleOCR для датчиков.
+"""
+
+import cv2
+import numpy as np
+from paddleocr import PaddleOCR
+import pytesseract
+
+
+# --------------------------
+# OCR титула (pytesseract)
+# --------------------------
+def ocr_title(img: np.ndarray) -> str:
+    """
+    OCR верхней области (титул мнемосхемы).
+    """
+    if img is None or img.size == 0:
+        return ""
+
+    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+    gray = cv2.convertScaleAbs(gray, alpha=2.0, beta=-40)
+
+    binary = cv2.adaptiveThreshold(
+        gray, 255,
+        cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
+        cv2.THRESH_BINARY_INV,
+        15, 9
+    )
+
+    text = pytesseract.image_to_string(binary, lang="eng", config="--psm 7").strip()
+    return text
+
+
+# --------------------------
+# OCR датчиков (PaddleOCR)
+# --------------------------
+paddle_ocr = PaddleOCR(
+    lang="en",
+    use_textline_orientation=True
+)
+
+
+def ocr_sensors(rois: list[np.ndarray]) -> list[dict]:
+    """
+    OCR областей сенсоров через PaddleOCR.predict().
+    Формат вывода:
+        [{"text": str, "score": float}, ...]
+    """
+    results = []
+
+    if not rois:
+        return []
+
+    try:
+        ocr_results = paddle_ocr.predict(rois)
+    except Exception as e:
+        print(f"⚠ Ошибка OCR.predict: {e}")
+        return [{"text": "?", "score": 0.0} for _ in rois]
+
+    for out in ocr_results:
+        text = out.get("rec_texts", ["?"])[0]
+        score = out.get("rec_scores", [0.0])[0]
+        results.append({"text": text, "score": float(score)})
+
+    return results
+

src/pipeline.py

@@ -0,0 +1,166 @@
+"""
+Главный pipeline проекта
+"""
+
+import cv2
+import numpy as np
+import pandas as pd
+from pathlib import Path
+from docx import Document
+
+from src.config_loader import load_config
+from src.ocr_utils import ocr_title, ocr_sensors
+
+
+# ---------------------------------------------------------
+# Итератор входных данных (СЦЕНАРИЙ A)
+# ---------------------------------------------------------
+def iter_input_images(input_dir: Path):
+    """
+    Универсальный вход:
+    - читаем ВСЕ png/jpg/bmp
+    - читаем ВСЕ изображения внутри DOCX
+    """
+    exts = {".png", ".jpg", ".jpeg", ".bmp"}
+
+    image_files = []
+    docx_files = []
+
+    for p in input_dir.rglob("*"):
+        if not p.is_file():
+            continue
+
+        suf = p.suffix.lower()
+        if suf in exts:
+            image_files.append(p)
+        elif suf == ".docx":
+            docx_files.append(p)
+
+    # ---------- PNG/JPG ----------
+    for path in sorted(image_files):
+        img = cv2.imread(str(path), cv2.IMREAD_COLOR)
+        if img is not None:
+            yield path.stem, img
+
+    # ---------- DOCX ----------
+    for docx_path in sorted(docx_files):
+        doc = Document(docx_path)
+
+        idx = 0
+        for rel in doc.part._rels.values():
+            if "image" not in rel.target_ref:
+                continue
+
+            idx += 1
+            blob = rel.target_part.blob
+            arr = np.frombuffer(blob, np.uint8)
+            img = cv2.imdecode(arr, cv2.IMREAD_COLOR)
+
+            if img is not None:
+                yield f"{docx_path.stem}_img{idx}", img
+
+    # Если не нашли ничего
+    if not image_files and not docx_files:
+        print("⚠ Папка не содержит ни изображений, ни DOCX.")
+
+
+
+# ---------------------------------------------------------
+# Извлечение текста титула + датчиков
+# ---------------------------------------------------------
+def extract_text(img: np.ndarray, color_ranges: dict) -> tuple[str, list[dict]]:
+    """
+    1. Обрезаем титул (верхняя область)
+    2. OCR титула
+    3. Извлекаем сенсоры по цветовым маскам
+    4. OCR сенсоров
+    """
+
+    # ---------- 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
+
+        hh = min(hh, 17)
+        rois.append(img[y:y + hh, x:x + ww])
+        positions.append((x, y, ww, hh))
+
+    # ---------- 3. OCR сенсоров ----------
+    sensors = []
+    if rois:
+        ocr_results = ocr_sensors(rois)
+
+        for (x, y, ww, hh), r in zip(positions, ocr_results):
+            sensors.append({
+                "text": r["text"],
+                "score": r["score"],
+                "x": x,
+                "y": y,
+                "w": ww,
+                "h": hh
+            })
+
+    return title_text, sensors
+
+
+# ---------------------------------------------------------
+# Основной pipeline → Excel
+# ---------------------------------------------------------
+def process_all_images_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 name, img in iter_input_images(input_dir):
+        title_text, sensors = extract_text(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)
+    pd.DataFrame.to_excel(df, excel_path, index=False, engine="openpyxl")
+
+    print(f"✅ Готово! Excel сохранён: {excel_path.resolve()}")
+
+
+# ---------------------------------------------------------
+# Запуск
+# ---------------------------------------------------------
+if __name__ == "__main__":
+    cfg = load_config("config.yaml")
+    process_all_images_to_excel(cfg)

tests/conftest.py

@@ -0,0 +1,8 @@
+import sys
+import os
+
+ROOT_DIR = os.path.dirname(os.path.dirname(__file__))
+SRC_DIR = os.path.join(ROOT_DIR, "src")
+
+sys.path.insert(0, SRC_DIR)
+

tests/test_extract_text.py

@@ -0,0 +1,72 @@
+import numpy as np
+import cv2
+import pytest
+from unittest.mock import patch
+
+from src.pipeline import extract_text   # поправь название модуля
+
+
+# ---------------------------------------------------------
+# ВСПОМОГАТЕЛЬНАЯ ФУНКЦИЯ: создаём тестовую картинку
+# ---------------------------------------------------------
+def make_test_image():
+    """
+    Создаёт искусственное изображение 200×300 + цветовой диапазон,
+    чтобы extract_text мог найти один "сенсор".
+    """
+    img = np.zeros((200, 300, 3), dtype=np.uint8)
+
+    # Рисуем титул (верхний левый угол) — просто белую полосу
+    img[0:45, 0:120] = (255, 255, 255)
+
+    # Рисуем сенсор в центре картинки
+    cv2.rectangle(img, (100, 80), (200, 97), (0, 255, 0), -1)
+
+    return img
+
+
+# ---------------------------------------------------------
+# ГЛАВНЫЙ ТЕСТ extract_text
+# ---------------------------------------------------------
+@patch("src.pipeline.ocr_title")
+@patch("src.pipeline.ocr_sensors")
+def test_extract_text(mock_ocr_sensors, mock_ocr_title):
+    img = make_test_image()
+
+    # Цветовой диапазон, по которому найдётся наш зелёный прямоугольник
+    color_ranges = {
+        "green": (
+            [50, 80, 80],   # нижняя граница HSV
+            [80, 255, 255]  # верхняя граница HSV
+        )
+    }
+
+    # --- подменяем OCR ---
+    mock_ocr_title.return_value = "TITLE_OK"
+
+    mock_ocr_sensors.return_value = [
+        {"text": "123", "score": 0.95},
+    ]
+
+    # --- вызываем тестируемую функцию ---
+    title, sensors = extract_text(img, color_ranges)
+
+    # --- ПРОВЕРКИ ---
+
+    # титул
+    assert title == "TITLE_OK"
+
+    # сенсоры
+    assert isinstance(sensors, list)
+    assert len(sensors) == 1
+
+    s = sensors[0]
+
+    assert s["text"] == "123"
+    assert s["score"] == 0.95
+
+    # координаты должны быть внутри изображения
+    assert 0 <= s["x"] <= img.shape[1]
+    assert 0 <= s["y"] <= img.shape[0]
+    assert s["w"] > 0
+    assert s["h"] > 0

tests/test_iter_input_images.py

@@ -0,0 +1,72 @@
+import numpy as np
+import pytest
+from unittest.mock import patch, MagicMock
+from pathlib import Path
+
+from  src.pipeline import iter_input_images   # замените под путь вашего файла
+
+@patch("src.pipeline.cv2.imread")
+def test_iter_input_images_with_images(mock_imread, tmp_path):
+    # Создаем временные файлы
+    img1 = tmp_path / "a.png"
+    img2 = tmp_path / "b.jpg"
+    img1.write_bytes(b"fake")
+    img2.write_bytes(b"fake")
+
+    # Мокаем результат чтения изображений
+    dummy_img = np.zeros((10, 10, 3), dtype=np.uint8)
+    mock_imread.return_value = dummy_img
+
+    results = list(iter_input_images(tmp_path))
+
+    assert len(results) == 2
+    assert results[0][0] == "a"     # stem
+    assert isinstance(results[0][1], np.ndarray)
+
+    assert results[1][0] == "b"
+    assert isinstance(results[1][1], np.ndarray)
+
+    assert mock_imread.call_count == 2
+
+@patch("src.pipeline.cv2.imdecode")
+@patch("src.pipeline.Document")
+def test_iter_input_images_with_docx(mock_Document, mock_imdecode, tmp_path):
+    docx_file = tmp_path / "test.docx"
+    docx_file.write_bytes(b"fake")
+
+    # Мокаем Document
+    mock_doc = MagicMock()
+    mock_Document.return_value = mock_doc
+
+    # Имитация doc.part._rels
+    mock_rel = MagicMock()
+    mock_rel.target_ref = "/media/image1.png"
+    mock_rel.target_part.blob = b"FAKE_IMAGE_DATA"
+
+    mock_doc.part._rels = {"r1": mock_rel}
+
+    # cv2.imdecode → возвращает искусственное изображение
+    dummy_img = np.ones((5, 5, 3), dtype=np.uint8)
+    mock_imdecode.return_value = dummy_img
+
+    results = list(iter_input_images(tmp_path))
+
+    assert len(results) == 1
+    name, img = results[0]
+
+    assert name == "test_img1"
+    assert isinstance(img, np.ndarray)
+    assert img.shape == (5, 5, 3)
+
+    mock_Document.assert_called_once()
+    mock_imdecode.assert_called_once()
+
+def test_iter_input_images_empty_folder(tmp_path, capsys):
+    results = list(iter_input_images(tmp_path))
+
+    # Функция ничего не вернёт
+    assert results == []
+
+    # Проверяем, что напечатано предупреждение
+    captured = capsys.readouterr()
+    assert "не содержит ни изображений" in captured.out.lower()

tests/test_ocr_utils.py

@@ -0,0 +1,50 @@
+import numpy as np
+import pytest
+from unittest.mock import patch
+
+from src.ocr_utils import ocr_title, ocr_sensors
+
+@patch("src.ocr_utils.pytesseract.image_to_string")
+def test_ocr_title_basic(mock_tesseract):
+    # Подделываем OCR-результат
+    mock_tesseract.return_value = "TEST_TITLE"
+
+    # Синтетическое изображение
+    img = np.zeros((50, 200, 3), dtype=np.uint8)
+
+    text = ocr_title(img)
+
+    assert text == "TEST_TITLE"
+    mock_tesseract.assert_called_once()
+
+@patch("src.ocr_utils.paddle_ocr.predict")
+def test_ocr_sensors_basic(mock_predict):
+    # Подменяем результат работы paddleocr.predict
+    mock_predict.return_value = [
+        {"rec_texts": ["123"], "rec_scores": [0.98]}
+    ]
+
+    # Одна синтетическая ROI
+    img = np.zeros((20, 60, 3), dtype=np.uint8)
+
+    out = ocr_sensors([img])
+
+    assert isinstance(out, list)
+    assert len(out) == 1
+    assert out[0]["text"] == "123"
+    assert out[0]["score"] == 0.98
+    mock_predict.assert_called_once()
+
+@patch("src.ocr_utils.paddle_ocr.predict")
+def test_ocr_sensors_empty_roi(mock_predict):
+    # paddleOCR должен корректно отработать
+    mock_predict.return_value = [
+        {"rec_texts": [""], "rec_scores": [0]}
+    ]
+
+    img = np.zeros((10, 10, 3), dtype=np.uint8)
+
+    out = ocr_sensors([img])
+
+    assert out[0]["text"] == ""
+    assert out[0]["score"] == 0

tests/test_process_all_images.py

@@ -0,0 +1,95 @@
+import pytest
+from unittest.mock import patch, MagicMock
+from pathlib import Path
+import pandas as pd
+import numpy as np
+
+from src.pipeline import process_all_images_to_excel   # заменить на реальный путь
+
+
+@patch("src.pipeline.extract_text")
+@patch("src.pipeline.iter_input_images")
+@patch("src.pipeline.pd.DataFrame.to_excel")
+def test_process_all_images_to_excel_ok(mock_to_excel, mock_iter, mock_extract, tmp_path):
+    # Подготовим конфиг
+    cfg = {
+        "paths": {
+            "input": str(tmp_path / "input"),
+            "output": str(tmp_path / "output")
+        },
+        "export": {
+            "excel_filename": "out.xlsx"
+        },
+        "colors": {"dummy_color": ([0,0,0],[255,255,255])}
+    }
+
+    # Создаём папки
+    (tmp_path / "input").mkdir()
+    (tmp_path / "output").mkdir()
+
+    # --- Мокаем iter_input_images ---
+    mock_iter.return_value = [
+        ("img1", np.zeros((10,10,3))),
+        ("img2", np.zeros((10,10,3))),
+    ]
+
+    # --- Мокаем extract_text ---
+    mock_extract.side_effect = [
+        ("TITLE1", [{"text": "S1", "score": 0.9}]),
+        ("TITLE2", [{"text": "S2", "score": 0.95}]),
+    ]
+
+    # --- Выполняем функцию ---
+    process_all_images_to_excel(cfg)
+
+    # --- ПРОВЕРКИ ---
+
+    # extract_text должен быть вызван 2 раза
+    assert mock_extract.call_count == 2
+
+    # Excel действительно формировался
+    mock_to_excel.assert_called_once()
+
+    # Проверяем, что путь правильный
+    called_path, called_kwargs = mock_to_excel.call_args[0][0], mock_to_excel.call_args[1]
+    assert called_kwargs["index"] is False
+    assert "openpyxl" in called_kwargs["engine"]
+
+    # Проверяем, что DataFrame содержит нужные строки
+    df_created: pd.DataFrame = mock_to_excel.call_args[0][0]  # аргумент-DataFrame
+
+    assert len(df_created) == 2
+    assert set(df_created.columns) == {"filename", "title", "sensor_name", "score"}
+
+    assert df_created.iloc[0]["filename"] == "img1"
+    assert df_created.iloc[0]["title"] == "TITLE1"
+    assert df_created.iloc[0]["sensor_name"] == "S1"
+    assert df_created.iloc[0]["score"] == 0.9
+
+
+@patch("src.pipeline.iter_input_images")
+@patch("src.pipeline.pd.DataFrame.to_excel")
+def test_process_all_images_no_data(mock_to_excel, mock_iter, tmp_path, capsys):
+    cfg = {
+        "paths": {
+            "input": str(tmp_path / "input"),
+            "output": str(tmp_path / "output")
+        },
+        "export": {"excel_filename": "out.xlsx"},
+        "colors": {}
+    }
+
+    (tmp_path / "input").mkdir()
+    (tmp_path / "output").mkdir()
+
+    # iter_input_images вернёт пустой список
+    mock_iter.return_value = []
+
+    process_all_images_to_excel(cfg)
+
+    # Excel не должен создаваться
+    mock_to_excel.assert_not_called()
+
+    # Проверяем вывод
+    captured = capsys.readouterr()
+    assert "нет данных" in captured.out.lower()