From 2d494893fc4da1a5ef0f4b67164617775df70d35 Mon Sep 17 00:00:00 2001 From: bright1205 Date: Sun, 28 Jan 2024 16:29:11 +0000 Subject: [PATCH] Refactor class DetectionsManager to handle empty detections_all array --- examples/traffic_analysis/inference_example.py | 8 ++++---- examples/traffic_analysis/ultralytics_example.py | 8 ++++---- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/examples/traffic_analysis/inference_example.py b/examples/traffic_analysis/inference_example.py index 1cba7ea3..5f06a710 100644 --- a/examples/traffic_analysis/inference_example.py +++ b/examples/traffic_analysis/inference_example.py @@ -48,10 +48,10 @@ class DetectionsManager: self.counts.setdefault(zone_out_id, {}) self.counts[zone_out_id].setdefault(zone_in_id, set()) self.counts[zone_out_id][zone_in_id].add(tracker_id) - - detections_all.class_id = np.vectorize( - lambda x: self.tracker_id_to_zone_id.get(x, -1) - )(detections_all.tracker_id) + if len(detections_all) > 0: + detections_all.class_id = np.vectorize( + lambda x: self.tracker_id_to_zone_id.get(x, -1) + )(detections_all.tracker_id) return detections_all[detections_all.class_id != -1] diff --git a/examples/traffic_analysis/ultralytics_example.py b/examples/traffic_analysis/ultralytics_example.py index 8cf3a0c4..11c8a45f 100644 --- a/examples/traffic_analysis/ultralytics_example.py +++ b/examples/traffic_analysis/ultralytics_example.py @@ -47,10 +47,10 @@ class DetectionsManager: self.counts.setdefault(zone_out_id, {}) self.counts[zone_out_id].setdefault(zone_in_id, set()) self.counts[zone_out_id][zone_in_id].add(tracker_id) - - detections_all.class_id = np.vectorize( - lambda x: self.tracker_id_to_zone_id.get(x, -1) - )(detections_all.tracker_id) + if len(detections_all) > 0: + detections_all.class_id = np.vectorize( + lambda x: self.tracker_id_to_zone_id.get(x, -1) + )(detections_all.tracker_id) return detections_all[detections_all.class_id != -1]