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