kadupul-git/BridgeDiseaseBackend-main/scripts/fix_model_texts.py

"""修复 model 表 disease_category / augmentation 及 detection 描述中的 ?? 乱码。"""
import os
import sys

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
if ROOT not in sys.path:
    sys.path.insert(0, ROOT)

from app import create_app
from app.config import mysql_uri_with_utf8mb4
from app.constants import DiseaseGrade
from app.models import Detection, Model, db

# 与 sql/init_db.sql 一致
MODEL_CATEGORIES = {
    'steel_corrosion.pt': '钢构件锈蚀',
    'steel_coating_peel_bubble.pt': '钢构件涂层剥脱/鼓包',
    'concrete_peeling_rebar.pt': '混凝土剥落露筋',
    'steel_crack.pt': '钢构件裂缝',
    'concrete_crack.pt': '混凝土裂缝',
    'concrete_weathering.pt': '混凝土风化',
    'road_pothole.pt': '路面坑凼',
}

AUGMENTATION = '随机点+颜色扭曲+高斯模糊'

DESCRIPTIONS = {
    DiseaseGrade.MILD: '隐患程度较轻，建议纳入日常巡检观察。',
    DiseaseGrade.MODERATE: '存在中等程度结构隐患，建议安排专项复核与养护。',
    DiseaseGrade.SEVERE: '隐患较为明显，应尽快组织检测评估并制定处置方案。',
    DiseaseGrade.CRITICAL: '隐患严重，需立即采取限载或封闭措施并启动应急处置。',
}


def main() -> None:
    uri = os.environ.get('SQLALCHEMY_DATABASE_URI')
    if uri:
        os.environ['SQLALCHEMY_DATABASE_URI'] = mysql_uri_with_utf8mb4(uri)

    app = create_app()
    with app.app_context():
        model_updated = 0
        for model in Model.query.all():
            category = MODEL_CATEGORIES.get(model.model_name)
            if not category:
                print(f'[skip] 未知模型: {model.model_name}')
                continue
            changed = False
            if model.disease_category != category:
                model.disease_category = category
                changed = True
            if model.augmentation != AUGMENTATION:
                model.augmentation = AUGMENTATION
                changed = True
            if changed:
                model_updated += 1
                print(f'[model] {model.model_name} -> {category}')

        det_updated = 0
        for det in Detection.query.all():
            model = Model.query.get(det.model_id)
            if not model:
                continue
            grade = det.disease_grade
            if isinstance(grade, str):
                grade = DiseaseGrade(grade.lower())
            desc = (
                f'【{model.disease_category}】{DESCRIPTIONS.get(grade, "暂无评估。")} '
                f'检出隐患 {det.disease_count} 处，覆盖面积约 {int(det.disease_area)} 像素。'
            )
            if det.disease_description != desc:
                det.disease_description = desc
                det_updated += 1

        db.session.commit()
        print(f'完成：更新模型 {model_updated} 条，更新检测描述 {det_updated} 条。')


if __name__ == '__main__':
    main()