|
|
--- |
|
|
license: mit |
|
|
tags: |
|
|
- tattoo-segmentation |
|
|
- image-segmentation |
|
|
- pytorch |
|
|
- unet |
|
|
- edge-detection |
|
|
--- |
|
|
|
|
|
# Deep Tattoo Segmentation v5/v7 |
|
|
|
|
|
Edge-Aware Attention U-Net for precise tattoo extraction with transparent background. |
|
|
|
|
|
## Models |
|
|
|
|
|
- `edge_aware_v3_clahe_best.pth`: **v5 Model** - Best performance (Val Dice: **90.50%**) |
|
|
- Edge-Aware Attention U-Net architecture |
|
|
- CLAHE preprocessing for lighting invariance |
|
|
- Trained on 24 manual labels + 165 auto-generated masks |
|
|
- Test-Time Augmentation (TTA) support |
|
|
|
|
|
- `edge_aware_v7_samrefiner_best.pth`: **v7 Model with SAMRefiner** (Val Dice: 74.25%) |
|
|
- Same architecture as v5 |
|
|
- Trained with 122 SAMRefiner-refined masks + 25 v2 fallback |
|
|
- SAMRefiner (ICLR 2025) for mask refinement |
|
|
- Better edge quality but lower validation score (due to v2-only validation set) |
|
|
|
|
|
- `edge_aware_improved_best.pth`: Base Model (Val Dice: 79.17%) |
|
|
- Foundation model before hybrid training |
|
|
|
|
|
## Key Features |
|
|
|
|
|
- **High Accuracy**: 90.50% Dice coefficient on validation set |
|
|
- **Edge Detection**: Specialized edge-aware attention mechanism |
|
|
- **Lighting Invariant**: CLAHE preprocessing handles various lighting conditions |
|
|
- **Transparent Output**: Extracts tattoos with alpha channel for transparent background |
|
|
- **Production Ready**: Optimized for inference with TTA |
|
|
|
|
|
## Usage |
|
|
|
|
|
```python |
|
|
import torch |
|
|
import cv2 |
|
|
import numpy as np |
|
|
from huggingface_hub import hf_hub_download |
|
|
|
|
|
# Download v5 model |
|
|
model_path = hf_hub_download( |
|
|
repo_id="jun710/deep-tattoo", |
|
|
filename="edge_aware_v3_clahe_best.pth" |
|
|
) |
|
|
|
|
|
# Load model |
|
|
checkpoint = torch.load(model_path, map_location='cpu') |
|
|
# Use with EdgeAwareAttentionUNet from the repository |
|
|
|
|
|
# Preprocess with CLAHE |
|
|
image = cv2.imread("tattoo.jpg") |
|
|
clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8)) |
|
|
# ... apply model and extract tattoo |
|
|
``` |
|
|
|
|
|
## Model Architecture |
|
|
|
|
|
- **Encoder**: EfficientNet-B3 backbone with edge detection branch |
|
|
- **Decoder**: Attention-based skip connections |
|
|
- **Output**: Binary segmentation mask (tattoo vs background) |
|
|
|
|
|
## Training Details |
|
|
|
|
|
- Image Size: 256x256 |
|
|
- Batch Size: 8 |
|
|
- Optimizer: Adam (lr=1e-4) |
|
|
- Loss: Boundary-Aware Loss (Dice + BCE + Edge) |
|
|
- Augmentation: Strong geometric + color transformations |
|
|
- Training Data: 189 images (24 manual + 165 auto-generated) |
|
|
|
|
|
## Performance |
|
|
|
|
|
- Validation Dice: 90.50% |
|
|
- Test Coverage: 52 diverse images |
|
|
- Success Rate: ~97% on typical tattoos |
|
|
- Limitation: Fine details on very delicate outlines (~3% cases) |
|
|
|
|
|
## Repository |
|
|
|
|
|
https://github.com/enjius/deep-tattoo |
|
|
|
