Update README.md

README.md

@@ -10,7 +10,8 @@ metrics:
 **Date:** 2025-09-17
 
 ## Overview
-A Capsule Network (CapsNet) implemented in TensorFlow/Keras to classify **four lung-disease categories** from masked chest X-ray images. The model uses routing-by-agreement and **margin loss**, and was trained with MLflow tracking. Input images are resized to **256×256×3**.
+A Capsule Network (CapsNet) implemented in TensorFlow/Keras to classify **four lung-disease categories** from masked chest X-ray images. 
+The model uses routing-by-agreement and **margin loss**, and was trained with MLflow tracking. Input images are resized to **256×256×3**.
 
 > ⚠️ **Not a medical device.** Outputs are for research/education. Clinician review is required before any clinical use.
 
@@ -20,7 +21,12 @@ A Capsule Network (CapsNet) implemented in TensorFlow/Keras to classify **four l
 - **Out-of-scope:** Direct clinical decision-making; deployment on patient data without formal validation and regulatory clearance.
 
 ## Model Details
-- **Architecture:** CapsNet with `PrimaryCaps` and `DigitCaps`; **routing iterations:** 3; **first Conv2D kernel size:** tuned over [5, 7, 9, 10, 11] after an exploratory sweep (3-epoch runs over [3,5,7,9,10,11]).  
+- **Architecture:** CapsNet with `PrimaryCaps` and `DigitCaps`
+  - built-in augmentation(disabled by default during inference) and rescaling layers:
+    - ´layers.RandomRotation(0.1)`, 'layers.RandomTranslation(height_factor=0.1, width_factor=0.1)', 'layers.RandomZoom(0.1)'
+    - 'Rescaling(1./255)'
+  - routing iterations: 3;
+  - first Conv2D kernel size: tuned over [5, 7, 9, 10, 11] after an exploratory sweep (3-epoch runs over [3,5,7,9,10,11]).  
 - **Loss:** `margin_loss` (capsule margin).  
 - **Optimizer:** `Adam` with a learning-rate scheduler (`lr_scheduler`).  
 - **Metrics:** accuracy.  
@@ -45,15 +51,17 @@ and
 https://www.kaggle.com/datasets/omkarmanohardalvi/lungs-disease-dataset-4-types (License: Unknown)
 
 ## Evaluation
-- **Protocol:** train on the balanced training set, validate with a held-out split; select first-layer kernel size from [5,7,9,10,11] based on validation performance; evaluate best checkpoint on the unseen test set.  
+- **Protocol:** train on the balanced training set, validate with a held-out split; select first-layer kernel size from [5,7,9,10,11] based on validation performance; evaluate best checkpoint on the external test set.  
 - **Reported metrics:**  Metrics for each model being trained are in `capsnet_training_metrics_all_runs.csv`
 
-## External test
->I excluded the lung opacity class from external tests because it often co-occurs with other diseases. This  makes it challenging to classify accurately.
+## External test and the winner model
+>I excluded the lung opacity class from external tests because it often co-occurs with other diseases. This class makes it challenging to classify accurately.
 However, the reported accuracy for the remaining classes is still quite representative.
+- krnl9 (best COVID F1, balanced performance) (the one here presented)-> Results CSV: `test_on_external_dataset_capsnet_lung_disease_classifier_krnl9.csv` 
+- krnl11 (highest accuracy, strongest Viral Pneumonia detection) 
+- optionally krnl10 (solid performance, close to the leaders)
 
 
- Results CSV: `test_on_external_dataset_capsnet_lung_disease_classifier_krnl9.csv` 
 
 ![alt text](image.png)
 
@@ -77,7 +85,7 @@ from tensorflow.keras.utils import load_img, img_to_array
 import numpy as np
 
 # Load trained Keras model
-model = tf.keras.models.load_model("path/to/model.keras",
+model = tf.keras.models.load_model("your/path/to/model.keras",
                                    custom_objects={"margin_loss": margin_loss,
                                                    "PrimaryCaps": PrimaryCaps,
                                                    "DigitCaps": DigitCaps,