diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 43acf5c8..e04d60d1 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -24,7 +24,7 @@ Before you contribute a new feature, consider submitting an Issue to discuss the ## How to Contribute Changes -First, fork this repository to your own GitHub account. Create a new branch that describes your changes (i.e. `line-counter-docs`). Push your changes to the branch on your fork and then submit a pull request to this repository. +First, fork this repository to your own GitHub account. Create a new branch that describes your changes (i.e. `line-counter-docs`). Push your changes to the branch on your fork and then submit a pull request to `develop` branch of this repository. When creating new functions, please ensure you have the following: @@ -32,6 +32,7 @@ When creating new functions, please ensure you have the following: 2. Unit tests for the function. 3. Examples in the documentation for the function. 4. Created an entry in our docs to autogenerate the documentation for the function. +5. Please share google colab with minimal code to test new feature or reproduce PR whenever it is possible. Please ensure that google colab can be accessed without any issue. All pull requests will be reviewed by the maintainers of the project. We will provide feedback and ask for changes if necessary. @@ -48,4 +49,4 @@ So far, **there is no types checking with mypy**. See [issue](https://github.com ## ๐Ÿงช tests -[`pytests`](https://docs.pytest.org/en/7.1.x/) is used to run our tests. \ No newline at end of file +[`pytests`](https://docs.pytest.org/en/7.1.x/) is used to run our tests. diff --git a/README.md b/README.md index 75b60a7c..9e7984b5 100644 --- a/README.md +++ b/README.md @@ -31,7 +31,7 @@ Pip install the supervision package in a pip install supervision[desktop] ``` -Read more about desktop, headless and local installation in our [guide](https://roboflow.github.io/supervision/). +Read more about desktop, headless, and local installation in our [guide](https://roboflow.github.io/supervision/). ## ๐Ÿ”ฅ quickstart @@ -52,7 +52,7 @@ Read more about desktop, headless and local installation in our [guide](https://
๐Ÿ‘‰ more detections utils -- Easily switch inference pipeline between supported object detection / instance segmentation models +- Easily switch inference pipeline between supported object detection/instance segmentation models ```python >>> import supervision as sv @@ -107,7 +107,7 @@ Read more about desktop, headless and local installation in our [guide](https://
๐Ÿ‘‰ more dataset utils -- Load object detection / instance segmentation datasets in one of supported formats +- Load object detection/instance segmentation datasets in one of the supported formats ```python >>> dataset = sv.DetectionDataset.from_yolo( @@ -138,7 +138,7 @@ Read more about desktop, headless and local installation in our [guide](https:// [ 20.154999, 347.825 , 416.125 , 915.895 ]], dtype=float32) ``` -- Split dataset for training, testing and validation +- Split dataset for training, testing, and validation ```python >>> train_dataset, test_dataset = dataset.split(split_ratio=0.7) @@ -148,7 +148,7 @@ Read more about desktop, headless and local installation in our [guide](https:// (700, 150, 150) ``` -- Merge multiple datasets together +- Merge multiple datasets ```python >>> ds_1 = sv.DetectionDataset(...) @@ -170,7 +170,7 @@ Read more about desktop, headless and local installation in our [guide](https:// ['cat', 'dog', 'person'] ``` -- Save object detection / instance segmentation datasets in one of supported formats +- Save object detection/instance segmentation datasets in one of the supported formats ```python >>> dataset.as_yolo( @@ -203,7 +203,7 @@ Read more about desktop, headless and local installation in our [guide](https:// ... ) ``` -- Load classification datasets in one of supported formats +- Load classification datasets in one of the supported formats ```python >>> cs = sv.ClassificationDataset.from_folder_structure( @@ -211,7 +211,7 @@ Read more about desktop, headless and local installation in our [guide](https:// ... ) ``` -- Save classification datasets in one of supported formats +- Save classification datasets in one of the supported formats ```python >>> cs.as_folder_structure( @@ -245,6 +245,30 @@ array([ ]) ``` +
+๐Ÿ‘‰ more metrics + +- Mean average precision (mAP) for object detection tasks. + + ```python + >>> import supervision as sv + + >>> dataset = sv.DetectionDataset.from_yolo(...) + + >>> def callback(image: np.ndarray) -> sv.Detections: + ... ... + + >>> mean_average_precision = sv.MeanAveragePrecision.benchmark( + ... dataset = dataset, + ... callback = callback + ... ) + + >>> mean_average_precision.map50_95 + 0.433 + ``` + +
+ ## ๐ŸŽฌ tutorials

diff --git a/docs/changelog.md b/docs/changelog.md index 3c3843bc..ded13202 100644 --- a/docs/changelog.md +++ b/docs/changelog.md @@ -1,3 +1,39 @@ +### 0.13.0 August 8, 2023 + +- Added [#236](https://github.com/roboflow/supervision/pull/236): support for mean average precision (mAP) for object detection models with [`sv.MeanAveragePrecision`](https://roboflow.github.io/supervision/metrics/detection/#meanaverageprecision). + +```python +>>> import supervision as sv +>>> from ultralytics import YOLO + +>>> dataset = sv.DetectionDataset.from_yolo(...) + +>>> model = YOLO(...) +>>> def callback(image: np.ndarray) -> sv.Detections: +... result = model(image)[0] +... return sv.Detections.from_yolov8(result) + +>>> mean_average_precision = sv.MeanAveragePrecision.benchmark( +... dataset = dataset, +... callback = callback +... ) + +>>> mean_average_precision.map50_95 +0.433 +``` + +- Added [#256](https://github.com/roboflow/supervision/pull/256): support for ByteTrack for object tracking with [`sv.ByteTrack`](https://roboflow.github.io/supervision/tracker/core/#bytetrack). + +- Added [#222](https://github.com/roboflow/supervision/pull/222): [`sv.Detections.from_ultralytics`](https://roboflow.github.io/supervision/detection/core/#supervision.detection.core.Detections.from_ultralytics) to enable seamless integration with [Ultralytics](https://github.com/ultralytics/ultralytics) framework. This will enable you to use `supervision` with all [models](https://docs.ultralytics.com/models/) that Ultralytics supports. + +!!! warning + + [`sv.Detections.from_yolov8`](https://roboflow.github.io/supervision/detection/core/#supervision.detection.core.Detections.from_yolov8) is now deprecated and will be removed with `supervision-0.15.0` release. + +- Added [#191](https://github.com/roboflow/supervision/pull/191): [`sv.Detections.from_paddledet`](https://roboflow.github.io/supervision/detection/core/#supervision.detection.core.Detections.from_paddledet) to enable seamless integration with [PaddleDetection](https://github.com/PaddlePaddle/PaddleDetection) framework. + +- Added [#245](https://github.com/roboflow/supervision/pull/245): support for loading PASCAL VOC segmentation datasets with [`sv.DetectionDataset.`](https://roboflow.github.io/supervision/dataset/core/#supervision.dataset.core.DetectionDataset.from_pascal_voc). + ### 0.12.0 July 24, 2023 !!! warning diff --git a/docs/metrics/detection.md b/docs/metrics/detection.md index 9e784464..feea7cce 100644 --- a/docs/metrics/detection.md +++ b/docs/metrics/detection.md @@ -6,3 +6,7 @@ ## ConfusionMatrix :::supervision.metrics.detection.ConfusionMatrix + +## MeanAveragePrecision + +:::supervision.metrics.detection.MeanAveragePrecision diff --git a/docs/tracker/core.md b/docs/tracker/core.md new file mode 100644 index 00000000..71673fbe --- /dev/null +++ b/docs/tracker/core.md @@ -0,0 +1,4 @@ +## ByteTrack + +:::supervision.tracker.byte_tracker.core.ByteTrack + diff --git a/mkdocs.yml b/mkdocs.yml index dd41adc9..d1977bfc 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -37,6 +37,8 @@ nav: - Utils: detection/utils.md - Tools: - Polygon Zone: detection/tools/polygon_zone.md + - Trackers: + - Core: tracker/core.md - Dataset: - Core: dataset/core.md - Metrics: diff --git a/poetry.lock b/poetry.lock index fa38d3ab..82b2dec4 100644 --- a/poetry.lock +++ b/poetry.lock @@ -2156,54 +2156,83 @@ files = [ [[package]] name = "pillow" -version = "8.4.0" +version = "9.5.0" description = "Python Imaging Library (Fork)" optional = false -python-versions = ">=3.6" +python-versions = ">=3.7" files = [ - {file = "Pillow-8.4.0-cp310-cp310-macosx_10_10_universal2.whl", hash = "sha256:81f8d5c81e483a9442d72d182e1fb6dcb9723f289a57e8030811bac9ea3fef8d"}, - {file = "Pillow-8.4.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:3f97cfb1e5a392d75dd8b9fd274d205404729923840ca94ca45a0af57e13dbe6"}, - {file = "Pillow-8.4.0-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:eb9fc393f3c61f9054e1ed26e6fe912c7321af2f41ff49d3f83d05bacf22cc78"}, - {file = "Pillow-8.4.0-cp310-cp310-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:d82cdb63100ef5eedb8391732375e6d05993b765f72cb34311fab92103314649"}, - {file = "Pillow-8.4.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:62cc1afda735a8d109007164714e73771b499768b9bb5afcbbee9d0ff374b43f"}, - {file = "Pillow-8.4.0-cp310-cp310-win32.whl", hash = "sha256:e3dacecfbeec9a33e932f00c6cd7996e62f53ad46fbe677577394aaa90ee419a"}, - {file = "Pillow-8.4.0-cp310-cp310-win_amd64.whl", hash = "sha256:620582db2a85b2df5f8a82ddeb52116560d7e5e6b055095f04ad828d1b0baa39"}, - {file = "Pillow-8.4.0-cp36-cp36m-macosx_10_10_x86_64.whl", hash = "sha256:1bc723b434fbc4ab50bb68e11e93ce5fb69866ad621e3c2c9bdb0cd70e345f55"}, - {file = "Pillow-8.4.0-cp36-cp36m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:72cbcfd54df6caf85cc35264c77ede902452d6df41166010262374155947460c"}, - {file = "Pillow-8.4.0-cp36-cp36m-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:70ad9e5c6cb9b8487280a02c0ad8a51581dcbbe8484ce058477692a27c151c0a"}, - {file = "Pillow-8.4.0-cp36-cp36m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:25a49dc2e2f74e65efaa32b153527fc5ac98508d502fa46e74fa4fd678ed6645"}, - {file = "Pillow-8.4.0-cp36-cp36m-win32.whl", hash = "sha256:93ce9e955cc95959df98505e4608ad98281fff037350d8c2671c9aa86bcf10a9"}, - {file = "Pillow-8.4.0-cp36-cp36m-win_amd64.whl", hash = "sha256:2e4440b8f00f504ee4b53fe30f4e381aae30b0568193be305256b1462216feff"}, - {file = "Pillow-8.4.0-cp37-cp37m-macosx_10_10_x86_64.whl", hash = "sha256:8c803ac3c28bbc53763e6825746f05cc407b20e4a69d0122e526a582e3b5e153"}, - {file = "Pillow-8.4.0-cp37-cp37m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:c8a17b5d948f4ceeceb66384727dde11b240736fddeda54ca740b9b8b1556b29"}, - {file = "Pillow-8.4.0-cp37-cp37m-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:1394a6ad5abc838c5cd8a92c5a07535648cdf6d09e8e2d6df916dfa9ea86ead8"}, - {file = "Pillow-8.4.0-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:792e5c12376594bfcb986ebf3855aa4b7c225754e9a9521298e460e92fb4a488"}, - {file = "Pillow-8.4.0-cp37-cp37m-win32.whl", hash = "sha256:d99ec152570e4196772e7a8e4ba5320d2d27bf22fdf11743dd882936ed64305b"}, - {file = "Pillow-8.4.0-cp37-cp37m-win_amd64.whl", hash = "sha256:7b7017b61bbcdd7f6363aeceb881e23c46583739cb69a3ab39cb384f6ec82e5b"}, - {file = "Pillow-8.4.0-cp38-cp38-macosx_10_10_x86_64.whl", hash = "sha256:d89363f02658e253dbd171f7c3716a5d340a24ee82d38aab9183f7fdf0cdca49"}, - {file = "Pillow-8.4.0-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:0a0956fdc5defc34462bb1c765ee88d933239f9a94bc37d132004775241a7585"}, - {file = "Pillow-8.4.0-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:5b7bb9de00197fb4261825c15551adf7605cf14a80badf1761d61e59da347779"}, - {file = "Pillow-8.4.0-cp38-cp38-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:72b9e656e340447f827885b8d7a15fc8c4e68d410dc2297ef6787eec0f0ea409"}, - {file = "Pillow-8.4.0-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a5a4532a12314149d8b4e4ad8ff09dde7427731fcfa5917ff16d0291f13609df"}, - {file = "Pillow-8.4.0-cp38-cp38-win32.whl", hash = "sha256:82aafa8d5eb68c8463b6e9baeb4f19043bb31fefc03eb7b216b51e6a9981ae09"}, - {file = "Pillow-8.4.0-cp38-cp38-win_amd64.whl", hash = "sha256:066f3999cb3b070a95c3652712cffa1a748cd02d60ad7b4e485c3748a04d9d76"}, - {file = "Pillow-8.4.0-cp39-cp39-macosx_10_10_x86_64.whl", hash = "sha256:5503c86916d27c2e101b7f71c2ae2cddba01a2cf55b8395b0255fd33fa4d1f1a"}, - {file = "Pillow-8.4.0-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:4acc0985ddf39d1bc969a9220b51d94ed51695d455c228d8ac29fcdb25810e6e"}, - {file = "Pillow-8.4.0-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:0b052a619a8bfcf26bd8b3f48f45283f9e977890263e4571f2393ed8898d331b"}, - {file = "Pillow-8.4.0-cp39-cp39-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:493cb4e415f44cd601fcec11c99836f707bb714ab03f5ed46ac25713baf0ff20"}, - {file = "Pillow-8.4.0-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:b8831cb7332eda5dc89b21a7bce7ef6ad305548820595033a4b03cf3091235ed"}, - {file = "Pillow-8.4.0-cp39-cp39-win32.whl", hash = "sha256:5e9ac5f66616b87d4da618a20ab0a38324dbe88d8a39b55be8964eb520021e02"}, - {file = "Pillow-8.4.0-cp39-cp39-win_amd64.whl", hash = "sha256:3eb1ce5f65908556c2d8685a8f0a6e989d887ec4057326f6c22b24e8a172c66b"}, - {file = "Pillow-8.4.0-pp36-pypy36_pp73-macosx_10_10_x86_64.whl", hash = "sha256:ddc4d832a0f0b4c52fff973a0d44b6c99839a9d016fe4e6a1cb8f3eea96479c2"}, - {file = "Pillow-8.4.0-pp36-pypy36_pp73-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:9a3e5ddc44c14042f0844b8cf7d2cd455f6cc80fd7f5eefbe657292cf601d9ad"}, - {file = "Pillow-8.4.0-pp36-pypy36_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c70e94281588ef053ae8998039610dbd71bc509e4acbc77ab59d7d2937b10698"}, - {file = "Pillow-8.4.0-pp37-pypy37_pp73-macosx_10_10_x86_64.whl", hash = "sha256:3862b7256046fcd950618ed22d1d60b842e3a40a48236a5498746f21189afbbc"}, - {file = "Pillow-8.4.0-pp37-pypy37_pp73-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:a4901622493f88b1a29bd30ec1a2f683782e57c3c16a2dbc7f2595ba01f639df"}, - {file = "Pillow-8.4.0-pp37-pypy37_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:84c471a734240653a0ec91dec0996696eea227eafe72a33bd06c92697728046b"}, - {file = "Pillow-8.4.0-pp37-pypy37_pp73-win_amd64.whl", hash = "sha256:244cf3b97802c34c41905d22810846802a3329ddcb93ccc432870243211c79fc"}, - {file = "Pillow-8.4.0.tar.gz", hash = "sha256:b8e2f83c56e141920c39464b852de3719dfbfb6e3c99a2d8da0edf4fb33176ed"}, + {file = "Pillow-9.5.0-cp310-cp310-macosx_10_10_x86_64.whl", hash = "sha256:ace6ca218308447b9077c14ea4ef381ba0b67ee78d64046b3f19cf4e1139ad16"}, + {file = "Pillow-9.5.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:d3d403753c9d5adc04d4694d35cf0391f0f3d57c8e0030aac09d7678fa8030aa"}, + {file = "Pillow-9.5.0-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:5ba1b81ee69573fe7124881762bb4cd2e4b6ed9dd28c9c60a632902fe8db8b38"}, + {file = "Pillow-9.5.0-cp310-cp310-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:fe7e1c262d3392afcf5071df9afa574544f28eac825284596ac6db56e6d11062"}, + {file = "Pillow-9.5.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:8f36397bf3f7d7c6a3abdea815ecf6fd14e7fcd4418ab24bae01008d8d8ca15e"}, + {file = "Pillow-9.5.0-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:252a03f1bdddce077eff2354c3861bf437c892fb1832f75ce813ee94347aa9b5"}, + {file = "Pillow-9.5.0-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:85ec677246533e27770b0de5cf0f9d6e4ec0c212a1f89dfc941b64b21226009d"}, + {file = "Pillow-9.5.0-cp310-cp310-musllinux_1_1_aarch64.whl", hash = "sha256:b416f03d37d27290cb93597335a2f85ed446731200705b22bb927405320de903"}, + {file = "Pillow-9.5.0-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:1781a624c229cb35a2ac31cc4a77e28cafc8900733a864870c49bfeedacd106a"}, + {file = "Pillow-9.5.0-cp310-cp310-win32.whl", hash = "sha256:8507eda3cd0608a1f94f58c64817e83ec12fa93a9436938b191b80d9e4c0fc44"}, + {file = "Pillow-9.5.0-cp310-cp310-win_amd64.whl", hash = "sha256:d3c6b54e304c60c4181da1c9dadf83e4a54fd266a99c70ba646a9baa626819eb"}, + {file = "Pillow-9.5.0-cp311-cp311-macosx_10_10_x86_64.whl", hash = "sha256:7ec6f6ce99dab90b52da21cf0dc519e21095e332ff3b399a357c187b1a5eee32"}, + {file = "Pillow-9.5.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:560737e70cb9c6255d6dcba3de6578a9e2ec4b573659943a5e7e4af13f298f5c"}, + {file = "Pillow-9.5.0-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:96e88745a55b88a7c64fa49bceff363a1a27d9a64e04019c2281049444a571e3"}, + {file = "Pillow-9.5.0-cp311-cp311-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:d9c206c29b46cfd343ea7cdfe1232443072bbb270d6a46f59c259460db76779a"}, + {file = "Pillow-9.5.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:cfcc2c53c06f2ccb8976fb5c71d448bdd0a07d26d8e07e321c103416444c7ad1"}, + {file = "Pillow-9.5.0-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:a0f9bb6c80e6efcde93ffc51256d5cfb2155ff8f78292f074f60f9e70b942d99"}, + {file = "Pillow-9.5.0-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:8d935f924bbab8f0a9a28404422da8af4904e36d5c33fc6f677e4c4485515625"}, + {file = "Pillow-9.5.0-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:fed1e1cf6a42577953abbe8e6cf2fe2f566daebde7c34724ec8803c4c0cda579"}, + {file = "Pillow-9.5.0-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:c1170d6b195555644f0616fd6ed929dfcf6333b8675fcca044ae5ab110ded296"}, + {file = "Pillow-9.5.0-cp311-cp311-win32.whl", hash = "sha256:54f7102ad31a3de5666827526e248c3530b3a33539dbda27c6843d19d72644ec"}, + {file = "Pillow-9.5.0-cp311-cp311-win_amd64.whl", hash = "sha256:cfa4561277f677ecf651e2b22dc43e8f5368b74a25a8f7d1d4a3a243e573f2d4"}, + {file = "Pillow-9.5.0-cp311-cp311-win_arm64.whl", hash = "sha256:965e4a05ef364e7b973dd17fc765f42233415974d773e82144c9bbaaaea5d089"}, + {file = "Pillow-9.5.0-cp312-cp312-win32.whl", hash = "sha256:22baf0c3cf0c7f26e82d6e1adf118027afb325e703922c8dfc1d5d0156bb2eeb"}, + {file = "Pillow-9.5.0-cp312-cp312-win_amd64.whl", hash = "sha256:432b975c009cf649420615388561c0ce7cc31ce9b2e374db659ee4f7d57a1f8b"}, + {file = "Pillow-9.5.0-cp37-cp37m-macosx_10_10_x86_64.whl", hash = "sha256:5d4ebf8e1db4441a55c509c4baa7a0587a0210f7cd25fcfe74dbbce7a4bd1906"}, + {file = "Pillow-9.5.0-cp37-cp37m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:375f6e5ee9620a271acb6820b3d1e94ffa8e741c0601db4c0c4d3cb0a9c224bf"}, + {file = "Pillow-9.5.0-cp37-cp37m-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:99eb6cafb6ba90e436684e08dad8be1637efb71c4f2180ee6b8f940739406e78"}, + {file = "Pillow-9.5.0-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:2dfaaf10b6172697b9bceb9a3bd7b951819d1ca339a5ef294d1f1ac6d7f63270"}, + {file = "Pillow-9.5.0-cp37-cp37m-manylinux_2_28_aarch64.whl", hash = "sha256:763782b2e03e45e2c77d7779875f4432e25121ef002a41829d8868700d119392"}, + {file = "Pillow-9.5.0-cp37-cp37m-manylinux_2_28_x86_64.whl", hash = "sha256:35f6e77122a0c0762268216315bf239cf52b88865bba522999dc38f1c52b9b47"}, + {file = "Pillow-9.5.0-cp37-cp37m-win32.whl", hash = "sha256:aca1c196f407ec7cf04dcbb15d19a43c507a81f7ffc45b690899d6a76ac9fda7"}, + {file = "Pillow-9.5.0-cp37-cp37m-win_amd64.whl", hash = "sha256:322724c0032af6692456cd6ed554bb85f8149214d97398bb80613b04e33769f6"}, + {file = "Pillow-9.5.0-cp38-cp38-macosx_10_10_x86_64.whl", hash = "sha256:a0aa9417994d91301056f3d0038af1199eb7adc86e646a36b9e050b06f526597"}, + {file = "Pillow-9.5.0-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:f8286396b351785801a976b1e85ea88e937712ee2c3ac653710a4a57a8da5d9c"}, + {file = "Pillow-9.5.0-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:c830a02caeb789633863b466b9de10c015bded434deb3ec87c768e53752ad22a"}, + {file = "Pillow-9.5.0-cp38-cp38-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:fbd359831c1657d69bb81f0db962905ee05e5e9451913b18b831febfe0519082"}, + {file = "Pillow-9.5.0-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f8fc330c3370a81bbf3f88557097d1ea26cd8b019d6433aa59f71195f5ddebbf"}, + {file = "Pillow-9.5.0-cp38-cp38-manylinux_2_28_aarch64.whl", hash = "sha256:7002d0797a3e4193c7cdee3198d7c14f92c0836d6b4a3f3046a64bd1ce8df2bf"}, + {file = "Pillow-9.5.0-cp38-cp38-manylinux_2_28_x86_64.whl", hash = "sha256:229e2c79c00e85989a34b5981a2b67aa079fd08c903f0aaead522a1d68d79e51"}, + {file = "Pillow-9.5.0-cp38-cp38-musllinux_1_1_aarch64.whl", hash = "sha256:9adf58f5d64e474bed00d69bcd86ec4bcaa4123bfa70a65ce72e424bfb88ed96"}, + {file = "Pillow-9.5.0-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:662da1f3f89a302cc22faa9f14a262c2e3951f9dbc9617609a47521c69dd9f8f"}, + {file = "Pillow-9.5.0-cp38-cp38-win32.whl", hash = "sha256:6608ff3bf781eee0cd14d0901a2b9cc3d3834516532e3bd673a0a204dc8615fc"}, + {file = "Pillow-9.5.0-cp38-cp38-win_amd64.whl", hash = "sha256:e49eb4e95ff6fd7c0c402508894b1ef0e01b99a44320ba7d8ecbabefddcc5569"}, + {file = "Pillow-9.5.0-cp39-cp39-macosx_10_10_x86_64.whl", hash = "sha256:482877592e927fd263028c105b36272398e3e1be3269efda09f6ba21fd83ec66"}, + {file = "Pillow-9.5.0-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:3ded42b9ad70e5f1754fb7c2e2d6465a9c842e41d178f262e08b8c85ed8a1d8e"}, + {file = "Pillow-9.5.0-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:c446d2245ba29820d405315083d55299a796695d747efceb5717a8b450324115"}, + {file = "Pillow-9.5.0-cp39-cp39-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:8aca1152d93dcc27dc55395604dcfc55bed5f25ef4c98716a928bacba90d33a3"}, + {file = "Pillow-9.5.0-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:608488bdcbdb4ba7837461442b90ea6f3079397ddc968c31265c1e056964f1ef"}, + {file = "Pillow-9.5.0-cp39-cp39-manylinux_2_28_aarch64.whl", hash = "sha256:60037a8db8750e474af7ffc9faa9b5859e6c6d0a50e55c45576bf28be7419705"}, + {file = "Pillow-9.5.0-cp39-cp39-manylinux_2_28_x86_64.whl", hash = "sha256:07999f5834bdc404c442146942a2ecadd1cb6292f5229f4ed3b31e0a108746b1"}, + {file = "Pillow-9.5.0-cp39-cp39-musllinux_1_1_aarch64.whl", hash = "sha256:a127ae76092974abfbfa38ca2d12cbeddcdeac0fb71f9627cc1135bedaf9d51a"}, + {file = "Pillow-9.5.0-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:489f8389261e5ed43ac8ff7b453162af39c3e8abd730af8363587ba64bb2e865"}, + {file = "Pillow-9.5.0-cp39-cp39-win32.whl", hash = "sha256:9b1af95c3a967bf1da94f253e56b6286b50af23392a886720f563c547e48e964"}, + {file = "Pillow-9.5.0-cp39-cp39-win_amd64.whl", hash = "sha256:77165c4a5e7d5a284f10a6efaa39a0ae8ba839da344f20b111d62cc932fa4e5d"}, + {file = "Pillow-9.5.0-pp38-pypy38_pp73-macosx_10_10_x86_64.whl", hash = "sha256:833b86a98e0ede388fa29363159c9b1a294b0905b5128baf01db683672f230f5"}, + {file = "Pillow-9.5.0-pp38-pypy38_pp73-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:aaf305d6d40bd9632198c766fb64f0c1a83ca5b667f16c1e79e1661ab5060140"}, + {file = "Pillow-9.5.0-pp38-pypy38_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:0852ddb76d85f127c135b6dd1f0bb88dbb9ee990d2cd9aa9e28526c93e794fba"}, + {file = "Pillow-9.5.0-pp38-pypy38_pp73-manylinux_2_28_x86_64.whl", hash = "sha256:91ec6fe47b5eb5a9968c79ad9ed78c342b1f97a091677ba0e012701add857829"}, + {file = "Pillow-9.5.0-pp38-pypy38_pp73-win_amd64.whl", hash = "sha256:cb841572862f629b99725ebaec3287fc6d275be9b14443ea746c1dd325053cbd"}, + {file = "Pillow-9.5.0-pp39-pypy39_pp73-macosx_10_10_x86_64.whl", hash = "sha256:c380b27d041209b849ed246b111b7c166ba36d7933ec6e41175fd15ab9eb1572"}, + {file = "Pillow-9.5.0-pp39-pypy39_pp73-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:7c9af5a3b406a50e313467e3565fc99929717f780164fe6fbb7704edba0cebbe"}, + {file = "Pillow-9.5.0-pp39-pypy39_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:5671583eab84af046a397d6d0ba25343c00cd50bce03787948e0fff01d4fd9b1"}, + {file = "Pillow-9.5.0-pp39-pypy39_pp73-manylinux_2_28_x86_64.whl", hash = "sha256:84a6f19ce086c1bf894644b43cd129702f781ba5751ca8572f08aa40ef0ab7b7"}, + {file = "Pillow-9.5.0-pp39-pypy39_pp73-win_amd64.whl", hash = "sha256:1e7723bd90ef94eda669a3c2c19d549874dd5badaeefabefd26053304abe5799"}, + {file = "Pillow-9.5.0.tar.gz", hash = "sha256:bf548479d336726d7a0eceb6e767e179fbde37833ae42794602631a070d630f1"}, ] +[package.extras] +docs = ["furo", "olefile", "sphinx (>=2.4)", "sphinx-copybutton", "sphinx-inline-tabs", "sphinx-removed-in", "sphinxext-opengraph"] +tests = ["check-manifest", "coverage", "defusedxml", "markdown2", "olefile", "packaging", "pyroma", "pytest", "pytest-cov", "pytest-timeout"] + [[package]] name = "pkginfo" version = "1.9.6" @@ -3049,6 +3078,44 @@ files = [ {file = "ruff-0.0.280.tar.gz", hash = "sha256:581c43e4ac5e5a7117ad7da2120d960a4a99e68ec4021ec3cd47fe1cf78f8380"}, ] +[[package]] +name = "scipy" +version = "1.10.1" +description = "Fundamental algorithms for scientific computing in Python" +optional = false +python-versions = "<3.12,>=3.8" +files = [ + {file = "scipy-1.10.1-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:e7354fd7527a4b0377ce55f286805b34e8c54b91be865bac273f527e1b839019"}, + {file = "scipy-1.10.1-cp310-cp310-macosx_12_0_arm64.whl", hash = "sha256:4b3f429188c66603a1a5c549fb414e4d3bdc2a24792e061ffbd607d3d75fd84e"}, + {file = "scipy-1.10.1-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1553b5dcddd64ba9a0d95355e63fe6c3fc303a8fd77c7bc91e77d61363f7433f"}, + {file = "scipy-1.10.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:4c0ff64b06b10e35215abce517252b375e580a6125fd5fdf6421b98efbefb2d2"}, + {file = "scipy-1.10.1-cp310-cp310-win_amd64.whl", hash = "sha256:fae8a7b898c42dffe3f7361c40d5952b6bf32d10c4569098d276b4c547905ee1"}, + {file = "scipy-1.10.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:0f1564ea217e82c1bbe75ddf7285ba0709ecd503f048cb1236ae9995f64217bd"}, + {file = "scipy-1.10.1-cp311-cp311-macosx_12_0_arm64.whl", hash = "sha256:d925fa1c81b772882aa55bcc10bf88324dadb66ff85d548c71515f6689c6dac5"}, + {file = "scipy-1.10.1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:aaea0a6be54462ec027de54fca511540980d1e9eea68b2d5c1dbfe084797be35"}, + {file = "scipy-1.10.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:15a35c4242ec5f292c3dd364a7c71a61be87a3d4ddcc693372813c0b73c9af1d"}, + {file = "scipy-1.10.1-cp311-cp311-win_amd64.whl", hash = "sha256:43b8e0bcb877faf0abfb613d51026cd5cc78918e9530e375727bf0625c82788f"}, + {file = "scipy-1.10.1-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:5678f88c68ea866ed9ebe3a989091088553ba12c6090244fdae3e467b1139c35"}, + {file = "scipy-1.10.1-cp38-cp38-macosx_12_0_arm64.whl", hash = "sha256:39becb03541f9e58243f4197584286e339029e8908c46f7221abeea4b749fa88"}, + {file = "scipy-1.10.1-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:bce5869c8d68cf383ce240e44c1d9ae7c06078a9396df68ce88a1230f93a30c1"}, + {file = "scipy-1.10.1-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:07c3457ce0b3ad5124f98a86533106b643dd811dd61b548e78cf4c8786652f6f"}, + {file = "scipy-1.10.1-cp38-cp38-win_amd64.whl", hash = "sha256:049a8bbf0ad95277ffba9b3b7d23e5369cc39e66406d60422c8cfef40ccc8415"}, + {file = "scipy-1.10.1-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:cd9f1027ff30d90618914a64ca9b1a77a431159df0e2a195d8a9e8a04c78abf9"}, + {file = "scipy-1.10.1-cp39-cp39-macosx_12_0_arm64.whl", hash = "sha256:79c8e5a6c6ffaf3a2262ef1be1e108a035cf4f05c14df56057b64acc5bebffb6"}, + {file = "scipy-1.10.1-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:51af417a000d2dbe1ec6c372dfe688e041a7084da4fdd350aeb139bd3fb55353"}, + {file = "scipy-1.10.1-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:1b4735d6c28aad3cdcf52117e0e91d6b39acd4272f3f5cd9907c24ee931ad601"}, + {file = "scipy-1.10.1-cp39-cp39-win_amd64.whl", hash = "sha256:7ff7f37b1bf4417baca958d254e8e2875d0cc23aaadbe65b3d5b3077b0eb23ea"}, + {file = "scipy-1.10.1.tar.gz", hash = "sha256:2cf9dfb80a7b4589ba4c40ce7588986d6d5cebc5457cad2c2880f6bc2d42f3a5"}, +] + +[package.dependencies] +numpy = ">=1.19.5,<1.27.0" + +[package.extras] +dev = ["click", "doit (>=0.36.0)", "flake8", "mypy", "pycodestyle", "pydevtool", "rich-click", "typing_extensions"] +doc = ["matplotlib (>2)", "numpydoc", "pydata-sphinx-theme (==0.9.0)", "sphinx (!=4.1.0)", "sphinx-design (>=0.2.0)"] +test = ["asv", "gmpy2", "mpmath", "pooch", "pytest", "pytest-cov", "pytest-timeout", "pytest-xdist", "scikit-umfpack", "threadpoolctl"] + [[package]] name = "secretstorage" version = "3.3.3" @@ -3443,4 +3510,4 @@ desktop = ["opencv-python"] [metadata] lock-version = "2.0" python-versions = ">=3.8,<3.12.0" -content-hash = "be27f05c8857580f327c9d6b89216524c9b8cacf662a24b1d5d4147e4a194f81" +content-hash = "fcefb11e72c40fa7defc13f4a9cb1fe24b9d3d171b64bb54fc263d8add9ef40f" diff --git a/pyproject.toml b/pyproject.toml index d091dd7f..09da872b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,10 +1,11 @@ [tool.poetry] name = "supervision" -version = "0.12.0" +version = "0.13.0" description = "A set of easy-to-use utils that will come in handy in any Computer Vision project" authors = ["Piotr Skalski "] maintainers = ["Piotr Skalski "] readme = "README.md" +license = "MIT" packages = [{include = "supervision"}] homepage = "https://github.com/roboflow/supervision" repository = "https://github.com/roboflow/supervision" @@ -14,7 +15,7 @@ keywords = ["machine-learning", "deep-learning", "vision", "ML", "DL", "AI", "YO classifiers=[ 'Intended Audience :: Developers', 'Intended Audience :: Science/Research', - 'License :: OSI Approved :: BSD License', + 'License :: OSI Approved :: MIT License', 'Programming Language :: Python :: 3', 'Programming Language :: Python :: 3.8', 'Programming Language :: Python :: 3.9', @@ -37,9 +38,10 @@ python = ">=3.8,<3.12.0" numpy = "^1.20.0" matplotlib = "^3.7.1" pyyaml = "^6.0" -pillow = "^8.4.0" +pillow = "^9.4.0" opencv-python = { version = "^4.8.0.74", optional = true } opencv-python-headless = "^4.8.0.74" +scipy = "^1.9.0" [tool.poetry.extras] diff --git a/supervision/__init__.py b/supervision/__init__.py index 9c0252fb..bb337009 100644 --- a/supervision/__init__.py +++ b/supervision/__init__.py @@ -30,7 +30,8 @@ from supervision.draw.color import Color, ColorPalette from supervision.draw.utils import draw_filled_rectangle, draw_polygon, draw_text from supervision.geometry.core import Point, Position, Rect from supervision.geometry.utils import get_polygon_center -from supervision.metrics.detection import ConfusionMatrix +from supervision.metrics.detection import ConfusionMatrix, MeanAveragePrecision +from supervision.tracker.byte_tracker.core import ByteTrack from supervision.utils.file import list_files_with_extensions from supervision.utils.image import ImageSink, crop from supervision.utils.notebook import plot_image, plot_images_grid diff --git a/supervision/dataset/core.py b/supervision/dataset/core.py index ce7e0edb..6e47d747 100644 --- a/supervision/dataset/core.py +++ b/supervision/dataset/core.py @@ -31,7 +31,6 @@ from supervision.dataset.utils import ( train_test_split, ) from supervision.detection.core import Detections -from supervision.utils.file import list_files_with_extensions @dataclass @@ -197,7 +196,10 @@ class DetectionDataset(BaseDataset): @classmethod def from_pascal_voc( - cls, images_directory_path: str, annotations_directory_path: str + cls, + images_directory_path: str, + annotations_directory_path: str, + force_masks: bool = False, ) -> DetectionDataset: """ Creates a Dataset instance from PASCAL VOC formatted data. @@ -205,6 +207,7 @@ class DetectionDataset(BaseDataset): Args: images_directory_path (str): The path to the directory containing the images. annotations_directory_path (str): The path to the directory containing the PASCAL VOC XML annotations. + force_masks (bool, optional): If True, forces masks to be loaded for all annotations, regardless of whether they are present. Returns: DetectionDataset: A DetectionDataset instance containing the loaded images and annotations. @@ -222,7 +225,7 @@ class DetectionDataset(BaseDataset): >>> project = rf.workspace(WORKSPACE_ID).project(PROJECT_ID) >>> dataset = project.version(PROJECT_VERSION).download("voc") - >>> ds = sv.DetectionDataset.from_yolo( + >>> ds = sv.DetectionDataset.from_pascal_voc( ... images_directory_path=f"{dataset.location}/train/images", ... annotations_directory_path=f"{dataset.location}/train/labels" ... ) @@ -231,34 +234,13 @@ class DetectionDataset(BaseDataset): ['dog', 'person'] ``` """ - image_paths = list_files_with_extensions( - directory=images_directory_path, extensions=["jpg", "jpeg", "png"] - ) - annotation_paths = list_files_with_extensions( - directory=annotations_directory_path, extensions=["xml"] + + classes, images, annotations = load_pascal_voc_annotations( + images_directory_path=images_directory_path, + annotations_directory_path=annotations_directory_path, + force_masks=force_masks, ) - raw_annotations: List[Tuple[str, Detections, List[str]]] = [ - load_pascal_voc_annotations(annotation_path=str(annotation_path)) - for annotation_path in annotation_paths - ] - - classes = [] - for annotation in raw_annotations: - classes.extend(annotation[2]) - classes = list(set(classes)) - - for annotation in raw_annotations: - class_id = [classes.index(class_name) for class_name in annotation[2]] - annotation[1].class_id = np.array(class_id) - - images = { - image_path.name: cv2.imread(str(image_path)) for image_path in image_paths - } - - annotations = { - image_name: detections for image_name, detections, _ in raw_annotations - } return DetectionDataset(classes=classes, images=images, annotations=annotations) @classmethod diff --git a/supervision/dataset/formats/coco.py b/supervision/dataset/formats/coco.py index a6066ead..35a8d0f3 100644 --- a/supervision/dataset/formats/coco.py +++ b/supervision/dataset/formats/coco.py @@ -124,7 +124,7 @@ def detections_to_coco_annotations( "category_id": int(class_id), "bbox": [xyxy[0], xyxy[1], box_width, box_height], "area": box_width * box_height, - "segmentation": polygon, + "segmentation": [polygon] if polygon else [], "iscrowd": 0, } coco_annotations.append(coco_annotation) diff --git a/supervision/dataset/formats/pascal_voc.py b/supervision/dataset/formats/pascal_voc.py index a3005eb9..4efb2917 100644 --- a/supervision/dataset/formats/pascal_voc.py +++ b/supervision/dataset/formats/pascal_voc.py @@ -1,12 +1,16 @@ -from typing import List, Optional, Tuple +import os +from pathlib import Path +from typing import Dict, List, Optional, Tuple from xml.dom.minidom import parseString from xml.etree.ElementTree import Element, SubElement, parse, tostring +import cv2 import numpy as np from supervision.dataset.utils import approximate_mask_with_polygons from supervision.detection.core import Detections -from supervision.detection.utils import polygon_to_xyxy +from supervision.detection.utils import polygon_to_mask, polygon_to_xyxy +from supervision.utils.file import list_files_with_extensions def object_to_pascal_voc( @@ -120,24 +124,91 @@ def detections_to_pascal_voc( def load_pascal_voc_annotations( - annotation_path: str, -) -> Tuple[str, Detections, List[str]]: + images_directory_path: str, + annotations_directory_path: str, + force_masks: bool = False, +) -> Tuple[List[str], Dict[str, np.ndarray], Dict[str, Detections]]: """ - Loads PASCAL VOC XML annotations and returns the image name, a Detections instance, and a list of class names. + Loads PASCAL VOC annotations and returns class names, images, and their corresponding detections. Args: - annotation_path (str): The path to the PASCAL VOC XML annotations file. + images_directory_path (str): The path to the directory containing the images. + annotations_directory_path (str): The path to the directory containing the PASCAL VOC annotation files. + force_masks (bool, optional): If True, forces masks to be loaded for all annotations, regardless of whether they are present. Returns: - Tuple[str, Detections, List[str]]: A tuple containing the image name, a Detections instance, and a list of class names of objects in the detections. + Tuple[List[str], Dict[str, np.ndarray], Dict[str, Detections]]: A tuple containing a list of class names, a dictionary with image names as keys and images as values, and a dictionary with image names as keys and corresponding Detections instances as values. """ - tree = parse(annotation_path) - root = tree.getroot() - image_name = root.find("filename").text + image_paths = list_files_with_extensions( + directory=images_directory_path, extensions=["jpg", "jpeg", "png"] + ) + classes = [] + images = {} + annotations = {} + + for image_path in image_paths: + image_name = Path(image_path).stem + image = cv2.imread(str(image_path)) + + annotation_path = os.path.join(annotations_directory_path, f"{image_name}.xml") + if not os.path.exists(annotation_path): + images[image_path.name] = image + annotations[image_path.name] = Detections.empty() + continue + + tree = parse(annotation_path) + root = tree.getroot() + + resolution_wh = (image.shape[1], image.shape[0]) + annotation, classes = detections_from_xml_obj( + root, classes, resolution_wh, force_masks + ) + + images[image_path.name] = image + annotations[image_path.name] = annotation + + return classes, images, annotations + + +def detections_from_xml_obj( + root: Element, classes: List[str], resolution_wh, force_masks: bool = False +) -> Tuple[Detections, List[str]]: + """ + Converts an XML object in Pascal VOC format to a Detections object. + Expected XML format: + + ... + + dog + + 48 + 240 + 195 + 371 + + + 48 + 240 + 195 + 240 + 195 + 371 + 48 + 371 + + + + + Returns: + Tuple[Detections, List[str]]: A tuple containing a Detections object and an updated list of class names, extended with the class names from the XML object. + """ xyxy = [] class_names = [] + masks = [] + with_masks = False + extended_classes = classes[:] for obj in root.findall("object"): class_name = obj.find("name").text class_names.append(class_name) @@ -150,7 +221,40 @@ def load_pascal_voc_annotations( xyxy.append([x1, y1, x2, y2]) - xyxy = np.array(xyxy) - detections = Detections(xyxy=xyxy) + with_masks = obj.find("polygon") is not None + with_masks = force_masks if force_masks else with_masks - return image_name, detections, class_names + for polygon in obj.findall("polygon"): + polygon_points = parse_polygon_points(polygon) + + mask_from_polygon = polygon_to_mask( + polygon=np.array(polygon_points), + resolution_wh=resolution_wh, + ) + masks.append(mask_from_polygon) + + xyxy = np.array(xyxy) if len(xyxy) > 0 else np.empty((0, 4)) + for k in set(class_names): + if k not in extended_classes: + extended_classes.append(k) + class_id = np.array( + [extended_classes.index(class_name) for class_name in class_names] + ) + + if with_masks: + annotation = Detections( + xyxy=xyxy, mask=np.array(masks).astype(bool), class_id=class_id + ) + else: + annotation = Detections(xyxy=xyxy, class_id=class_id) + return annotation, extended_classes + + +def parse_polygon_points(polygon: Element) -> List[List[int]]: + polygon_points = [] + coords = polygon.findall(".//*") + for i in range(0, len(coords), 2): + x = int(coords[i].text) + y = int(coords[i + 1].text) + polygon_points.append([x, y]) + return polygon_points diff --git a/supervision/dataset/formats/yolo.py b/supervision/dataset/formats/yolo.py index da070a0f..0d41f045 100644 --- a/supervision/dataset/formats/yolo.py +++ b/supervision/dataset/formats/yolo.py @@ -126,9 +126,6 @@ def load_yolo_annotations( image_paths = list_files_with_extensions( directory=images_directory_path, extensions=["jpg", "jpeg", "png"] ) - annotation_paths = list_files_with_extensions( - directory=annotations_directory_path, extensions=["txt"] - ) classes = _extract_class_names(file_path=data_yaml_path) images = {} diff --git a/supervision/detection/core.py b/supervision/detection/core.py index 9aab3630..4265c34c 100644 --- a/supervision/detection/core.py +++ b/supervision/detection/core.py @@ -7,12 +7,13 @@ import cv2 import numpy as np from supervision.detection.utils import ( - extract_yolov8_masks, + extract_ultralytics_masks, non_max_suppression, process_roboflow_result, xywh_to_xyxy, ) from supervision.geometry.core import Position +from supervision.utils.internal import deprecated def _validate_xyxy(xyxy: Any, n: int) -> None: @@ -26,7 +27,7 @@ def _validate_mask(mask: Any, n: int) -> None: isinstance(mask, np.ndarray) and len(mask.shape) == 3 and mask.shape[0] == n ) if not is_valid: - raise ValueError("mask must be 3d np.ndarray with (n, W, H) shape") + raise ValueError("mask must be 3d np.ndarray with (n, H, W) shape") def _validate_class_id(class_id: Any, n: int) -> None: @@ -59,7 +60,7 @@ class Detections: Data class containing information about the detections in a video frame. Attributes: xyxy (np.ndarray): An array of shape `(n, 4)` containing the bounding boxes coordinates in format `[x1, y1, x2, y2]` - mask: (Optional[np.ndarray]): An array of shape `(n, W, H)` containing the segmentation masks. + mask: (Optional[np.ndarray]): An array of shape `(n, H, W)` containing the segmentation masks. confidence (Optional[np.ndarray]): An array of shape `(n,)` containing the confidence scores of the detections. class_id (Optional[np.ndarray]): An array of shape `(n,)` containing the class ids of the detections. tracker_id (Optional[np.ndarray]): An array of shape `(n,)` containing the tracker ids of the detections. @@ -170,6 +171,9 @@ class Detections: ) @classmethod + @deprecated( + "This method is deprecated and removed in 0.15.0 release. Use sv.Detections.from_ultralytics() instead." + ) def from_yolov8(cls, yolov8_results) -> Detections: """ Creates a Detections instance from a [YOLOv8](https://github.com/ultralytics/ultralytics) inference result. @@ -196,7 +200,47 @@ class Detections: xyxy=yolov8_results.boxes.xyxy.cpu().numpy(), confidence=yolov8_results.boxes.conf.cpu().numpy(), class_id=yolov8_results.boxes.cls.cpu().numpy().astype(int), - mask=extract_yolov8_masks(yolov8_results), + mask=extract_ultralytics_masks(yolov8_results), + ) + + @classmethod + def from_ultralytics(cls, ultralytics_results) -> Detections: + """ + Creates a Detections instance from a [YOLOv8](https://github.com/ultralytics/ultralytics) inference result. + + Args: + ultralytics_results (ultralytics.yolo.engine.results.Results): The output Results instance from YOLOv8 + + Returns: + Detections: A new Detections object. + + Example: + ```python + >>> import cv2 + >>> from ultralytics import YOLO, FastSAM, SAM, RTDETR + >>> import supervision as sv + + >>> image = cv2.imread(SOURCE_IMAGE_PATH) + >>> model = YOLO('yolov8s.pt') + >>> model = SAM('sam_b.pt') + >>> model = SAM('mobile_sam.pt') + >>> model = FastSAM('FastSAM-s.pt') + >>> model = RTDETR('rtdetr-l.pt') + >>> # model inferences + >>> result = model(image)[0] + >>> # if tracker is enabled + >>> result = model.track(image)[0] + >>> detections = sv.Detections.from_ultralytics(result) + ``` + """ + return cls( + xyxy=ultralytics_results.boxes.xyxy.cpu().numpy(), + confidence=ultralytics_results.boxes.conf.cpu().numpy(), + class_id=ultralytics_results.boxes.cls.cpu().numpy().astype(int), + mask=extract_ultralytics_masks(ultralytics_results), + tracker_id=ultralytics_results.boxes.id.int().cpu().numpy() + if ultralytics_results.boxes.id is not None + else None, ) @classmethod @@ -235,10 +279,10 @@ class Detections: Also supported for [mmyolo](https://github.com/open-mmlab/mmyolo) Args: - mmdet_results (mmdet.structures.DetDataSample): The output Results instance from MMDetection + mmdet_results (mmdet.structures.DetDataSample): The output Results instance from MMDetection Returns: - Detections: A new Detections object. + Detections: A new Detections object. Example: ```python @@ -383,6 +427,42 @@ class Detections: return Detections(xyxy=xywh_to_xyxy(boxes_xywh=xywh), mask=mask) + @classmethod + def from_paddledet(cls, paddledet_result) -> Detections: + """ + Creates a Detections instance from [PaddleDetection](https://github.com/PaddlePaddle/PaddleDetection) inference result. + + Args: + paddledet_result (List[dict]): The output Results instance from SAM + + Returns: + Detections: A new Detections object. + + Example: + ```python + >>> import supervision as sv + >>> import paddle + >>> from ppdet.engine import Trainer + >>> from ppdet.core.workspace import load_config + + >>> weights = (...) + >>> config = (...) + + >>> cfg = load_config(config) + >>> trainer = Trainer(cfg, mode='test') + >>> trainer.load_weights(weights) + + >>> paddledet_result = trainer.predict([images])[0] + + >>> detections = sv.Detections.from_paddledet(paddledet_result=paddledet_result) + ``` + """ + return cls( + xyxy=paddledet_result["bbox"][:, 2:6], + confidence=paddledet_result["bbox"][:, 1], + class_id=paddledet_result["bbox"][:, 0].astype(int), + ) + @classmethod def empty(cls) -> Detections: """ diff --git a/supervision/detection/utils.py b/supervision/detection/utils.py index 63206ddc..71bd019c 100644 --- a/supervision/detection/utils.py +++ b/supervision/detection/utils.py @@ -17,7 +17,7 @@ def polygon_to_mask(polygon: np.ndarray, resolution_wh: Tuple[int, int]) -> np.n np.ndarray: The generated 2D mask, where the polygon is marked with `1`'s and the rest is filled with `0`'s. """ width, height = resolution_wh - mask = np.zeros((height, width), dtype=np.uint8) + mask = np.zeros((height, width)) cv2.fillPoly(mask, [polygon], color=1) return mask @@ -260,7 +260,7 @@ def approximate_polygon( return np.squeeze(approximated_points, axis=1) -def extract_yolov8_masks(yolov8_results) -> Optional[np.ndarray]: +def extract_ultralytics_masks(yolov8_results) -> Optional[np.ndarray]: if not yolov8_results.masks: return None @@ -288,7 +288,10 @@ def extract_yolov8_masks(yolov8_results) -> Optional[np.ndarray]: for i in range(masks.shape[0]): mask = masks[i] mask = mask[top:bottom, left:right] - mask = cv2.resize(mask, (orig_shape[1], orig_shape[0])) + + if mask.shape != orig_shape: + mask = cv2.resize(mask, (orig_shape[1], orig_shape[0])) + mask_maps.append(mask) return np.asarray(mask_maps, dtype=bool) diff --git a/supervision/metrics/detection.py b/supervision/metrics/detection.py index 5b5f710a..177cac24 100644 --- a/supervision/metrics/detection.py +++ b/supervision/metrics/detection.py @@ -12,6 +12,59 @@ from supervision.detection.core import Detections from supervision.detection.utils import box_iou_batch +def detections_to_tensor( + detections: Detections, with_confidence: bool = False +) -> np.ndarray: + """ + Convert Supervision Detections to numpy tensors for further computation + Args: + detections (sv.Detections): Detections/Targets in the format of sv.Detections + with_confidence (bool): Whether to include confidence in the tensor + Returns: + (np.ndarray): Detections as numpy tensors as in (xyxy, class_id, confidence) order + """ + if detections.class_id is None: + raise ValueError( + "ConfusionMatrix can only be calculated for Detections with class_id" + ) + + arrays_to_concat = [detections.xyxy, np.expand_dims(detections.class_id, 1)] + + if with_confidence: + if detections.confidence is None: + raise ValueError( + "ConfusionMatrix can only be calculated for Detections with confidence" + ) + arrays_to_concat.append(np.expand_dims(detections.confidence, 1)) + + return np.concatenate(arrays_to_concat, axis=1) + + +def validate_input_tensors(predictions: List[np.ndarray], targets: List[np.ndarray]): + """ + Checks for shape consistency of input tensors. + """ + if len(predictions) != len(targets): + raise ValueError( + f"Number of predictions ({len(predictions)}) and targets ({len(targets)}) must be equal." + ) + if len(predictions) > 0: + if not isinstance(predictions[0], np.ndarray) or not isinstance( + targets[0], np.ndarray + ): + raise ValueError( + f"Predictions and targets must be lists of numpy arrays. Got {type(predictions[0])} and {type(targets[0])} instead." + ) + if predictions[0].shape[1] != 6: + raise ValueError( + f"Predictions must have shape (N, 6). Got {predictions[0].shape} instead." + ) + if targets[0].shape[1] != 5: + raise ValueError( + f"Targets must have shape (N, 5). Got {targets[0].shape} instead." + ) + + @dataclass class ConfusionMatrix: """ @@ -85,11 +138,9 @@ class ConfusionMatrix: target_tensors = [] for prediction, target in zip(predictions, targets): prediction_tensors.append( - ConfusionMatrix.detections_to_tensor(prediction, with_confidence=True) - ) - target_tensors.append( - ConfusionMatrix.detections_to_tensor(target, with_confidence=False) + detections_to_tensor(prediction, with_confidence=True) ) + target_tensors.append(detections_to_tensor(target, with_confidence=False)) return cls.from_tensors( predictions=prediction_tensors, targets=target_tensors, @@ -98,26 +149,6 @@ class ConfusionMatrix: iou_threshold=iou_threshold, ) - @staticmethod - def detections_to_tensor( - detections: Detections, with_confidence: bool = False - ) -> np.ndarray: - if detections.class_id is None: - raise ValueError( - "ConfusionMatrix can only be calculated for Detections with class_id" - ) - - arrays_to_concat = [detections.xyxy, np.expand_dims(detections.class_id, 1)] - - if with_confidence: - if detections.confidence is None: - raise ValueError( - "ConfusionMatrix can only be calculated for Detections with confidence" - ) - arrays_to_concat.append(np.expand_dims(detections.confidence, 1)) - - return np.concatenate(arrays_to_concat, axis=1) - @classmethod def from_tensors( cls, @@ -184,7 +215,7 @@ class ConfusionMatrix: ]) ``` """ - cls._validate_input_tensors(predictions, targets) + validate_input_tensors(predictions, targets) num_classes = len(classes) matrix = np.zeros((num_classes + 1, num_classes + 1)) @@ -203,33 +234,6 @@ class ConfusionMatrix: iou_threshold=iou_threshold, ) - @classmethod - def _validate_input_tensors( - cls, predictions: List[np.ndarray], targets: List[np.ndarray] - ): - """ - Checks for shape consistency of input tensors. - """ - if len(predictions) != len(targets): - raise ValueError( - f"Number of predictions ({len(predictions)}) and targets ({len(targets)}) must be equal." - ) - if len(predictions) > 0: - if not isinstance(predictions[0], np.ndarray) or not isinstance( - targets[0], np.ndarray - ): - raise ValueError( - f"Predictions and targets must be lists of numpy arrays. Got {type(predictions[0])} and {type(targets[0])} instead." - ) - if predictions[0].shape[1] != 6: - raise ValueError( - f"Predictions must have shape (N, 6). Got {predictions[0].shape} instead." - ) - if targets[0].shape[1] != 5: - raise ValueError( - f"Targets must have shape (N, 5). Got {targets[0].shape} instead." - ) - @staticmethod def evaluate_detection_batch( predictions: np.ndarray, @@ -242,8 +246,8 @@ class ConfusionMatrix: Calculate confusion matrix for a batch of detections for a single image. Args: - predictions (List[np.ndarray]): Each element of the list describes a single image and has `shape = (M, 6)` where `M` is the number of detected objects. Each row is expected to be in `(x_min, y_min, x_max, y_max, class, conf)` format. - targets (List[np.ndarray]): Each element of the list describes a single image and has `shape = (N, 5)` where `N` is the number of ground-truth objects. Each row is expected to be in `(x_min, y_min, x_max, y_max, class)` format. + predictions (np.ndarray): Batch prediction. Describes a single image and has `shape = (M, 6)` where `M` is the number of detected objects. Each row is expected to be in `(x_min, y_min, x_max, y_max, class, conf)` format. + targets (np.ndarray): Batch target labels. Describes a single image and has `shape = (N, 5)` where `N` is the number of ground-truth objects. Each row is expected to be in `(x_min, y_min, x_max, y_max, class)` format. num_classes (int): Number of classes. conf_threshold (float): Detection confidence threshold between `0` and `1`. Detections with lower confidence will be excluded. iou_threshold (float): Detection iou threshold between `0` and `1`. Detections with lower iou will be classified as `FP`. @@ -319,7 +323,7 @@ class ConfusionMatrix: iou_threshold: float = 0.5, ) -> ConfusionMatrix: """ - Create confusion matrix from dataset and callback function. + Calculate confusion matrix from dataset and callback function. Args: dataset (DetectionDataset): Object detection dataset used for evaluation. @@ -456,3 +460,333 @@ class ConfusionMatrix: save_path, dpi=250, facecolor=fig.get_facecolor(), transparent=True ) return fig + + +@dataclass(frozen=True) +class MeanAveragePrecision: + """ + Mean Average Precision for object detection tasks. + + Attributes: + map50_95 (float): Mean Average Precision (mAP) calculated over IoU thresholds ranging from `0.50` to `0.95` with a step size of `0.05`. + map50 (float): Mean Average Precision (mAP) calculated specifically at an IoU threshold of `0.50`. + map75 (float): Mean Average Precision (mAP) calculated specifically at an IoU threshold of `0.75`. + per_class_ap50_95 (np.ndarray): Average Precision (AP) values calculated over IoU thresholds ranging from `0.50` to `0.95` with a step size of `0.05`, provided for each individual class. + """ + + map50_95: float + map50: float + map75: float + per_class_ap50_95: np.ndarray + + @classmethod + def from_detections( + cls, + predictions: List[Detections], + targets: List[Detections], + ) -> MeanAveragePrecision: + """ + Calculate mean average precision based on predicted and ground-truth detections. + + Args: + targets (List[Detections]): Detections objects from ground-truth. + predictions (List[Detections]): Detections objects predicted by the model. + Returns: + MeanAveragePrecision: New instance of ConfusionMatrix. + + Example: + ```python + >>> import supervision as sv + + >>> targets = [ + ... sv.Detections(...), + ... sv.Detections(...) + ... ] + + >>> predictions = [ + ... sv.Detections(...), + ... sv.Detections(...) + ... ] + + >>> mean_average_precision = sv.MeanAveragePrecision.from_detections( + ... predictions=predictions, + ... targets=target, + ... ) + + >>> mean_average_precison.map50_95 + 0.2899 + ``` + """ + prediction_tensors = [] + target_tensors = [] + for prediction, target in zip(predictions, targets): + prediction_tensors.append( + detections_to_tensor(prediction, with_confidence=True) + ) + target_tensors.append(detections_to_tensor(target, with_confidence=False)) + return cls.from_tensors( + predictions=prediction_tensors, + targets=target_tensors, + ) + + @classmethod + def benchmark( + cls, + dataset: DetectionDataset, + callback: Callable[[np.ndarray], Detections], + ) -> MeanAveragePrecision: + """ + Calculate mean average precision from dataset and callback function. + + Args: + dataset (DetectionDataset): Object detection dataset used for evaluation. + callback (Callable[[np.ndarray], Detections]): Function that takes an image as input and returns Detections object. + Returns: + MeanAveragePrecision: New instance of MeanAveragePrecision. + + Example: + ```python + >>> import supervision as sv + >>> from ultralytics import YOLO + + >>> dataset = sv.DetectionDataset.from_yolo(...) + + >>> model = YOLO(...) + >>> def callback(image: np.ndarray) -> sv.Detections: + ... result = model(image)[0] + ... return sv.Detections.from_yolov8(result) + + >>> mean_average_precision = sv.MeanAveragePrecision.benchmark( + ... dataset = dataset, + ... callback = callback + ... ) + + >>> mean_average_precision.map50_95 + 0.433 + ``` + """ + predictions, targets = [], [] + for img_name, img in dataset.images.items(): + predictions_batch = callback(img) + predictions.append(predictions_batch) + targets_batch = dataset.annotations[img_name] + targets.append(targets_batch) + return cls.from_detections( + predictions=predictions, + targets=targets, + ) + + @classmethod + def from_tensors( + cls, + predictions: List[np.ndarray], + targets: List[np.ndarray], + ) -> MeanAveragePrecision: + """ + Calculate Mean Average Precision based on predicted and ground-truth detections at different threshold. + + Args: + predictions (List[np.ndarray]): Each element of the list describes a single image and has `shape = (M, 6)` where `M` is the number of detected objects. Each row is expected to be in `(x_min, y_min, x_max, y_max, class, conf)` format. + targets (List[np.ndarray]): Each element of the list describes a single image and has `shape = (N, 5)` where `N` is the number of ground-truth objects. Each row is expected to be in `(x_min, y_min, x_max, y_max, class)` format. + Returns: + MeanAveragePrecision: New instance of MeanAveragePrecision. + + Example: + ```python + >>> import supervision as sv + + >>> targets = ( + ... [ + ... array( + ... [ + ... [0.0, 0.0, 3.0, 3.0, 1], + ... [2.0, 2.0, 5.0, 5.0, 1], + ... [6.0, 1.0, 8.0, 3.0, 2], + ... ] + ... ), + ... array([1.0, 1.0, 2.0, 2.0, 2]), + ... ] + ... ) + + >>> predictions = [ + ... array( + ... [ + ... [0.0, 0.0, 3.0, 3.0, 1, 0.9], + ... [0.1, 0.1, 3.0, 3.0, 0, 0.9], + ... [6.0, 1.0, 8.0, 3.0, 1, 0.8], + ... [1.0, 6.0, 2.0, 7.0, 1, 0.8], + ... ] + ... ), + ... array([[1.0, 1.0, 2.0, 2.0, 2, 0.8]]) + ... ] + + >>> mean_average_precison = sv.MeanAveragePrecision.from_tensors( + ... predictions=predictions, + ... targets=targets, + ... ) + + >>> mean_average_precison.map50_95 + 0.2899 + ``` + """ + validate_input_tensors(predictions, targets) + iou_thresholds = np.linspace(0.5, 0.95, 10) + stats = [] + + # Gather matching stats for predictions and targets + for true_objs, predicted_objs in zip(targets, predictions): + if predicted_objs.shape[0] == 0: + if true_objs.shape[0]: + stats.append( + ( + np.zeros((0, iou_thresholds.size), dtype=bool), + *np.zeros((2, 0)), + true_objs[:, 4], + ) + ) + continue + + if true_objs.shape[0]: + matches = cls._match_detection_batch( + predicted_objs, true_objs, iou_thresholds + ) + stats.append( + ( + matches, + predicted_objs[:, 5], + predicted_objs[:, 4], + true_objs[:, 4], + ) + ) + + # Compute average precisions if any matches exist + if stats: + concatenated_stats = [np.concatenate(items, 0) for items in zip(*stats)] + average_precisions = cls._average_precisions_per_class(*concatenated_stats) + map50 = average_precisions[:, 0].mean() + map75 = average_precisions[:, 5].mean() + map50_95 = average_precisions.mean() + else: + map50, map75, map50_95 = 0, 0, 0 + average_precisions = [] + + return cls( + map50_95=map50_95, + map50=map50, + map75=map75, + per_class_ap50_95=average_precisions, + ) + + @staticmethod + def compute_average_precision(recall: np.ndarray, precision: np.ndarray) -> float: + """ + Compute the average precision using 101-point interpolation (COCO), given the recall and precision curves. + + Args: + recall (np.ndarray): The recall curve. + precision (np.ndarray): The precision curve. + + Returns: + float: Average precision. + """ + extended_recall = np.concatenate(([0.0], recall, [1.0])) + extended_precision = np.concatenate(([1.0], precision, [0.0])) + max_accumulated_precision = np.flip( + np.maximum.accumulate(np.flip(extended_precision)) + ) + interpolated_recall_levels = np.linspace(0, 1, 101) + interpolated_precision = np.interp( + interpolated_recall_levels, extended_recall, max_accumulated_precision + ) + average_precision = np.trapz(interpolated_precision, interpolated_recall_levels) + return average_precision + + @staticmethod + def _match_detection_batch( + predictions: np.ndarray, targets: np.ndarray, iou_thresholds: np.ndarray + ) -> np.ndarray: + """ + Match predictions with target labels based on IoU levels. + + Args: + predictions (np.ndarray): Batch prediction. Describes a single image and has `shape = (M, 6)` where `M` is the number of detected objects. Each row is expected to be in `(x_min, y_min, x_max, y_max, class, conf)` format. + targets (np.ndarray): Batch target labels. Describes a single image and has `shape = (N, 5)` where `N` is the number of ground-truth objects. Each row is expected to be in `(x_min, y_min, x_max, y_max, class)` format. + iou_thresholds (np.ndarray): Array contains different IoU thresholds. + + Returns: + np.ndarray: Matched prediction with target labels result. + """ + num_predictions, num_iou_levels = predictions.shape[0], iou_thresholds.shape[0] + correct = np.zeros((num_predictions, num_iou_levels), dtype=bool) + iou = box_iou_batch(targets[:, :4], predictions[:, :4]) + correct_class = targets[:, 4:5] == predictions[:, 4] + + for i, iou_level in enumerate(iou_thresholds): + matched_indices = np.where((iou >= iou_level) & correct_class) + + if matched_indices[0].shape[0]: + combined_indices = np.stack(matched_indices, axis=1) + iou_values = iou[matched_indices][:, None] + matches = np.hstack([combined_indices, iou_values]) + + if matched_indices[0].shape[0] > 1: + matches = matches[matches[:, 2].argsort()[::-1]] + matches = matches[np.unique(matches[:, 1], return_index=True)[1]] + matches = matches[np.unique(matches[:, 0], return_index=True)[1]] + + correct[matches[:, 1].astype(int), i] = True + + return correct + + @staticmethod + def _average_precisions_per_class( + matches: np.ndarray, + prediction_confidence: np.ndarray, + prediction_class_ids: np.ndarray, + true_class_ids: np.ndarray, + eps: float = 1e-16, + ) -> np.ndarray: + """ + Compute the average precision, given the recall and precision curves. + Source: https://github.com/rafaelpadilla/Object-Detection-Metrics. + + Args: + matches (np.ndarray): True positives. + prediction_confidence (np.ndarray): Objectness value from 0-1. + prediction_class_ids (np.ndarray): Predicted object classes. + true_class_ids (np.ndarray): True object classes. + eps (float, optional): Small value to prevent division by zero. + + Returns: + np.ndarray: Average precision for different IoU levels. + """ + sorted_indices = np.argsort(-prediction_confidence) + matches = matches[sorted_indices] + prediction_class_ids = prediction_class_ids[sorted_indices] + + unique_classes, class_counts = np.unique(true_class_ids, return_counts=True) + num_classes = unique_classes.shape[0] + + average_precisions = np.zeros((num_classes, matches.shape[1])) + + for class_idx, class_id in enumerate(unique_classes): + is_class = prediction_class_ids == class_id + total_true = class_counts[class_idx] + total_prediction = is_class.sum() + + if total_prediction == 0 or total_true == 0: + continue + + false_positives = (1 - matches[is_class]).cumsum(0) + true_positives = matches[is_class].cumsum(0) + recall = true_positives / (total_true + eps) + precision = true_positives / (true_positives + false_positives) + + for iou_level_idx in range(matches.shape[1]): + average_precisions[ + class_idx, iou_level_idx + ] = MeanAveragePrecision.compute_average_precision( + recall[:, iou_level_idx], precision[:, iou_level_idx] + ) + + return average_precisions diff --git a/supervision/tracker/__init__.py b/supervision/tracker/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/supervision/tracker/byte_tracker/__init__.py b/supervision/tracker/byte_tracker/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/supervision/tracker/byte_tracker/basetrack.py b/supervision/tracker/byte_tracker/basetrack.py new file mode 100644 index 00000000..c8fadefa --- /dev/null +++ b/supervision/tracker/byte_tracker/basetrack.py @@ -0,0 +1,55 @@ +from collections import OrderedDict +from enum import Enum + +import numpy as np + + +class TrackState(Enum): + New = 0 + Tracked = 1 + Lost = 2 + Removed = 3 + + +class BaseTrack: + _count = 0 + + def __init__(self): + self.track_id = 0 + self.is_activated = False + self.state = TrackState.New + + self.history = OrderedDict() + self.features = [] + self.curr_feature = None + self.score = 0 + self.start_frame = 0 + self.frame_id = 0 + self.time_since_update = 0 + + # multi-camera + self.location = (np.inf, np.inf) + + @property + def end_frame(self) -> int: + return self.frame_id + + @staticmethod + def next_id() -> int: + BaseTrack._count += 1 + return BaseTrack._count + + def activate(self, *args): + raise NotImplementedError + + def predict(self): + raise NotImplementedError + + def update(self, *args, **kwargs): + raise NotImplementedError + + def mark_lost(self): + self.state = TrackState.Lost + + def mark_removed(self): + self.state = TrackState.Removed diff --git a/supervision/tracker/byte_tracker/core.py b/supervision/tracker/byte_tracker/core.py new file mode 100644 index 00000000..ad0eb8b3 --- /dev/null +++ b/supervision/tracker/byte_tracker/core.py @@ -0,0 +1,466 @@ +from typing import List, Tuple + +import numpy as np + +from supervision.detection.core import Detections +from supervision.tracker.byte_tracker import matching +from supervision.tracker.byte_tracker.basetrack import BaseTrack, TrackState +from supervision.tracker.byte_tracker.kalman_filter import KalmanFilter + + +class STrack(BaseTrack): + shared_kalman = KalmanFilter() + + def __init__(self, tlwh, score, class_ids): + # wait activate + self._tlwh = np.asarray(tlwh, dtype=np.float32) + self.kalman_filter = None + self.mean, self.covariance = None, None + self.is_activated = False + + self.score = score + self.class_ids = class_ids + self.tracklet_len = 0 + + def predict(self): + mean_state = self.mean.copy() + if self.state != TrackState.Tracked: + mean_state[7] = 0 + self.mean, self.covariance = self.kalman_filter.predict( + mean_state, self.covariance + ) + + @staticmethod + def multi_predict(stracks): + if len(stracks) > 0: + multi_mean = np.asarray([st.mean.copy() for st in stracks]) + multi_covariance = np.asarray([st.covariance for st in stracks]) + for i, st in enumerate(stracks): + if st.state != TrackState.Tracked: + multi_mean[i][7] = 0 + multi_mean, multi_covariance = STrack.shared_kalman.multi_predict( + multi_mean, multi_covariance + ) + for i, (mean, cov) in enumerate(zip(multi_mean, multi_covariance)): + stracks[i].mean = mean + stracks[i].covariance = cov + + def activate(self, kalman_filter, frame_id): + """Start a new tracklet""" + self.kalman_filter = kalman_filter + self.track_id = self.next_id() + self.mean, self.covariance = self.kalman_filter.initiate( + self.tlwh_to_xyah(self._tlwh) + ) + + self.tracklet_len = 0 + self.state = TrackState.Tracked + if frame_id == 1: + self.is_activated = True + self.frame_id = frame_id + self.start_frame = frame_id + + def re_activate(self, new_track, frame_id, new_id=False): + self.mean, self.covariance = self.kalman_filter.update( + self.mean, self.covariance, self.tlwh_to_xyah(new_track.tlwh) + ) + self.tracklet_len = 0 + self.state = TrackState.Tracked + self.is_activated = True + self.frame_id = frame_id + if new_id: + self.track_id = self.next_id() + self.score = new_track.score + + def update(self, new_track, frame_id): + """ + Update a matched track + :type new_track: STrack + :type frame_id: int + :type update_feature: bool + :return: + """ + self.frame_id = frame_id + self.tracklet_len += 1 + + new_tlwh = new_track.tlwh + self.mean, self.covariance = self.kalman_filter.update( + self.mean, self.covariance, self.tlwh_to_xyah(new_tlwh) + ) + self.state = TrackState.Tracked + self.is_activated = True + + self.score = new_track.score + + @property + def tlwh(self): + """Get current position in bounding box format `(top left x, top left y, + width, height)`. + """ + if self.mean is None: + return self._tlwh.copy() + ret = self.mean[:4].copy() + ret[2] *= ret[3] + ret[:2] -= ret[2:] / 2 + return ret + + @property + def tlbr(self): + """Convert bounding box to format `(min x, min y, max x, max y)`, i.e., + `(top left, bottom right)`. + """ + ret = self.tlwh.copy() + ret[2:] += ret[:2] + return ret + + @staticmethod + def tlwh_to_xyah(tlwh): + """Convert bounding box to format `(center x, center y, aspect ratio, + height)`, where the aspect ratio is `width / height`. + """ + ret = np.asarray(tlwh).copy() + ret[:2] += ret[2:] / 2 + ret[2] /= ret[3] + return ret + + def to_xyah(self): + return self.tlwh_to_xyah(self.tlwh) + + @staticmethod + def tlbr_to_tlwh(tlbr): + ret = np.asarray(tlbr).copy() + ret[2:] -= ret[:2] + return ret + + @staticmethod + def tlwh_to_tlbr(tlwh): + ret = np.asarray(tlwh).copy() + ret[2:] += ret[:2] + return ret + + def __repr__(self): + return "OT_{}_({}-{})".format(self.track_id, self.start_frame, self.end_frame) + + +def detections2boxes(detections: Detections) -> np.ndarray: + """ + Convert Supervision Detections to numpy tensors for further computation. + Args: + detections (Detections): Detections/Targets in the format of sv.Detections. + Returns: + (np.ndarray): Detections as numpy tensors as in `(x_min, y_min, x_max, y_max, confidence, class_id)` order. + """ + return np.hstack( + ( + detections.xyxy, + detections.confidence[:, np.newaxis], + detections.class_id[:, np.newaxis], + ) + ) + + +class ByteTrack: + """ + Initialize the ByteTrack object. + + Parameters: + track_thresh (float, optional): Detection confidence threshold for track activation. + track_buffer (int, optional): Number of frames to buffer when a track is lost. + match_thresh (float, optional): Threshold for matching tracks with detections. + frame_rate (int, optional): The frame rate of the video. + """ + + def __init__( + self, + track_thresh: float = 0.25, + track_buffer: int = 30, + match_thresh: float = 0.8, + frame_rate: int = 30, + ): + self.track_thresh = track_thresh + self.match_thresh = match_thresh + + self.frame_id = 0 + self.det_thresh = self.track_thresh + 0.1 + self.max_time_lost = int(frame_rate / 30.0 * track_buffer) + self.kalman_filter = KalmanFilter() + + self.tracked_tracks: List[STrack] = [] + self.lost_tracks: List[STrack] = [] + self.removed_tracks: List[STrack] = [] + + def update_with_detections(self, detections: Detections) -> Detections: + """ + Updates the tracker with the provided detections and returns the updated detection results. + + Parameters: + detections: The new detections to update with. + Returns: + Detection: The updated detection results that now include tracking IDs. + Example: + ```python + >>> import supervision as sv + >>> from ultralytics import YOLO + + >>> model = YOLO(...) + >>> byte_tracker = sv.ByteTrack() + >>> annotator = sv.BoxAnnotator() + + >>> def callback(frame: np.ndarray, index: int) -> np.ndarray: + ... results = model(frame)[0] + ... detections = sv.Detections.from_yolov8(results) + ... detections = byte_tracker.update_from_detections(detections=detections) + ... labels = [ + ... f"#{tracker_id} {model.model.names[class_id]} {confidence:0.2f}" + ... for _, _, confidence, class_id, tracker_id + ... in detections + ... ] + ... return annotator.annotate(scene=frame.copy(), detections=detections, labels=labels) + + >>> sv.process_video( + ... source_path='...', + ... target_path='...', + ... callback=callback + ... ) + ``` + """ + + tracks = self.update_with_tensors( + tensors=detections2boxes(detections=detections) + ) + detections = Detections.empty() + if len(tracks) > 0: + detections.xyxy = np.array( + [track.tlbr for track in tracks], dtype=np.float32 + ) + detections.class_id = np.array( + [int(t.class_ids) for t in tracks], dtype=int + ) + detections.tracker_id = np.array( + [int(t.track_id) for t in tracks], dtype=int + ) + detections.confidence = np.array( + [t.score for t in tracks], dtype=np.float32 + ) + + return detections + + def update_with_tensors(self, tensors: np.ndarray) -> List[STrack]: + """ + Updates the tracker with the provided tensors and returns the updated tracks. + + Parameters: + tensors: The new tensors to update with. + + Returns: + List[STrack]: Updated tracks. + """ + self.frame_id += 1 + activated_starcks = [] + refind_stracks = [] + lost_stracks = [] + removed_stracks = [] + + class_ids = tensors[:, 5] + scores = tensors[:, 4] + bboxes = tensors[:, :4] + + remain_inds = scores > self.track_thresh + inds_low = scores > 0.1 + inds_high = scores < self.track_thresh + + inds_second = np.logical_and(inds_low, inds_high) + dets_second = bboxes[inds_second] + dets = bboxes[remain_inds] + scores_keep = scores[remain_inds] + scores_second = scores[inds_second] + + class_ids_keep = class_ids[remain_inds] + class_ids_second = class_ids[inds_second] + + if len(dets) > 0: + """Detections""" + detections = [ + STrack(STrack.tlbr_to_tlwh(tlbr), s, c) + for (tlbr, s, c) in zip(dets, scores_keep, class_ids_keep) + ] + else: + detections = [] + + """ Add newly detected tracklets to tracked_stracks""" + unconfirmed = [] + tracked_stracks = [] # type: list[STrack] + for track in self.tracked_tracks: + if not track.is_activated: + unconfirmed.append(track) + else: + tracked_stracks.append(track) + + """ Step 2: First association, with high score detection boxes""" + strack_pool = joint_tracks(tracked_stracks, self.lost_tracks) + # Predict the current location with KF + STrack.multi_predict(strack_pool) + dists = matching.iou_distance(strack_pool, detections) + + dists = matching.fuse_score(dists, detections) + matches, u_track, u_detection = matching.linear_assignment( + dists, thresh=self.match_thresh + ) + + for itracked, idet in matches: + track = strack_pool[itracked] + det = detections[idet] + if track.state == TrackState.Tracked: + track.update(detections[idet], self.frame_id) + activated_starcks.append(track) + else: + track.re_activate(det, self.frame_id, new_id=False) + refind_stracks.append(track) + + """ Step 3: Second association, with low score detection boxes""" + # association the untrack to the low score detections + if len(dets_second) > 0: + """Detections""" + detections_second = [ + STrack(STrack.tlbr_to_tlwh(tlbr), s, c) + for (tlbr, s, c) in zip(dets_second, scores_second, class_ids_second) + ] + else: + detections_second = [] + r_tracked_stracks = [ + strack_pool[i] + for i in u_track + if strack_pool[i].state == TrackState.Tracked + ] + dists = matching.iou_distance(r_tracked_stracks, detections_second) + matches, u_track, u_detection_second = matching.linear_assignment( + dists, thresh=0.5 + ) + for itracked, idet in matches: + track = r_tracked_stracks[itracked] + det = detections_second[idet] + if track.state == TrackState.Tracked: + track.update(det, self.frame_id) + activated_starcks.append(track) + else: + track.re_activate(det, self.frame_id, new_id=False) + refind_stracks.append(track) + + for it in u_track: + track = r_tracked_stracks[it] + if not track.state == TrackState.Lost: + track.mark_lost() + lost_stracks.append(track) + + """Deal with unconfirmed tracks, usually tracks with only one beginning frame""" + detections = [detections[i] for i in u_detection] + dists = matching.iou_distance(unconfirmed, detections) + + dists = matching.fuse_score(dists, detections) + matches, u_unconfirmed, u_detection = matching.linear_assignment( + dists, thresh=0.7 + ) + for itracked, idet in matches: + unconfirmed[itracked].update(detections[idet], self.frame_id) + activated_starcks.append(unconfirmed[itracked]) + for it in u_unconfirmed: + track = unconfirmed[it] + track.mark_removed() + removed_stracks.append(track) + + """ Step 4: Init new stracks""" + for inew in u_detection: + track = detections[inew] + if track.score < self.det_thresh: + continue + track.activate(self.kalman_filter, self.frame_id) + activated_starcks.append(track) + """ Step 5: Update state""" + for track in self.lost_tracks: + if self.frame_id - track.end_frame > self.max_time_lost: + track.mark_removed() + removed_stracks.append(track) + + self.tracked_tracks = [ + t for t in self.tracked_tracks if t.state == TrackState.Tracked + ] + self.tracked_tracks = joint_tracks(self.tracked_tracks, activated_starcks) + self.tracked_tracks = joint_tracks(self.tracked_tracks, refind_stracks) + self.lost_tracks = sub_tracks(self.lost_tracks, self.tracked_tracks) + self.lost_tracks.extend(lost_stracks) + self.lost_tracks = sub_tracks(self.lost_tracks, self.removed_tracks) + self.removed_tracks.extend(removed_stracks) + self.tracked_tracks, self.lost_tracks = remove_duplicate_tracks( + self.tracked_tracks, self.lost_tracks + ) + output_stracks = [track for track in self.tracked_tracks if track.is_activated] + + return output_stracks + + +def joint_tracks( + track_list_a: List[STrack], track_list_b: List[STrack] +) -> List[STrack]: + """ + Joins two lists of tracks, ensuring that the resulting list does not + contain tracks with duplicate track_id values. + + Parameters: + track_list_a: First list of tracks (with track_id attribute). + track_list_b: Second list of tracks (with track_id attribute). + + Returns: + Combined list of tracks from track_list_a and track_list_b without duplicate track_id values. + """ + seen_track_ids = set() + result = [] + + for track in track_list_a + track_list_b: + if track.track_id not in seen_track_ids: + seen_track_ids.add(track.track_id) + result.append(track) + + return result + + +def sub_tracks(track_list_a: List, track_list_b: List) -> List[int]: + """ + Returns a list of tracks from track_list_a after removing any tracks + that share the same track_id with tracks in track_list_b. + + Parameters: + track_list_a: List of tracks (with track_id attribute). + track_list_b: List of tracks (with track_id attribute) to be subtracted from track_list_a. + Returns: + List of remaining tracks from track_list_a after subtraction. + """ + tracks = {track.track_id: track for track in track_list_a} + track_ids_b = {track.track_id for track in track_list_b} + + for track_id in track_ids_b: + tracks.pop(track_id, None) + + return list(tracks.values()) + + +def remove_duplicate_tracks(tracks_a: List, tracks_b: List) -> Tuple[List, List]: + pairwise_distance = matching.iou_distance(tracks_a, tracks_b) + matching_pairs = np.where(pairwise_distance < 0.15) + + duplicates_a, duplicates_b = set(), set() + for track_index_a, track_index_b in zip(*matching_pairs): + time_a = tracks_a[track_index_a].frame_id - tracks_a[track_index_a].start_frame + time_b = tracks_b[track_index_b].frame_id - tracks_b[track_index_b].start_frame + if time_a > time_b: + duplicates_b.add(track_index_b) + else: + duplicates_a.add(track_index_a) + + result_a = [ + track for index, track in enumerate(tracks_a) if index not in duplicates_a + ] + result_b = [ + track for index, track in enumerate(tracks_b) if index not in duplicates_b + ] + + return result_a, result_b diff --git a/supervision/tracker/byte_tracker/kalman_filter.py b/supervision/tracker/byte_tracker/kalman_filter.py new file mode 100644 index 00000000..f4b0b8c7 --- /dev/null +++ b/supervision/tracker/byte_tracker/kalman_filter.py @@ -0,0 +1,190 @@ +from typing import Tuple + +import numpy as np +import scipy.linalg + + +class KalmanFilter: + """ + A simple Kalman filter for tracking bounding boxes in image space. + + The 8-dimensional state space + + x, y, a, h, vx, vy, va, vh + + contains the bounding box center position (x, y), aspect ratio a, height h, + and their respective velocities. + + Object motion follows a constant velocity model. The bounding box location + (x, y, a, h) is taken as direct observation of the state space (linear + observation model). + """ + + def __init__(self): + ndim, dt = 4, 1.0 + + self._motion_mat = np.eye(2 * ndim, 2 * ndim) + for i in range(ndim): + self._motion_mat[i, ndim + i] = dt + self._update_mat = np.eye(ndim, 2 * ndim) + self._std_weight_position = 1.0 / 20 + self._std_weight_velocity = 1.0 / 160 + + def initiate(self, measurement: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: + """ + Create track from an unassociated measurement. + + Args: + measurement (ndarray): Bounding box coordinates (x, y, a, h) with center position (x, y), aspect ratio a, and height h. + + Returns: + Tuple[ndarray, ndarray]: Returns the mean vector (8 dimensional) and covariance matrix (8x8 dimensional) of the new track. Unobserved velocities are initialized to 0 mean. + """ + mean_pos = measurement + mean_vel = np.zeros_like(mean_pos) + mean = np.r_[mean_pos, mean_vel] + + std = [ + 2 * self._std_weight_position * measurement[3], + 2 * self._std_weight_position * measurement[3], + 1e-2, + 2 * self._std_weight_position * measurement[3], + 10 * self._std_weight_velocity * measurement[3], + 10 * self._std_weight_velocity * measurement[3], + 1e-5, + 10 * self._std_weight_velocity * measurement[3], + ] + covariance = np.diag(np.square(std)) + return mean, covariance + + def predict( + self, mean: np.ndarray, covariance: np.ndarray + ) -> Tuple[np.ndarray, np.ndarray]: + """ + Run Kalman filter prediction step. + + Args: + mean (ndarray): The 8 dimensional mean vector of the object state at the previous time step. + covariance (ndarray): The 8x8 dimensional covariance matrix of the object state at the previous time step. + + Returns: + Tuple[ndarray, ndarray]: Returns the mean vector and covariance matrix of the predicted state. Unobserved velocities are initialized to 0 mean. + """ + std_pos = [ + self._std_weight_position * mean[3], + self._std_weight_position * mean[3], + 1e-2, + self._std_weight_position * mean[3], + ] + std_vel = [ + self._std_weight_velocity * mean[3], + self._std_weight_velocity * mean[3], + 1e-5, + self._std_weight_velocity * mean[3], + ] + motion_cov = np.diag(np.square(np.r_[std_pos, std_vel])) + + mean = np.dot(mean, self._motion_mat.T) + covariance = ( + np.linalg.multi_dot((self._motion_mat, covariance, self._motion_mat.T)) + + motion_cov + ) + + return mean, covariance + + def project( + self, mean: np.ndarray, covariance: np.ndarray + ) -> Tuple[np.ndarray, np.ndarray]: + """ + Project state distribution to measurement space. + + Args: + mean (ndarray): The state's mean vector (8 dimensional array). + covariance (ndarray): The state's covariance matrix (8x8 dimensional). + + Returns: + Tuple[ndarray, ndarray]: Returns the projected mean and covariance matrix of the given state estimate. + """ + std = [ + self._std_weight_position * mean[3], + self._std_weight_position * mean[3], + 1e-1, + self._std_weight_position * mean[3], + ] + innovation_cov = np.diag(np.square(std)) + + mean = np.dot(self._update_mat, mean) + covariance = np.linalg.multi_dot( + (self._update_mat, covariance, self._update_mat.T) + ) + return mean, covariance + innovation_cov + + def multi_predict( + self, mean: np.ndarray, covariance: np.ndarray + ) -> Tuple[np.ndarray, np.ndarray]: + """ + Run Kalman filter prediction step (Vectorized version). + + Args: + mean (ndarray): The Nx8 dimensional mean matrix of the object states at the previous time step. + covariance (ndarray): The Nx8x8 dimensional covariance matrices of the object states at the previous time step. + + Returns: + Tuple[ndarray, ndarray]: Returns the mean vector and covariance matrix of the predicted state. Unobserved velocities are initialized to 0 mean. + """ + std_pos = [ + self._std_weight_position * mean[:, 3], + self._std_weight_position * mean[:, 3], + 1e-2 * np.ones_like(mean[:, 3]), + self._std_weight_position * mean[:, 3], + ] + std_vel = [ + self._std_weight_velocity * mean[:, 3], + self._std_weight_velocity * mean[:, 3], + 1e-5 * np.ones_like(mean[:, 3]), + self._std_weight_velocity * mean[:, 3], + ] + sqr = np.square(np.r_[std_pos, std_vel]).T + + motion_cov = [] + for i in range(len(mean)): + motion_cov.append(np.diag(sqr[i])) + motion_cov = np.asarray(motion_cov) + + mean = np.dot(mean, self._motion_mat.T) + left = np.dot(self._motion_mat, covariance).transpose((1, 0, 2)) + covariance = np.dot(left, self._motion_mat.T) + motion_cov + + return mean, covariance + + def update( + self, mean: np.ndarray, covariance: np.ndarray, measurement: np.ndarray + ) -> Tuple[np.ndarray, np.ndarray]: + """ + Run Kalman filter correction step. + + Args: + mean (ndarray): The predicted state's mean vector (8 dimensional). + covariance (ndarray): The state's covariance matrix (8x8 dimensional). + measurement (ndarray): The 4-dimensional measurement vector (x, y, a, h), where (x, y) is the center position, a the aspect ratio, and h the height of the bounding box. + + Returns: + Tuple[ndarray, ndarray]: Returns the measurement-corrected state distribution. + """ + projected_mean, projected_cov = self.project(mean, covariance) + + chol_factor, lower = scipy.linalg.cho_factor( + projected_cov, lower=True, check_finite=False + ) + kalman_gain = scipy.linalg.cho_solve( + (chol_factor, lower), + np.dot(covariance, self._update_mat.T).T, + check_finite=False, + ).T + innovation = measurement - projected_mean + + new_mean = mean + np.dot(innovation, kalman_gain.T) + new_covariance = covariance - np.linalg.multi_dot( + (kalman_gain, projected_cov, kalman_gain.T) + ) + return new_mean, new_covariance diff --git a/supervision/tracker/byte_tracker/matching.py b/supervision/tracker/byte_tracker/matching.py new file mode 100644 index 00000000..24abe224 --- /dev/null +++ b/supervision/tracker/byte_tracker/matching.py @@ -0,0 +1,64 @@ +from typing import List, Tuple + +import numpy as np +from scipy.optimize import linear_sum_assignment + +from supervision.detection.utils import box_iou_batch + + +def indices_to_matches( + cost_matrix: np.ndarray, indices: np.ndarray, thresh: float +) -> Tuple[np.ndarray, tuple, tuple]: + matched_cost = cost_matrix[tuple(zip(*indices))] + matched_mask = matched_cost <= thresh + + matches = indices[matched_mask] + unmatched_a = tuple(set(range(cost_matrix.shape[0])) - set(matches[:, 0])) + unmatched_b = tuple(set(range(cost_matrix.shape[1])) - set(matches[:, 1])) + return matches, unmatched_a, unmatched_b + + +def linear_assignment( + cost_matrix: np.ndarray, thresh: float +) -> [np.ndarray, Tuple[int], Tuple[int, int]]: + if cost_matrix.size == 0: + return ( + np.empty((0, 2), dtype=int), + tuple(range(cost_matrix.shape[0])), + tuple(range(cost_matrix.shape[1])), + ) + + cost_matrix[cost_matrix > thresh] = thresh + 1e-4 + row_ind, col_ind = linear_sum_assignment(cost_matrix) + indices = np.column_stack((row_ind, col_ind)) + + return indices_to_matches(cost_matrix, indices, thresh) + + +def iou_distance(atracks: List, btracks: List) -> np.ndarray: + if (len(atracks) > 0 and isinstance(atracks[0], np.ndarray)) or ( + len(btracks) > 0 and isinstance(btracks[0], np.ndarray) + ): + atlbrs = atracks + btlbrs = btracks + else: + atlbrs = [track.tlbr for track in atracks] + btlbrs = [track.tlbr for track in btracks] + + _ious = np.zeros((len(atlbrs), len(btlbrs)), dtype=np.float32) + if _ious.size != 0: + _ious = box_iou_batch(np.asarray(atlbrs), np.asarray(btlbrs)) + cost_matrix = 1 - _ious + + return cost_matrix + + +def fuse_score(cost_matrix: np.ndarray, detections: List) -> np.ndarray: + if cost_matrix.size == 0: + return cost_matrix + iou_sim = 1 - cost_matrix + det_scores = np.array([det.score for det in detections]) + det_scores = np.expand_dims(det_scores, axis=0).repeat(cost_matrix.shape[0], axis=0) + fuse_sim = iou_sim * det_scores + fuse_cost = 1 - fuse_sim + return fuse_cost diff --git a/supervision/utils/video.py b/supervision/utils/video.py index fcec53b1..dda68ca6 100644 --- a/supervision/utils/video.py +++ b/supervision/utils/video.py @@ -163,15 +163,15 @@ def process_video( Examples: ```python - >>> from supervision import process_video + >>> import supervision as sv - >>> def process_frame(scene: np.ndarray) -> np.ndarray: + >>> def callback(scene: np.ndarray, index: int) -> np.ndarray: ... ... >>> process_video( - ... source_path='source_video.mp4', - ... target_path='target_video.mp4', - ... callback=process_frame + ... source_path='...', + ... target_path='...', + ... callback=callback ... ) ``` """ diff --git a/test/dataset/formats/test_pascal_voc.py b/test/dataset/formats/test_pascal_voc.py new file mode 100644 index 00000000..fa3f1124 --- /dev/null +++ b/test/dataset/formats/test_pascal_voc.py @@ -0,0 +1,147 @@ +import xml.etree.ElementTree as ET +from contextlib import ExitStack as DoesNotRaise +from test.utils import mock_detections +from typing import List, Optional + +import numpy as np +import pytest + +from supervision.dataset.formats.pascal_voc import ( + detections_from_xml_obj, + object_to_pascal_voc, + parse_polygon_points, +) + + +def are_xml_elements_equal(elem1, elem2): + if ( + elem1.tag != elem2.tag + or elem1.attrib != elem2.attrib + or elem1.text != elem2.text + or len(elem1) != len(elem2) + ): + return False + + for child1, child2 in zip(elem1, elem2): + if not are_xml_elements_equal(child1, child2): + return False + + return True + + +@pytest.mark.parametrize( + "xyxy, name, polygon, expected_result, exception", + [ + ( + [0, 0, 10, 10], + "test", + None, + ET.fromstring( + """test001010""" + ), + DoesNotRaise(), + ), + ( + [0, 0, 10, 10], + "test", + [[0, 0], [10, 0], [10, 10], [0, 10]], + ET.fromstring( + """test001010001001010010""" + ), + DoesNotRaise(), + ), + ], +) +def test_object_to_pascal_voc( + xyxy: np.ndarray, + name: str, + polygon: Optional[np.ndarray], + expected_result, + exception: Exception, +): + with exception: + result = object_to_pascal_voc(xyxy=xyxy, name=name, polygon=polygon) + assert are_xml_elements_equal(result, expected_result) + + +@pytest.mark.parametrize( + "polygon_element, expected_result, exception", + [ + ( + ET.fromstring( + """001001010010""" + ), + [[0, 0], [10, 0], [10, 10], [0, 10]], + DoesNotRaise(), + ) + ], +) +def test_parse_polygon_points( + polygon_element, + expected_result: List[list], + exception, +): + with exception: + result = parse_polygon_points(polygon_element) + assert result == expected_result + + +ONE_CLASS_N_BBOX = """test001010test10102020""" + + +ONE_CLASS_ONE_BBOX = """test001010""" + + +N_CLASS_N_BBOX = """test001010test20303040test210102020""" + +NO_DETECTIONS = """""" + + +@pytest.mark.parametrize( + "xml_string, classes, resolution_wh, force_masks, expected_result, exception", + [ + ( + ONE_CLASS_ONE_BBOX, + ["test"], + (100, 100), + False, + mock_detections(np.array([[0, 0, 10, 10]]), None, [0]), + DoesNotRaise(), + ), + ( + ONE_CLASS_N_BBOX, + ["test"], + (100, 100), + False, + mock_detections(np.array([[0, 0, 10, 10], [10, 10, 20, 20]]), None, [0, 0]), + DoesNotRaise(), + ), + ( + N_CLASS_N_BBOX, + ["test", "test2"], + (100, 100), + False, + mock_detections( + np.array([[0, 0, 10, 10], [20, 30, 30, 40], [10, 10, 20, 20]]), + None, + [0, 0, 1], + ), + DoesNotRaise(), + ), + ( + NO_DETECTIONS, + [], + (100, 100), + False, + mock_detections(np.empty((0, 4)), None, []), + DoesNotRaise(), + ), + ], +) +def test_detections_from_xml_obj( + xml_string, classes, resolution_wh, force_masks, expected_result, exception +): + with exception: + root = ET.fromstring(xml_string) + result, _ = detections_from_xml_obj(root, classes, resolution_wh, force_masks) + assert result == expected_result diff --git a/test/metrics/test_detection.py b/test/metrics/test_detection.py index 0ecab3f0..e0a5f2ce 100644 --- a/test/metrics/test_detection.py +++ b/test/metrics/test_detection.py @@ -5,8 +5,8 @@ import numpy as np import pytest from supervision.detection.core import Detections -from supervision.metrics.detection import ConfusionMatrix -from test.utils import mock_detections +from supervision.metrics.detection import ConfusionMatrix, detections_to_tensor, MeanAveragePrecision +from test.utils import mock_detections, assert_almost_equal CLASSES = np.arange(80) NUM_CLASSES = len(CLASSES) @@ -167,7 +167,7 @@ def test_detections_to_tensor( exception: Exception ): with exception: - result = ConfusionMatrix.detections_to_tensor( + result = detections_to_tensor( detections=detections, with_confidence=with_confidence ) @@ -400,3 +400,43 @@ def test_drop_extra_matches( result = ConfusionMatrix._drop_extra_matches(matches) assert np.array_equal(result, expected_result) + + +@pytest.mark.parametrize( + 'recall, precision, expected_result, exception', + [ + ( + np.array([1.0]), + np.array([1.0]), + 1.0, + DoesNotRaise() + ), # perfect recall and precision + ( + np.array([0.0]), + np.array([0.0]), + 0.0, + DoesNotRaise() + ), # no recall and precision + ( + np.array([0.0, 0.2, 0.2, 0.8, 0.8, 1.0]), + np.array([0.7, 0.8, 0.4, 0.5, 0.1, 0.2]), + 0.5, + DoesNotRaise() + ), + ( + np.array([0.0, 0.5, 0.5, 1.0]), + np.array([0.75, 0.75, 0.75, 0.75]), + 0.75, + DoesNotRaise() + ) + ] +) +def test_compute_average_precision( + recall: np.ndarray, + precision: np.ndarray, + expected_result: float, + exception: Exception +) -> None: + with exception: + result = MeanAveragePrecision.compute_average_precision(recall=recall, precision=precision) + assert_almost_equal(result, expected_result, tolerance=0.01) diff --git a/test/utils.py b/test/utils.py index 20c294e5..a1d6907e 100644 --- a/test/utils.py +++ b/test/utils.py @@ -21,3 +21,7 @@ def mock_detections( if tracker_id is None else np.array(tracker_id, dtype=int), ) + + +def assert_almost_equal(actual, expected, tolerance=1e-5): + assert abs(actual - expected) < tolerance, f"Expected {expected}, but got {actual}."