unet学习资料

U-Net Image Segmentation - Learning Example

A hands-on U-Net implementation for learning deep learning image segmentation, using simulated cell microscopy data.

Project Structure

.
├── Unet.py                 # Model definition + training + visualization
├── predict.py              # Inference (auto-generates test data + runs prediction)
├── export.py               # One-click export: ONNX / RKNN / Allwinner guide
├── generate_test_data.py   # Generate independent test images
├── README.md
├── cell_data/              # Training data (auto-generated)
│   ├── images/
│   └── masks/
├── test_data/              # Test data (auto-generated by predict.py)
└── predict_results/        # Inference output images

Quick Start

1. Train the model

python Unet.py

This will:

• Generate 100 simulated cell images + masks (if not exist)
• Train U-Net for 20 epochs
• Save model weights to unet_model.pth
• Generate 5 visualization images:
  • unet_architecture.png - Network architecture diagram
  • unet_feature_maps.png - Encoder feature maps
  • unet_training_curves.png - Loss / Dice / IoU curves
  • unet_predictions.png - Prediction comparison
  • unet_metrics.png - Dice & IoU distribution

2. Generate test data

python generate_test_data.py

Creates 20 test images in test_data/ with a different random seed from training data.

3. Run inference

# One-click: auto-generate test data + load model + batch inference
python predict.py

# Or specify your own image / folder
python predict.py --image test_data/test_000.png
python predict.py --image test_data/

# Custom threshold (lower = more aggressive, higher = more conservative)
python predict.py --threshold 0.3

Results are saved to predict_results/. Each input image produces a xxx_result.png with three panels: input image, probability heatmap, and overlay.

Requirements

torch
torchvision
numpy
Pillow
scikit-learn
matplotlib
tqdm

Install all at once:

pip install torch torchvision numpy Pillow scikit-learn matplotlib tqdm

Export for Edge Deployment

python export.py                        # Export ONNX + try RKNN
python export.py --format onnx          # ONNX only
python export.py --format rknn          # ONNX + RKNN
python export.py --target rk3566        # Specify RK chip
python export.py --quantize             # INT8 quantization (faster on NPU)

Key Concepts

Tips

• No GPU? Reduce EPOCHS or IMG_SIZE in Unet.py for faster training.
• Want to use your own images? Put them in a folder and run predict.py --image your_folder/.
• The model expects RGB images of any size (auto-resized to 256x256 internally).