File size: 4,999 Bytes
7b0cda0 d249b39 7b0cda0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 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_demo 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) |