143 lines
4.3 KiB
Python
143 lines
4.3 KiB
Python
#!/usr/bin/env python3
|
|
# -*- coding: utf-8 -*-
|
|
"""
|
|
视频检测系统启动脚本
|
|
简化的用户界面，方便快速开始检测
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
from batch_detector import BatchVideoDetector
|
|
from config import config
|
|
|
|
def check_requirements():
|
|
"""检查运行环境"""
|
|
try:
|
|
import cv2
|
|
import numpy as np
|
|
from ultralytics import YOLO
|
|
print("✓ 所有依赖库已安装")
|
|
return True
|
|
except ImportError as e:
|
|
print(f"✗ 缺少依赖库: {e}")
|
|
print("请运行: pip install ultralytics opencv-python")
|
|
return False
|
|
|
|
def check_model():
|
|
"""检查模型文件"""
|
|
model_path = config.get_model_path()
|
|
if Path(model_path).exists():
|
|
print(f"✓ 模型文件存在: {model_path}")
|
|
return True
|
|
else:
|
|
print(f"✗ 模型文件不存在: {model_path}")
|
|
print("请检查模型路径配置")
|
|
return False
|
|
|
|
def setup_folders():
|
|
"""设置文件夹结构"""
|
|
config.create_folders()
|
|
|
|
input_folder = Path(config.INPUT_FOLDER)
|
|
output_folder = Path(config.OUTPUT_FOLDER)
|
|
|
|
print(f"✓ 输入文件夹: {input_folder.absolute()}")
|
|
print(f"✓ 输出文件夹: {output_folder.absolute()}")
|
|
|
|
if not any(input_folder.iterdir()) if input_folder.exists() else True:
|
|
print(f"⚠ 输入文件夹为空，请将视频文件放入: {input_folder.absolute()}")
|
|
return False
|
|
|
|
return True
|
|
|
|
def get_user_settings():
|
|
"""获取用户设置"""
|
|
print("\n=== 检测参数设置 ===")
|
|
|
|
# 置信度设置
|
|
while True:
|
|
try:
|
|
confidence_input = input(f"请输入置信度阈值 (0.1-0.9, 默认 {config.CONFIDENCE_THRESHOLD}): ").strip()
|
|
if not confidence_input:
|
|
confidence = config.CONFIDENCE_THRESHOLD
|
|
break
|
|
confidence = float(confidence_input)
|
|
if 0.1 <= confidence <= 0.9:
|
|
break
|
|
else:
|
|
print("置信度必须在 0.1 到 0.9 之间")
|
|
except ValueError:
|
|
print("请输入有效的数字")
|
|
|
|
# 输入文件夹
|
|
input_folder = input(f"输入视频文件夹 (默认 '{config.INPUT_FOLDER}'): ").strip()
|
|
if not input_folder:
|
|
input_folder = config.INPUT_FOLDER
|
|
|
|
# 输出文件夹
|
|
output_folder = input(f"输出图片文件夹 (默认 '{config.OUTPUT_FOLDER}'): ").strip()
|
|
if not output_folder:
|
|
output_folder = config.OUTPUT_FOLDER
|
|
|
|
return confidence, input_folder, output_folder
|
|
|
|
def main():
|
|
"""主函数"""
|
|
print("=" * 60)
|
|
print(" 道路损伤检测系统 - 视频批量处理")
|
|
print("=" * 60)
|
|
|
|
# 检查运行环境
|
|
print("\n1. 检查运行环境...")
|
|
if not check_requirements():
|
|
return
|
|
|
|
if not check_model():
|
|
return
|
|
|
|
# 设置文件夹
|
|
print("\n2. 设置文件夹...")
|
|
if not setup_folders():
|
|
input("\n请添加视频文件后重新运行程序。按回车键退出...")
|
|
return
|
|
|
|
# 获取用户设置
|
|
confidence, input_folder, output_folder = get_user_settings()
|
|
|
|
# 确认设置
|
|
print("\n=== 检测设置确认 ===")
|
|
print(f"模型文件: {config.get_model_path()}")
|
|
print(f"置信度阈值: {confidence}")
|
|
print(f"输入文件夹: {Path(input_folder).absolute()}")
|
|
print(f"输出文件夹: {Path(output_folder).absolute()}")
|
|
print(f"检测类别: {list(config.CLASS_NAMES_CN.values())}")
|
|
|
|
confirm = input("\n确认开始检测? (y/n): ").strip().lower()
|
|
if confirm not in ['y', 'yes', '是', '确认']:
|
|
print("检测已取消")
|
|
return
|
|
|
|
# 开始检测
|
|
print("\n=== 开始批量检测 ===")
|
|
try:
|
|
detector = BatchVideoDetector(
|
|
model_path=config.get_model_path(),
|
|
confidence=confidence
|
|
)
|
|
|
|
detector.process_all_videos(input_folder, output_folder)
|
|
|
|
print("\n检测完成! 请查看输出文件夹中的结果。")
|
|
|
|
except KeyboardInterrupt:
|
|
print("\n检测被用户中断")
|
|
except Exception as e:
|
|
print(f"\n检测过程中出现错误: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
|
|
input("\n按回车键退出...")
|
|
|
|
if __name__ == '__main__':
|
|
main() |