Add counts summary

ocean-data-factory-sweden · Nov 22, 2023 · 950559c · 950559c
1 parent 4945c15
commit 950559c
Showing 1 changed file with 9 additions and 8 deletions.
diff --git a/kso_utils/yolo_utils.py b/kso_utils/yolo_utils.py
@@ -1858,14 +1858,15 @@ def custom_class(x):
     logging.info(
         f"Saving tracking file to {str(Path(tracking_folder, 'tracking_clean.csv'))}"
     )
-    return (
-        total_df[
-            (total_df.max_frame_diff <= avg_diff_frames)
-            & (total_df.frame_length >= min_frames_length)
-        ]
-        .sort_index()
-        .to_csv(str(Path(tracking_folder, "tracking_clean.csv")))
-    )
+    filtered_df = total_df[
+        (total_df.max_frame_diff <= avg_diff_frames)
+        & (total_df.frame_length >= min_frames_length)
+    ].sort_index()
+    if len(names) > 0:
+        logging.info(filtered_df["species_name"].value_counts())
+    else:
+        logging.info(filtered_df["class_id"].value_counts())
+    return filtered_df.to_csv(str(Path(tracking_folder, "tracking_clean.csv")))
 
 
 def main():