from __future__ import annotations import copy import csv import hashlib import json import re import threading import time from datetime import datetime, timedelta from pathlib import Path from typing import Any, Callable from ..db import ensure_label_table, get_connection ROOT = Path(__file__).resolve().parents[3] LABEL_FILES = ( { "key": "vibration", "label": "振动异常段", "relativePath": "models/datas/reference/振动异常段.csv", }, { "key": "oil_pressure", "label": "润滑油压力小时阶段段", "relativePath": "models/datas/reference/润滑油压力小时阶段段.csv", }, { "key": "temperature", "label": "温度异常段", "relativePath": "_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 """ ProgressCallback = Callable[[int, str], None] def _now() -> str: return datetime.now().strftime("%Y-%m-%d %H:%M:%S") def read_csv(path: Path) -> list[dict[str, str]]: if not path.exists(): 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_for_file(source_type: str, path: Path, site_points: dict[str, str]) -> list[tuple[Any, ...]]: rows: list[tuple[Any, ...]] = [] source_file = str(path.relative_to(ROOT)) for row in read_csv(path): 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 if source_type == "vibration": side = (row.get("最强异常侧") or "").strip() metric = VIBRATION_SIDE_TO_METRIC.get(side) if not metric: continue rows.append( make_segment( source_type=source_type, 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=source_file, site_points=site_points, ) ) elif source_type == "oil_pressure": metric = "YSJ_5" if unit in {"7", "8", "9"} else "YSJ_1" rows.append( make_segment( source_type=source_type, unit=unit, metric=metric, start=start, end=end + timedelta(hours=1), label_type=(row.get("压力阶段") or "").strip(), label_subtype="润滑油压力", row=row, source_file=source_file, site_points=site_points, ) ) elif source_type == "temperature": metric = (row.get("阶段点位") or "").strip() if not re.fullmatch(r"YSJ_\d+", metric): continue rows.append( make_segment( source_type=source_type, 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=source_file, site_points=site_points, ) ) return [row for row in rows if row[17] in SEVERITY] def build_rows(site_points: dict[str, str]) -> list[tuple[Any, ...]]: rows: list[tuple[Any, ...]] = [] for item in LABEL_FILES: path = ROOT / item["relativePath"] rows.extend(build_rows_for_file(item["key"], path, site_points)) return rows 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 file_status() -> list[dict[str, Any]]: files: list[dict[str, Any]] = [] for item in LABEL_FILES: path = ROOT / item["relativePath"] exists = path.exists() files.append( { "key": item["key"], "label": item["label"], "relativePath": item["relativePath"], "exists": exists, "size": path.stat().st_size if exists else None, "modifiedAt": ( datetime.fromtimestamp(path.stat().st_mtime).strftime("%Y-%m-%d %H:%M:%S") if exists else None ), } ) return files def run_import(progress: ProgressCallback | None = None) -> dict[str, Any]: started = time.monotonic() ensure_label_table() if progress: progress(5, "准备") site_points = load_site_points() rows: list[tuple[Any, ...]] = [] by_source: dict[str, int] = {} skipped: list[str] = [] for index, item in enumerate(LABEL_FILES): path = ROOT / item["relativePath"] if progress: progress(15 + index * 18, f"读取{item['label']}") if not path.exists(): skipped.append(item["relativePath"]) continue file_rows = build_rows_for_file(item["key"], path, site_points) by_source[item["key"]] = len(file_rows) rows.extend(file_rows) if progress: progress(75, "写入数据库") if rows: with get_connection() as connection: with connection.cursor() as cursor: cursor.executemany(INSERT_SQL, rows) if progress: progress(90, "统计结果") with get_connection() as connection: with connection.cursor() as cursor: cursor.execute( "SELECT COUNT(*) AS n FROM compressor_label WHERE origin = 'algorithm' AND deleted_at IS NULL" ) active = int(cursor.fetchone()["n"]) duration_ms = int((time.monotonic() - started) * 1000) if progress: progress(100, "完成") return { "total": len(rows), "bySource": by_source, "skippedFiles": skipped, "activeLabels": active, "durationMs": duration_ms, } _job_lock = threading.Lock() _job_state: dict[str, Any] = { "running": False, "progress": 0, "stage": "空闲", "message": "", "startedAt": None, "finishedAt": None, "result": None, "error": None, } def job_status() -> dict[str, Any]: with _job_lock: return copy.deepcopy(_job_state) def _update_job(**kwargs: Any) -> None: with _job_lock: _job_state.update(kwargs) def _progress_callback(progress: int, stage: str) -> None: _update_job(progress=progress, stage=stage) def _run_job() -> None: try: result = run_import(progress=_progress_callback) _update_job( running=False, progress=100, stage="完成", result=result, finishedAt=_now(), error=None, ) except Exception as error: # noqa: BLE001 _update_job( running=False, stage="失败", error=str(error), finishedAt=_now(), ) def start_import() -> None: with _job_lock: if _job_state["running"]: raise ValueError("导入正在进行中") _job_state.update( { "running": True, "progress": 0, "stage": "准备", "message": "", "startedAt": _now(), "finishedAt": None, "result": None, "error": None, } ) threading.Thread(target=_run_job, daemon=True).start()