File size: 2,523 Bytes
c4a3bf8 2228b39 c4a3bf8 867ceb6 c4a3bf8 2228b39 c4a3bf8 2228b39 867ceb6 2228b39 c4a3bf8 2228b39 c4a3bf8 2228b39 c4a3bf8 2228b39 c4a3bf8 2228b39 c4a3bf8 2228b39 c4a3bf8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 |
---
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
|