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

```python
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](https://pytorch.org/)
- [TensorFlow/Keras](https://www.tensorflow.org/)