hongchen/door_detect/demo.py

"""
YOLOE 门目标检测脚本
基于 yoloe-26x-seg.pt 模型检测全景图中的门/门口目标
"""
import os
import argparse
from pathlib import Path
from tqdm import tqdm
from ultralytics import YOLOE


class DoorDetector:
    """门目标检测器，支持配置和批量处理"""

    # 门相关类别（精简版，覆盖主要语义）
    DOOR_CLASSES = [
        "door",         # 门（通用）
        "doorway",      # 门口
        "entrance",     # 入口
        "door frame",   # 门框
        "open door",    # 打开的门
        "closed door",  # 关闭的门
    ]

    def __init__(
        self,
        model_path: str = "yoloe-26x-seg.pt",
        classes: list = None,
        conf: float = 0.35,
        iou: float = 0.45,
        max_det: int = 50,
    ):
        """
        初始化检测器

        Args:
            model_path: YOLOE 模型路径
            classes: 检测类别列表
            conf: 置信度阈值
            iou: NMS IoU 阈值
            max_det: 最大检测数量
        """
        self.classes = classes or self.DOOR_CLASSES
        self.conf = conf
        self.iou = iou
        self.max_det = max_det

        print(f"加载模型：{model_path}")
        self.model = YOLOE(model_path)
        self.model.set_classes(self.classes)
        print(f"检测类别 ({len(self.classes)}): {self.classes}")

    def detect(self, img_path: str, save_path: str = None) -> dict:
        """
        检测单张图像

        Args:
            img_path: 输入图像路径
            save_path: 结果保存路径（可选）

        Returns:
            检测结果字典 {boxes, masks, scores, class_names}
        """
        if not os.path.exists(img_path):
            raise FileNotFoundError(f"图像不存在：{img_path}")

        results = self.model.predict(
            img_path,
            imgsz=(1024, 2048),
            conf=self.conf,
            iou=self.iou,
            max_det=self.max_det,
            augment=True,
            retina_masks=True,
            half=False,
            verbose=False,
        )

        result = results[0]

        # 提取检测结果
        detection_info = {
            "boxes": [],
            "scores": [],
            "class_names": [],
            "count": 0,
        }

        if result.boxes is not None:
            detection_info["boxes"] = result.boxes.xyxy.cpu().numpy().tolist()
            detection_info["scores"] = result.boxes.conf.cpu().numpy().tolist()
            detection_info["class_names"] = [
                self.model.names[int(cls)]
                for cls in result.boxes.cls.cpu().numpy()
            ]
            detection_info["count"] = len(detection_info["boxes"])

        # 保存可视化结果
        if save_path:
            os.makedirs(os.path.dirname(save_path), exist_ok=True)
            result.save(save_path)
            detection_info["save_path"] = save_path

        return detection_info

    def detect_batch(
        self,
        input_folder: str,
        output_folder: str,
        extensions: tuple = (".jpg", ".jpeg", ".png", ".bmp")
    ) -> list:
        """
        批量检测文件夹中的所有图像

        Args:
            input_folder: 输入文件夹路径
            output_folder: 输出文件夹路径
            extensions: 支持的图像扩展名

        Returns:
            所有检测结果的列表
        """
        input_path = Path(input_folder)
        output_path = Path(output_folder)
        output_path.mkdir(parents=True, exist_ok=True)

        # 获取所有图像文件
        img_files = [
            f for f in input_path.iterdir()
            if f.is_file() and f.suffix.lower() in extensions
        ]

        if not img_files:
            print(f"警告：在 '{input_folder}' 中未找到支持的图像文件")
            return []

        print(f"找到 {len(img_files)} 张图像，开始检测...")

        all_results = []
        for img_file in tqdm(img_files, desc="检测进度"):
            try:
                save_path = output_path / img_file.name
                result = self.detect(str(img_file), str(save_path))
                all_results.append({
                    "image": img_file.name,
                    **result
                })

                # 打印简要结果
                if result["count"] > 0:
                    tqdm.write(
                        f"  {img_file.name}: {result['count']} 个目标 "
                        f"(classes: {', '.join(set(result['class_names']))})"
                    )
                else:
                    tqdm.write(f"  {img_file.name}: 未检测到目标")

            except Exception as e:
                print(f"❌ 处理 {img_file.name} 时出错：{e}")
                all_results.append({
                    "image": img_file.name,
                    "error": str(e),
                    "count": 0
                })

        # 汇总统计
        total_detections = sum(r["count"] for r in all_results)
        images_with_detections = sum(1 for r in all_results if r["count"] > 0)

        print(f"\n{'='*50}")
        print(f"检测完成!")
        print(f"  总图像数：{len(img_files)}")
        print(f"  检测到目标的图像：{images_with_detections}")
        print(f"  总检测数量：{total_detections}")
        print(f"  结果保存至：{output_folder}")
        print(f"{'='*50}")

        return all_results


def main():
    parser = argparse.ArgumentParser(
        description="YOLOE 门目标检测工具",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument(
        "--input", "-i",
        type=str,
        default="image",
        help="输入图像文件夹 (默认：image)"
    )
    parser.add_argument(
        "--output", "-o",
        type=str,
        default="result",
        help="输出结果文件夹 (默认：result)"
    )
    parser.add_argument(
        "--model", "-m",
        type=str,
        default="yoloe-26x-seg.pt",
        help="YOLOE 模型路径 (默认：yoloe-26x-seg.pt)"
    )
    parser.add_argument(
        "--conf",
        type=float,
        default=0.35,
        help="置信度阈值 (默认：0.35)"
    )
    parser.add_argument(
        "--iou",
        type=float,
        default=0.45,
        help="NMS IoU 阈值 (默认：0.45)"
    )
    parser.add_argument(
        "--max-det",
        type=int,
        default=50,
        help="最大检测数量 (默认：50)"
    )
    parser.add_argument(
        "--classes",
        type=str,
        nargs="+",
        default=None,
        help="自定义检测类别列表 (可选)"
    )

    args = parser.parse_args()

    # 创建检测器
    detector = DoorDetector(
        model_path=args.model,
        classes=args.classes,
        conf=args.conf,
        iou=args.iou,
        max_det=args.max_det,
    )

    # 执行批量检测
    results = detector.detect_batch(args.input, args.output)

    return results


if __name__ == "__main__":
    main()