MVTec LOCO 前景分割工具

中文版 README | English README

本工具使用 U²-Net 为 MVTec LOCO 异常检测数据集生成二值前景掩码。

概述

mvtec_loco_fg_segmentation.py 脚本处理整个 MVTec LOCO 数据集，为所有图像生成二值前景掩码。它使用 U²-Net 模型执行显著性目标检测，并将概率图转换为二值掩码。

功能特点

完整数据集处理：处理所有类别（breakfast_box, screw_bag, juice_bottle, splicing_connectors, pushpins）
灵活的结构支持：处理 test 和 train 分割中的所有子目录（good, logical_anomalies, structural_anomalies）
二值掩码输出：生成标准的二值掩码（0/255），使用 L 模式（灰度）
参数可配置：可自定义阈值、类别、分割和处理选项
GPU/CPU 支持：自动检测并利用可用硬件

环境要求

环境设置

# 创建 conda 环境
conda create -n u2net python=3.8 -y
conda activate u2net

# 安装依赖
pip install torch==1.13.1+cu116 torchvision==0.14.1+cu116 --index-url https://download.pytorch.org/whl/cu116
pip install opencv-python scikit-image matplotlib numpy pillow

模型权重

下载 U²-Net 模型权重并放置到正确目录：

从 GoogleDrive 下载 u2net.pth（176.3 MB）
放置到：./saved_models/u2net/u2net.pth

数据集结构

确保您的 MVTec LOCO 数据集遵循以下结构：

mvtec_loco_anomaly_detection/
├── breakfast_box/
│   ├── test/
│   │   ├── good/
│   │   ├── logical_anomalies/
│   │   └── structural_anomalies/
│   └── train/
│       └── good/
├── screw_bag/
│   ├── test/
│   └── train/
└── ... (其他类别)

使用方法

基本用法

# 使用默认设置处理整个数据集
python mvtec_loco_fg_segmentation.py

# 显示帮助信息
python mvtec_loco_fg_segmentation.py -h

高级用法

# 指定自定义数据集和模型路径
python mvtec_loco_fg_segmentation.py \
    --dataset_path /path/to/mvtec_loco \
    --model_path /path/to/u2net.pth

# 仅处理特定类别
python mvtec_loco_fg_segmentation.py \
    --categories breakfast_box juice_bottle

# 仅处理测试集
python mvtec_loco_fg_segmentation.py \
    --splits test

# 使用不同的二值化阈值
python mvtec_loco_fg_segmentation.py \
    --threshold 0.3

# 自定义输出目录名称
python mvtec_loco_fg_segmentation.py \
    --output_dir custom_masks

# 使用多进程优化处理速度
python mvtec_loco_fg_segmentation.py \
    --num_workers 4 \
    --batch_size 4

命令行参数

参数	类型	默认值	说明
`--dataset_path`	str	`/root/hy-data/datasets/mvtec_loco_anomaly_detection`	MVTec LOCO 数据集根目录路径
`--model_path`	str	`./saved_models/u2net/u2net.pth`	U2NET 模型权重文件路径
`--output_dir`	str	`fg_mask`	生成掩码的输出目录名称
`--threshold`	float	`0.5`	二值掩码生成阈值
`--categories`	list	`所有5个类别`	要处理的类别
`--splits`	list	`['test', 'train']`	要处理的数据集分割
`--batch_size`	int	`1`	处理批次大小
`--num_workers`	int	`1`	数据加载工作进程数

输出结构

脚本生成以下结构的掩码：

mvtec_loco_anomaly_detection/
├── fg_mask/                    # 生成的掩码目录
│   ├── breakfast_box/
│   │   ├── test/
│   │   │   ├── good/
│   │   │   │   ├── 000.png     # 二值掩码（0/255 值）
│   │   │   │   ├── 001.png
│   │   │   │   └── ...
│   │   │   ├── logical_anomalies/
│   │   │   └── structural_anomalies/
│   │   └── train/
│   │       └── good/
│   └── ... (其他类别)
└── ... (原始数据集)

掩码属性

格式：PNG 图像
模式：L（灰度，单通道）
值：二值（背景为0，前景为255）
尺寸：与原始图像相同
阈值：可配置（默认0.5）

性能说明

推荐GPU：使用支持CUDA的GPU处理速度显著更快
内存使用：每张图像在处理期间需要约200MB GPU内存
处理时间：现代GPU上每张图像约2-3秒
总图像数：完整数据集中约5000+张图像

故障排除

常见问题

CUDA内存不足：减少批次大小或使用CPU处理
找不到模型：确保u2net.pth在正确目录中
数据集路径错误：验证MVTec LOCO数据集结构
权限错误：检查输出目录的写入权限

错误信息

ERROR: Dataset path not found：检查数据集路径和提取情况
ERROR: Model path not found：正确下载并放置u2net.pth
ERROR: Invalid categories：使用有效的类别名称

示例输出

脚本提供详细的进度信息：

配置信息:
  数据集路径: /root/hy-data/datasets/mvtec_loco_anomaly_detection
  模型路径: ./saved_models/u2net/u2net.pth
  输出目录: fg_mask
  二值化阈值: 0.5
  类别: ['breakfast_box', 'screw_bag', 'juice_bottle', 'splicing_connectors', 'pushpins']
  分割: ['test', 'train']

...加载 U2NET---
处理类别: breakfast_box
  处理 breakfast_box/test/good
    找到 102 张图像
    处理 1/102: 000.png
    处理 20/102: 019.png
    ...

引用

如果您在研究中使用此工具，请引用原始的 U²-Net 论文：

@InProceedings{Qin_2020_PR,
  title = {U2-Net: Going Deeper with Nested U-Structure for Salient Object Detection},
  author = {Qin, Xuebin and Zhang, Zichen and Huang, Chenyang and Dehghan, Masood and Zaiane, Osmar and Jagersand, Martin},
  journal = {Pattern Recognition},
  volume = {106},
  pages = {107404},
  year = {2020}
}

许可证

本工具扩展了原始 U²-Net 实现。请参考原始仓库的许可证信息。