You can deploy this model and start earning money today!

palmyz0/
Skin-Diseases-ViT-finetune-with-optuna

Name: Skin-Diseases-ViT-finetune-with-optuna
Availability: InStock

Views

No views yet

🩺 Skin Diseases Classification with Optimized Vision Transformers (ViT)

This model is a fine-tuned Vision Transformer (ViT-B/16) architecture designed to classify skin disease images into 5 categories. The model has been optimized using Optuna for hyperparameter tuning to achieve high accuracy and robustness on clinical images.

🎯 Model Performance

Accuracy: 93.31% (on Validation set)
Model Architecture: Vision Transformer (ViT-B/16) (Pre-trained on ImageNet)
Optimization Strategy: Bayesian Optimization via Optuna
Input Resolution: 224x224 pixels

📂 Class Labels

The model predicts one of the following 4 classes:

glioma
meningioma
notumor (Healthy)
pituitary

🛠️ How to Use (Python Code)

You can use this model directly with the Hugging Face huggingface_hub library or raw PyTorch.

Important: Images must be normalized using ImageNet statistics before inference.

python

import torch
import torch.nn as nn
from torchvision import transforms, models
from PIL import Image
from huggingface_hub import hf_hub_download

# 1. Define Model Architecture (ViT)
class VitModel(nn.Module):
    def __init__(self, num_classes=5, dropout_rate=0.0):
        super().__init__()
        # Load ViT-B/16 architecture
        self.backbone = models.vit_b_16(weights=None) # No need for weights when loading custom state_dict
        
        # Modify the head to match the training configuration
        input_dim = self.backbone.heads[0].in_features
        self.backbone.heads = nn.Sequential(
            nn.Linear(input_dim, 512),
            nn.LayerNorm(512),
            nn.GELU(),
            nn.Dropout(dropout_rate),
            nn.Linear(512, num_classes)
        )

    def forward(self, x):
        return self.backbone(x)

# 2. Configuration
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Replace with your actual Repo ID
REPO_ID = "Suphawit/skin-disease-vit-optuna" 
FILENAME = "final_best_skin_diseases_model.pth"

# 3. Load Model
model_path = hf_hub_download(repo_id=REPO_ID, filename=FILENAME)
model = VitModel(num_classes=5)
model.load_state_dict(torch.load(model_path, map_location=DEVICE))
model.to(DEVICE)
model.eval()

# 4. Predict
# Define Transform (Same as Validation)
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# Load Image
img_path = "path_to_your_skin_image.jpg"
img = Image.open(img_path).convert('RGB')
input_tensor = transform(img).unsqueeze(0).to(DEVICE)

# Inference
with torch.no_grad():
    output = model(input_tensor)
    probs = torch.nn.functional.softmax(output, dim=1)
    confidence, prediction = torch.max(probs, 1)
    
    print(f"Predicted Class Index: {prediction.item()}")
    print(f"Confidence: {confidence.item()*100:.2f}%")

"Try out the live model demo here: https://huggingface.co/spaces/palmyz0/brain-tumor-detection-demo"

👤 Author

Suphawit MeeSak

Junior AI Engineer / Junior Data Scientist

🔗 LinkedIn: Suphawit MeeSak
📧 Email: Suphawit11@icloud.com

palmyz0/Skin-Diseases-ViT-finetune-with-optuna