melanoma-segmentation

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:

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