CoolFace
Modelpublic

keanteng/swin-breast-cancer-classification-0728

sourceHugging Faceagpl-3.0updated 1y agoView on Hugging Face
0likes8downloads
Model Card

license: agpl-3.0 datasets:

  • —keanteng/miniddbs-jpeg base_model:
  • —microsoft/swin-base-patch4-window8-256 pipelinetag: image-classification libraryname: transformers tags:
  • —generative-ai
  • —medical-imaging
  • —swin-transformer
  • —breast-cancer
  • —classification ---

Breast Cancer Classification with Swin Transformer

This is a new version trained with some pipelines and parameters modification.

This repository contains a fine-tuned Swin Transformer model for breast cancer classification based on mammography images.

Due to the indistinguishable nature of the dataset various runs had been conducted to perform the original 3 classes classification according to the original DDSM dataset but the accuracy obtained is dismal (approx 67%) contrary to literature review of (>90%).

I have also explored dual input Swin Transformer using the Tumour Mask, however, similar dismal accuracy is obtained. We can look at the dataset and notice that the images all looks about the same except Normal. Thus, the detection strategy becomes detecting the presence of cancer by merging to Benign and Cancer images as a class against the Normal images.

With such approach, accuracy significant increases and achieve reliable performance.

Model Description

The model is based on the Swin Transformer Base architecture, fine-tuned on the Mini-DDBS-JPEG dataset for breast cancer classification. It uses a custom classification head consisting of a multi-layer perceptron with dropout for regularization.

Key Features

  • —Based on Swin Transformer architecture
  • —Input image size: 256x256 pixels
  • —Binary classification task (malignant vs benign)
  • —Mixed precision training for improved performance

Performance

The model was trained with class balancing techniques to handle data imbalance. Performance metrics on the test set:

MetricValue
Test Accuracy0.7672634271099744
Test Loss0.4335057186653547

For detailed performance metrics including precision, recall, and F1-score per class, please check the training notebook.

Usage

With Transformers Pipeline

python
from transformers import pipeline

classifier = pipeline("image-classification", model="keanteng/swin_v2_breast_cancer_classification")
result = classifier("path/to/mammogram.jpg")
print(result)
from transformers import AutoFeatureExtractor, AutoModelForImageClassification
from PIL import Image

# Load model and feature extractor
model = AutoModelForImageClassification.from_pretrained("keanteng/swin_v2_breast_cancer_classification")
feature_extractor = AutoFeatureExtractor.from_pretrained("keanteng/swin_v2_breast_cancer_classification")

# Prepare image
image = Image.open("path/to/mammogram.jpg").convert("RGB")
inputs = feature_extractor(images=image, return_tensors="pt")

# Get prediction
outputs = model(**inputs)
predicted_class_idx = outputs.logits.argmax(-1).item()
print(f"Predicted class: model.config.id2label[predicted_class_idx]")