Update README.md

README.md

@@ -49,11 +49,28 @@ This repo hosts baseline **YOLOv7** models trained on the **N-MARINE** dataset (
 > Supplementary + scripts: https://github.com/Pentaerythrittetranitrat/N-MARINE_dataset_supplementary
 
 ## Model list
-
-- `yolov7-nmarine-no-class-weights.pt` — best overall (recommended)
-- `yolov7-nmarine-class-weights.pt` — improves some rare classes but slightly reduces overall mAP
-
-Each model outputs **9 classes** (species IDs in `metadata/classes.csv` within the dataset/supplementary repo).
+The released models follow this directory structure, where each `datasetX` (X = 1–5) corresponds to one of the five cross-validation splits:
+
+Each `datasetX` directory contains two trained models:
+- `classweights/best.pt` — trained with inverse-frequency class weights.  
+- `no_classweights/best.pt` — trained without class weights (baseline, recommended).  
+
+The **dataset splits** used for training can be found in the supplementary repository:  
+[https://github.com/Pentaerythrittetranitrat/N-MARINE_dataset_supplementary](https://github.com/Pentaerythrittetranitrat/N-MARINE_dataset_supplementary)
+> **Note:** The models under `dataset5` correspond to the most balanced split across classes.
+
+Each model outputs **9 classes**. The **class order (index → class)** for all models is: 
+| Index | Class name           |
+|-------|----------------------|
+| 0     | Atlantic Cod         |
+| 1     | Roughhead Grenadier  |
+| 2     | Atlantic Halibut     |
+| 3     | Redfish Mentella     |
+| 4     | Thorny Skate         |
+| 5     | Striped Wolffish     |
+| 6     | Spinytail Skate      |
+| 7     | Whelk                |
+| 8     | Northern Wolffish    |
 
 ## Intended use
 
@@ -88,7 +105,7 @@ pip install -r requirements.txt
 
 # 2) Run inference
 python detect.py \
-  --weights /path/to/yolov7-nmarine-no-class-weights.pt \
+  --weights /path/to/N-MARINE_baseline_classifiers/n-marine_weights/dataset5/no_classweights/best.pt \
   --source /path/to/images_or_video \
   --img-size 640 \
   --conf-thres 0.25 \
@@ -97,7 +114,7 @@ python detect.py \
 ```
 ### Citation
 
-If you use these models or the dataset, please cite both:
+If you use the dataset, please cite:
 
 Dataset citation (plain text):
 Morris, C. J., Ayyagari, K. D., Porter, D., Nguyen, Q. K., Hanlon, J., & Whidden, C. (2025).
@@ -105,5 +122,7 @@ Newfoundland Marine Refuge Fish Classification Dataset (N-Marine).
 Government of Canada Open Data Portal.
 https://open.canada.ca/data/en/dataset/2ae46860-f82a-4127-bb1f-b02e36ef6a70
 
+If you use the models, please cite:
+
 Model citation (plain text):
 TODO: Add