README.md

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