t0m-R/vit-sem-classification

Vision Transformer for SEM Image Classification

This is a fine-tuned Vision Transformer (ViT-B/32) model for classifying Scanning Electron Microscopy (SEM) images into 10 distinct categories of nanostructures [1].

This model was developed as part of the NFFA-DI (Nano Foundries and Fine Analysis Digital Infrastructure) project, funded by the European Union's NextGenerationEU program.

Model Description

The model is based on the google/vit-base-patch32-224-in21k checkpoint and has been fine-tuned for a 10-class image classification task on SEM images. The 10 categories cover a wide range of nanostructures:

1. Porous Sponge
2. Patterned Surface
3. Particles
4. Films and Coated Surface
5. Powder
6. Tips
7. Nanowires
8. Biological
9. MEMS devices and electrodes
10. Fibres

How to Use

The following Python code shows how to load the model and its processor from the Hub and use it to classify a local SEM image.

from transformers import AutoImageProcessor, AutoModelForImageClassification
from PIL import Image
import torch

# Load the model and image processor from the Hub
model_name = "t0m-R/vit-sem-classification" 
image_processor = AutoImageProcessor.from_pretrained(model_name)
model = AutoModelForImageClassification.from_pretrained(model_name)

# Load and preprocess the image
image_path = "path/to/your/sem_image.jpg" 
try:
    image = Image.open(image_path).convert("RGB")

    # Prepare the image for the model
    inputs = image_processor(images=image, return_tensors="pt")

    # Run inference
    with torch.no_grad():
        logits = model(**inputs).logits
        predicted_label_id = logits.argmax(-1).item()
        predicted_label = model.config.id2label[predicted_label_id]

    print(f"Predicted Label: {predicted_label}")

except FileNotFoundError:
    print(f"Error: The file at {image_path} was not found.")

Training Data

This model was fine-tuned on the SEM Majority dataset, the first annotated set of scanning electron microscopy images for nanoscience.

The dataset consists of 25,537 SEM images manually classified into 10 categories. The classification labels were verified by a group of nanoscientists, and only images validated by the majority of the group were included in the dataset.

The dataset is publicly available at: https://doi.org/10.23728/b2share.e344a8afef08463a855ada08aadbf352

[1] Aversa, Rossella, et al. "The first annotated set of scanning electron microscopy images for nanoscience." Scientific data 5.1 (2018): 1-10.