"""Import algorithm stage labels into compressor_label. Phase 2: algorithm labels are written as segment records. Manual rows are never overwritten by a re-import. """ from __future__ import annotations import argparse import csv import hashlib import json import re from datetime import datetime, timedelta from pathlib import Path from typing import Any from app.db import ensure_label_table, get_connection ROOT = Path(__file__).resolve().parents[1] VIBRATION_CSV = ROOT / "models" / "datas" / "reference" / "振动异常段.csv" OIL_CSV = ROOT / "models" / "datas" / "reference" / "润滑油压力小时阶段段.csv" TEMPERATURE_CSV = ROOT / "_tmp_temperature" / "out" / "温度异常段.csv" SEVERITY = {"轻微": 1, "异常": 2, "严重异常": 3} VIBRATION_SIDE_TO_METRIC = {"联轴器端": "YSJ_10", "链轮端": "YSJ_11"} INSERT_SQL = """ INSERT INTO compressor_label ( source_type, origin, target_type, target_id, target_key, channel_key, segment_key, unit, device_code, device_part, device_point, metric_key, metric_label, start_time, end_time, algorithm_start, algorithm_end, label_type, label_subtype, severity, confirm_status, usable_for_train, source_file, source_row, algorithm_version, features_json ) VALUES ( %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s ) ON DUPLICATE KEY UPDATE label_type = IF(compressor_label.origin = 'manual', compressor_label.label_type, VALUES(label_type)), label_subtype = IF(compressor_label.origin = 'manual', compressor_label.label_subtype, VALUES(label_subtype)), severity = IF(compressor_label.origin = 'manual', compressor_label.severity, VALUES(severity)), confirm_status = IF(compressor_label.origin = 'manual', compressor_label.confirm_status, VALUES(confirm_status)), usable_for_train = IF(compressor_label.origin = 'manual', compressor_label.usable_for_train, VALUES(usable_for_train)), features_json = IF(compressor_label.origin = 'manual', compressor_label.features_json, VALUES(features_json)), updated_at = CURRENT_TIMESTAMP """ def read_csv(path: Path) -> list[dict[str, str]]: if not path.exists(): print(f"跳过：文件不存在 {path}") return [] with path.open("r", encoding="utf-8-sig", newline="") as handle: return list(csv.DictReader(handle)) def parse_unit(raw: str) -> str | None: match = re.match(r"\s*(\d+)\s*号机", raw or "") return match.group(1) if match else None def parse_time(raw: str) -> datetime | None: text = (raw or "").strip() if not text: return None for fmt in ("%Y-%m-%d %H:%M:%S", "%Y-%m-%d %H:%M", "%Y-%m-%d"): try: return datetime.strptime(text, fmt) except ValueError: continue return None def time_text(value: datetime) -> str: return value.strftime("%Y-%m-%d %H:%M:%S") def metric_label(site_points: dict[str, str], unit: str, metric: str) -> str: return site_points.get(f"YSJ{unit}_{metric.split('_')[-1]}", metric) def confirm_status(row: dict[str, str]) -> str: return "已确认" if (row.get("是否形成确认等级") or "").strip() == "是" else "未确认" def make_segment( *, source_type: str, unit: str, metric: str, start: datetime, end: datetime, label_type: str, label_subtype: str | None, row: dict[str, str], source_file: str, site_points: dict[str, str], ) -> tuple[Any, ...]: device_code = f"{unit}#" target_key = f"pks:{unit}:{device_code}:{metric}" source_row = (row.get("周期编号") or "").strip() segment_key = hashlib.md5( f"{source_type}|{unit}|{metric}|{source_row}|{time_text(start)}|{time_text(end)}".encode("utf-8") ).hexdigest() label = metric_label(site_points, unit, metric) return ( source_type, "algorithm", "pks", None, target_key, target_key, segment_key, unit, device_code, (row.get("机组") or "").strip(), label, metric, label, time_text(start), time_text(end), time_text(start), time_text(end), label_type, label_subtype, SEVERITY.get(label_type), confirm_status(row), 1, source_file, source_row, "csv-v1", json.dumps(row, ensure_ascii=False), ) def build_rows(site_points: dict[str, str]) -> list[tuple[Any, ...]]: rows: list[tuple[Any, ...]] = [] for row in read_csv(VIBRATION_CSV): unit = parse_unit(row.get("机组", "")) start = parse_time(row.get("开始小时", "")) end = parse_time(row.get("结束小时", "")) side = (row.get("最强异常侧") or "").strip() metric = VIBRATION_SIDE_TO_METRIC.get(side) if not unit or not start or not end or not metric: continue rows.append( make_segment( source_type="vibration", unit=unit, metric=metric, start=start, end=end + timedelta(hours=1), label_type=(row.get("振动阶段") or "").strip(), label_subtype=side or None, row=row, source_file=str(VIBRATION_CSV.relative_to(ROOT)), site_points=site_points, ) ) for row in read_csv(OIL_CSV): unit = parse_unit(row.get("机组", "")) start = parse_time(row.get("开始小时", "")) end = parse_time(row.get("结束小时", "")) if not unit or not start or not end: continue metric = "YSJ_5" if unit in {"7", "8", "9"} else "YSJ_1" rows.append( make_segment( source_type="oil_pressure", unit=unit, metric=metric, start=start, end=end + timedelta(hours=1), label_type=(row.get("压力阶段") or "").strip(), label_subtype="润滑油压力", row=row, source_file=str(OIL_CSV.relative_to(ROOT)), site_points=site_points, ) ) for row in read_csv(TEMPERATURE_CSV): unit = parse_unit(row.get("机组", "")) start = parse_time(row.get("开始小时", "")) end = parse_time(row.get("结束小时", "")) metric = (row.get("阶段点位") or "").strip() if not unit or not start or not end or not re.fullmatch(r"YSJ_\d+", metric): continue rows.append( make_segment( source_type="temperature", unit=unit, metric=metric, start=start, end=end + timedelta(hours=1), label_type=(row.get("温度阶段") or "").strip(), label_subtype=(row.get("阶段部位") or row.get("特征类型") or "").strip() or None, row=row, source_file=str(TEMPERATURE_CSV.relative_to(ROOT)), site_points=site_points, ) ) return [row for row in rows if row[17] in SEVERITY] def load_site_points() -> dict[str, str]: with get_connection() as connection: with connection.cursor() as cursor: cursor.execute("SELECT ItemName, ItemDescription FROM site_point") return {str(row["ItemName"]): str(row["ItemDescription"] or "") for row in cursor.fetchall()} def main() -> None: parser = argparse.ArgumentParser(description="导入算法阶段标签到 compressor_label") parser.add_argument("--dry-run", action="store_true", help="只统计，不写库") args = parser.parse_args() ensure_label_table() site_points = load_site_points() rows = build_rows(site_points) counts: dict[str, int] = {} for row in rows: counts[row[0]] = counts.get(row[0], 0) + 1 print(f"待导入 {len(rows)} 条：{counts}") if args.dry_run: return with get_connection() as connection: with connection.cursor() as cursor: cursor.executemany(INSERT_SQL, rows) cursor.execute( "SELECT source_type, COUNT(*) AS n FROM compressor_label " "WHERE origin = 'algorithm' GROUP BY source_type ORDER BY source_type" ) print("库内算法标签：", {row["source_type"]: int(row["n"]) for row in cursor.fetchall()}) if __name__ == "__main__": main()