import gradio as gr import pandas as pd from pathlib import Path import os import matplotlib.pyplot as plt # 导入所有模块中的核心函数和配置 from utils import load_images_from_folder_for_import, save_data_to_data_dir, \ load_training_data_for_management, \ add_new_image_entry, delete_image_entry, update_data_from_management_dataframe, \ _update_managed_display_from_index, _select_managed_image_for_edit, _navigate_managed_image, \ _process_managed_single_score_edit from train_predict import TrainingAndPredictionEngine from config import DEFAULT_EPOCHS, DEFAULT_LR, MODEL_SAVE_BASE_PATH, \ DEFAULT_DROPOUT_RATE, DEFAULT_WEIGHT_DECAY, DEFAULT_PCA_VARIANCE_RATIO, \ DEFAULT_OPTIMIZER, DEFAULT_LR_SCHEDULER, DEFAULT_SCHEDULER_PATIENCE, \ DEFAULT_SCHEDULER_FACTOR, DEFAULT_SCHEDULER_T_MAX, DEFAULT_LOSS_FUNCTION, \ DEFAULT_EARLY_STOPPING_PATIENCE, DEFAULT_BATCH_SIZE, DEFAULT_DATA_AUGMENTATION, DATA_DIR # --- UI辅助函数 (原始数据导入 Tab 专用) --- def _update_main_display_from_index(index, all_image_paths, all_scores): if not all_image_paths or not (0 <= index < len(all_image_paths)): return "", None, None preview_path = all_image_paths[index] score = all_scores[index] image_name = Path(preview_path).name return image_name, preview_path, score def _select_image_for_edit(evt: gr.SelectData, current_data_state): all_image_paths, all_scores, _ = current_data_state if all_image_paths and 0 <= evt.index < len(all_image_paths): current_data_state = (all_image_paths, all_scores, evt.index) selected_name, selected_path, selected_score = _update_main_display_from_index(evt.index, all_image_paths, all_scores) return selected_name, selected_path, selected_score, current_data_state return "", None, None, current_data_state def _navigate_image(direction, current_data_state): all_image_paths, all_scores, current_index = current_data_state if not all_image_paths: return "", None, None, current_data_state num_images = len(all_image_paths) if num_images == 0: return "", None, None, current_data_state new_index = current_index + direction if new_index < 0: new_index = num_images - 1 elif new_index >= num_images: new_index = 0 current_data_state = (all_image_paths, all_scores, new_index) selected_name, selected_path, selected_score = _update_main_display_from_index(new_index, all_image_paths, all_scores) return selected_name, selected_path, selected_score, current_data_state def _process_single_score_edit(entered_score, current_data_state): all_image_paths, all_scores, selected_index = current_data_state if entered_score is not None: final_score = max(0, min(100, round(entered_score))) else: final_score = 50 if selected_index != -1 and 0 <= selected_index < len(all_scores): all_scores[selected_index] = final_score current_data_state = (all_image_paths, all_scores, selected_index) dataframe_data = pd.DataFrame( [[Path(img_path_str).name, score] for img_path_str, score in zip(all_image_paths, all_scores)], columns=["文件名", "分数"]) return dataframe_data, current_data_state dataframe_data_current = pd.DataFrame() if all_image_paths and all_scores: dataframe_data_current = pd.DataFrame([[Path(p).name, s] for p, s in zip(all_image_paths, all_scores)], columns=["文件名", "分数"]) return dataframe_data_current, current_data_state def _create_static_error_plot(message): fig, ax = plt.subplots(figsize=(8, 4)) ax.text(0.5, 0.5, message, horizontalalignment='center', verticalalignment='center', transform=ax.transAxes, fontsize=12, color='red') ax.set_title("训练失败") ax.axis('off') plt.tight_layout() # 返回两个图表，因为 train_model 现在返回两个 # 创建一个空的MSE图 fig_mse, ax_mse = plt.subplots(figsize=(8, 4)) ax_mse.text(0.5, 0.5, "MSE图表：训练失败", horizontalalignment='center', verticalalignment='center', transform=ax_mse.transAxes, fontsize=12, color='red') ax_mse.axis('off') plt.tight_layout(fig_mse) return fig, fig_mse # --- 全局状态和初始化 --- training_engine = TrainingAndPredictionEngine() # 可选的基础CNN模型列表 BASE_CNN_MODELS = ["resnet50", "resnet18", "vgg16", "densenet121", "efficientnet_b0", "inception_v3"] OPTIMIZERS = ["Adam", "SGD", "AdamW"] LR_SCHEDULERS = ["None", "ReduceLROnPlateau", "CosineAnnealingLR"] LOSS_FUNCTIONS = ["MSELoss", "L1Loss", "SmoothL1Loss"] # --- Gradio UI界面定义 --- with gr.Blocks(title="图片分数编辑与模型训练") as demo: current_import_data_state = gr.State(([], [], -1)) current_managed_data_state = gr.State(([], [], -1)) gr.Markdown("## 图片分数编辑与模型训练系统") gr.Markdown("### 使用流程：") gr.Markdown( "1. **原始数据导入**: 从外部文件夹加载图片，解析文件名中的分数进行初始标注，并可预览、编辑。点击“保存数据”将图片复制到 `data` 目录并生成 `scores.txt`。") gr.Markdown( "2. **训练数据管理**: 浏览、编辑、添加、删除已保存到 `data` 目录的训练图片和分数。此处的更改会立即同步到文件。") gr.Markdown("3. **训练模型**: 选择模型类型和基础CNN模型，调整参数，开始训练。") gr.Markdown("4. **预测图片**: 上传图片，选择训练好的模型进行预测。") with gr.Tab("原始数据导入"): gr.Markdown("### 从外部文件夹导入图片并进行初始分数编辑") with gr.Row(): folder_input = gr.Textbox(label="输入图片文件夹路径", placeholder="请输入绝对路径，例如 D:\\images", value="") load_btn = gr.Button("加载图片", variant="primary", scale=0) import_status_text = gr.Textbox(label="操作状态", interactive=False, max_lines=2) with gr.Column(): gallery = gr.Gallery( label="图片列表", show_label=True, elem_id="gallery", height=600, preview=True, columns=5, rows=2, object_fit="contain" ) with gr.Group(): gr.Markdown("#### 当前图片信息与分数编辑") with gr.Row(): current_image_name_display = gr.Textbox( label="当前图片名称", interactive=False, show_label=True, placeholder="图片名称将在此处显示...", scale=3 ) current_score_input = gr.Number( label="当前分数 (0-100)", scale=1 ) confirm_score_btn = gr.Button("确定", variant="secondary", scale=0) with gr.Row(): prev_btn = gr.Button("⬅️ 上一张", scale=1) next_btn = gr.Button("下一张 ➡️", scale=1) current_image_preview = gr.Image( label="当前选中图片（预览）", height=300, show_download_button=False, container=False, interactive=False ) import_score_dataframe = gr.Dataframe( headers=["文件名", "分数"], datatype=["str", "number"], col_count=(2, "fixed"), interactive=True, label="所有图片分数表格 (可直接编辑)", value=[] ) save_import_data_btn = gr.Button("保存数据到训练目录", variant="secondary") with gr.Tab("训练数据管理"): gr.Markdown("### 管理已保存的训练图片和分数") gr.Markdown(f"当前训练数据目录: `{DATA_DIR}`") load_managed_data_btn = gr.Button("加载训练数据", variant="primary") managed_data_status_text = gr.Textbox(label="管理状态", interactive=False, max_lines=2) with gr.Column(): managed_gallery = gr.Gallery( label="已保存图片列表", show_label=True, elem_id="managed_gallery", height=400, preview=True, columns=5, rows=2, object_fit="contain", value=[] ) with gr.Group(): gr.Markdown("#### 当前图片信息与分数编辑 (训练数据)") with gr.Row(): managed_image_name_display = gr.Textbox( label="当前图片名称", interactive=False, show_label=True, placeholder="图片名称将在此处显示...", scale=3 ) managed_score_input = gr.Number( label="当前分数 (0-100)", scale=1 ) confirm_managed_score_btn = gr.Button("确定", variant="secondary", scale=0) with gr.Row(): prev_managed_btn = gr.Button("⬅️ 上一张", scale=1) next_managed_btn = gr.Button("下一张 ➡️", scale=1) managed_image_preview = gr.Image( label="当前选中图片（预览）", height=200, show_download_button=False, container=False, interactive=False ) managed_score_dataframe = gr.Dataframe( headers=["文件名", "分数"], datatype=["str", "number"], col_count=(2, "fixed"), interactive=True, label="所有图片分数表格 (可直接编辑)", value=[] ) with gr.Accordion("添加/删除图片条目", open=False): with gr.Row(): new_image_file_input = gr.File(label="上传新图片文件", type="filepath") new_image_name_input = gr.Textbox(label="或输入文件名 (例如：85.jpg)", placeholder="若上传文件可留空", scale=1) new_score_input = gr.Number(label="新分数 (0-100)", value=50) add_entry_btn = gr.Button("添加条目", variant="secondary", scale=0) with gr.Row(): delete_filename_input = gr.Textbox(label="要删除的文件名", placeholder="请输入完整文件名，例如 85.jpg") delete_entry_btn = gr.Button("删除条目", variant="stop", scale=0) with gr.Tab("训练模型"): gr.Markdown("## 模型训练") gr.Markdown( "1. **重要**: 确保在“原始数据导入”或“训练数据管理”标签页已加载并**保存**了数据。模型将使用 `data` 文件夹中的图片和分数。") gr.Markdown("2. 选择模型类型和基础CNN模型，调整参数，开始训练。") gr.Markdown("3. 点击“开始训练”按钮。训练过程中的损失曲线会实时显示。") with gr.Row(): model_type_selector = gr.Dropdown( ["深度学习", "端到端深度学习", "随机森林", "支持向量回归", "梯度提升回归", "堆叠回归", "K近邻", "线性回归"], label="选择模型类型", value="深度学习", interactive=True ) epochs_input = gr.Slider(1, 100, DEFAULT_EPOCHS, label="训练轮次", step=1) lr_input = gr.Number(DEFAULT_LR, label="学习率", precision=5) base_cnn_selector_train = gr.Dropdown( BASE_CNN_MODELS, label="选择基础CNN模型 (深度学习模式)", value="resnet50", interactive=True, visible=True ) with gr.Accordion("高级训练参数", open=False): with gr.Row(): batch_size_input = gr.Slider( 1, 128, DEFAULT_BATCH_SIZE, label="批量大小 (Batch Size)", step=1, interactive=True, visible=True ) dropout_rate_input = gr.Slider( 0.0, 0.9, DEFAULT_DROPOUT_RATE, label="Dropout 比率", step=0.05, interactive=True, visible=True ) with gr.Row(): weight_decay_input = gr.Number( DEFAULT_WEIGHT_DECAY, label="权重衰减 (L2正则化)", precision=7, interactive=True, visible=True ) pca_variance_ratio_input = gr.Slider( 0.7, 1.0, DEFAULT_PCA_VARIANCE_RATIO, label="PCA保留方差比例 (仅Sklearn)", step=0.01, interactive=True, visible=False ) with gr.Row(): optimizer_selector = gr.Dropdown( OPTIMIZERS, label="优化器", value=DEFAULT_OPTIMIZER, interactive=True, visible=True ) loss_function_selector = gr.Dropdown( LOSS_FUNCTIONS, label="损失函数", value=DEFAULT_LOSS_FUNCTION, interactive=True, visible=True ) with gr.Row(): lr_scheduler_selector = gr.Dropdown( LR_SCHEDULERS, label="学习率调度器", value=DEFAULT_LR_SCHEDULER, interactive=True, visible=True ) scheduler_patience_input = gr.Slider( 1, 20, DEFAULT_SCHEDULER_PATIENCE, label="调度器耐心 (ReduceLROnPlateau)", step=1, interactive=True, visible=True ) with gr.Row(): scheduler_factor_input = gr.Slider( 0.01, 0.5, DEFAULT_SCHEDULER_FACTOR, label="调度器因子 (ReduceLROnPlateau)", step=0.01, interactive=True, visible=True ) scheduler_t_max_input = gr.Slider( 1, 200, DEFAULT_SCHEDULER_T_MAX, label="调度器T_max (CosineAnnealingLR)", step=1, interactive=True, visible=True ) with gr.Row(): early_stopping_patience_input = gr.Slider( 1, 20, DEFAULT_EARLY_STOPPING_PATIENCE, label="早停耐心值", step=1, interactive=True, visible=True ) with gr.Row(): enable_augmentation_checkbox = gr.Checkbox( DEFAULT_DATA_AUGMENTATION, label="开启数据增强", interactive=True, visible=True ) train_start_btn = gr.Button("开始训练", variant="primary") train_status_text = gr.Textbox(label="训练状态", interactive=False) with gr.Row(): loss_plot_output = gr.Plot(label="训练与验证损失", scale=1) # 标签修改 metrics_plot_output = gr.Plot(label="验证MSE与MAE", scale=1) # 标签修改 with gr.Tab("预测图片"): gr.Markdown("## 图片预测") gr.Markdown("1. 选择之前训练过的模型类型。") gr.Markdown("2. 上传您要预测的图片。") gr.Markdown("3. 点击“预测”按钮获取分数。") with gr.Row(): predict_model_type_selector = gr.Dropdown( ["深度学习", "端到端深度学习", "随机森林", "支持向量回归", "梯度提升回归", "堆叠回归", "K近邻", "线性回归"], label="选择预测模型类型", value="深度学习", interactive=True ) image_for_predict = gr.Image(type="filepath", label="上传图片进行预测") predict_btn = gr.Button("预测", variant="primary") predicted_score_output = gr.Label(label="预测结果") base_cnn_selector_predict = gr.Dropdown( BASE_CNN_MODELS, label="选择基础CNN模型 (深度学习模式)", value="resnet50", interactive=True, visible=True ) # --- 事件绑定 --- # 1. 原始数据导入 Tab 的事件 load_btn.click( fn=load_images_from_folder_for_import, inputs=[folder_input], outputs=[ gallery, import_status_text, current_image_name_display, current_image_preview, current_score_input, current_import_data_state ] ).success( fn=lambda state_tuple_from_load: pd.DataFrame( [[Path(p).name, s] for p, s in zip(state_tuple_from_load[0], state_tuple_from_load[1])], columns=["文件名", "分数"] ), inputs=[current_import_data_state], outputs=import_score_dataframe ) gallery.select( fn=_select_image_for_edit, inputs=[current_import_data_state], outputs=[current_image_name_display, current_image_preview, current_score_input, current_import_data_state] ) prev_btn.click( fn=_navigate_image, inputs=[gr.State(-1), current_import_data_state], outputs=[current_image_name_display, current_image_preview, current_score_input, current_import_data_state] ) next_btn.click( fn=_navigate_image, inputs=[gr.State(1), current_import_data_state], outputs=[current_image_name_display, current_image_preview, current_score_input, current_import_data_state] ) confirm_score_btn.click( fn=_process_single_score_edit, inputs=[current_score_input, current_import_data_state], outputs=[import_score_dataframe, current_import_data_state] ) current_score_input.submit( fn=_process_single_score_edit, inputs=[current_score_input, current_import_data_state], outputs=[import_score_dataframe, current_import_data_state] ) import_score_dataframe.change( fn=lambda dataframe_data, state: (state[0], [row["分数"] for _, row in dataframe_data.iterrows()], state[2]), inputs=[import_score_dataframe, current_import_data_state], outputs=[current_import_data_state] ) save_import_data_btn.click( fn=save_data_to_data_dir, inputs=[current_import_data_state], outputs=[import_status_text, current_managed_data_state] ).then( fn=load_training_data_for_management, inputs=[], outputs=[managed_score_dataframe, managed_data_status_text, current_managed_data_state, managed_gallery, gr.State(current_managed_data_state.value[1])] ).then( fn=lambda all_image_paths, all_scores, index: _update_managed_display_from_index(index, all_image_paths, all_scores), inputs=[gr.State(current_managed_data_state.value[0]), gr.State(current_managed_data_state.value[1]), gr.State(current_managed_data_state.value[2])], outputs=[managed_image_name_display, managed_image_preview, managed_score_input, delete_filename_input] ) # 2. 训练数据管理 Tab 的事件 load_managed_data_btn.click( fn=load_training_data_for_management, inputs=[], outputs=[managed_score_dataframe, managed_data_status_text, current_managed_data_state, managed_gallery, gr.State(current_managed_data_state.value[1])] ).then( fn=lambda all_image_paths, all_scores, index: _update_managed_display_from_index(index, all_image_paths, all_scores), inputs=[gr.State(current_managed_data_state.value[0]), gr.State(current_managed_data_state.value[1]), gr.State(current_managed_data_state.value[2])], outputs=[managed_image_name_display, managed_image_preview, managed_score_input, delete_filename_input] ) managed_gallery.select( fn=_select_managed_image_for_edit, inputs=[current_managed_data_state], outputs=[managed_image_name_display, managed_image_preview, managed_score_input, delete_filename_input, current_managed_data_state] ) prev_managed_btn.click( fn=_navigate_managed_image, inputs=[gr.State(-1), current_managed_data_state], outputs=[managed_image_name_display, managed_image_preview, managed_score_input, delete_filename_input, current_managed_data_state] ) next_managed_btn.click( fn=_navigate_managed_image, inputs=[gr.State(1), current_managed_data_state], outputs=[managed_image_name_display, managed_image_preview, managed_score_input, delete_filename_input, current_managed_data_state] ) confirm_managed_score_btn.click( fn=_process_managed_single_score_edit, inputs=[managed_score_input, current_managed_data_state], outputs=[managed_data_status_text, managed_score_dataframe, managed_gallery, gr.State(current_managed_data_state.value[1]), current_managed_data_state] ) managed_score_input.submit( fn=_process_managed_single_score_edit, inputs=[managed_score_input, current_managed_data_state], outputs=[managed_data_status_text, managed_score_dataframe, managed_gallery, gr.State(current_managed_data_state.value[1]), current_managed_data_state] ) managed_score_dataframe.change( fn=update_data_from_management_dataframe, inputs=[managed_score_dataframe, current_managed_data_state], outputs=[managed_gallery, gr.State(current_managed_data_state.value[1]), current_managed_data_state] ).success( fn=lambda: "表格更新并保存成功！", inputs=[], outputs=[managed_data_status_text] ) add_entry_btn.click( fn=add_new_image_entry, inputs=[new_image_file_input, new_image_name_input, new_score_input, current_managed_data_state], outputs=[managed_data_status_text, managed_score_dataframe, new_image_file_input, managed_gallery, gr.State(current_managed_data_state.value[1]), current_managed_data_state] ).then( fn=lambda all_image_paths, all_scores, index: _update_managed_display_from_index(index, all_image_paths, all_scores), inputs=[gr.State(current_managed_data_state.value[0]), gr.State(current_managed_data_state.value[1]), gr.State(current_managed_data_state.value[2])], outputs=[managed_image_name_display, managed_image_preview, managed_score_input, delete_filename_input] ) delete_entry_btn.click( fn=delete_image_entry, inputs=[delete_filename_input, current_managed_data_state], outputs=[managed_data_status_text, managed_score_dataframe, delete_filename_input, managed_gallery, gr.State(current_managed_data_state.value[1]), current_managed_data_state] ).then( fn=lambda all_image_paths, all_scores, index: _update_managed_display_from_index(index, all_image_paths, all_scores), inputs=[gr.State(current_managed_data_state.value[0]), gr.State(current_managed_data_state.value[1]), gr.State(current_managed_data_state.value[2])], outputs=[managed_image_name_display, managed_image_preview, managed_score_input, delete_filename_input] ) # 3. 训练模型 Tab 的事件 def update_advanced_params_visibility(model_type): is_pytorch_model = model_type in ["深度学习", "端到端深度学习"] is_sklearn_model = model_type in ["随机森林", "支持向量回归", "梯度提升回归", "堆叠回归", "K近邻", "线性回归"] return { base_cnn_selector_train: gr.update(visible=is_pytorch_model or is_sklearn_model), lr_input: gr.update(visible=is_pytorch_model), epochs_input: gr.update(visible=is_pytorch_model), batch_size_input: gr.update(visible=is_pytorch_model), dropout_rate_input: gr.update(visible=is_pytorch_model), weight_decay_input: gr.update(visible=is_pytorch_model), optimizer_selector: gr.update(visible=is_pytorch_model), loss_function_selector: gr.update(visible=is_pytorch_model), lr_scheduler_selector: gr.update(visible=is_pytorch_model), scheduler_patience_input: gr.update( visible=is_pytorch_model and lr_scheduler_selector.value == "ReduceLROnPlateau"), scheduler_factor_input: gr.update( visible=is_pytorch_model and lr_scheduler_selector.value == "ReduceLROnPlateau"), scheduler_t_max_input: gr.update( visible=is_pytorch_model and lr_scheduler_selector.value == "CosineAnnealingLR"), early_stopping_patience_input: gr.update(visible=is_pytorch_model), enable_augmentation_checkbox: gr.update(visible=is_pytorch_model or is_sklearn_model), pca_variance_ratio_input: gr.update(visible=is_sklearn_model) } model_type_selector.change( fn=update_advanced_params_visibility, inputs=[model_type_selector], outputs=[ base_cnn_selector_train, lr_input, epochs_input, batch_size_input, dropout_rate_input, weight_decay_input, optimizer_selector, loss_function_selector, lr_scheduler_selector, scheduler_patience_input, scheduler_factor_input, scheduler_t_max_input, early_stopping_patience_input, enable_augmentation_checkbox, pca_variance_ratio_input ] ) lr_scheduler_selector.change( fn=lambda scheduler_name: { scheduler_patience_input: gr.update(visible=scheduler_name == "ReduceLROnPlateau"), scheduler_factor_input: gr.update(visible=scheduler_name == "ReduceLROnPlateau"), scheduler_t_max_input: gr.update(visible=scheduler_name == "CosineAnnealingLR") }, inputs=[lr_scheduler_selector], outputs=[ scheduler_patience_input, scheduler_factor_input, scheduler_t_max_input ] ) data_prep_success_flag_state = gr.State(False) train_start_btn.click( fn=training_engine.switch_model_type, inputs=[ model_type_selector, base_cnn_selector_train, dropout_rate_input, weight_decay_input, pca_variance_ratio_input, optimizer_selector, lr_scheduler_selector, scheduler_patience_input, scheduler_factor_input, scheduler_t_max_input, loss_function_selector, early_stopping_patience_input, batch_size_input, enable_augmentation_checkbox ], outputs=[train_status_text] ).then( fn=training_engine.prepare_data_for_training, inputs=[], outputs=[data_prep_success_flag_state, train_status_text] ).then( fn=lambda success_flag, epochs, lr: ( training_engine.train_model(epochs, lr) if success_flag else _create_static_error_plot("数据准备失败或无数据，无法训练") ), inputs=[data_prep_success_flag_state, epochs_input, lr_input], outputs=[loss_plot_output, metrics_plot_output] ) # 4. 预测图片 Tab 的事件 def update_predict_cnn_visibility(model_type): is_cnn_based = model_type in ["深度学习", "端到端深度学习", "随机森林", "支持向量回归", "梯度提升回归", "堆叠回归", "K近邻", "线性回归"] return gr.update(visible=is_cnn_based) predict_model_type_selector.change( fn=update_predict_cnn_visibility, inputs=[predict_model_type_selector], outputs=[base_cnn_selector_predict] ) predict_btn.click( fn=training_engine.predict_score, inputs=[image_for_predict, predict_model_type_selector, base_cnn_selector_predict], outputs=predicted_score_output ) # 启动 Gradio 应用 if __name__ == "__main__": Path(DATA_DIR).mkdir(exist_ok=True, parents=True) Path(MODEL_SAVE_BASE_PATH).parent.mkdir(exist_ok=True, parents=True) demo.launch( share=True, server_port=7860, show_error=True, debug=True )