Train and test on animal kingdom dataset

code/README.md

@@ -47,18 +47,33 @@ model_file = 'all_animals_no_pretrain_60.pth' # Path to the model file
 # Initialize the class
 model = Model()
 # Test the model on Protocol 1
-model.test(images_directory, mpii_json, model_file)
+model.test(images_directory, protocol='P1', model=model_file)
+# Test the model on Protocol 2
+model.test(images_directory, protocol='P2', model=model_file)
+# Test the model on birds class Protocol 3
+model.test(images_directory, protocol='P3', model=model_file, animal_class='bird')
+# Test the model on reptiles class Protocol 3
+model.test(images_directory, protocol='P3', model=model_file, animal_class='reptile')
+# Test the model on mammals class Protocol 3
+model.test(images_directory, protocol='P3', model=model_file, animal_class='mammal')
+# Test the model on fish class Protocol 3
+model.test(images_directory, protocol='P3', model=model_file, animal_class='fish')
+# Test the model on amphibian class Protocol 3
+model.test(images_directory, protocol='P3', model=model_file, animal_class='amphibian')
 ```
 You can also start to train your model or pretrain on top of ours
 ```python
 # train a VHR model
-train_json = '/content/AJA-pose/code/annotation/ak_P2/train.json' # labels for the train set
-valid_json = '/content/AJA-pose/code/annotation/ak_P2/test.json' # Labels for the validation set
+train_json = '' # labels for the train set (train.json)
+valid_json = '' # Labels for the validation set (test.json)
 model_file = '' # A pytorch model file to pretrain on.
 model.train(images_directory, train_json, valid_json, pretrained=model_file)
+
+# Train a model on a particular class e.g (Ampibian)
+model.train(images_directory, protocol='P3', animal_class='amphibian', model=model_file)
 ```
 
-### Results
+## Results
 A sanity check on our model.
 
 ![image3](https://github.com/Antony-gitau/AJA-pose/assets/88529649/266e526c-48aa-4401-b411-5f161a734c83)
@@ -70,7 +85,7 @@ Ground Truth
 Predictions
 ![image](https://github.com/Antony-gitau/AJA-pose/assets/88529649/efe360f8-3d5f-44b4-a396-a364096a2b4d)
 
-### Performance
+## Performance
 The performance of our model on the different animal classes is as shown below.
 
 | Animal Class | Samples |  | Head | Shoulder | Elbow | Wrist | Hip | Knee | Ankle | Mouth | Tail | Mean |
@@ -82,7 +97,10 @@ The performance of our model on the different animal classes is as shown below.
 | Amphibian | 1279 |  | 98.128 | 94.342 | 97.948 | 98.508 | 95.491 | 94.957 | 94.319 | 98.702 | 99.568 | 95.493 |
 
 The model performance on Protocol 1 and Protocol 2 is as shown below.
+
 | Protocol | Samples |  | Head | Shoulder | Elbow | Wrist | Hip | Knee | Ankle | Mouth | Tail | Mean |
 |----------|---------|--|------|----------|-------|-------|-----|------|-------|-------|------|------|
-| P1       |    6620     |  | 95.756 | 93.637 | 89.774 | 88.179 | 98.975 | 97.582 | 94.326 | 98.447 | 95.112 |  |
-| P2       |         |  | 91.538 | 85.291 | 84.662 | 85.587 | 90.457 | 88.097 | 85.239 | 96.723 | 83.925 | 89.553 |
+| P1       |    6620     |  | 94.230 | 91.054 | 90.806 | 90.920 | 94.414 | 93.233 | 92.094 | 96.867 | 92.346 | 93.073 |
+| P2       | 2883    |  | 88.683 | 75.815 | 80.223 | 81.136 | 85.568 | 83.840 | 82.028 | 94.799 | 72.506 | 83.711 |
+
+