Add complete U2Net project with HuggingFace preparation

- Add comprehensive .gitignore for Python projects
- Configure Git LFS for large files including:
- PyTorch model files (.pth): 338MB total
- Image datasets and results
- ZIP archives and media files
- Include all Python source code, documentation, and test data
- Ready for HuggingFace Hub deployment

🤖 Generated with [Claude Code](https://claude.ai/code)

Co-Authored-By: Claude <[email protected]>

.gitignore

@@ -0,0 +1,149 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# OS generated files
+.DS_Store
+.DS_Store?
+._*
+.Spotlight-V100
+.Trashes
+ehthumbs.db
+Thumbs.db
+
+# IDE files
+.vscode/
+.idea/
+*.swp
+*.swo
+*~
+
+# Temporary files
+*.tmp
+*.temp

LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

QUICKSTART.md

@@ -0,0 +1,76 @@
+# 🚀 Quick Start Guide
+
+## One-Command Setup
+
+### Method 1: Using our setup script
+```bash
+# Download the setup script
+curl -O https://huggingface.co/zhiqing0205/u2net-mvtec-loco-segmentation/raw/main/setup_project.py
+
+# Run setup (downloads everything automatically)
+python setup_project.py
+
+# Use the project
+cd u2net-mvtec-loco
+python mvtec_loco_fg_segmentation.py
+```
+
+### Method 2: Using HuggingFace CLI
+```bash
+# Install HuggingFace CLI
+pip install huggingface_hub
+
+# Download complete project (equivalent to git clone)
+huggingface-cli download zhiqing0205/u2net-mvtec-loco-segmentation \
+    --local-dir ./u2net-project --repo-type model
+
+# Use the project
+cd u2net-project
+python mvtec_loco_fg_segmentation.py
+```
+
+### Method 3: Using Python
+```bash
+# One-liner to download everything
+python -c "
+from huggingface_hub import snapshot_download
+snapshot_download('zhiqing0205/u2net-mvtec-loco-segmentation', local_dir='./u2net-project')
+print('Done! cd u2net-project && python mvtec_loco_fg_segmentation.py')
+"
+```
+
+## What Gets Downloaded
+
+✅ Complete source code  
+✅ Pre-trained model weights (u2net.pth - 169MB)  
+✅ Documentation (English + Chinese)  
+✅ Example scripts and utilities  
+✅ Ready to run immediately  
+
+## File Structure After Download
+```
+u2net-mvtec-loco/
+├── mvtec_loco_fg_segmentation.py  # Main script
+├── saved_models/
+│   └── u2net/
+│       └── u2net.pth             # Pre-trained model (169MB)
+├── model/                         # Model architecture
+├── data_loader.py                 # Data utilities
+├── README.md                      # English docs
+├── README_CN.md                   # Chinese docs
+└── ...
+```
+
+## Immediate Usage
+```bash
+# Process entire MVTec LOCO dataset
+python mvtec_loco_fg_segmentation.py
+
+# Process specific categories
+python mvtec_loco_fg_segmentation.py --categories breakfast_box
+
+# Custom threshold
+python mvtec_loco_fg_segmentation.py --threshold 0.3
+```
+
+That's it! 🎉

README.md

@@ -1,3 +1,215 @@
----
-license: apache-2.0
----
+# MVTec LOCO Foreground Segmentation Tool
+
+**[中文版 README](README_CN.md)** | **[English README](README.md)**
+
+This tool uses U²-Net for generating binary foreground masks from the MVTec LOCO anomaly detection dataset.
+
+## Overview
+
+The `mvtec_loco_fg_segmentation.py` script processes the entire MVTec LOCO dataset and generates binary foreground masks for all images. It uses the U²-Net model to perform salient object detection and converts the probability maps to binary masks.
+
+## Features
+
+- **Complete Dataset Processing**: Processes all categories (breakfast_box, screw_bag, juice_bottle, splicing_connectors, pushpins)
+- **Flexible Structure**: Handles both test and train splits with all subdirectories (good, logical_anomalies, structural_anomalies)
+- **Binary Mask Output**: Generates clean binary masks (0/255) in L mode (grayscale)
+- **Configurable Parameters**: Customizable threshold, categories, splits, and processing options
+- **GPU/CPU Support**: Automatic detection and utilization of available hardware
+
+## Requirements
+
+### Environment Setup
+```bash
+# Create conda environment
+conda create -n u2net python=3.8 -y
+conda activate u2net
+
+# Install dependencies
+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
+```
+
+### Model Weights
+
+**Option 1: Automatic Download (Recommended)**
+```bash
+# Install HuggingFace Hub
+pip install huggingface_hub
+
+# The model will be automatically downloaded when you run the script
+python mvtec_loco_fg_segmentation.py
+```
+
+**Option 2: Manual Download**
+- Download `u2net.pth` (176.3 MB) from [GoogleDrive](https://drive.google.com/file/d/1ao1ovG1Qtx4b7EoskHXmi2E9rp5CHLcZ/view)
+- Place it in: `./saved_models/u2net/u2net.pth`
+
+**Option 3: Download from HuggingFace**
+```bash
+# Download only the model
+python download_from_hf.py --model-only
+
+# Or download the complete repository
+python download_from_hf.py --complete-repo
+```
+
+### Dataset Structure
+Ensure your MVTec LOCO dataset follows this structure:
+```
+mvtec_loco_anomaly_detection/
+├── breakfast_box/
+│   ├── test/
+│   │   ├── good/
+│   │   ├── logical_anomalies/
+│   │   └── structural_anomalies/
+│   └── train/
+│       └── good/
+├── screw_bag/
+│   ├── test/
+│   └── train/
+└── ... (other categories)
+```
+
+## Usage
+
+### Basic Usage
+```bash
+# Process entire dataset with default settings
+python mvtec_loco_fg_segmentation.py
+
+# Show help
+python mvtec_loco_fg_segmentation.py -h
+```
+
+### Advanced Usage
+
+```bash
+# Specify custom dataset and model paths
+python mvtec_loco_fg_segmentation.py \
+    --dataset_path /path/to/mvtec_loco \
+    --model_path /path/to/u2net.pth
+
+# Process specific categories only
+python mvtec_loco_fg_segmentation.py \
+    --categories breakfast_box juice_bottle
+
+# Process only test split
+python mvtec_loco_fg_segmentation.py \
+    --splits test
+
+# Use different threshold for binary mask generation
+python mvtec_loco_fg_segmentation.py \
+    --threshold 0.3
+
+# Custom output directory name
+python mvtec_loco_fg_segmentation.py \
+    --output_dir custom_masks
+
+# Optimize processing with multiple workers
+python mvtec_loco_fg_segmentation.py \
+    --num_workers 4 \
+    --batch_size 4
+```
+
+## Command Line Arguments
+
+| Argument | Type | Default | Description |
+|----------|------|---------|-------------|
+| `--dataset_path` | str | `/root/hy-data/datasets/mvtec_loco_anomaly_detection` | Path to MVTec LOCO dataset root |
+| `--model_path` | str | `./saved_models/u2net/u2net.pth` | Path to U2NET model weights |
+| `--output_dir` | str | `fg_mask` | Output directory name for masks |
+| `--threshold` | float | `0.5` | Threshold for binary mask generation |
+| `--categories` | list | `all 5 categories` | Categories to process |
+| `--splits` | list | `['test', 'train']` | Dataset splits to process |
+| `--batch_size` | int | `1` | Batch size for processing |
+| `--num_workers` | int | `1` | Number of data loading workers |
+
+## Output Structure
+
+The script generates masks in the following structure:
+```
+mvtec_loco_anomaly_detection/
+├── fg_mask/                    # Generated masks directory
+│   ├── breakfast_box/
+│   │   ├── test/
+│   │   │   ├── good/
+│   │   │   │   ├── 000.png     # Binary mask (0/255 values)
+│   │   │   │   ├── 001.png
+│   │   │   │   └── ...
+│   │   │   ├── logical_anomalies/
+│   │   │   └── structural_anomalies/
+│   │   └── train/
+│   │       └── good/
+│   └── ... (other categories)
+└── ... (original dataset)
+```
+
+## Mask Properties
+
+- **Format**: PNG images
+- **Mode**: L (grayscale, single channel)
+- **Values**: Binary (0 for background, 255 for foreground)
+- **Size**: Same as original images
+- **Threshold**: Configurable (default 0.5)
+
+## Performance Notes
+
+- **GPU Recommended**: Processing is significantly faster with CUDA-enabled GPU
+- **Memory Usage**: Each image requires ~200MB GPU memory during processing
+- **Processing Time**: ~2-3 seconds per image on modern GPU
+- **Total Images**: ~5000+ images in complete dataset
+
+## Troubleshooting
+
+### Common Issues
+
+1. **CUDA Out of Memory**: Reduce batch size or use CPU processing
+2. **Model Not Found**: Ensure u2net.pth is in correct directory
+3. **Dataset Path Error**: Verify MVTec LOCO dataset structure
+4. **Permission Errors**: Check write permissions for output directory
+
+### Error Messages
+
+- `ERROR: Dataset path not found`: Check dataset path and extraction
+- `ERROR: Model path not found`: Download and place u2net.pth correctly
+- `ERROR: Invalid categories`: Use valid category names
+
+## Examples Output
+
+The script provides detailed progress information:
+```
+Configuration:
+  Dataset path: /root/hy-data/datasets/mvtec_loco_anomaly_detection
+  Model path: ./saved_models/u2net/u2net.pth
+  Output directory: fg_mask
+  Binary threshold: 0.5
+  Categories: ['breakfast_box', 'screw_bag', 'juice_bottle', 'splicing_connectors', 'pushpins']
+  Splits: ['test', 'train']
+
+...load U2NET---
+Processing category: breakfast_box
+  Processing breakfast_box/test/good
+    Found 102 images
+    Processing 1/102: 000.png
+    Processing 20/102: 019.png
+    ...
+```
+
+## Citation
+
+If you use this tool in your research, please cite the original U²-Net paper:
+
+```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}
+}
+```
+
+## License
+
+This tool extends the original U²-Net implementation. Please refer to the original repository for license information.

README_CN.md

@@ -0,0 +1,196 @@
+# 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 实现。请参考原始仓库的许可证信息。

README_HF.md

@@ -0,0 +1,141 @@
+---
+title: U²-Net MVTec LOCO Foreground Segmentation
+tags:
+- computer-vision
+- image-segmentation
+- anomaly-detection
+- u2net
+- mvtec-loco
+- pytorch
+license: apache-2.0
+language: en
+library_name: pytorch
+---
+
+# U²-Net MVTec LOCO Foreground Segmentation
+
+This repository contains a complete tool for generating binary foreground masks from the MVTec LOCO anomaly detection dataset using U²-Net.
+
+## 🚀 Quick Start
+
+### Installation
+```bash
+# Clone from HuggingFace
+git clone https://huggingface.co/zhiqing0205/u2net-mvtec-loco-segmentation
+cd u2net-mvtec-loco-segmentation
+
+# Install dependencies  
+pip install torch torchvision opencv-python scikit-image matplotlib numpy pillow huggingface_hub
+
+# Run segmentation (model auto-downloads)
+python mvtec_loco_fg_segmentation.py
+```
+
+### Download Options
+
+**Option 1: Auto-download (Recommended)**
+```python
+from download_from_hf import download_u2net_model
+
+# Download model only
+download_u2net_model()
+
+# Download complete repo
+from download_from_hf import download_complete_repo
+download_complete_repo()
+```
+
+**Option 2: Manual download**
+```bash
+python download_from_hf.py --model-only
+# or
+python download_from_hf.py --complete-repo
+```
+
+## 📁 Repository Contents
+
+```
+├── mvtec_loco_fg_segmentation.py  # Main segmentation script
+├── download_from_hf.py            # HuggingFace download utility
+├── model/                         # U2NET model architecture
+├── data_loader.py                 # Data loading utilities
+├── saved_models/
+│   └── u2net/
+│       └── u2net.pth             # Pre-trained U2NET weights (169MB)
+├── README.md                      # English documentation
+├── README_CN.md                   # Chinese documentation
+└── ...
+```
+
+## 🎯 Features
+
+- **Complete Dataset Processing**: All MVTec LOCO categories
+- **Binary Mask Output**: Standard 0/255 masks in grayscale
+- **GPU/CPU Support**: Automatic hardware detection
+- **Configurable Parameters**: Threshold, categories, splits
+- **Auto-download**: No manual model download needed
+
+## 💻 Usage
+
+### Basic Usage
+```bash
+python mvtec_loco_fg_segmentation.py
+```
+
+### Advanced Usage
+```bash
+# Custom parameters
+python mvtec_loco_fg_segmentation.py \
+    --threshold 0.3 \
+    --categories breakfast_box juice_bottle \
+    --splits test
+
+# Show all options
+python mvtec_loco_fg_segmentation.py -h
+```
+
+## 📊 Model Information
+
+- **Architecture**: U²-Net (U Square Net)
+- **Model Size**: 169MB
+- **Input Size**: 320×320 (auto-resized)
+- **Output**: Binary masks (0/255)
+- **Task**: Salient object detection → Foreground segmentation
+
+## 🏷️ Supported Categories
+
+- `breakfast_box`
+- `screw_bag` 
+- `juice_bottle`
+- `splicing_connectors`
+- `pushpins`
+
+## 📈 Performance
+
+- **GPU Processing**: ~2-3 seconds per image
+- **CPU Processing**: ~10-15 seconds per image  
+- **Memory Usage**: ~200MB GPU memory per image
+- **Total Dataset**: ~5000+ images
+
+## 📖 Citation
+
+```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}
+}
+```
+
+## 📜 License
+
+Apache-2.0 License
+
+## 🔗 Links
+
+- [Original U²-Net Paper](https://arxiv.org/pdf/2005.09007.pdf)
+- [MVTec LOCO Dataset](https://www.mvtec.com/company/research/datasets/mvtec-loco)
+- [GitHub Repository](https://github.com/NathanUA/U-2-Net)

README_original.md

@@ -0,0 +1,218 @@
+<p align="center">
+  <img width="320" height="320" src="figures/U2Net_Logo.png">
+  
+  <h1 align="center">U<sup>2</sup>-Net: U Square Net</h1>
+    
+</p>
+
+This is the official repo for our paper **U<sup>2</sup>-Net(U square net)** published in Pattern Recognition 2020:
+
+## [U<sup>2</sup>-Net: Going Deeper with Nested U-Structure for Salient Object Detection](https://arxiv.org/pdf/2005.09007.pdf)
+[Xuebin Qin](https://xuebinqin.github.io/), [Zichen Zhang](https://webdocs.cs.ualberta.ca/~zichen2/), [Chenyang Huang](https://chenyangh.com/), [Masood Dehghan](https://sites.google.com/view/masooddehghan), [Osmar R. Zaiane](http://webdocs.cs.ualberta.ca/~zaiane/) and [Martin Jagersand](https://webdocs.cs.ualberta.ca/~jag/)
+
+
+__Contact__: xuebin[at]ualberta[dot]ca
+
+## Updates !!!
+
+** (2022-Aug.-24) ** We are glad to announce that our U<sup>2</sup>-Net published in Pattern Recognition has been awarded the 2020 Pattern Recognition BEST PAPER AWARD !!!
+![u2net-best-paper](figures/u2net-best-paper.jpg)
+
+** (2022-Aug.-17) **
+Our U<sup>2</sup>-Net models are now available on [PlayTorch](https://playtorch.dev/), where you can build your own demo and run it on your Android/iOS phone. Try out this demo on [![PlayTorch Demo](https://github.com/facebookresearch/playtorch/blob/main/website/static/assets/playtorch_badge.svg)](https://playtorch.dev/snack/@playtorch/u2net/) and bring your ideas about U<sup>2</sup>-Net to truth in minutes!
+
+** (2022-Jul.-5)** Our new work **Highly Accurate Dichotomous Image Segmentation (DIS) [**Project Page**](https://xuebinqin.github.io/dis/index.html), [**Github**](https://github.com/xuebinqin/DIS) is accepted by ECCV 2022. Our code and dataset will be released before July 17th, 2022. Please be aware of our updates. 
+![ship-demo](figures/ship-demo.gif)
+![bg-removal](figures/bg-removal.gif)
+![view-move](figures/view-move.gif)
+![motor-demo](figures/motor-demo.gif)
+
+** (2022-Jun.-3)** Thank [**Adir Kol**](https://github.com/adirkol) for sharing the iOS App [**3D Photo Creator**](https://apps.apple.com/us/app/3d-photo-creator/id1619676262) based on our U<sup>2</sup>-Net.
+![portrait-ios-app](figures/3d-photo-re.jpg)
+
+** (2022-Mar.-31)** Thank [**Hikaru Tsuyumine**] for implementing the iOS App [**Portrait Drawing**](https://apps.apple.com/us/app/portrait-drawing/id1623269600) based on our U<sup>2</sup>-Net portrait generation model.
+![portrait-ios-app](figures/portrait-ios-app.jpg)
+
+** (2022-Apr.-12)** Thank [**Kevin Shah**](https://github.com/ioskevinshah) for providing us a great iOS App [**Lensto**](https://apps.apple.com/in/app/lensto-background-changer/id1574844033), ([**Demo Video**](https://www.youtube.com/shorts/jWwUiKZjfok)), based on U<sup>2</sup>-Net.
+![lensto](figures/lensto.png)
+
+** (2022-Mar.-31)** Our U<sup>2</sup>-Net model is also integrated by [**Hotpot.ai**](https://hotpot.ai/) for art design.
+![hotpot](figures/hotpot.png)
+
+** (2022-Mar-19)** Thank [**Kikedao**](https://github.com/Kikedao) for providing a fantastic webapp [**Silueta**](https://silueta.me/) based on U<sup>2</sup>-Net. More details can be found at [**https://github.com/xuebinqin/U-2-Net/issues/295**](https://github.com/xuebinqin/U-2-Net/issues/295).
+![silueta](figures/silueta.png) 
+
+** (2022-Mar-17)** Thank [**Ezaldeen Sahb**](https://github.com/Ezaldeen99/BackgroundRemoval) for implementing the iOS library for image background removal based on U<sup>2</sup>-Net, which will greatly facilitate the developing of mobile apps.
+![close-seg](figures/swift-u2net.jpeg) 
+
+<!-- ** (2022-Mar-10)** Thank [**Doron Adler**](https://github.com/Norod/U-2-Net-StyleTransfer) for training the awesome style transfer U<sup>2</sup>-Net.
+![style-trans](figures/style-trans.JPG)  -->
+
+** (2022-Mar-8)** Thank [**Levin Dabhi**](https://github.com/levindabhi/cloth-segmentation) for training the amazing clothes segmentation U<sup>2</sup>-Net.
+![close-seg](figures/close-seg.jpg) 
+
+** (2022-Mar-3)** Thank [**Renato Violin**](https://github.com/renatoviolin/bg-remove-augment) for providing an awesome webapp for image background removal and replacement based on our U<sup>2</sup>-Net.
+![bg-rm-aug](figures/bg-rm-aug.gif) 
+
+**(2021-Dec-21)** This [**blog**](https://rockyshikoku.medium.com/u2net-to-coreml-machine-learning-segmentation-on-iphone-eac0c721d67b) clearly describes the way of converting U<sup>2</sup>-Net to [**CoreML**](https://github.com/john-rocky/CoreML-Models) and running it on iphone. 
+
+**(2021-Nov-28)** Interesting Sky Segmentation models developed by [**xiongzhu**](https://github.com/xiongzhu666/Sky-Segmentation-and-Post-processing) using U<sup>2</sup>-Net. 
+
+![im_sky_segmentation](figures/sky-seg.png)
+
+**(2021-Nov-28)** Awesome image editing app [**Pixelmator pro**](https://www.pixelmator.com/pro/) uses U<sup>2</sup>-Net as one of its background removal models. 
+
+![im_sky_segmentation](figures/pixelmator.jpg)
+
+**(2021-Aug-24)** We played a bit more about fusing the orignal image and the generated portraits to composite different styles. You can <br/> 
+(1) Download this repo by
+```
+git clone https://github.com/NathanUA/U-2-Net.git
+```
+(2) Download the u2net_portrait.pth from [**GoogleDrive**](https://drive.google.com/file/d/1IG3HdpcRiDoWNookbncQjeaPN28t90yW/view?usp=sharing) or [**Baidu Pan(提取码：chgd)**](https://pan.baidu.com/s/1BYT5Ts6BxwpB8_l2sAyCkw)model and put it into the directory: ```./saved_models/u2net_portrait/```, <br/>
+(3) run the code by command 
+```
+python u2net_portrait_composite.py -s 20 -a 0.5
+```
+,where ``-s`` indicates the sigma of gaussian function for blurring the orignal image and ``-a`` denotes the alpha weights of the orignal image when fusing them. <br/>
+
+![im_portrait_composite](figures/im_composite.jpg)
+
+**(2021-July-16)** A new [background removal webapp](https://remove-background.net/) developed by Изатоп Василий. 
+
+![rm_bg](figures/rm_bg.JPG)
+
+**(2021-May-26)** Thank [**Dang Quoc Quy**](https://github.com/quyvsquy) for his [**Art Transfer APP**](https://play.google.com/store/apps/details?id=com.quyvsquy.arttransfer) built upon U<sup>2</sup>-Net.
+
+<!---![art_transfer](figures/art_transfer.JPG)--->
+
+**(2021-May-5)** Thank [**AK391**](https://github.com/AK391) for sharing his [**Gradio Web Demo of U<sup>2</sup>-Net**](https://gradio.app/hub/AK391/U-2-Net).
+
+![gradio_web_demo](figures/gradio_web_demo.jpg)
+
+
+**(2021-Apr-29)** Thanks [**Jonathan Benavides Vallejo**](https://www.linkedin.com/in/jonathanbv/) for releasing his App [**LensOCR: Extract Text & Image**](https://apps.apple.com/ch/app/lensocr-extract-text-image/id1549961729?l=en&mt=12), which uses U<sup>2</sup>-Net for extracting the image foreground.
+
+![LensOCR APP](figures/LensOCR.jpg)
+
+**(2021-Apr-18)** Thanks [**Andrea Scuderi**](https://www.linkedin.com/in/andreascuderi/) for releasing his App [**Clipping Camera**](https://apps.apple.com/us/app/clipping-camera/id1548192169?ign-mpt=uo%3D2), which is an U<sup>2</sup>-Net driven realtime camera app and "is able to detect relevant object from the scene and clip them to apply fancy filters". 
+
+![Clipping Camera APP](figures/clipping_camera.jpg)
+
+**(2021-Mar-17)** [**Dennis Bappert**](https://github.com/dennisbappert) re-trained the U<sup>2</sup>-Net model for [**human portrait matting**](https://github.com/dennisbappert/u-2-net-portrait). The results look very promising and he also provided the details of the training process and data generation(and augmentation) strategy, which are inspiring.
+
+**(2021-Mar-11)** Dr. Tim developed a [**video version rembg**](https://github.com/ecsplendid/rembg-greenscreen) for removing video backgrounds using U<sup>2</sup>-Net. The awesome demo results can be found on [**YouTube**](https://www.youtube.com/watch?v=4NjqR2vCV_k).
+
+**(2021-Mar-02)** We found some other interesting applications of our U<sup>2</sup>-Net including [**MOJO CUT**](https://play.google.com/store/apps/details?id=com.innoria.magicut&hl=en_CA&gl=US), [**Real-Time Background Removal on Iphone**](https://www.linkedin.com/feed/update/urn:li:activity:6752303661705170944/?updateEntityUrn=urn%3Ali%3Afs_feedUpdate%3A%28V2%2Curn%3Ali%3Aactivity%3A6752303661705170944%29), [**Video Background Removal**](https://nisargkapkar.hashnode.dev/image-and-video-background-removal-using-deep-learning), [**Another Online Portrait Generation Demo on AWS**](http://s3-website-hosting-u2net.s3-website-eu-west-1.amazonaws.com/), [**AI Scissor**](https://qooba.net/2020/09/11/ai-scissors-sharp-cut-with-neural-networks/).
+
+**(2021-Feb-15)** We just released an online demo [**http://profu.ai**](http://profu.ai) for the portrait generation. Please feel free to give it a try and provide any suggestions or comments. <br/>
+![Profuai](figures/profuai.png) <br/>
+
+**(2021-Feb-06)** Recently, some people asked the problem of using U<sup>2</sup>-Net for human segmentation, so we trained another example model for human segemntation based on [**Supervisely Person Dataset**](https://supervise.ly/explore/projects/supervisely-person-dataset-23304/datasets). <br/>
+
+(1) To run the human segmentation model, please first downlowd the [**u2net_human_seg.pth**](https://drive.google.com/file/d/1m_Kgs91b21gayc2XLW0ou8yugAIadWVP/view?usp=sharing) model weights into ``` ./saved_models/u2net_human_seg/```. <br/>
+(2) Prepare the to-be-segmented images into the corresponding directory, e.g. ```./test_data/test_human_images/```. <br/>
+(3) Run the inference by command: ```python u2net_human_seg_test.py``` and the results will be output into the corresponding dirctory, e.g. ```./test_data/u2net_test_human_images_results/```<br/>
+[**Notes: Due to the labeling accuracy of the Supervisely Person Dataset, the human segmentation model (u2net_human_seg.pth) here won't give you hair-level accuracy. But it should be more robust than u2net trained with DUTS-TR dataset on general human segmentation task. It can be used for human portrait segmentation, human body segmentation, etc.**](https://github.com/NathanUA/U-2-Net)<br/>
+
+![Human Image Segmentation](figures/human_seg.png) <br/>
+![Human Video](figures/human_seg_video.gif)
+![Human Video Results](figures/human_seg_results.gif)
+
+**(2020-Dec-28)** Some interesting applications and useful tools based on U<sup>2</sup>-Net: <br/>
+(1) [**Xiaolong Liu**](https://github.com/LiuXiaolong19920720) developed several very interesting applications based on U<sup>2</sup>-Net including [**Human Portrait Drawing**](https://www.cvpy.net/studio/cv/func/DeepLearning/sketch/sketch/page/)(As far as I know, Xiaolong is the first one who uses U<sup>2</sup>-Net for portrait generation), [**image matting**](https://www.cvpy.net/studio/cv/func/DeepLearning/matting/matting/page/) and [**so on**](https://www.cvpy.net/). <br/>
+(2) [**Vladimir Seregin**](https://github.com/peko/nn-lineart) developed an interesting tool, [**NN based lineart**](https://peko.github.io/nn-lineart/), for comparing the portrait results of U<sup>2</sup>-Net and that of another popular model, [**ArtLine**](https://github.com/vijishmadhavan/ArtLine), developed by [**Vijish Madhavan**](https://github.com/vijishmadhavan). <br/>
+(3) [**Daniel Gatis**](https://github.com/danielgatis/rembg) built a python tool, [**Rembg**](https://pypi.org/project/rembg/), for image backgrounds removal based on U<sup>2</sup>-Net. I think this tool will greatly facilitate the application of U<sup>2</sup>-Net in different fields. <br/>
+![REMBG](figures/rembg.png)
+
+**(2020-Nov-21)** Recently, we found an interesting application of U<sup>2</sup>-Net for [**human portrait drawing**](https://www.pythonf.cn/read/141098). Therefore, we trained another model for this task based on the [**APDrawingGAN dataset**](https://github.com/yiranran/APDrawingGAN).
+
+![Sample Results: Kids](figures/portrait_kids.png)
+
+![Sample Results: Ladies](figures/portrait_ladies.png)
+
+![Sample Results: Men](figures/portrait_men.png)
+
+### Usage for portrait generation
+1. Clone this repo to local
+```
+git clone https://github.com/NathanUA/U-2-Net.git
+```
+
+2. Download the u2net_portrait.pth from [**GoogleDrive**](https://drive.google.com/file/d/1IG3HdpcRiDoWNookbncQjeaPN28t90yW/view?usp=sharing) or [**Baidu Pan(提取码：chgd)**](https://pan.baidu.com/s/1BYT5Ts6BxwpB8_l2sAyCkw)model and put it into the directory: ```./saved_models/u2net_portrait/```.
+
+3. Run on the testing set. <br/>
+(1) Download the train and test set from [**APDrawingGAN**](https://github.com/yiranran/APDrawingGAN). These images and their ground truth are stitched side-by-side (512x1024). You need to split each of these images into two 512x512 images and put them into ```./test_data/test_portrait_images/portrait_im/```. You can also download the split testing set on [GoogleDrive](https://drive.google.com/file/d/1NkTsDDN8VO-JVik6VxXyV-3l2eo29KCk/view?usp=sharing). <br/>
+(2) Running the inference with command ```python u2net_portrait_test.py``` will ouptut the results into ```./test_data/test_portrait_images/portrait_results```. <br/>
+
+4. Run on your own dataset. <br/>
+(1) Prepare your images and put them into ```./test_data/test_portrait_images/your_portrait_im/```. [**To obtain enough details of the protrait, human head region in the input image should be close to or larger than 512x512. The head background should be relatively clear.**](https://github.com/NathanUA/U-2-Net) <br/>
+(2) Run the prediction by command ```python u2net_portrait_demo.py``` will outputs the results to ```./test_data/test_portrait_images/your_portrait_results/```. <br/>
+(3) The difference between ```python u2net_portrait_demo.py``` and ```python u2net_portrait_test.py``` is that we added a simple [**face detection**](https://opencv-python-tutroals.readthedocs.io/en/latest/py_tutorials/py_objdetect/py_face_detection/py_face_detection.html) step before the portrait generation in ```u2net_portrait_demo.py```.  Because the testing set of APDrawingGAN are normalized and cropped to 512x512 for including only heads of humans, while our own dataset may varies with different resolutions and contents. Therefore, the code ```python u2net_portrait_demo.py``` will detect the biggest face from the given image and then crop, pad and resize the ROI to 512x512 for feeding to the network. The following figure shows how to take your own photos for generating high quality portraits.
+
+**(2020-Sep-13)** Our U<sup>2</sup>-Net based model is the **6th** in [**MICCAI 2020 Thyroid Nodule Segmentation Challenge**](https://tn-scui2020.grand-challenge.org/Resultannouncement/).
+
+**(2020-May-18)** The official paper of our **U<sup>2</sup>-Net (U square net)** ([**PDF in elsevier**(free until July 5 2020)](https://www.sciencedirect.com/science/article/pii/S0031320320302077?dgcid=author), [**PDF in arxiv**](http://arxiv.org/abs/2005.09007)) is now available. If you are not able to access that, please feel free to drop me an email.
+
+**(2020-May-16)** We fixed the upsampling issue of the network. Now, the model should be able to handle **arbitrary input size**. (Tips: This modification is to facilitate the retraining of U<sup>2</sup>-Net on your own datasets. When using our pre-trained model on SOD datasets, please keep the input size as 320x320 to guarantee the performance.)
+
+**(2020-May-16)** We highly appreciate **Cyril Diagne** for building this fantastic AR project: [**AR Copy and Paste**](https://github.com/cyrildiagne/ar-cutpaste) using our **U<sup>2</sup>-Net** (Qin *et al*, PR 2020) and [**BASNet**](https://github.com/NathanUA/BASNet)(Qin *et al*, CVPR 2019). The [**demo video**](https://twitter.com/cyrildiagne/status/1256916982764646402) in twitter has achieved over **5M** views, which is phenomenal and shows us more application possibilities of SOD.
+
+## U<sup>2</sup>-Net Results (176.3 MB)
+
+![U<sup>2</sup>-Net Results](figures/u2netqual.png)
+
+
+## Our previous work: [BASNet (CVPR 2019)](https://github.com/NathanUA/BASNet)
+
+## Required libraries
+
+Python 3.6  
+numpy 1.15.2  
+scikit-image 0.14.0  
+python-opencv
+PIL 5.2.0  
+PyTorch 0.4.0  
+torchvision 0.2.1  
+glob  
+
+## Usage for salient object detection
+1. Clone this repo
+```
+git clone https://github.com/NathanUA/U-2-Net.git
+```
+2. Download the pre-trained model u2net.pth (176.3 MB) from [**GoogleDrive**](https://drive.google.com/file/d/1ao1ovG1Qtx4b7EoskHXmi2E9rp5CHLcZ/view?usp=sharing) or [**Baidu Pan 提取码: pf9k**](https://pan.baidu.com/s/1WjwyEwDiaUjBbx_QxcXBwQ) or u2netp.pth (4.7 MB) from [**GoogleDrive**](https://drive.google.com/file/d/1rbSTGKAE-MTxBYHd-51l2hMOQPT_7EPy/view?usp=sharing) or [**Baidu Pan 提取码: 8xsi**](https://pan.baidu.com/s/10tW12OlecRpE696z8FxdNQ) and put it into the dirctory './saved_models/u2net/' and './saved_models/u2netp/'
+
+3.  Cd to the directory 'U-2-Net', run the train or inference process by command: ```python u2net_train.py```
+or ```python u2net_test.py``` respectively. The 'model_name' in both files can be changed to 'u2net' or 'u2netp' for using different models.  
+
+ We also provide the predicted saliency maps ([u2net results](https://drive.google.com/file/d/1mZFWlS4WygWh1eVI8vK2Ad9LrPq4Hp5v/view?usp=sharing),[u2netp results](https://drive.google.com/file/d/1j2pU7vyhOO30C2S_FJuRdmAmMt3-xmjD/view?usp=sharing)) for datasets SOD, ECSSD, DUT-OMRON, PASCAL-S, HKU-IS and DUTS-TE.
+
+
+## U<sup>2</sup>-Net Architecture
+
+![U<sup>2</sup>-Net architecture](figures/U2NETPR.png)
+
+
+## Quantitative Comparison
+
+![Quantitative Comparison](figures/quan_1.png)
+
+![Quantitative Comparison](figures/quan_2.png)
+
+
+## Qualitative Comparison
+
+![Qualitative Comparison](figures/qual.png?raw=true)
+
+## Citation
+```
+@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}
+}
+```

data_loader.py

@@ -0,0 +1,266 @@
+# data loader
+from __future__ import print_function, division
+import glob
+import torch
+from skimage import io, transform, color
+import numpy as np
+import random
+import math
+import matplotlib.pyplot as plt
+from torch.utils.data import Dataset, DataLoader
+from torchvision import transforms, utils
+from PIL import Image
+
+#==========================dataset load==========================
+class RescaleT(object):
+
+	def __init__(self,output_size):
+		assert isinstance(output_size,(int,tuple))
+		self.output_size = output_size
+
+	def __call__(self,sample):
+		imidx, image, label = sample['imidx'], sample['image'],sample['label']
+
+		h, w = image.shape[:2]
+
+		if isinstance(self.output_size,int):
+			if h > w:
+				new_h, new_w = self.output_size*h/w,self.output_size
+			else:
+				new_h, new_w = self.output_size,self.output_size*w/h
+		else:
+			new_h, new_w = self.output_size
+
+		new_h, new_w = int(new_h), int(new_w)
+
+		# #resize the image to new_h x new_w and convert image from range [0,255] to [0,1]
+		# img = transform.resize(image,(new_h,new_w),mode='constant')
+		# lbl = transform.resize(label,(new_h,new_w),mode='constant', order=0, preserve_range=True)
+
+		img = transform.resize(image,(self.output_size,self.output_size),mode='constant')
+		lbl = transform.resize(label,(self.output_size,self.output_size),mode='constant', order=0, preserve_range=True)
+
+		return {'imidx':imidx, 'image':img,'label':lbl}
+
+class Rescale(object):
+
+	def __init__(self,output_size):
+		assert isinstance(output_size,(int,tuple))
+		self.output_size = output_size
+
+	def __call__(self,sample):
+		imidx, image, label = sample['imidx'], sample['image'],sample['label']
+
+		if random.random() >= 0.5:
+			image = image[::-1]
+			label = label[::-1]
+
+		h, w = image.shape[:2]
+
+		if isinstance(self.output_size,int):
+			if h > w:
+				new_h, new_w = self.output_size*h/w,self.output_size
+			else:
+				new_h, new_w = self.output_size,self.output_size*w/h
+		else:
+			new_h, new_w = self.output_size
+
+		new_h, new_w = int(new_h), int(new_w)
+
+		# #resize the image to new_h x new_w and convert image from range [0,255] to [0,1]
+		img = transform.resize(image,(new_h,new_w),mode='constant')
+		lbl = transform.resize(label,(new_h,new_w),mode='constant', order=0, preserve_range=True)
+
+		return {'imidx':imidx, 'image':img,'label':lbl}
+
+class RandomCrop(object):
+
+	def __init__(self,output_size):
+		assert isinstance(output_size, (int, tuple))
+		if isinstance(output_size, int):
+			self.output_size = (output_size, output_size)
+		else:
+			assert len(output_size) == 2
+			self.output_size = output_size
+	def __call__(self,sample):
+		imidx, image, label = sample['imidx'], sample['image'], sample['label']
+
+		if random.random() >= 0.5:
+			image = image[::-1]
+			label = label[::-1]
+
+		h, w = image.shape[:2]
+		new_h, new_w = self.output_size
+
+		top = np.random.randint(0, h - new_h)
+		left = np.random.randint(0, w - new_w)
+
+		image = image[top: top + new_h, left: left + new_w]
+		label = label[top: top + new_h, left: left + new_w]
+
+		return {'imidx':imidx,'image':image, 'label':label}
+
+class ToTensor(object):
+	"""Convert ndarrays in sample to Tensors."""
+
+	def __call__(self, sample):
+
+		imidx, image, label = sample['imidx'], sample['image'], sample['label']
+
+		tmpImg = np.zeros((image.shape[0],image.shape[1],3))
+		tmpLbl = np.zeros(label.shape)
+
+		image = image/np.max(image)
+		if(np.max(label)<1e-6):
+			label = label
+		else:
+			label = label/np.max(label)
+
+		if image.shape[2]==1:
+			tmpImg[:,:,0] = (image[:,:,0]-0.485)/0.229
+			tmpImg[:,:,1] = (image[:,:,0]-0.485)/0.229
+			tmpImg[:,:,2] = (image[:,:,0]-0.485)/0.229
+		else:
+			tmpImg[:,:,0] = (image[:,:,0]-0.485)/0.229
+			tmpImg[:,:,1] = (image[:,:,1]-0.456)/0.224
+			tmpImg[:,:,2] = (image[:,:,2]-0.406)/0.225
+
+		tmpLbl[:,:,0] = label[:,:,0]
+
+
+		tmpImg = tmpImg.transpose((2, 0, 1))
+		tmpLbl = label.transpose((2, 0, 1))
+
+		return {'imidx':torch.from_numpy(imidx), 'image': torch.from_numpy(tmpImg), 'label': torch.from_numpy(tmpLbl)}
+
+class ToTensorLab(object):
+	"""Convert ndarrays in sample to Tensors."""
+	def __init__(self,flag=0):
+		self.flag = flag
+
+	def __call__(self, sample):
+
+		imidx, image, label =sample['imidx'], sample['image'], sample['label']
+
+		tmpLbl = np.zeros(label.shape)
+
+		if(np.max(label)<1e-6):
+			label = label
+		else:
+			label = label/np.max(label)
+
+		# change the color space
+		if self.flag == 2: # with rgb and Lab colors
+			tmpImg = np.zeros((image.shape[0],image.shape[1],6))
+			tmpImgt = np.zeros((image.shape[0],image.shape[1],3))
+			if image.shape[2]==1:
+				tmpImgt[:,:,0] = image[:,:,0]
+				tmpImgt[:,:,1] = image[:,:,0]
+				tmpImgt[:,:,2] = image[:,:,0]
+			else:
+				tmpImgt = image
+			tmpImgtl = color.rgb2lab(tmpImgt)
+
+			# nomalize image to range [0,1]
+			tmpImg[:,:,0] = (tmpImgt[:,:,0]-np.min(tmpImgt[:,:,0]))/(np.max(tmpImgt[:,:,0])-np.min(tmpImgt[:,:,0]))
+			tmpImg[:,:,1] = (tmpImgt[:,:,1]-np.min(tmpImgt[:,:,1]))/(np.max(tmpImgt[:,:,1])-np.min(tmpImgt[:,:,1]))
+			tmpImg[:,:,2] = (tmpImgt[:,:,2]-np.min(tmpImgt[:,:,2]))/(np.max(tmpImgt[:,:,2])-np.min(tmpImgt[:,:,2]))
+			tmpImg[:,:,3] = (tmpImgtl[:,:,0]-np.min(tmpImgtl[:,:,0]))/(np.max(tmpImgtl[:,:,0])-np.min(tmpImgtl[:,:,0]))
+			tmpImg[:,:,4] = (tmpImgtl[:,:,1]-np.min(tmpImgtl[:,:,1]))/(np.max(tmpImgtl[:,:,1])-np.min(tmpImgtl[:,:,1]))
+			tmpImg[:,:,5] = (tmpImgtl[:,:,2]-np.min(tmpImgtl[:,:,2]))/(np.max(tmpImgtl[:,:,2])-np.min(tmpImgtl[:,:,2]))
+
+			# tmpImg = tmpImg/(np.max(tmpImg)-np.min(tmpImg))
+
+			tmpImg[:,:,0] = (tmpImg[:,:,0]-np.mean(tmpImg[:,:,0]))/np.std(tmpImg[:,:,0])
+			tmpImg[:,:,1] = (tmpImg[:,:,1]-np.mean(tmpImg[:,:,1]))/np.std(tmpImg[:,:,1])
+			tmpImg[:,:,2] = (tmpImg[:,:,2]-np.mean(tmpImg[:,:,2]))/np.std(tmpImg[:,:,2])
+			tmpImg[:,:,3] = (tmpImg[:,:,3]-np.mean(tmpImg[:,:,3]))/np.std(tmpImg[:,:,3])
+			tmpImg[:,:,4] = (tmpImg[:,:,4]-np.mean(tmpImg[:,:,4]))/np.std(tmpImg[:,:,4])
+			tmpImg[:,:,5] = (tmpImg[:,:,5]-np.mean(tmpImg[:,:,5]))/np.std(tmpImg[:,:,5])
+
+		elif self.flag == 1: #with Lab color
+			tmpImg = np.zeros((image.shape[0],image.shape[1],3))
+
+			if image.shape[2]==1:
+				tmpImg[:,:,0] = image[:,:,0]
+				tmpImg[:,:,1] = image[:,:,0]
+				tmpImg[:,:,2] = image[:,:,0]
+			else:
+				tmpImg = image
+
+			tmpImg = color.rgb2lab(tmpImg)
+
+			# tmpImg = tmpImg/(np.max(tmpImg)-np.min(tmpImg))
+
+			tmpImg[:,:,0] = (tmpImg[:,:,0]-np.min(tmpImg[:,:,0]))/(np.max(tmpImg[:,:,0])-np.min(tmpImg[:,:,0]))
+			tmpImg[:,:,1] = (tmpImg[:,:,1]-np.min(tmpImg[:,:,1]))/(np.max(tmpImg[:,:,1])-np.min(tmpImg[:,:,1]))
+			tmpImg[:,:,2] = (tmpImg[:,:,2]-np.min(tmpImg[:,:,2]))/(np.max(tmpImg[:,:,2])-np.min(tmpImg[:,:,2]))
+
+			tmpImg[:,:,0] = (tmpImg[:,:,0]-np.mean(tmpImg[:,:,0]))/np.std(tmpImg[:,:,0])
+			tmpImg[:,:,1] = (tmpImg[:,:,1]-np.mean(tmpImg[:,:,1]))/np.std(tmpImg[:,:,1])
+			tmpImg[:,:,2] = (tmpImg[:,:,2]-np.mean(tmpImg[:,:,2]))/np.std(tmpImg[:,:,2])
+
+		else: # with rgb color
+			tmpImg = np.zeros((image.shape[0],image.shape[1],3))
+			image = image/np.max(image)
+			if image.shape[2]==1:
+				tmpImg[:,:,0] = (image[:,:,0]-0.485)/0.229
+				tmpImg[:,:,1] = (image[:,:,0]-0.485)/0.229
+				tmpImg[:,:,2] = (image[:,:,0]-0.485)/0.229
+			else:
+				tmpImg[:,:,0] = (image[:,:,0]-0.485)/0.229
+				tmpImg[:,:,1] = (image[:,:,1]-0.456)/0.224
+				tmpImg[:,:,2] = (image[:,:,2]-0.406)/0.225
+
+		tmpLbl[:,:,0] = label[:,:,0]
+
+
+		tmpImg = tmpImg.transpose((2, 0, 1))
+		tmpLbl = label.transpose((2, 0, 1))
+
+		return {'imidx':torch.from_numpy(imidx), 'image': torch.from_numpy(tmpImg), 'label': torch.from_numpy(tmpLbl)}
+
+class SalObjDataset(Dataset):
+	def __init__(self,img_name_list,lbl_name_list,transform=None):
+		# self.root_dir = root_dir
+		# self.image_name_list = glob.glob(image_dir+'*.png')
+		# self.label_name_list = glob.glob(label_dir+'*.png')
+		self.image_name_list = img_name_list
+		self.label_name_list = lbl_name_list
+		self.transform = transform
+
+	def __len__(self):
+		return len(self.image_name_list)
+
+	def __getitem__(self,idx):
+
+		# image = Image.open(self.image_name_list[idx])#io.imread(self.image_name_list[idx])
+		# label = Image.open(self.label_name_list[idx])#io.imread(self.label_name_list[idx])
+
+		image = io.imread(self.image_name_list[idx])
+		imname = self.image_name_list[idx]
+		imidx = np.array([idx])
+
+		if(0==len(self.label_name_list)):
+			label_3 = np.zeros(image.shape)
+		else:
+			label_3 = io.imread(self.label_name_list[idx])
+
+		label = np.zeros(label_3.shape[0:2])
+		if(3==len(label_3.shape)):
+			label = label_3[:,:,0]
+		elif(2==len(label_3.shape)):
+			label = label_3
+
+		if(3==len(image.shape) and 2==len(label.shape)):
+			label = label[:,:,np.newaxis]
+		elif(2==len(image.shape) and 2==len(label.shape)):
+			image = image[:,:,np.newaxis]
+			label = label[:,:,np.newaxis]
+
+		sample = {'imidx':imidx, 'image':image, 'label':label}
+
+		if self.transform:
+			sample = self.transform(sample)
+
+		return sample

download_from_hf.py

@@ -0,0 +1,96 @@
+#!/usr/bin/env python3
+"""
+HuggingFace model download utility for U2NET MVTec LOCO segmentation
+"""
+
+import os
+from huggingface_hub import hf_hub_download, snapshot_download
+from pathlib import Path
+
+def download_u2net_model(model_path="./saved_models/u2net/u2net.pth", 
+                         repo_id="zhiqing0205/u2net-mvtec-loco-segmentation",
+                         force_download=False):
+    """
+    Download U2NET model from HuggingFace Hub
+    
+    Args:
+        model_path: Local path to save the model
+        repo_id: HuggingFace repository ID  
+        force_download: Force re-download even if file exists
+    """
+    
+    # Check if model already exists
+    if os.path.exists(model_path) and not force_download:
+        print(f"Model already exists at {model_path}")
+        return model_path
+    
+    print(f"Downloading U2NET model from HuggingFace: {repo_id}")
+    
+    try:
+        # Ensure directory exists
+        os.makedirs(os.path.dirname(model_path), exist_ok=True)
+        
+        # Download specific model file
+        downloaded_path = hf_hub_download(
+            repo_id=repo_id,
+            filename="saved_models/u2net/u2net.pth",
+            local_dir=".",
+            local_dir_use_symlinks=False
+        )
+        
+        print(f"Model downloaded successfully to: {downloaded_path}")
+        return downloaded_path
+        
+    except Exception as e:
+        print(f"Error downloading model: {e}")
+        print("Please download manually from:")
+        print("https://drive.google.com/file/d/1ao1ovG1Qtx4b7EoskHXmi2E9rp5CHLcZ/view")
+        return None
+
+def download_complete_repo(local_dir="./u2net-mvtec-loco",
+                          repo_id="zhiqing0205/u2net-mvtec-loco-segmentation"):
+    """
+    Download complete repository from HuggingFace Hub
+    
+    Args:
+        local_dir: Local directory to save the repo
+        repo_id: HuggingFace repository ID
+    """
+    
+    print(f"Downloading complete repository: {repo_id}")
+    
+    try:
+        # Download entire repository
+        snapshot_download(
+            repo_id=repo_id,
+            local_dir=local_dir,
+            local_dir_use_symlinks=False
+        )
+        
+        print(f"Repository downloaded successfully to: {local_dir}")
+        return local_dir
+        
+    except Exception as e:
+        print(f"Error downloading repository: {e}")
+        return None
+
+if __name__ == "__main__":
+    import argparse
+    
+    parser = argparse.ArgumentParser(description="Download U2NET model or complete repo from HuggingFace")
+    parser.add_argument("--model-only", action="store_true", 
+                       help="Download only the model file")
+    parser.add_argument("--complete-repo", action="store_true",
+                       help="Download the complete repository") 
+    parser.add_argument("--repo-id", type=str, 
+                       default="zhiqing0205/u2net-mvtec-loco-segmentation",
+                       help="HuggingFace repository ID")
+    parser.add_argument("--force", action="store_true",
+                       help="Force download even if files exist")
+    
+    args = parser.parse_args()
+    
+    if args.complete_repo:
+        download_complete_repo(repo_id=args.repo_id)
+    else:
+        download_u2net_model(repo_id=args.repo_id, force_download=args.force)

fg_mask_mvtec_loco.zip

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:46d1b9829d323c7388a42841fb7f7759dd82027714697129b22ade369b98c714
+size 27964370

gradio/demo.py

@@ -0,0 +1,37 @@
+import cv2
+import paddlehub as hub
+import gradio as gr
+import torch
+
+# Images
+torch.hub.download_url_to_file('https://cdn.pixabay.com/photo/2018/08/12/16/59/ara-3601194_1280.jpg', 'parrot.jpg')
+torch.hub.download_url_to_file('https://cdn.pixabay.com/photo/2016/10/21/14/46/fox-1758183_1280.jpg', 'fox.jpg')
+
+model = hub.Module(name='U2Net')
+
+def infer(img):
+  result = model.Segmentation(
+      images=[cv2.imread(img.name)],
+      paths=None,
+      batch_size=1,
+      input_size=320,
+      output_dir='output',
+      visualization=True)
+  return result[0]['front'][:,:,::-1], result[0]['mask']
+
+inputs = gr.inputs.Image(type='file', label="Original Image")
+outputs = [
+           gr.outputs.Image(type="numpy",label="Front"),
+           gr.outputs.Image(type="numpy",label="Mask")
+           ]
+
+title = "U^2-Net"
+description = "demo for U^2-Net. To use it, simply upload your image, or click one of the examples to load them. Read more at the links below."
+article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2005.09007'>U^2-Net: Going Deeper with Nested U-Structure for Salient Object Detection</a> | <a href='https://github.com/xuebinqin/U-2-Net'>Github Repo</a></p>"
+
+examples = [
+  ['fox.jpg'],
+  ['parrot.jpg']
+]
+
+gr.Interface(infer, inputs, outputs, title=title, description=description, article=article, examples=examples).launch()

model/__init__.py

@@ -0,0 +1,2 @@
+from .u2net import U2NET
+from .u2net import U2NETP