Upload README.md with huggingface_hub

README.md

@@ -4,36 +4,79 @@ tags:
 - tattoo-segmentation
 - image-segmentation
 - pytorch
+- unet
+- edge-detection
 ---
 
-# Deep Tattoo Segmentation
+# Deep Tattoo Segmentation v5
 
-2-Stage tattoo extraction model using Attention U-Net.
+Edge-Aware Attention U-Net for precise tattoo extraction with transparent background.
 
 ## Models
 
-- `best_model_pytorch.pth`: Stage 1 - Skin detection (Val Dice: 93.08%)
-- `tattoo_stage2_best.pth`: Stage 2 v1 - Tattoo extraction (Val Dice: 75.77%)
-- `tattoo_stage2_v2_best.pth`: Stage 2 v2 - Fine-line tattoo detection (Val Dice: 73.63%)
-- `tattoo_final_best.pth`: Final model trained on 9,000 images (Best checkpoint)
-- `tattoo_final_last.pth`: Final model trained on 9,000 images (Last checkpoint)
+- `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_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 model
+# Download v5 model
 model_path = hf_hub_download(
     repo_id="jun710/deep-tattoo",
-    filename="best_model_pytorch.pth"
+    filename="edge_aware_v3_clahe_best.pth"
 )
 
 # Load model
-model = torch.load(model_path, map_location='cpu')
+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