# MVTec LOCO 前景分割工具

**[中文版 README](README_CN.md)** | **[English README](README.md)**

本工具使用 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 支持**：自动检测并利用可用硬件

## 环境要求

### 环境设置
```bash
# 创建 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](https://drive.google.com/file/d/1ao1ovG1Qtx4b7EoskHXmi2E9rp5CHLcZ/view) 下载 `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/
└── ... (其他类别)
```

## 使用方法

### 基本用法
```bash
# 使用默认设置处理整个数据集
python mvtec_loco_fg_segmentation.py

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

### 高级用法

```bash
# 指定自定义数据集和模型路径
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+张图像

## 故障排除

### 常见问题

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

### 错误信息

- `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 论文：

```bibtex
@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 实现。请参考原始仓库的许可证信息。