Update app.py

app.py

@@ -8,7 +8,7 @@ from mouse_tracker import MouseTrackerAnalyzer
 # 全局变量
 analyzer = None
 video_file_path = None
-model_file_path = None
+model_file_path = "weights/fst-v1.2-n.onnx"  # 直接指定模型文件路径
 total_frames = 0
 output_path = None
 
@@ -41,14 +41,14 @@ def select_video(video_file):
     global video_file_path, total_frames
     
     if not video_file:
-        return None, "Please select a video file", gr.Slider(minimum=0, maximum=0, value=0), gr.Slider(minimum=0, maximum=0, value=0)
+        return None, "请选择视频文件", gr.Slider(minimum=0, maximum=0, value=0), gr.Slider(minimum=0, maximum=0, value=0)
         
     video_file_path = video_file
     
     # 获取视频总帧数
     cap = cv2.VideoCapture(video_file_path)
     if not cap.isOpened():
-        return None, "Cannot open video file", gr.Slider(minimum=0, maximum=0, value=0), gr.Slider(minimum=0, maximum=0, value=0)
+        return None, "无法打开视频文件", gr.Slider(minimum=0, maximum=0, value=0), gr.Slider(minimum=0, maximum=0, value=0)
         
     total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
     
@@ -57,7 +57,7 @@ def select_video(video_file):
     cap.release()
     
     if not ret:
-        return None, "Cannot read video frame", gr.Slider(minimum=0, maximum=0, value=0), gr.Slider(minimum=0, maximum=0, value=0)
+        return None, "无法读取视频帧", gr.Slider(minimum=0, maximum=0, value=0), gr.Slider(minimum=0, maximum=0, value=0)
     
     # 转为RGB
     first_frame_rgb = cv2.cvtColor(first_frame, cv2.COLOR_BGR2RGB)
@@ -66,39 +66,30 @@ def select_video(video_file):
     start_slider = gr.Slider(minimum=0, maximum=total_frames-1, value=0, step=1)
     end_slider = gr.Slider(minimum=0, maximum=total_frames-1, value=total_frames-1, step=1)
     
-    return first_frame_rgb, f"Video loaded successfully, total frames: {total_frames}", start_slider, end_slider
-
-# 选择模型文件
-def select_model(model_file):
-    global model_file_path
-    
-    if model_file is None:
-        return "Please select a model file"
-        
-    model_file_path = model_file
-    return f"Model selected: {os.path.basename(model_file_path)}"
+    model_status = f"使用模型: {os.path.basename(model_file_path)}"
+    return first_frame_rgb, f"视频加载成功，总帧数: {total_frames}. {model_status}", start_slider, end_slider
 
 # 预览帧
 def preview_frame(video_file, frame_num):
     if not video_file:
-        return None, "Please select a video first"
+        return None, "请先选择视频文件"
         
     # 从视频提取帧
     frame = extract_frame(video_file, frame_num)
     if frame is None:
-        return None, "Cannot read specified frame"
+        return None, "无法读取指定帧"
     
-    return frame, f"Frame {frame_num}"
+    return frame, f"帧 {frame_num}"
 
 # 开始分析
-def start_analysis(video, model, conf, iou, max_det, start_frame, end_frame, threshold):
-    global analyzer, output_path
+def start_analysis(video, conf, iou, max_det, start_frame, end_frame, threshold):
+    global analyzer, output_path, model_file_path
     
-    if not video or not model:
-        return None, None, "Please select a video and model file"
+    if not video:
+        return None, None, "请选择视频文件"
         
     if start_frame >= end_frame:
-        return None, None, "Start frame must be less than end frame"
+        return None, None, "起始帧必须小于结束帧"
     
     # 创建输出路径
     video_name = os.path.splitext(os.path.basename(video))[0]
@@ -108,7 +99,7 @@ def start_analysis(video, model, conf, iou, max_det, start_frame, end_frame, thr
     try:
         # 创建分析器
         analyzer = MouseTrackerAnalyzer(
-            model_path=model, 
+            model_path=model_file_path, 
             conf=conf, 
             iou=iou, 
             max_det=max_det, 
@@ -118,11 +109,12 @@ def start_analysis(video, model, conf, iou, max_det, start_frame, end_frame, thr
         
         # 处理视频的进度回调
         def progress_update(progress, frame, results):
-            print(f"Processing: {progress}%, Objects detected: {len(results)}")
+            print(f"处理进度: {progress}%, 检测到对象数: {len(results)}")
             
-        print(f"Processing video: {video}")
-        print(f"Output path: {output_path}")
-        print(f"Parameters: conf={conf}, iou={iou}, max_det={max_det}, threshold={threshold}")
+        print(f"处理视频: {video}")
+        print(f"输出路径: {output_path}")
+        print(f"参数: conf={conf}, iou={iou}, max_det={max_det}, threshold={threshold}")
+        print(f"使用模型: {model_file_path}")
         
         # 提取视频帧数范围并分析
         results = analyzer.process_video(
@@ -134,123 +126,124 @@ def start_analysis(video, model, conf, iou, max_det, start_frame, end_frame, thr
         )
         
         # 保存结果到CSV
-        print(f"Saving results to CSV: {csv_path}")
+        print(f"保存结果到CSV: {csv_path}")
         analyzer.save_results(csv_path)
-        print(f"Results saved to CSV with {len(analyzer.results)} frames of data")
+        print(f"结果已保存到CSV，共 {len(analyzer.results)} 帧数据")
         
         # 生成分析图表
-        print("Generating time series plot...")
+        print("生成时间序列图...")
         if len(analyzer.results) == 0:
-            print("WARNING: No results available for plotting!")
+            print("警告: 没有可用于绘图的结果!")
             plot_path = None
         else:
             plot_path = analyzer.generate_time_series_plot()
             if plot_path and os.path.exists(plot_path):
-                print(f"Plot generated and saved to: {plot_path}, size: {os.path.getsize(plot_path)/1024:.2f}KB")
+                print(f"图表已生成并保存到: {plot_path}, 大小: {os.path.getsize(plot_path)/1024:.2f}KB")
             else:
-                print(f"Failed to generate plot or plot file does not exist!")
+                print(f"生成图表失败或图表文件不存在!")
                 plot_path = None
         
         # 检查输出文件是否存在
         if os.path.exists(output_path):
             file_size = os.path.getsize(output_path) / (1024 * 1024)  # MB
-            print(f"Output video size: {file_size:.2f}MB")
+            print(f"输出视频大小: {file_size:.2f}MB")
             
             # 处理debug帧
             debug_frame_path = os.path.join(os.path.dirname(output_path), "debug_frame.jpg")
             if os.path.exists(debug_frame_path):
-                print(f"Debug frame saved at: {debug_frame_path}")
+                print(f"调试帧保存在: {debug_frame_path}")
         
         if plot_path and os.path.exists(plot_path):
-            print(f"Plot file exists at: {plot_path}, size: {os.path.getsize(plot_path)/1024:.2f}KB")
+            print(f"图表文件存在于: {plot_path}, 大小: {os.path.getsize(plot_path)/1024:.2f}KB")
             
         # 确保返回正确的文件路径
-        status_message = "Analysis complete. "
+        status_message = "分析完成。"
         
         if os.path.exists(output_path):
-            status_message += f"Video saved."
+            status_message += f"视频已保存。"
         else:
-            status_message += "WARNING: Output video not found. "
+            status_message += "警告: 未找到输出视频。"
             
         if plot_path and os.path.exists(plot_path):
-            status_message += f" Time series plot generated."
+            status_message += f" 时间序列图已生成。"
         else:
-            status_message += " WARNING: Failed to generate time series plot."
+            status_message += " 警告: 生成时间序列图失败。"
             
-        status_message += f" Results saved to: {csv_path}"
+        status_message += f" 结果已保存到: {csv_path}"
         
         return output_path, plot_path, status_message
     except Exception as e:
         import traceback
         traceback.print_exc()
-        return None, None, f"Processing error: {str(e)}"
+        return None, None, f"处理错误: {str(e)}"
 
 # 创建Gradio界面
 def create_interface():
-    with gr.Blocks(title="Mouse Struggle Analysis - Object Tracking") as app:
-        gr.Markdown("# Mouse Forced Swim Test Struggle Analysis (Object Tracking)")
+    with gr.Blocks(title="鼠强迫游泳挣扎度分析 - 对象跟踪") as app:
+        gr.Markdown("# 鼠强迫游泳测试挣扎度分析 (对象跟踪)")
         
         with gr.Row():
             with gr.Column(scale=1):
-                # 视频和模型选择
-                video_input = gr.Video(label="Input Video")
-                model_input = gr.File(label="Model File (.pt format recommended)")
+                # 只保留视频选择，移除模型选择
+                video_input = gr.Video(label="输入视频")
+                
+                # 显示当前使用的模型
+                model_info = gr.Textbox(label="模型信息", value=f"使用模型: {os.path.basename(model_file_path)}", interactive=False)
                 
                 # 参数设置
                 with gr.Row():
-                    conf = gr.Slider(minimum=0.1, maximum=0.9, value=0.25, step=0.05, label="Confidence Threshold")
-                    iou = gr.Slider(minimum=0.1, maximum=0.9, value=0.45, step=0.05, label="IoU Threshold")
+                    conf = gr.Slider(minimum=0.1, maximum=0.9, value=0.25, step=0.05, label="置信度阈值")
+                    iou = gr.Slider(minimum=0.1, maximum=0.9, value=0.45, step=0.05, label="IoU阈值")
                 
                 with gr.Row():
-                    max_det = gr.Slider(minimum=1, maximum=50, value=20, step=1, label="Max Detections")
-                    threshold = gr.Slider(minimum=0, maximum=1, value=0.3, step=0.01, label="Struggle Threshold")
+                    max_det = gr.Slider(minimum=1, maximum=50, value=20, step=1, label="最大检测数")
+                    threshold = gr.Slider(minimum=0, maximum=1, value=0.3, step=0.01, label="挣扎阈值")
                 
                 # 帧选择
-                start_frame = gr.Slider(minimum=0, maximum=999999, value=0, step=1, label="Start Frame")
-                end_frame = gr.Slider(minimum=0, maximum=999999, value=999999, step=1, label="End Frame")
+                start_frame = gr.Slider(minimum=0, maximum=999999, value=0, step=1, label="起始帧")
+                end_frame = gr.Slider(minimum=0, maximum=999999, value=999999, step=1, label="结束帧")
                 
                 # 预览按钮
-                preview_btn = gr.Button("Preview Frame")
+                preview_btn = gr.Button("预览帧")
                 
                 # 开始分析
-                start_btn = gr.Button("Start Analysis", variant="primary")
+                start_btn = gr.Button("开始分析", variant="primary")
                 
             with gr.Column(scale=2):
                 # 显示区域
-                with gr.Tab("Preview"):
+                with gr.Tab("预览"):
                     # 图像预览
-                    preview_image = gr.Image(label="Preview Image", type="numpy", height=400)
-                    status_text = gr.Textbox(label="Status", interactive=False)
+                    preview_image = gr.Image(label="预览图像", type="numpy", height=400)
+                    status_text = gr.Textbox(label="状态", interactive=False)
                     gr.Markdown("""
-                    ### Instructions:
-                    1. Select a video and model file (.pt format segmentation model like yolov8n-seg.pt recommended)
-                    2. Adjust parameters
-                       - Confidence Threshold: Minimum confidence for object detection, lower values detect more potential objects
-                       - IoU Threshold: For filtering overlapping detections
-                       - Max Detections: Maximum number of objects to detect per frame
-                       - Struggle Threshold: Minimum score to classify as struggle state
-                    3. Set frame range
-                    4. Click "Start Analysis" button
+                    ### 使用说明:
+                    1. 选择一个视频文件
+                    2. 调整参数
+                       - 置信度阈值: 对象检测的最低置信度，较低的值会检测更多潜在对象
+                       - IoU阈值: 用于过滤重叠检测
+                       - 最大检测数: 每帧检测的最大对象数
+                       - 挣扎阈值: 分类为挣扎状态的最低分数
+                    3. 设置帧范围
+                    4. 点击"开始分析"按钮
                     
-                    The system will automatically track mice and analyze their struggle behavior, no need to manually define regions
+                    系统将自动跟踪小鼠并分析其挣扎行为，无需手动定义区域
                     """)
                 
-                with gr.Tab("Results"):
+                with gr.Tab("结果"):
                     with gr.Row():
-                        output_video = gr.Video(label="Analysis Result Video")
-                        result_plot = gr.Image(label="Struggle Score Time Series")
+                        output_video = gr.Video(label="分析结果视频")
+                        result_plot = gr.Image(label="挣扎分数时间序列")
                     
-                    result_status = gr.Textbox(label="Analysis Status", interactive=False)
+                    result_status = gr.Textbox(label="分析状态", interactive=False)
         
         # 绑定事件
         video_input.change(select_video, inputs=[video_input], outputs=[preview_image, status_text, start_frame, end_frame])
-        model_input.change(select_model, inputs=[model_input], outputs=[status_text])
         
         preview_btn.click(preview_frame, inputs=[video_input, start_frame], outputs=[preview_image, status_text])
         
         start_btn.click(
             start_analysis, 
-            inputs=[video_input, model_input, conf, iou, max_det, start_frame, end_frame, threshold], 
+            inputs=[video_input, conf, iou, max_det, start_frame, end_frame, threshold], 
             outputs=[output_video, result_plot, result_status]
         )
     
@@ -258,6 +251,19 @@ def create_interface():
 
 # 启动应用
 if __name__ == "__main__":
+    # 清除可能干扰的代理设置
+    if 'http_proxy' in os.environ:
+        del os.environ['http_proxy']
+    if 'https_proxy' in os.environ:
+        del os.environ['https_proxy']
+    if 'all_proxy' in os.environ:
+        del os.environ['all_proxy']
+    
+    # 检查模型文件是否存在
+    if not os.path.exists(model_file_path):
+        print(f"警告: 模型文件 {model_file_path} 不存在!")
+    else:
+        print(f"使用模型: {model_file_path}")
     
     app = create_interface()
     # 使用简化的启动配置