[PolygonZone] - allow per class counting - #817
karanjakhar wants to merge 8 commits into
Conversation
SkalskiP
added
enhancement
|
Hi @karanjakhar 👋🏻 Thanks a lot for your interest in Supervision! I see you implemented the current Would you be willing to extend the logic and implement |
|
@SkalskiP Yes. I will add that also. |
|
Awesome! Let me know once you'll be ready for review. |
|
Hi @SkalskiP , I have made the changes. Please review. |
|
Hi @SkalskiP |
|
Hi @SkalskiP |
|
@karanjakhar could you pls check the conflicts? 🦩 |
There was a problem hiding this comment.
Pull request overview
This PR adds per-class counting functionality to the PolygonZone class, enabling tracking of object counts by class ID both inside and outside the zone. This enhancement addresses issue #791 and provides more granular detection statistics.
- Added four new dictionary attributes to track per-class counts for objects inside/outside the zone in current and cumulative totals
- Modified the trigger method to populate per-class count dictionaries while maintaining backward compatibility
- Added comprehensive test suite for the new per-class counting feature
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 5 comments.
| File | Description |
|---|---|
| supervision/detection/tools/polygon_zone.py | Adds four new dictionary attributes (class_in_current_count, class_out_current_count, class_in_total_count, class_out_total_count) and implements per-class counting logic in the trigger method |
| test/detection/test_polygon_zone.py | Adds new test file with parametrized tests covering empty detections, single detection inside zone, and single detection outside zone scenarios |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| @pytest.mark.parametrize( | ||
| "polygon, detections, expected_result, current_count, class_in_current_count,\ | ||
| class_out_current_count, class_in_total_count, class_out_total_count", | ||
| [ | ||
| ( | ||
| np.array([[0, 0], [10, 0], [10, 10], [0, 10]]), | ||
| mock_detections(xyxy=np.empty((0, 4)), class_id=[]), | ||
| np.array([]), | ||
| 0, | ||
| {}, | ||
| {}, | ||
| {}, | ||
| {}, | ||
| ), # empty detections | ||
| ( | ||
| np.array([[0, 0], [10, 0], [10, 10], [0, 10]]), | ||
| mock_detections(xyxy=[[5, 5, 10, 10]], class_id=[2], confidence=[0.9]), | ||
| np.array([True]), | ||
| 1, | ||
| {2: 1}, | ||
| {}, | ||
| {2: 1}, | ||
| {}, | ||
| ), # single detection in zone | ||
| ( | ||
| np.array([[0, 0], [10, 0], [10, 10], [0, 10]]), | ||
| mock_detections(xyxy=[[25, 25, 55, 55]], class_id=[1], confidence=[0.9]), | ||
| np.array([False]), | ||
| 0, | ||
| {}, | ||
| {1: 1}, | ||
| {}, | ||
| {1: 1}, | ||
| ), # single detection outside zone | ||
| ], | ||
| ) |
There was a problem hiding this comment.
The test suite is missing test cases for important scenarios including: multiple detections with the same class_id (to verify counting), multiple detections with different class_ids in the same frame, and detections without class_id (None). These edge cases should be tested to ensure robustness.
| for in_class_id in detections.class_id[is_in_zone]: | ||
| if in_class_id in self.class_in_current_count: | ||
| self.class_in_current_count[in_class_id] += 1 | ||
| else: | ||
| self.class_in_current_count[in_class_id] = 1 | ||
|
|
||
| if in_class_id in self.class_in_total_count: | ||
| self.class_in_total_count[in_class_id] += 1 | ||
| else: | ||
| self.class_in_total_count[in_class_id] = 1 | ||
|
|
||
| for out_class_id in detections.class_id[~is_in_zone]: | ||
| if out_class_id in self.class_out_current_count: | ||
| self.class_out_current_count[out_class_id] += 1 | ||
| else: | ||
| self.class_out_current_count[out_class_id] = 1 | ||
|
|
||
| if out_class_id in self.class_out_total_count: | ||
| self.class_out_total_count[out_class_id] += 1 | ||
| else: | ||
| self.class_out_total_count[out_class_id] = 1 | ||
|
|
There was a problem hiding this comment.
The trigger method doesn't handle the case where detections.class_id is None. Since class_id is an Optional field in the Detections class, attempting to index it with is_in_zone when it's None will raise a TypeError. Consider adding a check to ensure class_id is not None before accessing it, or documenting that class_id is required for per-class counting functionality.
| for in_class_id in detections.class_id[is_in_zone]: | |
| if in_class_id in self.class_in_current_count: | |
| self.class_in_current_count[in_class_id] += 1 | |
| else: | |
| self.class_in_current_count[in_class_id] = 1 | |
| if in_class_id in self.class_in_total_count: | |
| self.class_in_total_count[in_class_id] += 1 | |
| else: | |
| self.class_in_total_count[in_class_id] = 1 | |
| for out_class_id in detections.class_id[~is_in_zone]: | |
| if out_class_id in self.class_out_current_count: | |
| self.class_out_current_count[out_class_id] += 1 | |
| else: | |
| self.class_out_current_count[out_class_id] = 1 | |
| if out_class_id in self.class_out_total_count: | |
| self.class_out_total_count[out_class_id] += 1 | |
| else: | |
| self.class_out_total_count[out_class_id] = 1 | |
| if detections.class_id is not None: | |
| for in_class_id in detections.class_id[is_in_zone]: | |
| if in_class_id in self.class_in_current_count: | |
| self.class_in_current_count[in_class_id] += 1 | |
| else: | |
| self.class_in_current_count[in_class_id] = 1 | |
| if in_class_id in self.class_in_total_count: | |
| self.class_in_total_count[in_class_id] += 1 | |
| else: | |
| self.class_in_total_count[in_class_id] = 1 | |
| for out_class_id in detections.class_id[~is_in_zone]: | |
| if out_class_id in self.class_out_current_count: | |
| self.class_out_current_count[out_class_id] += 1 | |
| else: | |
| self.class_out_current_count[out_class_id] = 1 | |
| if out_class_id in self.class_out_total_count: | |
| self.class_out_total_count[out_class_id] += 1 | |
| else: | |
| self.class_out_total_count[out_class_id] = 1 |
| class_in_current_count (Dict[int, int]): A dictionary storing per-class | ||
| counts of objects inside the zone for current frame | ||
| class_out_current_count (Dict[int, int]): A dictionary storing per-class | ||
| counts of objects outside the zone for current frame | ||
| class_in_total_count (Dict[int, int]): A dictionary storing per-class counts | ||
| of objects inside the zone for all frames | ||
| class_out_total_count (Dict[int, int]): A dictionary storing per-class counts | ||
| of objects outside the zone for all frames |
There was a problem hiding this comment.
The docstring references Dict[int, int] type but Dict is not imported in the module. Add 'Dict' to the imports from typing module to ensure proper type hint documentation.
| for in_class_id in detections.class_id[is_in_zone]: | ||
| if in_class_id in self.class_in_current_count: | ||
| self.class_in_current_count[in_class_id] += 1 | ||
| else: | ||
| self.class_in_current_count[in_class_id] = 1 | ||
|
|
||
| if in_class_id in self.class_in_total_count: | ||
| self.class_in_total_count[in_class_id] += 1 | ||
| else: | ||
| self.class_in_total_count[in_class_id] = 1 | ||
|
|
||
| for out_class_id in detections.class_id[~is_in_zone]: | ||
| if out_class_id in self.class_out_current_count: | ||
| self.class_out_current_count[out_class_id] += 1 | ||
| else: | ||
| self.class_out_current_count[out_class_id] = 1 | ||
|
|
||
| if out_class_id in self.class_out_total_count: | ||
| self.class_out_total_count[out_class_id] += 1 | ||
| else: | ||
| self.class_out_total_count[out_class_id] = 1 | ||
|
|
There was a problem hiding this comment.
The per-class counting logic uses Python loops which can be inefficient for large numbers of detections. Consider using numpy operations like numpy.unique with return_counts=True to count class occurrences more efficiently. For example: unique_classes, counts = np.unique(detections.class_id[is_in_zone], return_counts=True) followed by dict(zip(unique_classes, counts)).
| for in_class_id in detections.class_id[is_in_zone]: | |
| if in_class_id in self.class_in_current_count: | |
| self.class_in_current_count[in_class_id] += 1 | |
| else: | |
| self.class_in_current_count[in_class_id] = 1 | |
| if in_class_id in self.class_in_total_count: | |
| self.class_in_total_count[in_class_id] += 1 | |
| else: | |
| self.class_in_total_count[in_class_id] = 1 | |
| for out_class_id in detections.class_id[~is_in_zone]: | |
| if out_class_id in self.class_out_current_count: | |
| self.class_out_current_count[out_class_id] += 1 | |
| else: | |
| self.class_out_current_count[out_class_id] = 1 | |
| if out_class_id in self.class_out_total_count: | |
| self.class_out_total_count[out_class_id] += 1 | |
| else: | |
| self.class_out_total_count[out_class_id] = 1 | |
| # count in-zone classes using vectorized numpy operations | |
| in_class_ids = detections.class_id[is_in_zone] | |
| if in_class_ids.size > 0: | |
| unique_in, counts_in = np.unique(in_class_ids, return_counts=True) | |
| for class_id, count in zip(unique_in, counts_in): | |
| self.class_in_current_count[class_id] = int(count) | |
| self.class_in_total_count[class_id] = ( | |
| self.class_in_total_count.get(class_id, 0) + int(count) | |
| ) | |
| # count out-of-zone classes using vectorized numpy operations | |
| out_class_ids = detections.class_id[~is_in_zone] | |
| if out_class_ids.size > 0: | |
| unique_out, counts_out = np.unique(out_class_ids, return_counts=True) | |
| for class_id, count in zip(unique_out, counts_out): | |
| self.class_out_current_count[class_id] = int(count) | |
| self.class_out_total_count[class_id] = ( | |
| self.class_out_total_count.get(class_id, 0) + int(count) | |
| ) |
| def test_trigger_detections_single_frame( | ||
| polygon: np.ndarray, | ||
| detections: Detections, | ||
| expected_result: np.ndarray, | ||
| current_count: int, | ||
| class_in_current_count: Dict[int, int], | ||
| class_out_current_count: Dict[int, int], | ||
| class_in_total_count: Dict[int, int], | ||
| class_out_total_count: Dict[int, int], | ||
| ) -> None: | ||
| zone = PolygonZone(polygon, (100, 100)) | ||
| result = zone.trigger(detections) | ||
|
|
||
| assert np.array_equal(result, expected_result) | ||
| assert zone.current_count == current_count | ||
| assert zone.class_in_current_count == class_in_current_count | ||
| assert zone.class_out_current_count == class_out_current_count | ||
| assert zone.class_in_total_count == class_in_total_count | ||
| assert zone.class_out_total_count == class_out_total_count |
There was a problem hiding this comment.
The tests only verify single-frame behavior but don't test the cumulative nature of class_in_total_count and class_out_total_count across multiple trigger calls. Consider adding a test case that calls trigger multiple times to verify that total counts accumulate correctly while current counts reset.
Description
Resolves #791
Type of change
How has this change been tested, please provide a testcase or example of how you tested the change?
https://colab.research.google.com/drive/1RQqqMvQfyWv4fYg5vUVJfK_xt8lSQ5PQ?usp=sharing