Melanoma Segmentation (U-Net + EfficientNet-B0 encoder)
U-Net segmentation model with an EfficientNet-B0 encoder, trained to segment melanoma lesions in dermoscopic images. Part of a 3-stage melanoma analysis pipeline (detection → segmentation → classification) — see melanoma-pipeline for the full inference chain.
Model Details
- Architecture: U-Net with EfficientNet-B0 encoder (ImageNet-pretrained backbone), via
segmentation_models_pytorch - Task: Binary semantic segmentation (lesion vs. background)
- Framework: PyTorch
- Input size: 256×256, normalized to [0, 1] (simple /255 scaling, no ImageNet mean/std normalization)
- Training data: ISIC 2018 Task 1 (lesion segmentation subset)
- File:
unet_melanoma.pth(state_dict)
Training Details
- Loss: Dice Loss (
smp.losses.DiceLoss(mode='binary')) - Optimizer: Adam, learning rate 0.001
- Epochs: 5
- Batch size: 8
- Train/val split: first 2,000 images for training, remainder for validation (sequential split, not shuffled)
Intended Use
This model is intended for research and educational purposes — as a component in an experimental melanoma analysis pipeline. It is not a certified medical device and must not be used for clinical diagnosis without validation by qualified medical professionals and regulatory approval.
How to Use
import torch
import segmentation_models_pytorch as smp
from huggingface_hub import hf_hub_download
# Download weights
weights_path = hf_hub_download(
repo_id="Ai-Adam-Six-Sigma/melanoma-segmentation",
filename="unet_melanoma.pth"
)
model = smp.Unet(
encoder_name="efficientnet-b0",
encoder_weights=None,
in_channels=3,
classes=1,
)
model.load_state_dict(torch.load(weights_path, map_location="cpu"))
model.eval()
# Preprocessing: resize to 256x256, scale to [0, 1] (no ImageNet mean/std normalization)
from PIL import Image
import numpy as np
img = Image.open("lesion.jpg").convert("RGB").resize((256, 256))
input_tensor = torch.tensor(np.array(img)).permute(2, 0, 1).float().unsqueeze(0) / 255.0
# Inference
with torch.no_grad():
mask_logits = model(input_tensor)
mask = torch.sigmoid(mask_logits) > 0.5
Training Data
Trained on the ISIC 2018 Task 1 lesion segmentation dataset (dermoscopic images with binary lesion masks). Split: first 2,000 images for training, remainder for validation (sequential, not shuffled).
Metrics
Trained for 5 epochs with binary Dice Loss:
| Epoch | Train Loss | Val Loss |
|---|---|---|
| 1 | 0.2058 | 0.1386 |
| 2 | 0.1346 | 0.1583 |
| 3 | 0.1195 | 0.1126 |
| 4 | 0.1091 | 0.1340 |
| 5 | 0.0972 | 0.1194 |
Since the loss is Dice Loss (1 - Dice coefficient), the final epoch corresponds to a validation Dice coefficient of ≈ 0.88. Note validation loss fluctuates across epochs (best at epoch 3: 0.1126 → Dice ≈ 0.89) rather than decreasing monotonically, likely due to the small dataset split and no learning rate scheduling.
Limitations
- Trained on dermoscopic images only; performance on smartphone/clinical photos is not validated.
- ISIC 2018 dataset may not represent all skin tones and lesion types equally.
- Not validated for clinical use.
License
Apache 2.0