| --- |
| license: mit |
| language: |
| - en |
| tags: |
| - semantic-segmentation |
| - pytorch |
| - unet |
| - resnet34 |
| - materials-science |
| - microscopy |
| - sem |
| - computer-vision |
| - segmentation-models-pytorch |
| datasets: |
| - safdar/sem-ni-wc-metal-matrix-composites |
| metrics: |
| - mean_iou |
| - dice |
| library_name: segmentation_models_pytorch |
| pipeline_tag: image-segmentation |
| --- |
| |
| # SEM Microstructure Semantic Segmentation |
|
|
| UNet with a pretrained ResNet-34 encoder for pixel-wise segmentation of Scanning Electron Microscopy (SEM) images of additively manufactured Ni-WC metal matrix composites. Trained to identify five microstructural phases at pixel level. |
|
|
| ## Model Details |
|
|
| | Property | Detail | |
| |------------------|---------------------------------------------| |
| | Architecture | UNet + ResNet-34 encoder | |
| | Encoder weights | ImageNet pretrained | |
| | Input channels | 1 (grayscale SEM image) | |
| | Output classes | 5 (pixel-wise segmentation) | |
| | Framework | PyTorch + segmentation_models_pytorch | |
| | Training device | Apple MPS (M-series) | |
|
|
| ## Classes |
|
|
| | ID | Phase | |
| |----|----------------| |
| | 0 | Matrix | |
| | 1 | Carbide | |
| | 2 | Void | |
| | 3 | Reprecipitate | |
| | 4 | Dilution Zone | |
|
|
| ## Test Performance |
|
|
| Evaluated on 54 held-out test images, scored at image level to avoid batch-size bias. |
|
|
| | Metric | Mean | Std | Min | Max | |
| |--------|-------|-------|-------|-------| |
| | mIoU | 0.872 | 0.088 | 0.723 | 0.958 | |
| | mDice | 0.912 | 0.079 | 0.759 | 0.978 | |
| | mBF1 | 0.728 | 0.027 | 0.678 | 0.773 | |
|
|
| ### Per-Class Performance |
|
|
| | Class | IoU | Dice | |
| |---------------|-------|-------| |
| | Matrix | 0.939 | 0.969 | |
| | Carbide | 0.753 | 0.859 | |
| | Void | 0.976 | 0.988 | |
| | Reprecipitate | 0.891 | 0.942 | |
| | Dilution Zone | 0.881 | 0.937 | |
|
|
| ## Files |
|
|
| | File | Description | |
| |----------------------------|------------------------------------| |
| | `best_model.pth` | Best pretrained encoder checkpoint | |
| | `ScratchUNet_best_unet.pth`| Best scratch-built UNet checkpoint | |
|
|
| ## How to Use |
|
|
| ### Install dependencies |
| ```bash |
| pip install torch segmentation-models-pytorch huggingface_hub |
| ``` |
|
|
| ### Load and run inference |
| ```python |
| import torch |
| from huggingface_hub import hf_hub_download |
| import segmentation_models_pytorch as smp |
| |
| # Download checkpoint |
| ckpt_path = hf_hub_download( |
| repo_id="imranlabs/sem-microstructure-segmentation", |
| filename="best_model.pth" |
| ) |
| |
| # Rebuild architecture |
| model = smp.Unet( |
| encoder_name = "resnet34", |
| encoder_weights = None, # weights loaded from checkpoint |
| in_channels = 1, |
| classes = 5, |
| activation = None, |
| ) |
| |
| # Load weights |
| ckpt = torch.load(ckpt_path, map_location="cpu") |
| model.load_state_dict(ckpt["model_state"]) |
| model.eval() |
| |
| # Inference — input: (1, 1, H, W) float32 tensor normalised to [0, 1] |
| with torch.no_grad(): |
| logits = model(image_tensor) # (1, 5, H, W) |
| preds = torch.argmax(logits, dim=1) # (1, H, W) class labels |
| ``` |
|
|
| ## Training Details |
|
|
| - **Loss:** 0.6 × Weighted CrossEntropy + 0.4 × DiceLoss |
| - **Optimizer:** AdamW (weight decay 1e-5) |
| - **Scheduler:** ReduceLROnPlateau |
| - **Strategy:** Encoder frozen for first 4 epochs, then full fine-tuning |
| - **Early stopping:** Monitored validation mIoU |
| - **Best checkpoint:** Epoch 12, validation mIoU = 0.876 |
|
|
| ## Dataset |
|
|
| Safdar, M. (2025). *Scanning Electron Microscopy (SEM) Dataset of Additively Manufactured Ni-WC Metal Matrix Composites for Semantic Segmentation* (Version 1). Zenodo. |
| https://doi.org/10.5281/zenodo.17315241 |
|
|
| ## Links |
|
|
| - GitHub: [sem-microstructure-segmentation](https://github.com/imranlabs/sem-microstructure-segmentation) |
|
|
| ## Author |
|
|
| **Imran Khan** — Physics PhD · SEM Metrology Engineer · Computer Vision / AI |
|
|
