Save predictions as JSON list

Previously was being stored as Python list (eg. with single quotes strings)
ultralytics · Jun 4, 2023 · fc3be9e · fc3be9e
1 parent e37d96a
commit fc3be9e
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/classify/predict.py b/classify/predict.py
@@ -30,6 +30,7 @@
 
 import argparse
 import csv
+import json
 import io
 import os
 import platform
@@ -152,7 +153,7 @@ def run(
                 top_pred_index = top5i[0]
                 csv_output.writerow({
                     'path': path,
-                    'top_5_predicted': [(names[j], prob[j].item()) for j in top5i],
+                    'top_5_predicted': json.dumps([(names[j], prob[j].item()) for j in top5i]),
                     'label': names[top_pred_index],
                     'confidence': f'{prob[top_pred_index]:.2f}'
                 })