Update README.md

README.md

@@ -1,14 +1,16 @@
 # ModelAuditor Pre-trained Models
 
-Pre-trained ResNet50 models for medical image classification, used with the [ModelAuditor](https://github.com/lukaskuhndkfz/ModelAuditor) framework for AI-powered model auditing and robustness evaluation.
+Pre-trained models for medical image classification, used with the [ModelAuditor](https://github.com/lukaskuhndkfz/ModelAuditor) framework for AI-powered model auditing and robustness evaluation.
 
 ## Models
 
-| Model | Domain | Task | Input Size |
-|-------|--------|------|------------|
-| `camelyon17_resnet50_1_224.pt` | Pathology | Tumor detection in lymph node sections | 224x224 |
-| `chexpert_resnet50_1_224.pt` | Radiology | Chest X-ray classification | 224x224 |
-| `ham10000_resnet50_1_224.pt` | Dermatology | Skin lesion classification (melanoma vs. benign keratosis) | 224x224 |
+| Model | Architecture | Domain | Task | Input Size |
+|-------|-------------|--------|------|------------|
+| `camelyon17_resnet50_1_224.pt` | ResNet50 | Pathology | Tumor detection in lymph node sections | 224x224 |
+| `chexpert_resnet50_1_224.pt` | ResNet50 | Radiology | Chest X-ray classification | 224x224 |
+| `ham10000_resnet50_1_224.pt` | ResNet50 | Dermatology | Skin lesion classification (melanoma vs. benign keratosis) | 224x224 |
+| `cifar10.pth` | ResNet50 | General | CIFAR-10 image classification | 32x32 |
+| `DermaMNIST_resnet18.pth` | ResNet18 | Dermatology | Skin lesion classification (7 classes) | 28x28 |
 
 ## Usage
 
@@ -23,9 +25,10 @@ huggingface-cli download lukaskuhndkfz/ModelAuditor --local-dir models
 # Or download individually
 huggingface-cli download lukaskuhndkfz/ModelAuditor ham10000_resnet50_1_224.pt --local-dir models
 ```
-Use with ModelAuditor
 
-```
+### Use with ModelAuditor
+
+```bash
 git clone https://github.com/lukaskuhndkfz/ModelAuditor
 cd ModelAuditor
 pip install -e ".[medical]"
@@ -34,8 +37,9 @@ pip install -e ".[medical]"
 python main.py --model resnet50 --dataset ham10000 --weights models/ham10000_resnet50_1_224.pt
 ```
 
-Load in PyTorch
-```
+### Load in PyTorch
+
+```python
 import torch
 from torchvision.models import resnet50
 
@@ -44,8 +48,25 @@ model.load_state_dict(torch.load("ham10000_resnet50_1_224.pt", map_location="cpu
 model.eval()
 ```
 
-Datasets
+For DermaMNIST (ResNet18):
+
+```python
+import torch
+from torchvision.models import resnet18
+
+model = resnet18(num_classes=7)
+model.load_state_dict(torch.load("DermaMNIST_resnet18.pth", map_location="cpu"))
+model.eval()
+```
+
+## Training
+
+Training scripts for all ResNet50 models are available in the [ModelAuditor repository](https://github.com/lukaskuhndkfz/ModelAuditor) under `training/`.
+
+## Datasets
 
 - Camelyon17: https://wilds.stanford.edu/datasets/#camelyon17
 - CheXpert: https://stanfordmlgroup.github.io/competitions/chexpert/
-- HAM10000: https://dataverse.harvard.edu/dataset.xhtml?persistentId=doi:10.7910/DVN/DBW86T
+- HAM10000: https://dataverse.harvard.edu/dataset.xhtml?persistentId=doi:10.7910/DVN/DBW86T
+- CIFAR-10: https://www.cs.toronto.edu/~kriz/cifar.html
+- DermaMNIST: https://medmnist.com/