Shadow0482/ResNet50-APTOS-DR-ONNX

base_model: ResNet50 tags:

• image-classification
• diabetic-retinopathy
• onnx license: mit

ResNet50-APTOS-DR-ONNX Model

This repository contains a ResNet50 model, originally trained for Diabetic Retinopathy (DR) detection on the APTOS dataset, exported to ONNX format for efficient inference.

Model Overview

• Architecture: ResNet50
• Task: Diabetic Retinopathy Classification (5 classes: No DR, Mild DR, Moderate DR, Severe DR, Proliferative DR)
• Format: ONNX (Opset 18)

Usage (ONNX Inference)

To use this model for inference, you will need the onnxruntime library. Below is a basic example:

import onnxruntime as ort
import numpy as np
from PIL import Image
from torchvision import transforms

ONNX_MODEL_PATH = "mithu-vit.onnx" # Path to the downloaded ONNX model
CLASSES = ["No DR", "Mild DR", "Moderate DR", "Severe DR", "Proliferative DR"]

# Image preprocessing (matching the training pipeline)
preprocess = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

def predict_image(image_path):
    img = Image.open(image_path).convert('RGB')
    input_tensor = preprocess(img)
    input_numpy = input_tensor.unsqueeze(0).numpy() # Add batch dimension

    session = ort.InferenceSession(ONNX_MODEL_PATH)
    input_name = session.get_inputs()[0].name
    output_name = session.get_outputs()[0].name

    outputs = session.run([output_name], {input_name: input_numpy})
    logits = outputs[0][0]
    probs = np.exp(logits) / np.sum(np.exp(logits))
    pred_index = np.argmax(probs)

    print(f"Predicted Class: {CLASSES[pred_index]} (Class {pred_index})")
    print(f"Confidence: {probs[pred_index] * 100:.2f}%")
    print("All Probabilities:")
    for i, p in enumerate(probs):
        print(f"  {CLASSES[i]}: {p*100:.2f}%")

# Example usage:
# predict_image("path/to/your/image.jpg")

Fine-tuning

The original model was trained using PyTorch. If you wish to fine-tune this model on a custom dataset or for a slightly different task, you can use the original PyTorch weights (if available) or adapt the ONNX model for further training in a suitable framework.

Steps for fine-tuning generally involve:

1. Load the pre-trained model: Start with the original PyTorch model or a version compatible with transfer learning.
2. Prepare your dataset: Ensure your images are properly labeled and preprocessed (resized to 224x224, normalized with ImageNet stats).
3. Modify the head: Replace the final classification layer to match the number of classes in your new dataset.
4. Define optimizer and loss function: Choose appropriate settings for your fine-tuning task.
5. Train: Fine-tune the model, typically with a lower learning rate than initial training, focusing on training the new head and potentially unfreezing earlier layers for more granular adjustments.
6. Export to ONNX: After fine-tuning, export your updated model to ONNX format following similar steps to the original export process.

Recommended Frameworks for Fine-tuning:

• PyTorch
• TensorFlow/Keras