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
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)