Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
1bddf26f34 | ||
|
|
73ed3f5f5e | ||
|
|
88e7fb76d2 | ||
|
|
09c5407347 | ||
|
|
bfcf6469d4 | ||
|
|
ec72f81749 | ||
|
|
3a52f7b082 | ||
|
|
5efb3e2df1 | ||
|
|
6e33cb04ca | ||
|
|
2455ef7cfe | ||
|
|
1d0ffeba72 | ||
|
|
55d14cf33e | ||
|
|
a66d835a57 | ||
|
|
440be5d416 | ||
|
|
4fd1ba4dd3 | ||
|
|
61108048e1 | ||
|
|
a0b64342ca | ||
|
|
7cf593b55d | ||
|
|
1c1936e411 | ||
|
|
331ccb8dde | ||
|
|
7366ba9036 | ||
|
|
9d5e53030a | ||
|
|
d67fa3ae9b | ||
|
|
10ad27d726 | ||
|
|
8511bbb762 | ||
|
|
14388e56a2 | ||
|
|
1eb74ed88a | ||
|
|
87d830d3a1 | ||
|
|
d7a5624d60 | ||
|
|
d3c71c2119 | ||
|
|
9ad5e50ae9 | ||
|
|
67fa66c7e7 | ||
|
|
492f033882 | ||
|
|
e911516180 | ||
|
|
ed30ef68af | ||
|
|
fefd273a0f | ||
|
|
8746dd3d0e | ||
|
|
7cb57c9cc9 | ||
|
|
ddc036ea3a | ||
|
|
c83959bb16 | ||
|
|
4f513ce275 | ||
|
|
6f30c87612 | ||
|
|
93a8110908 | ||
|
|
410b7eb4fc | ||
|
|
e6e1397f56 | ||
|
|
0ca9a3ecdf | ||
|
|
c5016bbd8f | ||
|
|
9fbd0ddf62 | ||
|
|
a630870968 | ||
|
|
06416607cc | ||
|
|
1b449294c4 | ||
|
|
36f95d060d | ||
|
|
07c344e050 | ||
|
|
f3bdb6efa2 | ||
|
|
f4f9e7ddc7 | ||
|
|
53f2b4fd45 | ||
|
|
c7582f7c07 | ||
|
|
3e8906d830 | ||
|
|
6521f9f4ed | ||
|
|
560a35b75f | ||
|
|
405a4c73f9 | ||
|
|
87fd10c366 | ||
|
|
8141cb357f | ||
|
|
87fb336eb6 | ||
|
|
28d279bec3 | ||
|
|
1b0bcea5a6 | ||
|
|
507e000bca | ||
|
|
f7bf766a60 | ||
|
|
1e0f18f792 | ||
|
|
be5c1b1aa1 | ||
|
|
c35a7714d2 | ||
|
|
a8cc5a68bb | ||
|
|
de4056192e | ||
|
|
d337a03d28 | ||
|
|
00b67588d8 | ||
|
|
e10e0f21df | ||
|
|
4fb6760d2c | ||
|
|
de7ed5e115 | ||
|
|
d929b86c9f | ||
|
|
0b1627f835 | ||
|
|
20b1b21774 | ||
|
|
e5eaf26137 | ||
|
|
be3e67949c | ||
|
|
2587fdfefd | ||
|
|
08a88a8c0c | ||
|
|
181c0b833b | ||
|
|
11601e3f2d | ||
|
|
1b8ff617c3 | ||
|
|
5bc83bf6cc | ||
|
|
dd80753b2b | ||
|
|
b2bb69bebb | ||
|
|
9c081b0e11 | ||
|
|
7f613c6d27 | ||
|
|
9736eba9e6 | ||
|
|
53dfeac19b | ||
|
|
5c25e58e39 | ||
|
|
52949f2fd1 | ||
|
|
56c511b5d8 | ||
|
|
62014186a2 | ||
|
|
6cd7cf0209 | ||
|
|
a39c92016f | ||
|
|
4e88e26d91 | ||
|
|
484e53324e | ||
|
|
e5c4d8e3b9 | ||
|
|
52ef3fadcb | ||
|
|
472ddd8d0e | ||
|
|
90b23430d2 | ||
|
|
e91f0cb496 | ||
|
|
fa6238dbc4 | ||
|
|
1a7306c4ae | ||
|
|
c6b8fbed31 | ||
|
|
b9598f54b6 | ||
|
|
21e5d5139c | ||
|
|
8f583404ce | ||
|
|
bb42753e23 | ||
|
|
ea0ec29c81 | ||
|
|
3d4326c79b | ||
|
|
ca212b70c8 | ||
|
|
0008b33eb1 | ||
|
|
77190f287c | ||
|
|
121d4f31bb | ||
|
|
7c97e87e79 | ||
|
|
9640f1d756 | ||
|
|
2419d42224 | ||
|
|
c6517de143 | ||
|
|
2706304a73 | ||
|
|
8538c26a22 | ||
|
|
a0c60fe043 | ||
|
|
01d8b0316c | ||
|
|
6356a8ac3d | ||
|
|
bbd9a955f8 | ||
|
|
7b2bcf00ad | ||
|
|
775a3f2795 | ||
|
|
5dae02e93e | ||
|
|
666427df08 | ||
|
|
78411b14b5 | ||
|
|
69ae5c6cd9 | ||
|
|
94d08bc9b5 | ||
|
|
693465d74d | ||
|
|
f82f0fade6 | ||
|
|
474a4db248 | ||
|
|
89a9307855 | ||
|
|
5f94b6c23a | ||
|
|
18c972e422 | ||
|
|
426c4498df | ||
|
|
a3596584c7 | ||
|
|
d8787080ec | ||
|
|
c9e10b8a04 | ||
|
|
597e12d723 | ||
|
|
bb2d2c9f80 | ||
|
|
c2a29d6577 | ||
|
|
7a6da63a69 | ||
|
|
796c17f085 | ||
|
|
ddd400ed24 | ||
|
|
38b6ac9935 | ||
|
|
4a450ce2ed | ||
|
|
9c40f6f4aa | ||
|
|
f6766cb116 | ||
|
|
52f5d9ef63 | ||
|
|
02117dd8ae | ||
|
|
bddd32dbb5 | ||
|
|
07c4472eac | ||
|
|
11541fe271 | ||
|
|
6d7de56c25 | ||
|
|
85de9b255b | ||
|
|
a3a19d93bc | ||
|
|
1eaccd3e87 | ||
|
|
eb1155cc17 | ||
|
|
beb3a2ae9f | ||
|
|
b4efe69d1e | ||
|
|
3d6b5e75a8 | ||
|
|
a7edf46717 | ||
|
|
cd6a7fd186 | ||
|
|
82729a61b0 | ||
|
|
30a51c72d6 | ||
|
|
6b0be9c55b | ||
|
|
cf2c8cb364 | ||
|
|
7f8d091f50 | ||
|
|
36e03f3707 | ||
|
|
f8baad3065 | ||
|
|
18f7cad081 | ||
|
|
bebb427a76 | ||
|
|
6f03b7ffd0 | ||
|
|
06f41e83dd | ||
|
|
2ec33dbe9a | ||
|
|
b86f8ccdb8 | ||
|
|
7de50e6b0b | ||
|
|
0f889935fd | ||
|
|
61f972db79 | ||
|
|
4b419fa19e | ||
|
|
738d00cc89 | ||
|
|
9c53bc16ac | ||
|
|
8a84463aaa | ||
|
|
9ca1204a4c | ||
|
|
6a2203c07a | ||
|
|
04df220737 | ||
|
|
981bcc71ea | ||
|
|
1d42992e6b | ||
|
|
1c0434c6d8 | ||
|
|
b40eb5b47f | ||
|
|
98f12b60ea | ||
|
|
2c988ac269 | ||
|
|
f53c51ac8d | ||
|
|
7426d52c4a | ||
|
|
ad26e13704 | ||
|
|
d8aecc4fe7 | ||
|
|
7dbf279477 | ||
|
|
d59f9c4c74 | ||
|
|
0d991009f7 | ||
|
|
3a6de392fc | ||
|
|
5bd78d11cc | ||
|
|
cd8541f948 | ||
|
|
77a1faa7de | ||
|
|
b47bc808f7 | ||
|
|
16f3a61431 | ||
|
|
a051be0920 | ||
|
|
7d449abee2 | ||
|
|
12ec48105b | ||
|
|
afc240e6be | ||
|
|
c2b9745fe5 | ||
|
|
4c5c15fcf4 | ||
|
|
47268c893c | ||
|
|
512f98f074 | ||
|
|
1f0b50645c | ||
|
|
3c5f5f7741 | ||
|
|
01956a2735 | ||
|
|
a2b8cf50d6 | ||
|
|
bfe65013ab | ||
|
|
824b2f6dc2 | ||
|
|
8f70513edc | ||
|
|
9167101e92 | ||
|
|
b064356280 | ||
|
|
4e5f87686f | ||
|
|
7bd6a107ce | ||
|
|
a71d8c254a | ||
|
|
755365a92d | ||
|
|
8c70197734 | ||
|
|
735de7de7e | ||
|
|
dacce7aa5f | ||
|
|
a5c6fe57f6 | ||
|
|
b69d264ffe | ||
|
|
92413d3cc2 | ||
|
|
5f314cf21f | ||
|
|
bc4ba7f67c | ||
|
|
47385f5a97 | ||
|
|
234f6fa43b | ||
|
|
26402386aa | ||
|
|
6f427b9ca4 | ||
|
|
6a40516371 | ||
|
|
1e1ce3c060 | ||
|
|
8fb8432e8b | ||
|
|
0ed404602f | ||
|
|
d81d187da6 | ||
|
|
b90f9a1677 | ||
|
|
67a81afcfa | ||
|
|
f0bc628eaf | ||
|
|
f38c438c3c | ||
|
|
d37c62d0cf | ||
|
|
8b764218eb | ||
|
|
4f9d8c02c7 | ||
|
|
4f79d29934 | ||
|
|
d86623600a | ||
|
|
475b2ceeea | ||
|
|
b1325132d7 | ||
|
|
3d3b3e6436 | ||
|
|
95661f3bb1 | ||
|
|
e6ab865a0c | ||
|
|
63669e587a | ||
|
|
2166bc23bc | ||
|
|
867dd78f58 | ||
|
|
255b34a183 | ||
|
|
f7b59870ed | ||
|
|
7948234f94 | ||
|
|
b5e83fa953 | ||
|
|
de536a3f8e | ||
|
|
6741ae642a | ||
|
|
3834e5cdf9 | ||
|
|
aa77c396f7 | ||
|
|
bee6dec22c | ||
|
|
db2eede55b | ||
|
|
ba880402ad | ||
|
|
62e097b559 | ||
|
|
451bb5e513 | ||
|
|
330167fc73 | ||
|
|
1c23e806ba | ||
|
|
2c31f8ec94 | ||
|
|
a6c73f70ec | ||
|
|
13ee396b20 | ||
|
|
0f222375e2 | ||
|
|
a63f65fff3 | ||
|
|
94f15a6e50 | ||
|
|
525e1b00ec | ||
|
|
7f4662cf3d | ||
|
|
3a31ea2986 | ||
|
|
e0ca4ca5e7 | ||
|
|
0fcc17ee11 | ||
|
|
52cbe684d4 | ||
|
|
fce6e660d9 | ||
|
|
c6f2bf27c1 | ||
|
|
523d76f86d | ||
|
|
1ba1b67b47 | ||
|
|
a45cc865dd | ||
|
|
4942d3d817 | ||
|
|
514f61b885 | ||
|
|
fc154e34fe | ||
|
|
6f0f3bc3bc | ||
|
|
b40e2f228a | ||
|
|
21062b22a4 | ||
|
|
f9d3aaa67e | ||
|
|
58156fe61e | ||
|
|
d170cbeda0 | ||
|
|
e131f8b197 | ||
|
|
f372fb69e9 | ||
|
|
9f8291405a | ||
|
|
da9d50b4e2 | ||
|
|
8f437107b8 | ||
|
|
2a9a1f05bb | ||
|
|
414e0327bc | ||
|
|
f2d2827eee | ||
|
|
8eaa42abc6 | ||
|
|
aa50052ff5 | ||
|
|
a59e4dfe73 | ||
|
|
d6d4760fce | ||
|
|
f2bbb7c38c | ||
|
|
fa44884ea7 | ||
|
|
6fdabf0a96 | ||
|
|
26289e3743 | ||
|
|
e2d1fe7695 | ||
|
|
36d3e94a76 | ||
|
|
c0c292ef82 | ||
|
|
9462a3058d | ||
|
|
77c2afe559 | ||
|
|
95c8576aaf | ||
|
|
12b3def032 | ||
|
|
e8f616ba2f | ||
|
|
f8005b268b | ||
|
|
e4ef57ec2d | ||
|
|
1642432541 | ||
|
|
1febd715e2 | ||
|
|
9a9d0c5259 | ||
|
|
446ae23d99 | ||
|
|
a181826527 | ||
|
|
0a1212f23d | ||
|
|
b0e0d074af | ||
|
|
27e8e2d91b | ||
|
|
f4c17f09c8 | ||
|
|
b219961f6e | ||
|
|
1645ecf4bc | ||
|
|
d4001a5bfa | ||
|
|
70fe357a4e | ||
|
|
4b621c3c8d | ||
|
|
277a9cb6b9 | ||
|
|
6dde70eb0c | ||
|
|
a91118322b | ||
|
|
86b33ed558 | ||
|
|
3723eae50c | ||
|
|
d7e08d42eb | ||
|
|
af82b4e2d2 | ||
|
|
9fef53f8a3 | ||
|
|
8d8703a1c6 | ||
|
|
84a59a7e01 | ||
|
|
dec982200b | ||
|
|
da6a324c1f | ||
|
|
7e391d0861 | ||
|
|
c2018fa1ba | ||
|
|
6f7a50f41b | ||
|
|
204a04a85f | ||
|
|
5727fb5cb2 | ||
|
|
a9dce586c2 | ||
|
|
57e17c34a0 | ||
|
|
36cdcf1b12 | ||
|
|
3c64871c87 | ||
|
|
7ec6bcd916 | ||
|
|
da13366d0a | ||
|
|
771a47e15d | ||
|
|
fe00da70bf | ||
|
|
7067741310 | ||
|
|
117ba2fcc0 | ||
|
|
f4dc515df1 | ||
|
|
0754cb75bd | ||
|
|
44a174f889 | ||
|
|
e9fbcf5dcb | ||
|
|
9b6b74dc77 | ||
|
|
e249ea8644 | ||
|
|
64c446ae44 | ||
|
|
dfd26545f8 | ||
|
|
a28755bbdb | ||
|
|
06d0646fb8 | ||
|
|
8afd54a1cc | ||
|
|
cddab1cbe6 | ||
|
|
8662671e6b | ||
|
|
c7c8e67ffd | ||
|
|
b2b75f1105 | ||
|
|
2f726ca795 | ||
|
|
40a6fe1b5f | ||
|
|
0cfad8407c | ||
|
|
1d28b8ae26 | ||
|
|
cc2e08ac72 | ||
|
|
ba03fefbd8 | ||
|
|
4cd9f05fab | ||
|
|
8066125f03 | ||
|
|
aa772dde8d | ||
|
|
6587f39b0f | ||
|
|
0f95c5d8b3 | ||
|
|
800e39d64c | ||
|
|
926f935bad | ||
|
|
d3da6afa07 | ||
|
|
97ae155b1e | ||
|
|
24f013757f | ||
|
|
532245640b | ||
|
|
864764e665 | ||
|
|
ca348401d5 | ||
|
|
6def59444d | ||
|
|
bc0a38500d | ||
|
|
300ebea8ec | ||
|
|
5cc7da289e | ||
|
|
5edd58d6de | ||
|
|
5729f8ffc3 | ||
|
|
7a07517750 | ||
|
|
699fc01a0c | ||
|
|
43339e380a | ||
|
|
89a33cda8a | ||
|
|
e84302b43f | ||
|
|
10d3e3d48f | ||
|
|
3c59618152 | ||
|
|
e882c398b7 | ||
|
|
93ab21dda7 | ||
|
|
effa1415b5 | ||
|
|
e9e075acaf | ||
|
|
61e9584f5a | ||
|
|
2a640d3ab6 | ||
|
|
2f478f4d06 | ||
|
|
0c20830be5 | ||
|
|
6c51c4a1b4 | ||
|
|
d6883caadc | ||
|
|
f1f2728e0d | ||
|
|
b6b7663ea7 | ||
|
|
3d074fa213 | ||
|
|
9fd267acb2 | ||
|
|
727451b382 | ||
|
|
1b5e6db196 | ||
|
|
5f432fbc93 | ||
|
|
ea4638f5f5 | ||
|
|
fc4bb27aa8 | ||
|
|
ace6ce363f | ||
|
|
4b38634c63 | ||
|
|
985608146c | ||
|
|
f0cef99de7 | ||
|
|
04b7badcbe | ||
|
|
54524a4a1f | ||
|
|
3f55e7a79a | ||
|
|
7a782b1ac3 | ||
|
|
6016338a4f | ||
|
|
8d131ad3e8 | ||
|
|
ae55e9c908 | ||
|
|
5cb3c603cd | ||
|
|
56e0b77211 | ||
|
|
bc9fe6f690 | ||
|
|
fb211a6160 | ||
|
|
37b7224380 | ||
|
|
314454c7a7 | ||
|
|
ac39bdd487 | ||
|
|
fc1e899575 | ||
|
|
fccbe6baa2 | ||
|
|
d21154e882 | ||
|
|
1b3af74756 | ||
|
|
dc1083711a | ||
|
|
8df5c5b0c4 | ||
|
|
abbc0e44b0 | ||
|
|
950744111a | ||
|
|
c4d296fbc6 | ||
|
|
9c4555a800 | ||
|
|
20954a2996 | ||
|
|
22d82c306f | ||
|
|
399121b65e | ||
|
|
1b9c65bdf8 | ||
|
|
3a8c4d63ea | ||
|
|
bfe058c307 | ||
|
|
12ce5feab7 | ||
|
|
186bfc44fa | ||
|
|
2c0905ba28 | ||
|
|
fd717ea749 | ||
|
|
40853350fd | ||
|
|
fdd8127293 | ||
|
|
21a480b483 | ||
|
|
1b013f85e5 | ||
|
|
36b48def89 | ||
|
|
b7f72b79de | ||
|
|
141428310a | ||
|
|
8f39df2fc0 | ||
|
|
5107dd359f | ||
|
|
4b423fd402 | ||
|
|
d480ff5be2 | ||
|
|
2b876bdd64 | ||
|
|
476be527e5 | ||
|
|
f78ac90464 | ||
|
|
27efbc044d | ||
|
|
2a2c259964 | ||
|
|
d4ea295f09 | ||
|
|
9c5c296fb6 | ||
|
|
85c4d5a644 | ||
|
|
ad06937653 | ||
|
|
e01e495c92 | ||
|
|
6784c78b4f | ||
|
|
4375e19245 | ||
|
|
a75a92708c | ||
|
|
c83e57c638 | ||
|
|
fcde1cea9c | ||
|
|
8614984f8f | ||
|
|
c705d50e74 | ||
|
|
d5b8a3aa1a | ||
|
|
cc81ae723b | ||
|
|
f7a4a788ed | ||
|
|
35d9ee7bad | ||
|
|
abbde887dc | ||
|
|
4f5f87acb8 | ||
|
|
3f78150a1d | ||
|
|
183ba1e809 | ||
|
|
597c776c39 | ||
|
|
5291706bc1 | ||
|
|
a50e47e7c0 | ||
|
|
c9d85477f1 | ||
|
|
19271f362e | ||
|
|
9e1e9fc2fc | ||
|
|
629744dee1 | ||
|
|
c9d2b52900 | ||
|
|
43cca52a58 | ||
|
|
e1ee80cb37 | ||
|
|
6a8f3fc03b | ||
|
|
d9f6c5c73c | ||
|
|
c4b33fe1d6 | ||
|
|
65a93f8287 | ||
|
|
1bcf8346dc | ||
|
|
aba6ff57a9 | ||
|
|
b5e1342038 | ||
|
|
ebfaee22eb | ||
|
|
9c062aac29 | ||
|
|
64d223b76d | ||
|
|
21968e53c6 | ||
|
|
9b77b2a1c6 | ||
|
|
3e4cb6df42 | ||
|
|
5fdbb70c60 | ||
|
|
9aee232d04 | ||
|
|
9adc953780 | ||
|
|
4ace93f9ab | ||
|
|
c1d3ab2944 | ||
|
|
2072eeac12 | ||
|
|
c74efc4908 | ||
|
|
08f399cf19 | ||
|
|
17020b23de | ||
|
|
cd929525f6 | ||
|
|
bb18de8dce | ||
|
|
dec96aec4d | ||
|
|
ee45c684d0 | ||
|
|
ea517642f0 | ||
|
|
cdc08c0455 | ||
|
|
a0c808cfdf | ||
|
|
f146e88ceb | ||
|
|
66d29015d6 | ||
|
|
46cb82797e | ||
|
|
194ff0aaad | ||
|
|
a579ed46c6 | ||
|
|
ddabe04409 | ||
|
|
8918783702 | ||
|
|
6adb0f019d | ||
|
|
84ae12768a | ||
|
|
9cbec9fc94 | ||
|
|
ed12daf77d | ||
|
|
84cfa9609b | ||
|
|
f2eccae88e | ||
|
|
55c8295755 | ||
|
|
7849033da3 | ||
|
|
2e443c80c7 | ||
|
|
006caa0b9b | ||
|
|
34e45bed5e | ||
|
|
0d39de898a | ||
|
|
782e7e5a39 | ||
|
|
b40265db9e | ||
|
|
de1ceff065 | ||
|
|
5bf78e5505 | ||
|
|
f63336a824 | ||
|
|
4d82460358 | ||
|
|
540efc80ce | ||
|
|
beea21a97d | ||
|
|
484e3d4afd | ||
|
|
ab9728cca2 | ||
|
|
769b6b967c | ||
|
|
253113fa88 | ||
|
|
3b2902c87c | ||
|
|
9cff624e15 | ||
|
|
584c4fef82 | ||
|
|
282aa2c5d3 | ||
|
|
453d4df068 | ||
|
|
4d4b38b292 | ||
|
|
13e8afa996 | ||
|
|
dd4cae328b | ||
|
|
ed7103851a | ||
|
|
ca40b001dc | ||
|
|
de3ccc5c5d | ||
|
|
1b6d02898c | ||
|
|
5ac92bc43b | ||
|
|
86d9293770 | ||
|
|
0e943e673e | ||
|
|
234e376fca | ||
|
|
54b7cbeb70 | ||
|
|
d000292eb2 | ||
|
|
b551e6d410 | ||
|
|
138c271af0 | ||
|
|
d22940ce30 | ||
|
|
06da0e70d2 | ||
|
|
5d8b378f35 | ||
|
|
a4830389de | ||
|
|
d9a2efc64b | ||
|
|
f91b844b7d | ||
|
|
dc45d47e7e | ||
|
|
b5dedf09cd | ||
|
|
eb1ae1b790 | ||
|
|
f2f8a564e0 | ||
|
|
140a219091 | ||
|
|
86af9b9dce | ||
|
|
4366e4480f | ||
|
|
b66e9c8c32 | ||
|
|
d0d5626156 | ||
|
|
c661bf250d | ||
|
|
544f9cb0e2 | ||
|
|
84928a2da9 |
@@ -0,0 +1 @@
|
||||
*.ipynb linguist-vendored
|
||||
@@ -62,4 +62,4 @@ body:
|
||||
description: >
|
||||
(Optional) We encourage you to submit a [Pull Request](https://github.com/roboflow/supervision/pulls) (PR) to help improve Supervision for everyone, especially if you have a good understanding of how to implement a fix or feature.
|
||||
options:
|
||||
- label: Yes I'd like to help by submitting a PR!
|
||||
- label: Yes I'd like to help by submitting a PR!
|
||||
|
||||
@@ -46,4 +46,4 @@ body:
|
||||
description: >
|
||||
(Optional) We encourage you to submit a [Pull Request](https://github.com/roboflow/supervision/pulls) (PR) to help improve Supervision for everyone, especially if you have a good understanding of how to implement a fix or feature.
|
||||
options:
|
||||
- label: Yes I'd like to help by submitting a PR!
|
||||
- label: Yes I'd like to help by submitting a PR!
|
||||
|
||||
@@ -30,4 +30,4 @@ body:
|
||||
- type: textarea
|
||||
attributes:
|
||||
label: Additional
|
||||
description: Anything else you would like to share?
|
||||
description: Anything else you would like to share?
|
||||
|
||||
@@ -5,14 +5,15 @@ on:
|
||||
branches:
|
||||
- master
|
||||
- main
|
||||
- develop
|
||||
jobs:
|
||||
deploy:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v2
|
||||
- uses: actions/setup-python@v2
|
||||
- uses: actions/checkout@v3
|
||||
- uses: actions/setup-python@v4
|
||||
with:
|
||||
python-version: 3.x
|
||||
- run: pip install mkdocs-material
|
||||
- run: pip install "mkdocstrings[python]"
|
||||
- run: mkdocs gh-deploy --force
|
||||
- run: mkdocs gh-deploy --force
|
||||
|
||||
@@ -0,0 +1,41 @@
|
||||
name: Supervision Test Releases to PyPi
|
||||
on:
|
||||
push:
|
||||
tags:
|
||||
- '[0-9]+.[0-9]+[0-9]+.[0-9]+a[0-9]'
|
||||
- '[0-9]+.[0-9]+[0-9]+.[0-9]+b[0-9]'
|
||||
- '[0-9]+.[0-9]+[0-9]+.[0-9]+rc[0-9]'
|
||||
|
||||
# Allows you to run this workflow manually from the Actions tab
|
||||
workflow_dispatch:
|
||||
|
||||
jobs:
|
||||
build-n-publish:
|
||||
name: Build and publish to PyPI
|
||||
runs-on: ubuntu-latest
|
||||
|
||||
steps:
|
||||
- name: Checkout source
|
||||
uses: actions/checkout@v3
|
||||
|
||||
- name: 🐍 Set up Python 3.8 environment for build
|
||||
uses: actions/setup-python@v4
|
||||
with:
|
||||
python-version: "3.8"
|
||||
|
||||
- name: 🏗️ Build source and wheel distributions
|
||||
run: |
|
||||
python -m pip install --upgrade build twine
|
||||
python -m build
|
||||
twine check --strict dist/*
|
||||
- name: 🚀 Publish distribution to PyPI
|
||||
uses: pypa/gh-action-pypi-publish@release/v1
|
||||
with:
|
||||
user: ${{ secrets.PYPI_USERNAME }}
|
||||
password: ${{ secrets.PYPI_PASSWORD }}
|
||||
- name: 🚀 Publish to Test-PyPi
|
||||
uses: pypa/gh-action-pypi-publish@release/v1
|
||||
with:
|
||||
repository-url: https://test.pypi.org/legacy/
|
||||
user: ${{ secrets.PYPI_TEST_USERNAME }}
|
||||
password: ${{ secrets.PYPI_TEST_PASSWORD }}
|
||||
@@ -1,8 +1,13 @@
|
||||
name: Publish WorkFlow
|
||||
|
||||
name: Supervision Releases to PyPi
|
||||
on:
|
||||
release:
|
||||
types: [created]
|
||||
push:
|
||||
tags:
|
||||
- '[0-9]+.[0-9]+[0-9]+.[0-9]'
|
||||
- '[0-9]+.[0-9]+[0-9]+.[0-9]'
|
||||
- '[0-9]+.[0-9]+[0-9]+.[0-9]'
|
||||
|
||||
# Allows you to run this workflow manually from the Actions tab
|
||||
workflow_dispatch:
|
||||
|
||||
jobs:
|
||||
build:
|
||||
@@ -16,17 +21,23 @@ jobs:
|
||||
with:
|
||||
ref: ${{ github.head_ref }}
|
||||
- name: 🐍 Set up Python ${{ matrix.python-version }}
|
||||
uses: actions/setup-python@v2
|
||||
uses: actions/setup-python@v4
|
||||
with:
|
||||
python-version: ${{ matrix.python-version }}
|
||||
- name: 🦾 Install dependencies
|
||||
|
||||
- name: 🏗️ Build source and wheel distributions
|
||||
run: |
|
||||
python -m pip install --upgrade pip
|
||||
pip install ".[dev]"
|
||||
python -m pip install --upgrade build twine
|
||||
python -m build
|
||||
twine check --strict dist/*
|
||||
- name: 🚀 Publish to PyPi
|
||||
env:
|
||||
PYPI_USERNAME: ${{ secrets.PYPI_USERNAME }}
|
||||
PYPI_PASSWORD: ${{ secrets.PYPI_PASSWORD }}
|
||||
PYPI_TEST_PASSWORD: ${{ secrets.PYPI_TEST_PASSWORD }}
|
||||
run: |
|
||||
make publish -e PYPI_USERNAME=$PYPI_USERNAME -e PYPI_PASSWORD=$PYPI_PASSWORD -e PYPI_TEST_PASSWORD=$PYPI_TEST_PASSWORD
|
||||
uses: pypa/gh-action-pypi-publish@release/v1
|
||||
with:
|
||||
user: ${{ secrets.PYPI_USERNAME }}
|
||||
password: ${{ secrets.PYPI_PASSWORD }}
|
||||
- name: 🚀 Publish to Test-PyPi
|
||||
uses: pypa/gh-action-pypi-publish@release/v1
|
||||
with:
|
||||
repository-url: https://test.pypi.org/legacy/
|
||||
user: ${{ secrets.PYPI_TEST_USERNAME }}
|
||||
password: ${{ secrets.PYPI_TEST_PASSWORD }}
|
||||
|
||||
@@ -2,21 +2,19 @@ name: Test WorkFlow
|
||||
|
||||
on:
|
||||
pull_request:
|
||||
branches: [main]
|
||||
branches: [main, develop]
|
||||
|
||||
jobs:
|
||||
build:
|
||||
runs-on: ubuntu-latest
|
||||
strategy:
|
||||
matrix:
|
||||
python-version: [3.7, 3.8, 3.9, '3.10']
|
||||
python-version: ["3.8", "3.9", "3.10","3.11"]
|
||||
steps:
|
||||
- name: 🛎️ Checkout
|
||||
uses: actions/checkout@v3
|
||||
with:
|
||||
ref: ${{ github.head_ref }}
|
||||
- name: 🐍 Set up Python ${{ matrix.python-version }}
|
||||
uses: actions/setup-python@v2
|
||||
uses: actions/setup-python@v4
|
||||
with:
|
||||
python-version: ${{ matrix.python-version }}
|
||||
- name: 🦾 Install dependencies
|
||||
@@ -24,11 +22,5 @@ jobs:
|
||||
python -m pip install --upgrade pip
|
||||
pip install -e .
|
||||
pip install pytest
|
||||
pip install isort
|
||||
pip install flake8
|
||||
pip install "black==22.3.0"
|
||||
- name: 🧹 Lint with flake8
|
||||
run: |
|
||||
make check_code_quality
|
||||
- name: 🧪 Test
|
||||
run: "python -m pytest ./test"
|
||||
run: "python -m pytest ./test"
|
||||
|
||||
@@ -15,4 +15,4 @@ jobs:
|
||||
with:
|
||||
repo-token: ${{ secrets.GITHUB_TOKEN }}
|
||||
issue-message: "Hello there, thank you for opening an Issue ! 🙏🏻 The team was notified and they will get back to you asap."
|
||||
pr-message: "Hello there, thank you for opening an PR ! 🙏🏻 The team was notified and they will get back to you asap."
|
||||
pr-message: "Hello there, thank you for opening an PR ! 🙏🏻 The team was notified and they will get back to you asap."
|
||||
|
||||
@@ -54,6 +54,7 @@ coverage.xml
|
||||
*.py,cover
|
||||
.hypothesis/
|
||||
.pytest_cache/
|
||||
.ruff_cache/
|
||||
|
||||
# Translations
|
||||
*.mo
|
||||
|
||||
@@ -0,0 +1,80 @@
|
||||
|
||||
ci:
|
||||
autofix_prs: true
|
||||
autoupdate_schedule: weekly
|
||||
autofix_commit_msg: "fix(pre_commit): 🎨 auto format pre-commit hooks"
|
||||
autoupdate_commit_msg: "chore(pre_commit): ⬆ pre_commit autoupdate"
|
||||
|
||||
repos:
|
||||
- repo: https://github.com/pre-commit/pre-commit-hooks
|
||||
rev: v4.4.0
|
||||
hooks:
|
||||
- id: end-of-file-fixer
|
||||
- id: trailing-whitespace
|
||||
- id: check-yaml
|
||||
- id: check-docstring-first
|
||||
- id: check-executables-have-shebangs
|
||||
- id: check-toml
|
||||
- id: check-case-conflict
|
||||
- id: check-added-large-files
|
||||
args: ['--maxkb=2048']
|
||||
exclude: ^logo/
|
||||
- id: detect-private-key
|
||||
- id: forbid-new-submodules
|
||||
- id: pretty-format-json
|
||||
args: ['--autofix', '--no-sort-keys', '--indent=4']
|
||||
- id: end-of-file-fixer
|
||||
- id: mixed-line-ending
|
||||
|
||||
# - repo: https://github.com/asottile/pyupgrade
|
||||
# rev: v3.9.0
|
||||
# hooks:
|
||||
# - id: pyupgrade
|
||||
# name: Upgrade code
|
||||
# args: [--py38-plus]
|
||||
|
||||
- repo: https://github.com/PyCQA/isort
|
||||
rev: 5.12.0
|
||||
hooks:
|
||||
- id: isort
|
||||
name: Sort imports
|
||||
|
||||
- repo: https://github.com/PyCQA/flake8
|
||||
rev: 6.1.0
|
||||
hooks:
|
||||
- id: flake8
|
||||
name: Flake8 Checks
|
||||
entry: flake8
|
||||
additional_dependencies: [Flake8-pyproject]
|
||||
|
||||
|
||||
- repo: https://github.com/PyCQA/bandit
|
||||
rev: '1.7.5'
|
||||
hooks:
|
||||
- id: bandit
|
||||
args: ["-c", "pyproject.toml"]
|
||||
additional_dependencies: ["bandit[toml]"]
|
||||
|
||||
- repo: https://github.com/pycqa/isort
|
||||
rev: 5.12.0
|
||||
hooks:
|
||||
- id: isort
|
||||
name: isort (python)
|
||||
- id: isort
|
||||
name: isort (cython)
|
||||
types: [cython]
|
||||
- id: isort
|
||||
name: isort (pyi)
|
||||
types: [pyi]
|
||||
|
||||
|
||||
- repo: https://github.com/psf/black
|
||||
rev: 23.9.1
|
||||
hooks:
|
||||
- id: black
|
||||
|
||||
- repo: https://github.com/astral-sh/ruff-pre-commit
|
||||
rev: v0.0.292
|
||||
hooks:
|
||||
- id: ruff
|
||||
args: [--fix, --exit-non-zero-on-fix]
|
||||
+19
-8
@@ -1,12 +1,23 @@
|
||||
# This CITATION.cff file was generated with cffinit.
|
||||
# Visit https://bit.ly/cffinit to generate yours today!
|
||||
|
||||
cff-version: 1.2.0
|
||||
author: Roboflow
|
||||
message: If you use this software, please cite it as below.
|
||||
title: supervision
|
||||
version: 0.1.0
|
||||
date-released: 2023-01-19
|
||||
license: MIT
|
||||
repository-code: https://github.com/roboflow/supervision
|
||||
title: Supervision
|
||||
message: >-
|
||||
If you use this software, please cite it using the
|
||||
metadata from this file.
|
||||
type: software
|
||||
authors:
|
||||
- given-names: Roboflow
|
||||
email: support@roboflow.com
|
||||
repository-code: 'https://github.com/roboflow/supervision'
|
||||
url: 'https://roboflow.github.io/supervision/'
|
||||
abstract: >-
|
||||
supervision features a range of utilities for use in
|
||||
computer vision projects, from detections processing and
|
||||
filtering to confusion matrix calcuation.
|
||||
keywords:
|
||||
- computer vision
|
||||
- image processing
|
||||
- video processing
|
||||
- video processing
|
||||
license: MIT
|
||||
|
||||
@@ -0,0 +1,134 @@
|
||||
|
||||
# Contributor Covenant Code of Conduct
|
||||
|
||||
## Our Pledge
|
||||
|
||||
We as members, contributors, and leaders pledge to make participation in our
|
||||
community a harassment-free experience for everyone, regardless of age, body
|
||||
size, visible or invisible disability, ethnicity, sex characteristics, gender
|
||||
identity and expression, level of experience, education, socio-economic status,
|
||||
nationality, personal appearance, race, caste, color, religion, or sexual
|
||||
identity and orientation.
|
||||
|
||||
We pledge to act and interact in ways that contribute to an open, welcoming,
|
||||
diverse, inclusive, and healthy community.
|
||||
|
||||
## Our Standards
|
||||
|
||||
Examples of behavior that contributes to a positive environment for our
|
||||
community include:
|
||||
|
||||
* Demonstrating empathy and kindness toward other people
|
||||
* Being respectful of differing opinions, viewpoints, and experiences
|
||||
* Giving and gracefully accepting constructive feedback
|
||||
* Accepting responsibility and apologizing to those affected by our mistakes,
|
||||
and learning from the experience
|
||||
* Focusing on what is best not just for us as individuals, but for the overall
|
||||
community
|
||||
|
||||
Examples of unacceptable behavior include:
|
||||
|
||||
* The use of sexualized language or imagery, and sexual attention or advances of
|
||||
any kind
|
||||
* Trolling, insulting or derogatory comments, and personal or political attacks
|
||||
* Public or private harassment
|
||||
* Publishing others' private information, such as a physical or email address,
|
||||
without their explicit permission
|
||||
* Other conduct which could reasonably be considered inappropriate in a
|
||||
professional setting
|
||||
|
||||
## Enforcement Responsibilities
|
||||
|
||||
Community leaders are responsible for clarifying and enforcing our standards of
|
||||
acceptable behavior and will take appropriate and fair corrective action in
|
||||
response to any behavior that they deem inappropriate, threatening, offensive,
|
||||
or harmful.
|
||||
|
||||
Community leaders have the right and responsibility to remove, edit, or reject
|
||||
comments, commits, code, wiki edits, issues, and other contributions that are
|
||||
not aligned to this Code of Conduct, and will communicate reasons for moderation
|
||||
decisions when appropriate.
|
||||
|
||||
## Scope
|
||||
|
||||
This Code of Conduct applies within all community spaces, and also applies when
|
||||
an individual is officially representing the community in public spaces.
|
||||
Examples of representing our community include using an official e-mail address,
|
||||
posting via an official social media account, or acting as an appointed
|
||||
representative at an online or offline event.
|
||||
|
||||
## Enforcement
|
||||
|
||||
Instances of abusive, harassing, or otherwise unacceptable behavior may be
|
||||
reported to the community leaders responsible for enforcement at
|
||||
community-reports@roboflow.com.
|
||||
|
||||
All complaints will be reviewed and investigated promptly and fairly.
|
||||
|
||||
All community leaders are obligated to respect the privacy and security of the
|
||||
reporter of any incident.
|
||||
|
||||
## Enforcement Guidelines
|
||||
|
||||
Community leaders will follow these Community Impact Guidelines in determining
|
||||
the consequences for any action they deem in violation of this Code of Conduct:
|
||||
|
||||
### 1. Correction
|
||||
|
||||
**Community Impact**: Use of inappropriate language or other behavior deemed
|
||||
unprofessional or unwelcome in the community.
|
||||
|
||||
**Consequence**: A private, written warning from community leaders, providing
|
||||
clarity around the nature of the violation and an explanation of why the
|
||||
behavior was inappropriate. A public apology may be requested.
|
||||
|
||||
### 2. Warning
|
||||
|
||||
**Community Impact**: A violation through a single incident or series of
|
||||
actions.
|
||||
|
||||
**Consequence**: A warning with consequences for continued behavior. No
|
||||
interaction with the people involved, including unsolicited interaction with
|
||||
those enforcing the Code of Conduct, for a specified period of time. This
|
||||
includes avoiding interactions in community spaces as well as external channels
|
||||
like social media. Violating these terms may lead to a temporary or permanent
|
||||
ban.
|
||||
|
||||
### 3. Temporary Ban
|
||||
|
||||
**Community Impact**: A serious violation of community standards, including
|
||||
sustained inappropriate behavior.
|
||||
|
||||
**Consequence**: A temporary ban from any sort of interaction or public
|
||||
communication with the community for a specified period of time. No public or
|
||||
private interaction with the people involved, including unsolicited interaction
|
||||
with those enforcing the Code of Conduct, is allowed during this period.
|
||||
Violating these terms may lead to a permanent ban.
|
||||
|
||||
### 4. Permanent Ban
|
||||
|
||||
**Community Impact**: Demonstrating a pattern of violation of community
|
||||
standards, including sustained inappropriate behavior, harassment of an
|
||||
individual, or aggression toward or disparagement of classes of individuals.
|
||||
|
||||
**Consequence**: A permanent ban from any sort of public interaction within the
|
||||
community.
|
||||
|
||||
## Attribution
|
||||
|
||||
This Code of Conduct is adapted from the [Contributor Covenant][homepage],
|
||||
version 2.1, available at
|
||||
[https://www.contributor-covenant.org/version/2/1/code_of_conduct.html][v2.1].
|
||||
|
||||
Community Impact Guidelines were inspired by
|
||||
[Mozilla's code of conduct enforcement ladder][Mozilla CoC].
|
||||
|
||||
For answers to common questions about this code of conduct, see the FAQ at
|
||||
[https://www.contributor-covenant.org/faq][FAQ]. Translations are available at
|
||||
[https://www.contributor-covenant.org/translations][translations].
|
||||
|
||||
[homepage]: https://www.contributor-covenant.org
|
||||
[v2.1]: https://www.contributor-covenant.org/version/2/1/code_of_conduct.html
|
||||
[Mozilla CoC]: https://github.com/mozilla/diversity
|
||||
[FAQ]: https://www.contributor-covenant.org/faq
|
||||
[translations]: https://www.contributor-covenant.org/translations
|
||||
+80
-9
@@ -14,7 +14,7 @@ We welcome contributions to:
|
||||
4. Submit a request for a new feature.
|
||||
5. Improve our test coverage.
|
||||
|
||||
### Contributing Features
|
||||
### Contributing Features ✨
|
||||
|
||||
Supervision is designed to provide generic utilities to solve problems. Thus, we focus on contributions that can have an impact on a wide range of projects.
|
||||
|
||||
@@ -24,7 +24,72 @@ 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. Click "fork" in the top corner of the `supervision` repository to get started:
|
||||
|
||||

|
||||
|
||||

|
||||
|
||||
Then, run `git clone` to download the project code to your computer.
|
||||
|
||||
Move to a new branch using the `git checkout` command:
|
||||
|
||||
```bash
|
||||
git checkout -b <your_branch_name>
|
||||
```
|
||||
|
||||
The name you choose for your branch should describe the change you want to make (i.e. `line-counter-docs`).
|
||||
|
||||
Make any changes you want to the project code, then run the following commands to commit your changes:
|
||||
|
||||
```bash
|
||||
git add .
|
||||
git commit -m "Your commit message"
|
||||
git push -u origin main
|
||||
```
|
||||
|
||||
## 🎨 Code quality
|
||||
|
||||
### Pre-commit tool
|
||||
|
||||
This project utilizes the [pre-commit](https://pre-commit.com/) tool to maintain code quality and consistency. Before submitting a pull request or making any commits, it is important to run the pre-commit tool to ensure that your changes meet the project's guidelines.
|
||||
|
||||
Furthermore, we have integrated a pre-commit GitHub Action into our workflow. This means that with every pull request opened, the pre-commit checks will be automatically enforced, streamlining the code review process and ensuring that all contributions adhere to our quality standards.
|
||||
|
||||
To run the pre-commit tool, follow these steps:
|
||||
|
||||
1. Install pre-commit by running the following command: `poetry install`. It will not only install pre-commit but also install all the deps and dev-deps of project
|
||||
|
||||
2. Once pre-commit is installed, navigate to the project's root directory.
|
||||
|
||||
3. Run the command `pre-commit run --all-files`. This will execute the pre-commit hooks configured for this project against the modified files. If any issues are found, the pre-commit tool will provide feedback on how to resolve them. Make the necessary changes and re-run the pre-commit command until all issues are resolved.
|
||||
|
||||
4. You can also install pre-commit as a git hook by execute `pre-commit install`. Every time you made `git commit` pre-commit run automatically for you.
|
||||
|
||||
### Docstrings
|
||||
|
||||
All new functions and classes in `supervision` should include docstrings. This is a prerequisite for any new functions and classes to be added to the library.
|
||||
|
||||
`supervision` adheres to the [Google Python docstring style](https://google.github.io/styleguide/pyguide.html#383-functions-and-methods). Please refer to the style guide while writing docstrings for your contribution.
|
||||
|
||||
### Type checking
|
||||
|
||||
So far, **there is no type checking with mypy**. See [issue](https://github.com/roboflow-ai/template-python/issues/4).
|
||||
|
||||
|
||||
Then, go back to your fork of the `supervision` repository, click "Pull Requests", and click "New Pull Request".
|
||||
|
||||

|
||||
|
||||
Make sure the `base` branch is `develop` before submitting your PR.
|
||||
|
||||
On the next page, review your changes then click "Create pull request":
|
||||
|
||||

|
||||
|
||||
Next, write a description for your pull request, and click "Create pull request" again to submit it for review:
|
||||
|
||||

|
||||
|
||||
When creating new functions, please ensure you have the following:
|
||||
|
||||
@@ -32,20 +97,26 @@ 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 a 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.
|
||||
|
||||
When you submit your Pull Request, you will be asked to sign a Contributor License Agreement (CLA) by the `cla-assistant` GitHub bot. We can only respond to PRs from contributors who have signed the project CLA.
|
||||
|
||||
All pull requests will be reviewed by the maintainers of the project. We will provide feedback and ask for changes if necessary.
|
||||
|
||||
PRs must pass all tests and linting requirements before they can be merged.
|
||||
|
||||
## 🧹 code quality
|
||||
## 📝 documentation
|
||||
|
||||
We provide two handy commands inside the `Makefile`, namely:
|
||||
The `supervision` documentation is stored in a folder called `docs`. The project documentation is built using `mkdocs`.
|
||||
|
||||
- `make style` to format the code
|
||||
- `make check_code_quality` to check code quality (PEP8 basically)
|
||||
To run the documentation, install the project requirements with `poetry install dev`. Then, run `mkdocs serve` to start the documentation server.
|
||||
|
||||
So far, **there is no types checking with mypy**. See [issue](https://github.com/roboflow-ai/template-python/issues/4).
|
||||
You can learn more about mkdocs on the [mkdocs website](https://www.mkdocs.org/).
|
||||
|
||||
## 🧪 tests
|
||||
## 🧪 tests
|
||||
|
||||
[`pytests`](https://docs.pytest.org/en/7.1.x/) is used to run our tests.
|
||||
[`pytests`](https://docs.pytest.org/en/7.1.x/) is used to run our tests.
|
||||
|
||||
## 📄 license
|
||||
|
||||
By contributing, you agree that your contributions will be licensed under an [MIT license](https://github.com/roboflow/supervision/blob/develop/LICENSE.md).
|
||||
|
||||
@@ -12,11 +12,11 @@ check_code_quality:
|
||||
isort --check-only --profile black $(check_dirs)
|
||||
# stop the build if there are Python syntax errors or undefined names
|
||||
flake8 $(check_dirs) --count --select=E9,F63,F7,F82 --show-source --statistics
|
||||
# exit-zero treats all errors as warnings. E203 for black, E501 for docstring, W503 for line breaks before logical operators
|
||||
# exit-zero treats all errors as warnings. E203 for black, E501 for docstring, W503 for line breaks before logical operators
|
||||
flake8 $(check_dirs) --count --max-line-length=88 --exit-zero --ignore=D --extend-ignore=E203,E501,W503 --statistics
|
||||
|
||||
|
||||
publish:
|
||||
python setup.py sdist bdist_wheel
|
||||
twine upload -r testpypi dist/* -u ${PYPI_USERNAME} -p ${PYPI_TEST_PASSWORD} --verbose
|
||||
poetry build
|
||||
twine upload -r testpypi dist/* -u ${PYPI_USERNAME} -p ${PYPI_TEST_PASSWORD} --verbose
|
||||
twine check dist/*
|
||||
twine upload dist/* -u ${PYPI_USERNAME} -p ${PYPI_PASSWORD} --verbose
|
||||
twine upload dist/* -u ${PYPI_USERNAME} -p ${PYPI_PASSWORD} --verbose
|
||||
|
||||
@@ -10,11 +10,15 @@
|
||||
|
||||
<br>
|
||||
|
||||
[](https://badge.fury.io/py/supervision)
|
||||
[](https://pypistats.org/packages/supervision)
|
||||
[](https://github.com/roboflow/supervision/blob/main/LICENSE.md)
|
||||
[](https://badge.fury.io/py/supervision)
|
||||
[](https://colab.research.google.com/github/roboflow-ai/notebooks/blob/main/notebooks/how-to-detect-and-count-objects-in-polygon-zone.ipynb)
|
||||
[notebooks](https://github.com/roboflow/notebooks) | [inference](https://github.com/roboflow/inference) | [autodistill](https://github.com/autodistill/autodistill) | [collect](https://github.com/roboflow/roboflow-collect)
|
||||
|
||||
<br>
|
||||
|
||||
[](https://badge.fury.io/py/supervision)
|
||||
[](https://pypistats.org/packages/supervision)
|
||||
[](https://github.com/roboflow/supervision/blob/main/LICENSE.md)
|
||||
[](https://badge.fury.io/py/supervision)
|
||||
[](https://colab.research.google.com/github/roboflow/supervision/blob/main/demo.ipynb)
|
||||
|
||||
</div>
|
||||
|
||||
@@ -25,29 +29,13 @@
|
||||
## 💻 install
|
||||
|
||||
Pip install the supervision package in a
|
||||
[**3.11>=Python>=3.7**](https://www.python.org/) environment.
|
||||
[**3.11>=Python>=3.8**](https://www.python.org/) environment.
|
||||
|
||||
```bash
|
||||
pip install supervision
|
||||
pip install supervision[desktop]
|
||||
```
|
||||
|
||||
<details close>
|
||||
<summary>👉 install from source</summary>
|
||||
|
||||
```bash
|
||||
# clone repository and navigate to root directory
|
||||
git clone https://github.com/roboflow/supervision.git
|
||||
cd supervision
|
||||
|
||||
# setup python environment and activate it
|
||||
python3 -m venv venv
|
||||
source venv/bin/activate
|
||||
|
||||
# install
|
||||
pip install -e ".[dev]"
|
||||
```
|
||||
|
||||
</details>
|
||||
Read more about desktop, headless, and local installation in our [guide](https://roboflow.github.io/supervision/).
|
||||
|
||||
## 🔥 quickstart
|
||||
|
||||
@@ -59,7 +47,7 @@ pip install -e ".[dev]"
|
||||
|
||||
>>> model = YOLO('yolov8s.pt')
|
||||
>>> result = model(IMAGE)[0]
|
||||
>>> detections = sv.Detections.from_yolov8(result)
|
||||
>>> detections = sv.Detections.from_ultralytics(result)
|
||||
|
||||
>>> len(detections)
|
||||
5
|
||||
@@ -67,9 +55,9 @@ pip install -e ".[dev]"
|
||||
|
||||
<details close>
|
||||
<summary>👉 more detections utils</summary>
|
||||
|
||||
- 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
|
||||
>>> from segment_anything import sam_model_registry, SamAutomaticMaskGenerator
|
||||
@@ -79,17 +67,17 @@ pip install -e ".[dev]"
|
||||
>>> sam_result = mask_generator.generate(IMAGE)
|
||||
>>> detections = sv.Detections.from_sam(sam_result=sam_result)
|
||||
```
|
||||
|
||||
|
||||
- [Advanced filtering](https://roboflow.github.io/supervision/quickstart/detections/)
|
||||
|
||||
|
||||
```python
|
||||
>>> detections = detections[detections.class_id == 0]
|
||||
>>> detections = detections[detections.confidence > 0.5]
|
||||
>>> detections = detections[detections.area > 1000]
|
||||
```
|
||||
|
||||
|
||||
- Image annotation
|
||||
|
||||
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
@@ -99,7 +87,7 @@ pip install -e ".[dev]"
|
||||
... detections=detections
|
||||
... )
|
||||
```
|
||||
|
||||
|
||||
</details>
|
||||
|
||||
### [datasets processing](https://roboflow.github.io/supervision/dataset/core/)
|
||||
@@ -123,7 +111,7 @@ pip install -e ".[dev]"
|
||||
<details close>
|
||||
<summary>👉 more dataset utils</summary>
|
||||
|
||||
- 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(
|
||||
@@ -136,8 +124,13 @@ pip install -e ".[dev]"
|
||||
... images_directory_path='...',
|
||||
... annotations_directory_path='...'
|
||||
... )
|
||||
|
||||
>>> dataset = sv.DetectionDataset.from_coco(
|
||||
... images_directory_path='...',
|
||||
... annotations_path='...'
|
||||
... )
|
||||
```
|
||||
|
||||
|
||||
- Loop over dataset entries
|
||||
|
||||
```python
|
||||
@@ -148,19 +141,41 @@ pip install -e ".[dev]"
|
||||
[155. , 497. , 404. , 833.5 ],
|
||||
[ 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)
|
||||
>>> test_dataset, valid_dataset = test_dataset.split(split_ratio=0.5)
|
||||
|
||||
|
||||
>>> len(train_dataset), len(test_dataset), len(valid_dataset)
|
||||
(700, 150, 150)
|
||||
```
|
||||
|
||||
- Save object detection / instance segmentation datasets in one of supported formats
|
||||
|
||||
|
||||
- Merge multiple datasets
|
||||
|
||||
```python
|
||||
>>> ds_1 = sv.DetectionDataset(...)
|
||||
>>> len(ds_1)
|
||||
100
|
||||
>>> ds_1.classes
|
||||
['dog', 'person']
|
||||
|
||||
>>> ds_2 = sv.DetectionDataset(...)
|
||||
>>> len(ds_2)
|
||||
200
|
||||
>>> ds_2.classes
|
||||
['cat']
|
||||
|
||||
>>> ds_merged = sv.DetectionDataset.merge([ds_1, ds_2])
|
||||
>>> len(ds_merged)
|
||||
300
|
||||
>>> ds_merged.classes
|
||||
['cat', 'dog', 'person']
|
||||
```
|
||||
|
||||
- Save object detection/instance segmentation datasets in one of the supported formats
|
||||
|
||||
```python
|
||||
>>> dataset.as_yolo(
|
||||
... images_directory_path='...',
|
||||
@@ -172,10 +187,15 @@ pip install -e ".[dev]"
|
||||
... images_directory_path='...',
|
||||
... annotations_directory_path='...'
|
||||
... )
|
||||
|
||||
>>> dataset.as_coco(
|
||||
... images_directory_path='...',
|
||||
... annotations_path='...'
|
||||
... )
|
||||
```
|
||||
|
||||
- Convert labels between suppoted formats
|
||||
|
||||
|
||||
- Convert labels between supported formats
|
||||
|
||||
```python
|
||||
>>> sv.DetectionDataset.from_yolo(
|
||||
... images_directory_path='...',
|
||||
@@ -187,27 +207,99 @@ pip install -e ".[dev]"
|
||||
... )
|
||||
```
|
||||
|
||||
- Load classification datasets in one of the supported formats
|
||||
|
||||
```python
|
||||
>>> cs = sv.ClassificationDataset.from_folder_structure(
|
||||
... root_directory_path='...'
|
||||
... )
|
||||
```
|
||||
|
||||
- Save classification datasets in one of the supported formats
|
||||
|
||||
```python
|
||||
>>> cs.as_folder_structure(
|
||||
... root_directory_path='...'
|
||||
... )
|
||||
```
|
||||
|
||||
</details>
|
||||
|
||||
### [model evaluation](https://roboflow.github.io/supervision/metrics/detection/)
|
||||
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> dataset = sv.DetectionDataset.from_yolo(...)
|
||||
|
||||
>>> def callback(image: np.ndarray) -> sv.Detections:
|
||||
... ...
|
||||
|
||||
>>> confusion_matrix = sv.ConfusionMatrix.benchmark(
|
||||
... dataset = dataset,
|
||||
... callback = callback
|
||||
... )
|
||||
|
||||
>>> confusion_matrix.matrix
|
||||
array([
|
||||
[0., 0., 0., 0.],
|
||||
[0., 1., 0., 1.],
|
||||
[0., 1., 1., 0.],
|
||||
[1., 1., 0., 0.]
|
||||
])
|
||||
```
|
||||
|
||||
<details close>
|
||||
<summary>👉 more metrics</summary>
|
||||
|
||||
- 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
|
||||
```
|
||||
|
||||
</details>
|
||||
|
||||
## 🎬 tutorials
|
||||
|
||||
<p align="left">
|
||||
<a href="https://youtu.be/oEQYStnF2l8" title="Accelerate Image Annotation with SAM and Grounding DINO"><img src="https://github.com/SkalskiP/SkalskiP/assets/26109316/ae1ca38e-40b7-4b35-8582-e8ea5de3806e" alt="Accelerate Image Annotation with SAM and Grounding DINO" width="300px" align="left" /></a>
|
||||
<a href="https://youtu.be/oEQYStnF2l8" title="Accelerate Image Annotation with SAM and Grounding DINO"><strong>Accelerate Image Annotation with SAM and Grounding DINO</strong></a>
|
||||
<div><strong>Created: 20 Apr 2023</strong> | <strong>Updated: 20 Apr 2023</strong></div>
|
||||
<br/> Discover how to speed up your image annotation process using Grounding DINO and Segment Anything Model (SAM). Learn how to convert object detection datasets into instance segmentation datasets, and see the potential of using these models to automatically annotate your datasets for real-time detectors like YOLOv8... </p>
|
||||
<a href="https://youtu.be/4Q3ut7vqD5o" title="Traffic Analysis with YOLOv8 and ByteTrack - Vehicle Detection and Tracking"><img src="https://github.com/roboflow/supervision/assets/26109316/54afdf1c-218c-4451-8f12-627fb85f1682" alt="Traffic Analysis with YOLOv8 and ByteTrack - Vehicle Detection and Tracking" width="300px" align="left" /></a>
|
||||
<a href="https://youtu.be/4Q3ut7vqD5o" title="Traffic Analysis with YOLOv8 and ByteTrack - Vehicle Detection and Tracking"><strong>Traffic Analysis with YOLOv8 and ByteTrack - Vehicle Detection and Tracking</strong></a>
|
||||
<div><strong>Created: 6 Sep 2023</strong> | <strong>Updated: 6 Sep 2023</strong></div>
|
||||
<br/> In this video, we explore real-time traffic analysis using YOLOv8 and ByteTrack to detect and track vehicles on aerial images. Harnessing the power of Python and Supervision, we delve deep into assigning cars to specific entry zones and understanding their direction of movement. By visualizing their paths, we gain insights into traffic flow across bustling roundabouts... </p>
|
||||
|
||||
<br/>
|
||||
<br/>
|
||||
|
||||
<p align="left">
|
||||
<a href="https://youtu.be/oEQYStnF2l8" title="SAM - Segment Anything Model by Meta AI: Complete Guide"><img src="https://github.com/SkalskiP/SkalskiP/assets/26109316/6913ff11-53c6-4341-8d90-eaff3023c3fd" alt="SAM - Segment Anything Model by Meta AI: Complete Guide" width="300px" align="left" /></a>
|
||||
<a href="https://youtu.be/oEQYStnF2l8" title="SAM - Segment Anything Model by Meta AI: Complete Guide"><strong>SAM - Segment Anything Model by Meta AI: Complete Guide</strong></a>
|
||||
<a href="https://youtu.be/D-D6ZmadzPE" title="SAM - Segment Anything Model by Meta AI: Complete Guide"><img src="https://github.com/SkalskiP/SkalskiP/assets/26109316/6913ff11-53c6-4341-8d90-eaff3023c3fd" alt="SAM - Segment Anything Model by Meta AI: Complete Guide" width="300px" align="left" /></a>
|
||||
<a href="https://youtu.be/D-D6ZmadzPE" title="SAM - Segment Anything Model by Meta AI: Complete Guide"><strong>SAM - Segment Anything Model by Meta AI: Complete Guide</strong></a>
|
||||
<div><strong>Created: 11 Apr 2023</strong> | <strong>Updated: 11 Apr 2023</strong></div>
|
||||
<br/> Discover the incredible potential of Meta AI's Segment Anything Model (SAM)! We dive into SAM, an efficient and promptable model for image segmentation, which has revolutionized computer vision tasks. With over 1 billion masks on 11M licensed and privacy-respecting images, SAM's zero-shot performance is often competitive with or even superior to prior fully supervised results... </p>
|
||||
|
||||
## 💜 built with supervision
|
||||
|
||||
Did you build something cool using supervision? [Let us know!](https://github.com/roboflow/supervision/discussions/categories/built-with-supervision)
|
||||
|
||||
https://user-images.githubusercontent.com/26109316/207858600-ee862b22-0353-440b-ad85-caa0c4777904.mp4
|
||||
|
||||
https://github.com/roboflow/supervision/assets/26109316/c9436828-9fbf-4c25-ae8c-60e9c81b3900
|
||||
|
||||
## 📚 documentation
|
||||
|
||||
Curious how Supervision can help you solve problems on your project? Visit our [documentation](https://roboflow.github.io/supervision) page!
|
||||
Visit our [documentation](https://roboflow.github.io/supervision) page to learn how supervision can help you build computer vision applications faster and more reliably.
|
||||
|
||||
## 🏆 contribution
|
||||
|
||||
|
||||
Vendored
+1325
File diff suppressed because one or more lines are too long
@@ -0,0 +1 @@
|
||||
supervision.roboflow.com
|
||||
@@ -0,0 +1,199 @@
|
||||
=== "BoundingBox"
|
||||
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> image = ...
|
||||
>>> detections = sv.Detections(...)
|
||||
|
||||
>>> bounding_box_annotator = sv.BoundingBoxAnnotator()
|
||||
>>> annotated_frame = bounding_box_annotator.annotate(
|
||||
... scene=image.copy(),
|
||||
... detections=detections
|
||||
... )
|
||||
```
|
||||
|
||||
<div class="result" markdown>
|
||||
|
||||
{ align=center width="800" }
|
||||
|
||||
</div>
|
||||
|
||||
=== "Mask"
|
||||
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> image = ...
|
||||
>>> detections = sv.Detections(...)
|
||||
|
||||
>>> mask_annotator = sv.MaskAnnotator()
|
||||
>>> annotated_frame = mask_annotator.annotate(
|
||||
... scene=image.copy(),
|
||||
... detections=detections
|
||||
... )
|
||||
```
|
||||
|
||||
<div class="result" markdown>
|
||||
|
||||
{ align=center width="800" }
|
||||
|
||||
</div>
|
||||
|
||||
=== "Ellipse"
|
||||
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> image = ...
|
||||
>>> detections = sv.Detections(...)
|
||||
|
||||
>>> ellipse_annotator = sv.EllipseAnnotator()
|
||||
>>> annotated_frame = ellipse_annotator.annotate(
|
||||
... scene=image.copy(),
|
||||
... detections=detections
|
||||
... )
|
||||
```
|
||||
|
||||
<div class="result" markdown>
|
||||
|
||||
{ align=center width="800" }
|
||||
|
||||
</div>
|
||||
|
||||
=== "BoxCorner"
|
||||
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> image = ...
|
||||
>>> detections = sv.Detections(...)
|
||||
|
||||
>>> corner_annotator = sv.BoxCornerAnnotator()
|
||||
>>> annotated_frame = corner_annotator.annotate(
|
||||
... scene=image.copy(),
|
||||
... detections=detections
|
||||
... )
|
||||
```
|
||||
|
||||
<div class="result" markdown>
|
||||
|
||||
{ align=center width="800" }
|
||||
|
||||
</div>
|
||||
|
||||
=== "Circle"
|
||||
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> image = ...
|
||||
>>> detections = sv.Detections(...)
|
||||
|
||||
>>> circle_annotator = sv.CircleAnnotator()
|
||||
>>> annotated_frame = circle_annotator.annotate(
|
||||
... scene=image.copy(),
|
||||
... detections=detections
|
||||
... )
|
||||
```
|
||||
|
||||
<div class="result" markdown>
|
||||
|
||||
{ align=center width="800" }
|
||||
|
||||
</div>
|
||||
|
||||
=== "Label"
|
||||
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> image = ...
|
||||
>>> detections = sv.Detections(...)
|
||||
|
||||
>>> label_annotator = sv.LabelAnnotator(text_position=sv.Position.CENTER)
|
||||
>>> annotated_frame = label_annotator.annotate(
|
||||
... scene=image.copy(),
|
||||
... detections=detections
|
||||
... )
|
||||
```
|
||||
|
||||
<div class="result" markdown>
|
||||
|
||||
{ align=center width="800" }
|
||||
|
||||
</div>
|
||||
|
||||
=== "Blur"
|
||||
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> image = ...
|
||||
>>> detections = sv.Detections(...)
|
||||
|
||||
>>> blur_annotator = sv.BlurAnnotator()
|
||||
>>> annotated_frame = blur_annotator.annotate(
|
||||
... scene=image.copy(),
|
||||
... detections=detections
|
||||
... )
|
||||
```
|
||||
|
||||
<div class="result" markdown>
|
||||
|
||||
{ align=center width="800" }
|
||||
|
||||
</div>
|
||||
|
||||
=== "Trace"
|
||||
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> image = ...
|
||||
>>> detections = sv.Detections(...)
|
||||
|
||||
>>> trace_annotator = sv.TraceAnnotator()
|
||||
>>> annotated_frame = trace_annotator.annotate(
|
||||
... scene=image.copy(),
|
||||
... detections=detections
|
||||
... )
|
||||
```
|
||||
|
||||
<div class="result" markdown>
|
||||
|
||||
{ align=center width="800" }
|
||||
|
||||
</div>
|
||||
|
||||
## BoundingBoxAnnotator
|
||||
|
||||
:::supervision.annotators.core.BoundingBoxAnnotator
|
||||
|
||||
## MaskAnnotator
|
||||
|
||||
:::supervision.annotators.core.MaskAnnotator
|
||||
|
||||
## EllipseAnnotator
|
||||
|
||||
:::supervision.annotators.core.EllipseAnnotator
|
||||
|
||||
## BoxCornerAnnotator
|
||||
|
||||
:::supervision.annotators.core.BoxCornerAnnotator
|
||||
|
||||
## CircleAnnotator
|
||||
|
||||
:::supervision.annotators.core.CircleAnnotator
|
||||
|
||||
## LabelAnnotator
|
||||
|
||||
:::supervision.annotators.core.LabelAnnotator
|
||||
|
||||
## TraceAnnotator
|
||||
|
||||
:::supervision.annotators.core.TraceAnnotator
|
||||
|
||||
## BlurAnnotator
|
||||
|
||||
:::supervision.annotators.core.BlurAnnotator
|
||||
+268
-74
@@ -1,6 +1,233 @@
|
||||
### 0.15.0 <small>October 5, 2023</small>
|
||||
|
||||
- Added [#170](https://github.com/roboflow/supervision/pull/170): [`sv.BoundingBoxAnnotator`](https://supervision.roboflow.com/annotators/#supervision.annotators.core.BoundingBoxAnnotator) allowing to annotate images and videos with bounding boxes.
|
||||
|
||||
- Added [#170](https://github.com/roboflow/supervision/pull/170): [`sv.BoxCornerAnnotator `](https://supervision.roboflow.com/annotators/#supervision.annotators.core.BoxCornerAnnotator) allowing to annotate images and videos with just bounding box corners.
|
||||
|
||||
- Added [#170](https://github.com/roboflow/supervision/pull/170): [`sv.MaskAnnotator`](https://supervision.roboflow.com/annotators/#supervision.annotators.core.MaskAnnotator) allowing to annotate images and videos with segmentation masks.
|
||||
|
||||
- Added [#170](https://github.com/roboflow/supervision/pull/170): [`sv.EllipseAnnotator`](https://supervision.roboflow.com/annotators/#supervision.annotators.core.EllipseAnnotator) allowing to annotate images and videos with ellipses (sports game style).
|
||||
|
||||
- Added [#386](https://github.com/roboflow/supervision/pull/386): [`sv.CircleAnnotator`](https://supervision.roboflow.com/annotators/#supervision.annotators.core.CircleAnnotator) allowing to annotate images and videos with circles.
|
||||
|
||||
- Added [#354](https://github.com/roboflow/supervision/pull/354): [`sv.TraceAnnotator`](https://supervision.roboflow.com/annotators/#supervision.annotators.core.TraceAnnotator) allowing to draw path of moving objects on videos.
|
||||
|
||||
- Added [#405](https://github.com/roboflow/supervision/pull/405): [`sv.BlurAnnotator`](https://supervision.roboflow.com/annotators/#supervision.annotators.core.BlurAnnotator) allowing to blur objects on images and videos.
|
||||
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> image = ...
|
||||
>>> detections = sv.Detections(...)
|
||||
|
||||
>>> bounding_box_annotator = sv.BoundingBoxAnnotator()
|
||||
>>> annotated_frame = bounding_box_annotator.annotate(
|
||||
... scene=image.copy(),
|
||||
... detections=detections
|
||||
... )
|
||||
```
|
||||
|
||||
- Added [#354](https://github.com/roboflow/supervision/pull/354): Supervision usage [example](https://github.com/roboflow/supervision/tree/develop/examples/traffic_analysis). You can now learn how to perform traffic flow analysis with Supervision.
|
||||
|
||||
- Changed [#399](https://github.com/roboflow/supervision/pull/399): [`sv.Detections.from_roboflow`](https://supervision.roboflow.com/detection/core/#supervision.detection.core.Detections.from_roboflow) now does not require `class_list` to be specified. The `class_id` value can be extracted directly from the [inference](https://github.com/roboflow/inference) response.
|
||||
|
||||
- Changed [#381](https://github.com/roboflow/supervision/pull/381): [`sv.VideoSink`](https://supervision.roboflow.com/utils/video/#videosink) now allows to customize the output codec.
|
||||
|
||||
- Changed [#361](https://github.com/roboflow/supervision/pull/361): [`sv.InferenceSlicer`](https://supervision.roboflow.com/detection/tools/inference_slicer/#supervision.detection.tools.inference_slicer.InferenceSlicer) can now operate in multithreading mode.
|
||||
|
||||
- Fixed [#348](https://github.com/roboflow/supervision/pull/348): [`sv.Detections.from_deepsparse`](https://supervision.roboflow.com/detection/core/#supervision.detection.core.Detections.from_deepsparse) to allow processing empty [deepsparse](https://github.com/neuralmagic/deepsparse) result object.
|
||||
|
||||
### 0.14.0 <small>August 31, 2023</small>
|
||||
|
||||
- Added [#282](https://github.com/roboflow/supervision/pull/282): support for SAHI inference technique with [`sv.InferenceSlicer`](https://supervision.roboflow.com/detection/tools/inference_slicer).
|
||||
|
||||
```python
|
||||
>>> import cv2
|
||||
>>> import supervision as sv
|
||||
>>> from ultralytics import YOLO
|
||||
|
||||
>>> image = cv2.imread(SOURCE_IMAGE_PATH)
|
||||
>>> model = YOLO(...)
|
||||
|
||||
>>> def callback(image_slice: np.ndarray) -> sv.Detections:
|
||||
... result = model(image_slice)[0]
|
||||
... return sv.Detections.from_ultralytics(result)
|
||||
|
||||
>>> slicer = sv.InferenceSlicer(callback = callback)
|
||||
|
||||
>>> detections = slicer(image)
|
||||
```
|
||||
|
||||
- Added [#297](https://github.com/roboflow/supervision/pull/297): [`Detections.from_deepsparse`](https://roboflow.github.io/supervision/detection/core/#supervision.detection.core.Detections.from_deepsparse) to enable seamless integration with [DeepSparse](https://github.com/neuralmagic/deepsparse) framework.
|
||||
|
||||
- Added [#281](https://github.com/roboflow/supervision/pull/281): [`sv.Classifications.from_ultralytics`](https://supervision.roboflow.com/classification/core/#supervision.classification.core.Classifications.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) and [sv.Classifications.from_yolov8](https://supervision.roboflow.com/classification/core/#supervision.classification.core.Classifications.from_yolov8) are now deprecated and will be removed with supervision-0.16.0 release.
|
||||
|
||||
- Added [#341](https://github.com/roboflow/supervision/pull/341): First supervision usage example script showing how to detect and track objects on video using YOLOv8 + Supervision.
|
||||
|
||||
- Changed [#296](https://github.com/roboflow/supervision/pull/296): [`sv.ClassificationDataset`](https://supervision.roboflow.com/dataset/core/#supervision.dataset.core.ClassificationDataset) and [`sv.DetectionDataset`](https://supervision.roboflow.com/dataset/core/#supervision.dataset.core.DetectionDataset) now use image path (not image name) as dataset keys.
|
||||
|
||||
- Fixed [#300](https://github.com/roboflow/supervision/pull/300): [`Detections.from_roboflow`](https://supervision.roboflow.com/detection/core/#supervision.detection.core.Detections.from_roboflow) to filter out polygons with less than 3 points.
|
||||
|
||||
### 0.13.0 <small>August 8, 2023</small>
|
||||
|
||||
- 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 <small>July 24, 2023</small>
|
||||
|
||||
!!! warning
|
||||
|
||||
With the `supervision-0.12.0` release, we are terminating official support for Python 3.7.
|
||||
|
||||
- Added [#177](https://github.com/roboflow/supervision/pull/177): initial support for object detection model benchmarking with [`sv.ConfusionMatrix`](https://roboflow.github.io/supervision/metrics/detection/#confusionmatrix).
|
||||
|
||||
```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)
|
||||
|
||||
>>> confusion_matrix = sv.ConfusionMatrix.benchmark(
|
||||
... dataset = dataset,
|
||||
... callback = callback
|
||||
... )
|
||||
|
||||
>>> confusion_matrix.matrix
|
||||
array([
|
||||
[0., 0., 0., 0.],
|
||||
[0., 1., 0., 1.],
|
||||
[0., 1., 1., 0.],
|
||||
[1., 1., 0., 0.]
|
||||
])
|
||||
```
|
||||
|
||||
- Added [#173](https://github.com/roboflow/supervision/pull/173): [`Detections.from_mmdetection`](https://roboflow.github.io/supervision/detection/core/#supervision.detection.core.Detections.from_mmdetection) to enable seamless integration with [MMDetection](https://github.com/open-mmlab/mmdetection) framework.
|
||||
|
||||
- Added [#130](https://github.com/roboflow/supervision/issues/130): ability to [install](https://roboflow.github.io/supervision/) package in `headless` or `desktop` mode.
|
||||
|
||||
- Changed [#180](https://github.com/roboflow/supervision/pull/180): packing method from `setup.py` to `pyproject.toml`.
|
||||
|
||||
- Fixed [#188](https://github.com/roboflow/supervision/issues/188): [`sv.DetectionDataset.from_cooc`](https://roboflow.github.io/supervision/dataset/core/#supervision.dataset.core.DetectionDataset.from_coco) can't be loaded when there are images without annotations.
|
||||
|
||||
- Fixed [#226](https://github.com/roboflow/supervision/issues/226): [`sv.DetectionDataset.from_yolo`](https://roboflow.github.io/supervision/dataset/core/#supervision.dataset.core.DetectionDataset.from_yolo) can't load background instances.
|
||||
|
||||
### 0.11.1 <small>June 29, 2023</small>
|
||||
|
||||
- Fix [#165](https://github.com/roboflow/supervision/pull/165): [`as_folder_structure`](https://roboflow.github.io/supervision/dataset/core/#supervision.dataset.core.ClassificationDataset.as_folder_structure) fails to save [`sv.ClassificationDataset`](https://roboflow.github.io/supervision/dataset/core/#classificationdataset) when it is result of inference.
|
||||
|
||||
### 0.11.0 <small>June 28, 2023</small>
|
||||
|
||||
- Added [#150](https://github.com/roboflow/supervision/pull/150): ability to load and save [`sv.DetectionDataset`](https://roboflow.github.io/supervision/dataset/core/#detectiondataset) in COCO format using [`as_coco`](https://roboflow.github.io/supervision/dataset/core/#supervision.dataset.core.DetectionDataset.as_coco) and [`from_coco`](https://roboflow.github.io/supervision/dataset/core/#supervision.dataset.core.DetectionDataset.from_coco) methods.
|
||||
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> ds = sv.DetectionDataset.from_coco(
|
||||
... images_directory_path='...',
|
||||
... annotations_path='...'
|
||||
... )
|
||||
|
||||
>>> ds.as_coco(
|
||||
... images_directory_path='...',
|
||||
... annotations_path='...'
|
||||
... )
|
||||
```
|
||||
|
||||
- Added [#158](https://github.com/roboflow/supervision/pull/158): ability to marge multiple [`sv.DetectionDataset`](https://roboflow.github.io/supervision/dataset/core/#detectiondataset) together using [`merge`](https://roboflow.github.io/supervision/dataset/core/#supervision.dataset.core.DetectionDataset.merge) method.
|
||||
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> ds_1 = sv.DetectionDataset(...)
|
||||
>>> len(ds_1)
|
||||
100
|
||||
>>> ds_1.classes
|
||||
['dog', 'person']
|
||||
|
||||
>>> ds_2 = sv.DetectionDataset(...)
|
||||
>>> len(ds_2)
|
||||
200
|
||||
>>> ds_2.classes
|
||||
['cat']
|
||||
|
||||
>>> ds_merged = sv.DetectionDataset.merge([ds_1, ds_2])
|
||||
>>> len(ds_merged)
|
||||
300
|
||||
>>> ds_merged.classes
|
||||
['cat', 'dog', 'person']
|
||||
```
|
||||
|
||||
- Added [#162](https://github.com/roboflow/supervision/pull/162): additional `start` and `end` arguments to [`sv.get_video_frames_generator`](https://roboflow.github.io/supervision/utils/video/#get_video_frames_generator) allowing to generate frames only for a selected part of the video.
|
||||
|
||||
- Fix [#157](https://github.com/roboflow/supervision/pull/157): incorrect loading of YOLO dataset class names from `data.yaml`.
|
||||
|
||||
### 0.10.0 <small>June 14, 2023</small>
|
||||
|
||||
- Added [#125](https://github.com/roboflow/supervision/pull/125): ability to load and save [`sv.ClassificationDataset`](https://roboflow.github.io/supervision/dataset/core/#classificationdataset) in a folder structure format.
|
||||
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> cs = sv.ClassificationDataset.from_folder_structure(
|
||||
... root_directory_path='...'
|
||||
... )
|
||||
|
||||
>>> cs.as_folder_structure(
|
||||
... root_directory_path='...'
|
||||
... )
|
||||
```
|
||||
|
||||
- Added [#125](https://github.com/roboflow/supervision/pull/125): support for [`sv.ClassificationDataset.split`](https://roboflow.github.io/supervision/dataset/core/#supervision.dataset.core.ClassificationDataset.split) allowing to divide `sv.ClassificationDataset` into two parts.
|
||||
|
||||
- Added [#110](https://github.com/roboflow/supervision/pull/110): ability to extract masks from Roboflow API results using [`sv.Detections.from_roboflow`](https://roboflow.github.io/supervision/detection/core/#supervision.detection.core.Detections.from_roboflow).
|
||||
|
||||
- Added [commit hash](https://github.com/roboflow/supervision/commit/d000292eb2f2342544e0947b65528082e60fb8d6): Supervision Quickstart [notebook](https://colab.research.google.com/github/roboflow/supervision/blob/main/demo.ipynb) where you can learn more about Detection, Dataset and Video APIs.
|
||||
|
||||
- Changed [#135](https://github.com/roboflow/supervision/pull/135): `sv.get_video_frames_generator` documentation to better describe actual behavior.
|
||||
|
||||
### 0.9.0 <small>June 7, 2023</small>
|
||||
|
||||
- Added [[#118](https://github.com/roboflow/supervision/pull/118)]: ability to select [`sv.Detections`](https://roboflow.github.io/supervision/detection/core/#supervision.detection.core.Detections.__getitem__) by index, list of indexes or slice. Here is an example illustrating the new selection methods.
|
||||
- Added [#118](https://github.com/roboflow/supervision/pull/118): ability to select [`sv.Detections`](https://roboflow.github.io/supervision/detection/core/#supervision.detection.core.Detections.__getitem__) by index, list of indexes or slice. Here is an example illustrating the new selection methods.
|
||||
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
@@ -14,44 +241,11 @@
|
||||
2
|
||||
```
|
||||
|
||||
- Added [[#101](https://github.com/roboflow/supervision/pull/101)]: ability to extract masks from YOLOv8 result using [`sv.Detections.from_yolov8`](https://roboflow.github.io/supervision/detection/core/#supervision.detection.core.Detections.from_yolov8). Here is an example illustrating how to extract boolean masks from the result of the YOLOv8 model inference.
|
||||
- Added [#101](https://github.com/roboflow/supervision/pull/101): ability to extract masks from YOLOv8 result using [`sv.Detections.from_yolov8`](https://roboflow.github.io/supervision/detection/core/#supervision.detection.core.Detections.from_yolov8). Here is an example illustrating how to extract boolean masks from the result of the YOLOv8 model inference.
|
||||
|
||||
```python
|
||||
>>> import cv2
|
||||
>>> from ultralytics import YOLO
|
||||
>>> import supervision as sv
|
||||
- Added [#122](https://github.com/roboflow/supervision/pull/122): ability to crop image using [`sv.crop`](https://roboflow.github.io/supervision/utils/image/#crop). Here is an example showing how to get a separate crop for each detection in `sv.Detections`.
|
||||
|
||||
>>> image = cv2.imread(...)
|
||||
>>> image.shape
|
||||
(640, 640, 3)
|
||||
|
||||
>>> model = YOLO('yolov8s-seg.pt')
|
||||
>>> result = model(image)[0]
|
||||
>>> detections = sv.Detections.from_yolov8(result)
|
||||
>>> detections.mask.shape
|
||||
(2, 640, 640)
|
||||
```
|
||||
|
||||
- Added [[#122](https://github.com/roboflow/supervision/pull/122)]: ability to crop image using [`sv.crop`](https://roboflow.github.io/supervision/utils/image/#crop). Here is an example showing how to get a separate crop for each detection in `sv.Detections`.
|
||||
|
||||
```python
|
||||
>>> import cv2
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> image = cv2.imread(...)
|
||||
>>> detections = sv.Detections(...)
|
||||
>>> len(detections)
|
||||
2
|
||||
>>> crops = [
|
||||
... sv.crop(image=image, xyxy=xyxy)
|
||||
... for xyxy
|
||||
... in detections.xyxy
|
||||
... ]
|
||||
>>> len(crops)
|
||||
2
|
||||
```
|
||||
|
||||
- Added [[#120](https://github.com/roboflow/supervision/pull/120)]: ability to conveniently save multiple images into directory using [`sv.ImageSink`](https://roboflow.github.io/supervision/utils/image/#imagesink). Here is an example showing how to save every tenth video frame as a separate image.
|
||||
- Added [#120](https://github.com/roboflow/supervision/pull/120): ability to conveniently save multiple images into directory using [`sv.ImageSink`](https://roboflow.github.io/supervision/utils/image/#imagesink). Here is an example showing how to save every tenth video frame as a separate image.
|
||||
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
@@ -61,12 +255,12 @@
|
||||
... sink.save_image(image=image)
|
||||
```
|
||||
|
||||
- Fixed [[#106](https://github.com/roboflow/supervision/issues/106)]: inconvenient handling of [`sv.PolygonZone`](https://roboflow.github.io/supervision/detection/tools/polygon_zone/#polygonzone) coordinates. Now `sv.PolygonZone` accepts coordinates in the form of `[[x1, y1], [x2, y2], ...]` that can be both integers and floats.
|
||||
- Fixed [#106](https://github.com/roboflow/supervision/issues/106): inconvenient handling of [`sv.PolygonZone`](https://roboflow.github.io/supervision/detection/tools/polygon_zone/#polygonzone) coordinates. Now `sv.PolygonZone` accepts coordinates in the form of `[[x1, y1], [x2, y2], ...]` that can be both integers and floats.
|
||||
|
||||
### 0.8.0 <small>May 17, 2023</small>
|
||||
|
||||
- Added [[#100](https://github.com/roboflow/supervision/pull/100)]: support for dataset inheritance. The current `Dataset` got renamed to `DetectionDataset`. Now [`DetectionDataset`](https://roboflow.github.io/supervision/dataset/core/#detectiondataset) inherits from `BaseDataset`. This change was made to enforce the future consistency of APIs of different types of computer vision datasets.
|
||||
- Added [[#100](https://github.com/roboflow/supervision/pull/100)]: ability to save datasets in YOLO format using [`DetectionDataset.as_yolo`](https://roboflow.github.io/supervision/dataset/core/#supervision.dataset.core.DetectionDataset.as_yolo).
|
||||
- Added [#100](https://github.com/roboflow/supervision/pull/100): support for dataset inheritance. The current `Dataset` got renamed to `DetectionDataset`. Now [`DetectionDataset`](https://roboflow.github.io/supervision/dataset/core/#detectiondataset) inherits from `BaseDataset`. This change was made to enforce the future consistency of APIs of different types of computer vision datasets.
|
||||
- Added [#100](https://github.com/roboflow/supervision/pull/100): ability to save datasets in YOLO format using [`DetectionDataset.as_yolo`](https://roboflow.github.io/supervision/dataset/core/#supervision.dataset.core.DetectionDataset.as_yolo).
|
||||
|
||||
```python
|
||||
>>> import roboflow
|
||||
@@ -90,7 +284,7 @@
|
||||
['dog', 'person']
|
||||
```
|
||||
|
||||
- Added [[#102](https://github.com/roboflow/supervision/pull/103)]: support for [`DetectionDataset.split`](https://roboflow.github.io/supervision/dataset/core/#supervision.dataset.core.DetectionDataset.split) allowing to divide `DetectionDataset` into two parts.
|
||||
- Added [#102](https://github.com/roboflow/supervision/pull/103): support for [`DetectionDataset.split`](https://roboflow.github.io/supervision/dataset/core/#supervision.dataset.core.DetectionDataset.split) allowing to divide `DetectionDataset` into two parts.
|
||||
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
@@ -102,27 +296,27 @@
|
||||
(700, 300)
|
||||
```
|
||||
|
||||
- Changed [[#100](https://github.com/roboflow/supervision/pull/100)]: default value of `approximation_percentage` parameter from `0.75` to `0.0` in `DetectionDataset.as_yolo` and `DetectionDataset.as_pascal_voc`.
|
||||
- Changed [#100](https://github.com/roboflow/supervision/pull/100): default value of `approximation_percentage` parameter from `0.75` to `0.0` in `DetectionDataset.as_yolo` and `DetectionDataset.as_pascal_voc`.
|
||||
|
||||
### 0.7.0 <small>May 11, 2023</small>
|
||||
|
||||
- Added [[#91](https://github.com/roboflow/supervision/pull/91)]: `Detections.from_yolo_nas` to enable seamless integration with [YOLO-NAS](https://github.com/Deci-AI/super-gradients/blob/master/YOLONAS.md) model.
|
||||
- Added [[#86](https://github.com/roboflow/supervision/pull/86)]: ability to load datasets in YOLO format using `Dataset.from_yolo`.
|
||||
- Added [[#84](https://github.com/roboflow/supervision/pull/84)]: `Detections.merge` to merge multiple `Detections` objects together.
|
||||
- Fixed [[#81](https://github.com/roboflow/supervision/pull/81)]: `LineZoneAnnotator.annotate` does not return annotated frame.
|
||||
- Changed [[#44](https://github.com/roboflow/supervision/pull/44)]: `LineZoneAnnotator.annotate` to allow for custom text for the in and out tags.
|
||||
- Added [#91](https://github.com/roboflow/supervision/pull/91): `Detections.from_yolo_nas` to enable seamless integration with [YOLO-NAS](https://github.com/Deci-AI/super-gradients/blob/master/YOLONAS.md) model.
|
||||
- Added [#86](https://github.com/roboflow/supervision/pull/86): ability to load datasets in YOLO format using `Dataset.from_yolo`.
|
||||
- Added [#84](https://github.com/roboflow/supervision/pull/84): `Detections.merge` to merge multiple `Detections` objects together.
|
||||
- Fixed [#81](https://github.com/roboflow/supervision/pull/81): `LineZoneAnnotator.annotate` does not return annotated frame.
|
||||
- Changed [#44](https://github.com/roboflow/supervision/pull/44): `LineZoneAnnotator.annotate` to allow for custom text for the in and out tags.
|
||||
|
||||
### 0.6.0 <small>April 19, 2023</small>
|
||||
|
||||
- Added [[#71](https://github.com/roboflow/supervision/pull/71)]: initial `Dataset` support and ability to save `Detections` in Pascal VOC XML format.
|
||||
- Added [[#71](https://github.com/roboflow/supervision/pull/71)]: new `mask_to_polygons`, `filter_polygons_by_area`, `polygon_to_xyxy` and `approximate_polygon` utilities.
|
||||
- Added [[#72](https://github.com/roboflow/supervision/pull/72)]: ability to load Pascal VOC XML **object detections** dataset as `Dataset`.
|
||||
- Changed [[#70](https://github.com/roboflow/supervision/pull/70)]: order of `Detections` attributes to make it consistent with order of objects in `__iter__` tuple.
|
||||
- Changed [[#71](https://github.com/roboflow/supervision/pull/71)]: `generate_2d_mask` to `polygon_to_mask`.
|
||||
- Added [#71](https://github.com/roboflow/supervision/pull/71): initial `Dataset` support and ability to save `Detections` in Pascal VOC XML format.
|
||||
- Added [#71](https://github.com/roboflow/supervision/pull/71): new `mask_to_polygons`, `filter_polygons_by_area`, `polygon_to_xyxy` and `approximate_polygon` utilities.
|
||||
- Added [#72](https://github.com/roboflow/supervision/pull/72): ability to load Pascal VOC XML **object detections** dataset as `Dataset`.
|
||||
- Changed [#70](https://github.com/roboflow/supervision/pull/70): order of `Detections` attributes to make it consistent with order of objects in `__iter__` tuple.
|
||||
- Changed [#71](https://github.com/roboflow/supervision/pull/71): `generate_2d_mask` to `polygon_to_mask`.
|
||||
|
||||
### 0.5.2 <small>April 13, 2023</small>
|
||||
|
||||
- Fixed [[#63](https://github.com/roboflow/supervision/pull/63)]: `LineZone.trigger` function expects 4 values instead of 5.
|
||||
- Fixed [#63](https://github.com/roboflow/supervision/pull/63): `LineZone.trigger` function expects 4 values instead of 5.
|
||||
|
||||
### 0.5.1 <small>April 12, 2023</small>
|
||||
|
||||
@@ -131,41 +325,41 @@
|
||||
|
||||
### 0.5.0 <small>April 10, 2023</small>
|
||||
|
||||
- Added [[#58](https://github.com/roboflow/supervision/pull/58)]: `Detections.mask` to enable segmentation support.
|
||||
- Added [[#58](https://github.com/roboflow/supervision/pull/58)]: `MaskAnnotator` to allow easy `Detections.mask` annotation.
|
||||
- Added [[#58](https://github.com/roboflow/supervision/pull/58)]: `Detections.from_sam` to enable native Segment Anything Model (SAM) support.
|
||||
- Changed [[#58](https://github.com/roboflow/supervision/pull/58)]: `Detections.area` behaviour to work not only with boxes but also with masks.
|
||||
- Added [#58](https://github.com/roboflow/supervision/pull/58): `Detections.mask` to enable segmentation support.
|
||||
- Added [#58](https://github.com/roboflow/supervision/pull/58): `MaskAnnotator` to allow easy `Detections.mask` annotation.
|
||||
- Added [#58](https://github.com/roboflow/supervision/pull/58): `Detections.from_sam` to enable native Segment Anything Model (SAM) support.
|
||||
- Changed [#58](https://github.com/roboflow/supervision/pull/58): `Detections.area` behaviour to work not only with boxes but also with masks.
|
||||
|
||||
### 0.4.0 <small>April 5, 2023</small>
|
||||
### 0.4.0 <small>April 5, 2023</small>
|
||||
|
||||
- Added [[#46](https://github.com/roboflow/supervision/discussions/48)]: `Detections.empty` to allow easy creation of empty `Detections` objects.
|
||||
- Added [[#56](https://github.com/roboflow/supervision/pull/56)]: `Detections.from_roboflow` to allow easy creation of `Detections` objects from Roboflow API inference results.
|
||||
- Added [[#56](https://github.com/roboflow/supervision/pull/56)]: `plot_images_grid` to allow easy plotting of multiple images on single plot.
|
||||
- Added [[#56](https://github.com/roboflow/supervision/pull/56)]: initial support for Pascal VOC XML format with `detections_to_voc_xml` method.
|
||||
- Changed [[#56](https://github.com/roboflow/supervision/pull/56)]: `show_frame_in_notebook` refactored and renamed to `plot_image`.
|
||||
- Added [#46](https://github.com/roboflow/supervision/discussions/48): `Detections.empty` to allow easy creation of empty `Detections` objects.
|
||||
- Added [#56](https://github.com/roboflow/supervision/pull/56): `Detections.from_roboflow` to allow easy creation of `Detections` objects from Roboflow API inference results.
|
||||
- Added [#56](https://github.com/roboflow/supervision/pull/56): `plot_images_grid` to allow easy plotting of multiple images on single plot.
|
||||
- Added [#56](https://github.com/roboflow/supervision/pull/56): initial support for Pascal VOC XML format with `detections_to_voc_xml` method.
|
||||
- Changed [#56](https://github.com/roboflow/supervision/pull/56): `show_frame_in_notebook` refactored and renamed to `plot_image`.
|
||||
|
||||
### 0.3.2 <small>March 23, 2023</small>
|
||||
### 0.3.2 <small>March 23, 2023</small>
|
||||
|
||||
- Changed [[#50](https://github.com/roboflow/supervision/issues/50)]: Allow `Detections.class_id` to be `None`.
|
||||
- Changed [#50](https://github.com/roboflow/supervision/issues/50): Allow `Detections.class_id` to be `None`.
|
||||
|
||||
### 0.3.1 <small>March 6, 2023</small>
|
||||
### 0.3.1 <small>March 6, 2023</small>
|
||||
|
||||
- Fixed [[#41](https://github.com/roboflow/supervision/issues/41)]: `PolygonZone` throws an exception when the object touches the bottom edge of the image.
|
||||
- Fixed [[#42](https://github.com/roboflow/supervision/issues/42)]: `Detections.wth_nms` method throws an exception when `Detections` is empty.
|
||||
- Changed [[#36](https://github.com/roboflow/supervision/pull/36)]: `Detections.wth_nms` support class agnostic and non-class agnostic case.
|
||||
- Fixed [#41](https://github.com/roboflow/supervision/issues/41): `PolygonZone` throws an exception when the object touches the bottom edge of the image.
|
||||
- Fixed [#42](https://github.com/roboflow/supervision/issues/42): `Detections.wth_nms` method throws an exception when `Detections` is empty.
|
||||
- Changed [#36](https://github.com/roboflow/supervision/pull/36): `Detections.wth_nms` support class agnostic and non-class agnostic case.
|
||||
|
||||
### 0.3.0 <small>March 6, 2023</small>
|
||||
### 0.3.0 <small>March 6, 2023</small>
|
||||
|
||||
- Changed: Allow `Detections.confidence` to be `None`.
|
||||
- Added: `Detections.from_transformers` and `Detections.from_detectron2` to enable seamless integration with Transformers and Detectron2 models.
|
||||
- Added: `Detections.from_transformers` and `Detections.from_detectron2` to enable seamless integration with Transformers and Detectron2 models.
|
||||
- Added: `Detections.area` to dynamically calculate bounding box area.
|
||||
- Added: `Detections.wth_nms` to filter out double detections with NMS. Initial - only class agnostic - implementation.
|
||||
- Added: `Detections.wth_nms` to filter out double detections with NMS. Initial - only class agnostic - implementation.
|
||||
|
||||
### 0.2.0 <small>February 2, 2023</small>
|
||||
### 0.2.0 <small>February 2, 2023</small>
|
||||
|
||||
- Added: Advanced `Detections` filtering with pandas-like API.
|
||||
- Added: `Detections.from_yolov5` and `Detections.from_yolov8` to enable seamless integration with YOLOv5 and YOLOv8 models.
|
||||
|
||||
### 0.1.0 <small>January 19, 2023</small>
|
||||
### 0.1.0 <small>January 19, 2023</small>
|
||||
|
||||
Say hello to Supervision 👋
|
||||
|
||||
@@ -0,0 +1,3 @@
|
||||
## Classifications
|
||||
|
||||
:::supervision.classification.core.Classifications
|
||||
@@ -1,8 +0,0 @@
|
||||
!!! warning
|
||||
|
||||
Dataset API is still fluid and may change. If you use Dataset API in your project until further notice, freeze the
|
||||
`supervision` version in your `requirements.txt` or `setup.py`.
|
||||
|
||||
## DetectionDataset
|
||||
|
||||
:::supervision.dataset.core.DetectionDataset
|
||||
@@ -0,0 +1,12 @@
|
||||
!!! warning
|
||||
|
||||
Dataset API is still fluid and may change. If you use Dataset API in your project until further notice, freeze the
|
||||
`supervision` version in your `requirements.txt` or `setup.py`.
|
||||
|
||||
## DetectionDataset
|
||||
|
||||
:::supervision.dataset.core.DetectionDataset
|
||||
|
||||
## ClassificationDataset
|
||||
|
||||
:::supervision.dataset.core.ClassificationDataset
|
||||
@@ -1,7 +0,0 @@
|
||||
## BoxAnnotator
|
||||
|
||||
:::supervision.detection.annotate.BoxAnnotator
|
||||
|
||||
## MaskAnnotator
|
||||
|
||||
:::supervision.detection.annotate.MaskAnnotator
|
||||
@@ -0,0 +1,3 @@
|
||||
## InferenceSlicer
|
||||
|
||||
:::supervision.detection.tools.inference_slicer.InferenceSlicer
|
||||
@@ -4,4 +4,4 @@
|
||||
|
||||
## PolygonZoneAnnotator
|
||||
|
||||
:::supervision.detection.tools.polygon_zone.PolygonZoneAnnotator
|
||||
:::supervision.detection.tools.polygon_zone.PolygonZoneAnnotator
|
||||
|
||||
@@ -24,4 +24,4 @@
|
||||
|
||||
## filter_polygons_by_area
|
||||
|
||||
:::supervision.detection.utils.filter_polygons_by_area
|
||||
:::supervision.detection.utils.filter_polygons_by_area
|
||||
|
||||
+1
-1
@@ -16,4 +16,4 @@
|
||||
|
||||
## draw_text
|
||||
|
||||
:::supervision.draw.utils.draw_text
|
||||
:::supervision.draw.utils.draw_text
|
||||
|
||||
@@ -0,0 +1,3 @@
|
||||
## Position
|
||||
|
||||
:::supervision.geometry.core.Position
|
||||
+54
-18
@@ -9,29 +9,65 @@
|
||||
</p>
|
||||
</div>
|
||||
|
||||
## 👋 Welcome
|
||||
## 👋 Hello
|
||||
|
||||
Supervision is a set of easy-to-use utilities that will come in handy in any computer vision project.
|
||||
We write your reusable computer vision tools. Whether you need to load your dataset from your hard drive, draw detections on an image or video, or count how many detections are in a zone. You can count on us!
|
||||
|
||||
**Supervision** is still in
|
||||
pre-release stage 🚧 Keep your eyes open for potential bugs and be aware that at this stage our API is still fluid and may change.
|
||||
## 💻 Install
|
||||
|
||||
## 💻 How to Install
|
||||
You can install `supervision` with pip in a
|
||||
[**3.11>=Python>=3.8**](https://www.python.org/) environment.
|
||||
|
||||
You can install `supervision` with pip in a
|
||||
[**3.11>=Python>=3.7**](https://www.python.org/) environment.
|
||||
!!! example "pip install (recommended)"
|
||||
|
||||
!!! example "Pip install method (recommended)"
|
||||
=== "headless"
|
||||
The headless installation of `supervision` is designed for environments where graphical user interfaces (GUI) are not needed, making it more lightweight and suitable for server-side applications.
|
||||
|
||||
```bash
|
||||
pip install supervision
|
||||
```
|
||||
```bash
|
||||
pip install supervision
|
||||
```
|
||||
|
||||
!!! example "Git clone method (for development)"
|
||||
=== "desktop"
|
||||
If you require the full version of `supervision` with GUI support you can install the desktop version. This version includes the GUI components of OpenCV, allowing you to display images and videos on the screen.
|
||||
|
||||
```bash
|
||||
git https://github.com/roboflow/supervision.git
|
||||
cd supervision
|
||||
pip install -e '.[dev]'
|
||||
```
|
||||
See contributing section to know more about contributing to the project
|
||||
```bash
|
||||
pip install supervision[desktop]
|
||||
```
|
||||
|
||||
!!! example "git clone (for development)"
|
||||
|
||||
=== "virtualenv"
|
||||
|
||||
```bash
|
||||
# clone repository and navigate to root directory
|
||||
git clone https://github.com/roboflow/supervision.git
|
||||
cd supervision
|
||||
|
||||
# setup python environment and activate it
|
||||
python3 -m venv venv
|
||||
source venv/bin/activate
|
||||
|
||||
# headless install
|
||||
pip install -e "."
|
||||
|
||||
# desktop install
|
||||
pip install -e ".[desktop]"
|
||||
```
|
||||
|
||||
=== "poetry"
|
||||
|
||||
```bash
|
||||
# clone repository and navigate to root directory
|
||||
git clone https://github.com/roboflow/supervision.git
|
||||
cd supervision
|
||||
|
||||
# setup python environment and activate it
|
||||
poetry env use python 3.10
|
||||
poetry shell
|
||||
|
||||
# headless install
|
||||
poetry install
|
||||
|
||||
# desktop install
|
||||
poetry install --extras "desktop"
|
||||
```
|
||||
|
||||
@@ -0,0 +1,12 @@
|
||||
!!! warning
|
||||
|
||||
Evaluation API is still fluid and may change. If you use Evaluation API in your project until further notice, freeze the
|
||||
`supervision` version in your `requirements.txt` or `setup.py`.
|
||||
|
||||
## ConfusionMatrix
|
||||
|
||||
:::supervision.metrics.detection.ConfusionMatrix
|
||||
|
||||
## MeanAveragePrecision
|
||||
|
||||
:::supervision.metrics.detection.MeanAveragePrecision
|
||||
@@ -1,10 +1,10 @@
|
||||
## advanced filtering
|
||||
|
||||
|
||||
The advanced filtering capabilities of the `Detections` class offer users a versatile and efficient way to narrow down
|
||||
and refine object detections. This section outlines various filtering methods, including filtering by specific class
|
||||
or a set of classes, confidence, object area, bounding box area, relative area, box dimensions, and designated zones.
|
||||
Each method is demonstrated with concise code examples to provide users with a clear understanding of how to implement
|
||||
The advanced filtering capabilities of the `Detections` class offer users a versatile and efficient way to narrow down
|
||||
and refine object detections. This section outlines various filtering methods, including filtering by specific class
|
||||
or a set of classes, confidence, object area, bounding box area, relative area, box dimensions, and designated zones.
|
||||
Each method is demonstrated with concise code examples to provide users with a clear understanding of how to implement
|
||||
the filters in their applications.
|
||||
|
||||
### by specific class
|
||||
@@ -44,14 +44,14 @@ Allows you to select detections that belong only to one selected class.
|
||||
|
||||
### by set of classes
|
||||
|
||||
Allows you to select detections that belong only to selected set of classes.
|
||||
Allows you to select detections that belong only to selected set of classes.
|
||||
|
||||
=== "After"
|
||||
|
||||
```python
|
||||
import numpy as np
|
||||
import supervision as sv
|
||||
|
||||
|
||||
selected_classes = [0, 2, 3]
|
||||
detections = sv.Detections(...)
|
||||
detections = detections[np.isin(detections.class_id, selected_classes)]
|
||||
@@ -68,7 +68,7 @@ Allows you to select detections that belong only to selected set of classes.
|
||||
```python
|
||||
import numpy as np
|
||||
import supervision as sv
|
||||
|
||||
|
||||
class_id = [0, 2, 3]
|
||||
detections = sv.Detections(...)
|
||||
detections = detections[np.isin(detections.class_id, class_id)]
|
||||
@@ -116,8 +116,8 @@ Allows you to select detections with specific confidence value, for example high
|
||||
|
||||
### by area
|
||||
|
||||
Allows you to select detections based on their size. We define the area as the number of pixels occupied by the
|
||||
detection in the image. In the example below, we have sifted out the detections that are too small.
|
||||
Allows you to select detections based on their size. We define the area as the number of pixels occupied by the
|
||||
detection in the image. In the example below, we have sifted out the detections that are too small.
|
||||
|
||||
=== "After"
|
||||
|
||||
@@ -151,9 +151,9 @@ detection in the image. In the example below, we have sifted out the detections
|
||||
|
||||
### by relative area
|
||||
|
||||
Allows you to select detections based on their size in relation to the size of whole image. Sometimes the concept of
|
||||
detection size changes depending on the image. Detection occupying 10000 square px can be large on a 1280x720 image
|
||||
but small on a 3840x2160 image. In such cases, we can filter out detections based on the percentage of the image area
|
||||
Allows you to select detections based on their size in relation to the size of whole image. Sometimes the concept of
|
||||
detection size changes depending on the image. Detection occupying 10000 square px can be large on a 1280x720 image
|
||||
but small on a 3840x2160 image. In such cases, we can filter out detections based on the percentage of the image area
|
||||
occupied by them. In the example below, we remove too large detections.
|
||||
|
||||
=== "After"
|
||||
@@ -164,7 +164,7 @@ occupied by them. In the example below, we remove too large detections.
|
||||
image = ...
|
||||
height, width, channels = image.shape
|
||||
image_area = height * width
|
||||
|
||||
|
||||
detections = sv.Detections(...)
|
||||
detections = detections[(detections.area / image_area) < 0.8]
|
||||
```
|
||||
@@ -183,7 +183,7 @@ occupied by them. In the example below, we remove too large detections.
|
||||
image = ...
|
||||
height, width, channels = image.shape
|
||||
image_area = height * width
|
||||
|
||||
|
||||
detections = sv.Detections(...)
|
||||
detections = detections[(detections.area / image_area) < 0.8]
|
||||
```
|
||||
@@ -196,8 +196,8 @@ occupied by them. In the example below, we remove too large detections.
|
||||
|
||||
### by box dimensions
|
||||
|
||||
Allows you to select detections based on their dimensions. The size of the bounding box, as well as its coordinates,
|
||||
can be criteria for rejecting detection. Implementing such filtering requires a bit of custom code but is relatively
|
||||
Allows you to select detections based on their dimensions. The size of the bounding box, as well as its coordinates,
|
||||
can be criteria for rejecting detection. Implementing such filtering requires a bit of custom code but is relatively
|
||||
simple and fast.
|
||||
|
||||
=== "After"
|
||||
@@ -236,7 +236,7 @@ simple and fast.
|
||||
|
||||
### by `PolygonZone`
|
||||
|
||||
Allows you to use `Detections` in combination with `PolygonZone` to weed out bounding boxes that are in and out of the
|
||||
Allows you to use `Detections` in combination with `PolygonZone` to weed out bounding boxes that are in and out of the
|
||||
zone. In the example below you can see how to filter out all detections located in the lower part of the image.
|
||||
|
||||
=== "After"
|
||||
@@ -309,4 +309,4 @@ zone. In the example below you can see how to filter out all detections located
|
||||
|
||||
{ align=center width="800" }
|
||||
|
||||
</div>
|
||||
</div>
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
:root {
|
||||
--md-primary-fg-color: #8315F9;
|
||||
--md-accent-fg-color: #00FFCE;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,3 @@
|
||||
## ByteTrack
|
||||
|
||||
:::supervision.tracker.byte_tracker.core.ByteTrack
|
||||
+1
-1
@@ -4,4 +4,4 @@
|
||||
|
||||
## crop
|
||||
|
||||
:::supervision.utils.image.crop
|
||||
:::supervision.utils.image.crop_image
|
||||
|
||||
@@ -4,4 +4,4 @@
|
||||
|
||||
## plot_images_grid
|
||||
|
||||
:::supervision.utils.notebook.plot_images_grid
|
||||
:::supervision.utils.notebook.plot_images_grid
|
||||
|
||||
+1
-1
@@ -12,4 +12,4 @@
|
||||
|
||||
## process_video
|
||||
|
||||
:::supervision.utils.video.process_video
|
||||
:::supervision.utils.video.process_video
|
||||
|
||||
@@ -0,0 +1,20 @@
|
||||
## 👋 hello
|
||||
|
||||
This script provides functionality for processing videos using YOLOv8 for object
|
||||
detection and Supervision for tracking and annotation.
|
||||
|
||||
## 💻 install
|
||||
|
||||
```bash
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
## ⚙️ parameters
|
||||
|
||||
| parameter | required | description |
|
||||
|:-------------------------|:--------:|:----------------------------------------------------------------------------------|
|
||||
| `--source_weights_path` | ✓ | Path to the source weights file for YOLOv8. |
|
||||
| `--source_video_path` | ✓ | Path to the source video file to be processed. |
|
||||
| `--target_video_path` | ✓ | Path to the target video file (output). |
|
||||
| `--confidence_threshold` | ✗ | Confidence threshold for YOLO model detection. Default is 0.3. |
|
||||
| `--iou_threshold` | ✗ | IOU (Intersection over Union) threshold for YOLO model detection. Default is 0.7. |
|
||||
@@ -0,0 +1,3 @@
|
||||
supervision
|
||||
tqdm
|
||||
ultralytics
|
||||
@@ -0,0 +1,83 @@
|
||||
import argparse
|
||||
|
||||
from tqdm import tqdm
|
||||
from ultralytics import YOLO
|
||||
|
||||
import supervision as sv
|
||||
|
||||
|
||||
def process_video(
|
||||
source_weights_path: str,
|
||||
source_video_path: str,
|
||||
target_video_path: str,
|
||||
confidence_threshold: float = 0.3,
|
||||
iou_threshold: float = 0.7,
|
||||
) -> None:
|
||||
model = YOLO(source_weights_path)
|
||||
|
||||
tracker = sv.ByteTrack()
|
||||
box_annotator = sv.BoundingBoxAnnotator()
|
||||
label_annotator = sv.LabelAnnotator(classes=model.names)
|
||||
frame_generator = sv.get_video_frames_generator(source_path=source_video_path)
|
||||
video_info = sv.VideoInfo.from_video_path(video_path=source_video_path)
|
||||
|
||||
with sv.VideoSink(target_path=target_video_path, video_info=video_info) as sink:
|
||||
for frame in tqdm(frame_generator, total=video_info.total_frames):
|
||||
results = model(
|
||||
frame, verbose=False, conf=confidence_threshold, iou=iou_threshold
|
||||
)[0]
|
||||
detections = sv.Detections.from_ultralytics(results)
|
||||
detections = tracker.update_with_detections(detections)
|
||||
|
||||
annotated_frame = box_annotator.annotate(
|
||||
scene=frame.copy(), detections=detections
|
||||
)
|
||||
|
||||
annotated_labeled_frame = label_annotator.annotate(
|
||||
scene=annotated_frame, detections=detections
|
||||
)
|
||||
|
||||
sink.write_frame(frame=annotated_labeled_frame)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Video Processing with YOLO and ByteTrack"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--source_weights_path",
|
||||
required=True,
|
||||
help="Path to the source weights file",
|
||||
type=str,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--source_video_path",
|
||||
required=True,
|
||||
help="Path to the source video file",
|
||||
type=str,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--target_video_path",
|
||||
required=True,
|
||||
help="Path to the target video file (output)",
|
||||
type=str,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--confidence_threshold",
|
||||
default=0.3,
|
||||
help="Confidence threshold for the model",
|
||||
type=float,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--iou_threshold", default=0.7, help="IOU threshold for the model", type=float
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
process_video(
|
||||
source_weights_path=args.source_weights_path,
|
||||
source_video_path=args.source_video_path,
|
||||
target_video_path=args.target_video_path,
|
||||
confidence_threshold=args.confidence_threshold,
|
||||
iou_threshold=args.iou_threshold,
|
||||
)
|
||||
@@ -0,0 +1 @@
|
||||
data/
|
||||
@@ -0,0 +1,44 @@
|
||||
## 👋 hello
|
||||
|
||||
This script performs traffic flow analysis using YOLOv8, an object-detection method and ByteTrack, a simple yet effective online multi-object tracking method. It uses the supervision package for multiple tasks such as tracking, annotations, etc.
|
||||
|
||||
https://github.com/roboflow/supervision/assets/26109316/c9436828-9fbf-4c25-ae8c-60e9c81b3900
|
||||
|
||||
## 💻 install
|
||||
|
||||
- clone repository and navigate to example directory
|
||||
|
||||
```bash
|
||||
git clone https://github.com/roboflow/supervision.git
|
||||
cd supervision/examples/traffic_analysis
|
||||
```
|
||||
|
||||
- setup python environment and activate it [optional]
|
||||
|
||||
```bash
|
||||
python3 -m venv venv
|
||||
source venv/bin/activate
|
||||
```
|
||||
|
||||
- install required dependencies
|
||||
|
||||
```bash
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
- download `traffic_analysis.pt` and `traffic_analysis.mov` files
|
||||
|
||||
```bash
|
||||
./setup.sh
|
||||
```
|
||||
|
||||
## ⚙️ run
|
||||
|
||||
```bash
|
||||
python script.py \
|
||||
--source_weights_path data/traffic_analysis.pt \
|
||||
--source_video_path data/traffic_analysis.mov \
|
||||
--confidence_threshold 0.3 \
|
||||
--iou_threshold 0.5 \
|
||||
--target_video_path data/traffic_analysis_result.mov
|
||||
```
|
||||
@@ -0,0 +1,4 @@
|
||||
supervision>=0.15.0rc1
|
||||
tqdm
|
||||
ultralytics
|
||||
gdown
|
||||
@@ -0,0 +1,219 @@
|
||||
import argparse
|
||||
from typing import Dict, List, Set, Tuple
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
from tqdm import tqdm
|
||||
from ultralytics import YOLO
|
||||
|
||||
import supervision as sv
|
||||
|
||||
COLORS = sv.ColorPalette.default()
|
||||
|
||||
ZONE_IN_POLYGONS = [
|
||||
np.array([[592, 282], [900, 282], [900, 82], [592, 82]]),
|
||||
np.array([[950, 860], [1250, 860], [1250, 1060], [950, 1060]]),
|
||||
np.array([[592, 582], [592, 860], [392, 860], [392, 582]]),
|
||||
np.array([[1250, 282], [1250, 530], [1450, 530], [1450, 282]]),
|
||||
]
|
||||
|
||||
ZONE_OUT_POLYGONS = [
|
||||
np.array([[950, 282], [1250, 282], [1250, 82], [950, 82]]),
|
||||
np.array([[592, 860], [900, 860], [900, 1060], [592, 1060]]),
|
||||
np.array([[592, 282], [592, 550], [392, 550], [392, 282]]),
|
||||
np.array([[1250, 860], [1250, 560], [1450, 560], [1450, 860]]),
|
||||
]
|
||||
|
||||
|
||||
class DetectionsManager:
|
||||
def __init__(self) -> None:
|
||||
self.tracker_id_to_zone_id: Dict[int, int] = {}
|
||||
self.counts: Dict[int, Dict[int, Set[int]]] = {}
|
||||
|
||||
def update(
|
||||
self,
|
||||
detections_all: sv.Detections,
|
||||
detections_in_zones: List[sv.Detections],
|
||||
detections_out_zones: List[sv.Detections],
|
||||
) -> sv.Detections:
|
||||
for zone_in_id, detections_in_zone in enumerate(detections_in_zones):
|
||||
for tracker_id in detections_in_zone.tracker_id:
|
||||
self.tracker_id_to_zone_id.setdefault(tracker_id, zone_in_id)
|
||||
|
||||
for zone_out_id, detections_out_zone in enumerate(detections_out_zones):
|
||||
for tracker_id in detections_out_zone.tracker_id:
|
||||
if tracker_id in self.tracker_id_to_zone_id:
|
||||
zone_in_id = self.tracker_id_to_zone_id[tracker_id]
|
||||
self.counts.setdefault(zone_out_id, {})
|
||||
self.counts[zone_out_id].setdefault(zone_in_id, set())
|
||||
self.counts[zone_out_id][zone_in_id].add(tracker_id)
|
||||
|
||||
detections_all.class_id = np.vectorize(
|
||||
lambda x: self.tracker_id_to_zone_id.get(x, -1)
|
||||
)(detections_all.tracker_id)
|
||||
return detections_all[detections_all.class_id != -1]
|
||||
|
||||
|
||||
def initiate_polygon_zones(
|
||||
polygons: List[np.ndarray],
|
||||
frame_resolution_wh: Tuple[int, int],
|
||||
triggering_position: sv.Position = sv.Position.CENTER,
|
||||
) -> List[sv.PolygonZone]:
|
||||
return [
|
||||
sv.PolygonZone(
|
||||
polygon=polygon,
|
||||
frame_resolution_wh=frame_resolution_wh,
|
||||
triggering_position=triggering_position,
|
||||
)
|
||||
for polygon in polygons
|
||||
]
|
||||
|
||||
|
||||
class VideoProcessor:
|
||||
def __init__(
|
||||
self,
|
||||
source_weights_path: str,
|
||||
source_video_path: str,
|
||||
target_video_path: str = None,
|
||||
confidence_threshold: float = 0.3,
|
||||
iou_threshold: float = 0.7,
|
||||
) -> None:
|
||||
self.conf_threshold = confidence_threshold
|
||||
self.iou_threshold = iou_threshold
|
||||
self.source_video_path = source_video_path
|
||||
self.target_video_path = target_video_path
|
||||
|
||||
self.model = YOLO(source_weights_path)
|
||||
self.tracker = sv.ByteTrack()
|
||||
|
||||
self.video_info = sv.VideoInfo.from_video_path(source_video_path)
|
||||
self.zones_in = initiate_polygon_zones(
|
||||
ZONE_IN_POLYGONS, self.video_info.resolution_wh, sv.Position.CENTER
|
||||
)
|
||||
self.zones_out = initiate_polygon_zones(
|
||||
ZONE_OUT_POLYGONS, self.video_info.resolution_wh, sv.Position.CENTER
|
||||
)
|
||||
|
||||
self.box_annotator = sv.BoxAnnotator(color=COLORS)
|
||||
self.trace_annotator = sv.TraceAnnotator(
|
||||
color=COLORS, position=sv.Position.CENTER, trace_length=100, thickness=2
|
||||
)
|
||||
self.detections_manager = DetectionsManager()
|
||||
|
||||
def process_video(self):
|
||||
frame_generator = sv.get_video_frames_generator(
|
||||
source_path=self.source_video_path
|
||||
)
|
||||
|
||||
if self.target_video_path:
|
||||
with sv.VideoSink(self.target_video_path, self.video_info) as sink:
|
||||
for frame in tqdm(frame_generator, total=self.video_info.total_frames):
|
||||
annotated_frame = self.process_frame(frame)
|
||||
sink.write_frame(annotated_frame)
|
||||
else:
|
||||
for frame in tqdm(frame_generator, total=self.video_info.total_frames):
|
||||
annotated_frame = self.process_frame(frame)
|
||||
cv2.imshow("Processed Video", annotated_frame)
|
||||
if cv2.waitKey(1) & 0xFF == ord("q"):
|
||||
break
|
||||
cv2.destroyAllWindows()
|
||||
|
||||
def annotate_frame(
|
||||
self, frame: np.ndarray, detections: sv.Detections
|
||||
) -> np.ndarray:
|
||||
annotated_frame = frame.copy()
|
||||
for i, (zone_in, zone_out) in enumerate(zip(self.zones_in, self.zones_out)):
|
||||
annotated_frame = sv.draw_polygon(
|
||||
annotated_frame, zone_in.polygon, COLORS.colors[i]
|
||||
)
|
||||
annotated_frame = sv.draw_polygon(
|
||||
annotated_frame, zone_out.polygon, COLORS.colors[i]
|
||||
)
|
||||
|
||||
labels = [f"#{tracker_id}" for tracker_id in detections.tracker_id]
|
||||
annotated_frame = self.trace_annotator.annotate(annotated_frame, detections)
|
||||
annotated_frame = self.box_annotator.annotate(
|
||||
annotated_frame, detections, labels
|
||||
)
|
||||
|
||||
for zone_out_id, zone_out in enumerate(self.zones_out):
|
||||
zone_center = sv.get_polygon_center(polygon=zone_out.polygon)
|
||||
if zone_out_id in self.detections_manager.counts:
|
||||
counts = self.detections_manager.counts[zone_out_id]
|
||||
for i, zone_in_id in enumerate(counts):
|
||||
count = len(self.detections_manager.counts[zone_out_id][zone_in_id])
|
||||
text_anchor = sv.Point(x=zone_center.x, y=zone_center.y + 40 * i)
|
||||
annotated_frame = sv.draw_text(
|
||||
scene=annotated_frame,
|
||||
text=str(count),
|
||||
text_anchor=text_anchor,
|
||||
background_color=COLORS.colors[zone_in_id],
|
||||
)
|
||||
|
||||
return annotated_frame
|
||||
|
||||
def process_frame(self, frame: np.ndarray) -> np.ndarray:
|
||||
results = self.model(
|
||||
frame, verbose=False, conf=self.conf_threshold, iou=self.iou_threshold
|
||||
)[0]
|
||||
detections = sv.Detections.from_ultralytics(results)
|
||||
detections.class_id = np.zeros(len(detections))
|
||||
detections = self.tracker.update_with_detections(detections)
|
||||
|
||||
detections_in_zones = []
|
||||
detections_out_zones = []
|
||||
|
||||
for i, (zone_in, zone_out) in enumerate(zip(self.zones_in, self.zones_out)):
|
||||
detections_in_zone = detections[zone_in.trigger(detections=detections)]
|
||||
detections_in_zones.append(detections_in_zone)
|
||||
detections_out_zone = detections[zone_out.trigger(detections=detections)]
|
||||
detections_out_zones.append(detections_out_zone)
|
||||
|
||||
detections = self.detections_manager.update(
|
||||
detections, detections_in_zones, detections_out_zones
|
||||
)
|
||||
return self.annotate_frame(frame, detections)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Traffic Flow Analysis with YOLO and ByteTrack"
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--source_weights_path",
|
||||
required=True,
|
||||
help="Path to the source weights file",
|
||||
type=str,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--source_video_path",
|
||||
required=True,
|
||||
help="Path to the source video file",
|
||||
type=str,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--target_video_path",
|
||||
default=None,
|
||||
help="Path to the target video file (output)",
|
||||
type=str,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--confidence_threshold",
|
||||
default=0.3,
|
||||
help="Confidence threshold for the model",
|
||||
type=float,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--iou_threshold", default=0.7, help="IOU threshold for the model", type=float
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
processor = VideoProcessor(
|
||||
source_weights_path=args.source_weights_path,
|
||||
source_video_path=args.source_video_path,
|
||||
target_video_path=args.target_video_path,
|
||||
confidence_threshold=args.confidence_threshold,
|
||||
iou_threshold=args.iou_threshold,
|
||||
)
|
||||
processor.process_video()
|
||||
Executable
+17
@@ -0,0 +1,17 @@
|
||||
#!/bin/bash
|
||||
|
||||
# Get the directory where the script is located
|
||||
DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"
|
||||
|
||||
# Check if 'data' directory does not exist and then create it
|
||||
if [[ ! -e $DIR/data ]]; then
|
||||
mkdir "$DIR/data"
|
||||
else
|
||||
echo "'data' directory already exists."
|
||||
fi
|
||||
|
||||
# Download the traffic_analysis.mov file from Google Drive
|
||||
gdown -O "$DIR/data/traffic_analysis.mov" "https://drive.google.com/uc?id=1qadBd7lgpediafCpL_yedGjQPk-FLK-W"
|
||||
|
||||
# Download the traffic_analysis.pt file from Google Drive
|
||||
gdown -O "$DIR/data/traffic_analysis.pt" "https://drive.google.com/uc?id=1y-IfToCjRXa3ZdC1JpnKRopC7mcQW-5z"
|
||||
+22
-17
@@ -28,22 +28,27 @@ nav:
|
||||
- Home: index.md
|
||||
- Quickstart:
|
||||
- Detections: quickstart/detections.md
|
||||
- API reference:
|
||||
- Detections:
|
||||
- Core: detection/core.md
|
||||
- Annotate: detection/annotate.md
|
||||
- Utils: detection/utils.md
|
||||
- Tools:
|
||||
- Polygon Zone: detection/tools/polygon_zone.md
|
||||
- Dataset:
|
||||
- Core: dataset/core.md
|
||||
- Draw:
|
||||
- Utils: draw/utils.md
|
||||
- Utils:
|
||||
- Video: utils/video.md
|
||||
- Image: utils/image.md
|
||||
- Notebook: utils/notebook.md
|
||||
- File: utils/file.md
|
||||
- Classifications:
|
||||
- Core: classification/core.md
|
||||
- Detections:
|
||||
- Core: detection/core.md
|
||||
- Annotate: detection/annotate.md
|
||||
- Utils: detection/utils.md
|
||||
- Tools:
|
||||
- Polygon Zone: detection/tools/polygon_zone.md
|
||||
- Inference Slicer: detection/tools/inference_slicer.md
|
||||
- Annotators: annotators.md
|
||||
- Trackers: trackers.md
|
||||
- Datasets: datasets.md
|
||||
- Metrics:
|
||||
- Object Detection: metrics/detection.md
|
||||
- Draw:
|
||||
- Utils: draw/utils.md
|
||||
- Utils:
|
||||
- Video: utils/video.md
|
||||
- Image: utils/image.md
|
||||
- Notebook: utils/notebook.md
|
||||
- File: utils/file.md
|
||||
- Changelog: changelog.md
|
||||
|
||||
theme:
|
||||
@@ -81,4 +86,4 @@ markdown_extensions:
|
||||
- pymdownx.tabbed:
|
||||
alternate_style: true
|
||||
- toc:
|
||||
permalink: true
|
||||
permalink: true
|
||||
|
||||
Generated
+3523
File diff suppressed because it is too large
Load Diff
+183
@@ -0,0 +1,183 @@
|
||||
[tool.poetry]
|
||||
name = "supervision"
|
||||
version = "0.15.0"
|
||||
description = "A set of easy-to-use utils that will come in handy in any Computer Vision project"
|
||||
authors = ["Piotr Skalski <piotr.skalski92@gmail.com>"]
|
||||
maintainers = ["Piotr Skalski <piotr.skalski92@gmail.com>"]
|
||||
readme = "README.md"
|
||||
license = "MIT"
|
||||
packages = [{include = "supervision"}]
|
||||
homepage = "https://github.com/roboflow/supervision"
|
||||
repository = "https://github.com/roboflow/supervision"
|
||||
documentation = "https://github.com/roboflow/supervision/blob/main/README.md"
|
||||
keywords = ["machine-learning", "deep-learning", "vision", "ML", "DL", "AI", "YOLOv5", "YOLOv8", "Roboflow"]
|
||||
|
||||
classifiers=[
|
||||
'Development Status :: 4 - Beta',
|
||||
'Intended Audience :: Developers',
|
||||
'Intended Audience :: Education',
|
||||
'Intended Audience :: Science/Research',
|
||||
'License :: OSI Approved :: MIT License',
|
||||
'Programming Language :: Python :: 3',
|
||||
'Programming Language :: Python :: 3.8',
|
||||
'Programming Language :: Python :: 3.9',
|
||||
'Programming Language :: Python :: 3.10',
|
||||
'Programming Language :: Python :: 3.11',
|
||||
'Programming Language :: Python :: 3 :: Only',
|
||||
'Topic :: Software Development',
|
||||
'Topic :: Scientific/Engineering',
|
||||
"Topic :: Scientific/Engineering :: Artificial Intelligence",
|
||||
"Topic :: Scientific/Engineering :: Image Recognition",
|
||||
'Typing :: Typed',
|
||||
'Operating System :: Microsoft :: Windows',
|
||||
'Operating System :: POSIX :: Linux',
|
||||
'Operating System :: MacOS'
|
||||
]
|
||||
|
||||
[tool.poetry.dependencies]
|
||||
python = ">=3.8,<3.12.0"
|
||||
numpy = "^1.20.0"
|
||||
matplotlib = "^3.7.1"
|
||||
pyyaml = "^6.0"
|
||||
pillow = ">=9.4,<11.0"
|
||||
opencv-python = { version = "^4.8.0.74", optional = true }
|
||||
opencv-python-headless = "^4.8.0.74"
|
||||
scipy = "^1.9.0"
|
||||
|
||||
|
||||
[tool.poetry.extras]
|
||||
desktop = ["opencv-python"]
|
||||
|
||||
|
||||
[tool.poetry.group.dev.dependencies]
|
||||
twine = "^4.0.2"
|
||||
pytest = "^7.2.2"
|
||||
wheel = "^0.40.0"
|
||||
notebook = "^6.5.3"
|
||||
build = "^0.10.0"
|
||||
ruff = "^0.0.280"
|
||||
isort = "^5.12.0"
|
||||
black = "^23.7.0"
|
||||
mypy = "^1.4.1"
|
||||
pre-commit = "^3.3.3"
|
||||
flake8 = { version = "*", python = ">=3.8.1,<3.12.0" }
|
||||
|
||||
|
||||
[tool.poetry.group.docs.dependencies]
|
||||
mkdocs-material = "^9.1.4"
|
||||
mkdocstrings = {extras = ["python"], version = "^0.20.0"}
|
||||
|
||||
[tool.flake8]
|
||||
exclude = ".venv"
|
||||
max-complexity = 10
|
||||
max-line-length = 88
|
||||
extend-ignore = """
|
||||
W503,
|
||||
E203,
|
||||
E701,
|
||||
C901,
|
||||
"""
|
||||
per-file-ignores = """
|
||||
__init__.py: F401
|
||||
"""
|
||||
|
||||
[tool.isort]
|
||||
line_length = 88
|
||||
profile = "black"
|
||||
|
||||
[tool.bandit]
|
||||
target = ["test", "supervision"]
|
||||
tests = ["B201", "B301"]
|
||||
|
||||
[tool.autoflake]
|
||||
check = true
|
||||
imports = ["cv2", "supervision"]
|
||||
|
||||
|
||||
[tool.black]
|
||||
target-version = ["py38"]
|
||||
line-length = 88
|
||||
include = '\.pyi?$'
|
||||
exclude = '''
|
||||
/(
|
||||
\.git
|
||||
| \.hg
|
||||
| \.mypy_cache
|
||||
| \.tox
|
||||
| \.venv
|
||||
| _build
|
||||
| buck-out
|
||||
| build
|
||||
| dist
|
||||
| docs
|
||||
)/
|
||||
'''
|
||||
|
||||
[tool.ruff]
|
||||
target-version = "py38"
|
||||
# Enable pycodestyle (`E`) and Pyflakes (`F`) codes by default.
|
||||
select = ["E", "F"]
|
||||
ignore = []
|
||||
|
||||
# Allow autofix for all enabled rules (when `--fix`) is provided.
|
||||
fixable = ["A", "B", "C", "D", "E", "F", "G", "I", "N", "Q", "S", "T", "W", "ANN", "ARG", "BLE", "COM", "DJ", "DTZ", "EM", "ERA", "EXE", "FBT", "ICN", "INP", "ISC", "NPY", "PD", "PGH", "PIE", "PL", "PT", "PTH", "PYI", "RET", "RSE", "RUF", "SIM", "SLF", "TCH", "TID", "TRY", "UP", "YTT"]
|
||||
unfixable = []
|
||||
|
||||
# Exclude a variety of commonly ignored directories.
|
||||
exclude = [
|
||||
".bzr",
|
||||
".direnv",
|
||||
".eggs",
|
||||
".git",
|
||||
".git-rewrite",
|
||||
".hg",
|
||||
".mypy_cache",
|
||||
".nox",
|
||||
".pants.d",
|
||||
".pytype",
|
||||
".ruff_cache",
|
||||
".svn",
|
||||
".tox",
|
||||
".venv",
|
||||
"__pypackages__",
|
||||
"_build",
|
||||
"buck-out",
|
||||
"build",
|
||||
"dist",
|
||||
"node_modules",
|
||||
"venv",
|
||||
"yarn-error.log",
|
||||
"yarn.lock",
|
||||
"docs",
|
||||
]
|
||||
|
||||
# Same as Black.
|
||||
line-length = 88
|
||||
|
||||
# Allow unused variables when underscore-prefixed.
|
||||
dummy-variable-rgx = "^(_+|(_+[a-zA-Z0-9_]*[a-zA-Z0-9]+?))$"
|
||||
|
||||
[tool.ruff.flake8-quotes]
|
||||
inline-quotes = "double"
|
||||
multiline-quotes = "double"
|
||||
docstring-quotes = "double"
|
||||
|
||||
[tool.ruff.pydocstyle]
|
||||
convention = "google"
|
||||
|
||||
[tool.ruff.per-file-ignores]
|
||||
"__init__.py" = ["E402","F401"]
|
||||
|
||||
[tool.ruff.pylint]
|
||||
max-args = 20
|
||||
|
||||
|
||||
[tool.setuptools]
|
||||
include-package-data = false
|
||||
|
||||
[tool.setuptools.packages.find]
|
||||
exclude = ["docs*", "test*"]
|
||||
|
||||
[build-system]
|
||||
requires = ["poetry-core"]
|
||||
build-backend = "poetry.core.masonry.api"
|
||||
@@ -1,69 +0,0 @@
|
||||
import setuptools
|
||||
from setuptools import find_packages
|
||||
import re
|
||||
from pathlib import Path
|
||||
|
||||
FILE = Path(__file__).resolve()
|
||||
PARENT = FILE.parent # root directory
|
||||
README = (PARENT / "README.md").read_text(encoding="utf-8")
|
||||
|
||||
|
||||
def get_version():
|
||||
file = PARENT / 'supervision/__init__.py'
|
||||
return re.search(r'^__version__ = [\'"]([^\'"]*)[\'"]', file.read_text(encoding="utf-8"), re.M)[1]
|
||||
|
||||
|
||||
setuptools.setup(
|
||||
name='supervision',
|
||||
version=get_version(),
|
||||
author='Piotr Skalski',
|
||||
author_email='piotr.skalski92@gmail.com',
|
||||
license='MIT',
|
||||
description='A set of easy-to-use utils that will come in handy in any Computer Vision project',
|
||||
long_description=README,
|
||||
long_description_content_type='text/markdown',
|
||||
url='https://github.com/roboflow/supervision',
|
||||
install_requires=[
|
||||
'numpy>=1.20.0',
|
||||
'opencv-python',
|
||||
'matplotlib',
|
||||
'pyyaml'
|
||||
],
|
||||
packages=find_packages(exclude=("tests",)),
|
||||
extras_require={
|
||||
'dev': [
|
||||
'flake8',
|
||||
'black==22.3.0',
|
||||
'isort',
|
||||
'twine',
|
||||
'pytest',
|
||||
'wheel',
|
||||
'notebook',
|
||||
'mkdocs-material',
|
||||
'mkdocstrings[python]'
|
||||
],
|
||||
},
|
||||
classifiers=[
|
||||
'Intended Audience :: Developers',
|
||||
'Intended Audience :: Science/Research',
|
||||
'License :: OSI Approved :: BSD License',
|
||||
'Programming Language :: Python :: 3',
|
||||
'Programming Language :: Python :: 3.7',
|
||||
'Programming Language :: Python :: 3.8',
|
||||
'Programming Language :: Python :: 3.9',
|
||||
'Programming Language :: Python :: 3.10',
|
||||
'Programming Language :: Python :: 3.11',
|
||||
'Programming Language :: Python :: 3 :: Only',
|
||||
'Topic :: Software Development',
|
||||
'Topic :: Scientific/Engineering',
|
||||
"Topic :: Scientific/Engineering :: Artificial Intelligence",
|
||||
"Topic :: Scientific/Engineering :: Image Recognition",
|
||||
'Typing :: Typed',
|
||||
'Operating System :: Microsoft :: Windows',
|
||||
'Operating System :: POSIX',
|
||||
'Operating System :: Unix',
|
||||
'Operating System :: MacOS'
|
||||
],
|
||||
keywords="machine-learning, deep-learning, vision, ML, DL, AI, YOLOv5, YOLOv8, Roboflow",
|
||||
python_requires='>=3.7',
|
||||
)
|
||||
+28
-4
@@ -1,9 +1,31 @@
|
||||
__version__ = "0.9.0"
|
||||
import importlib.metadata as importlib_metadata
|
||||
|
||||
from supervision.dataset.core import BaseDataset, DetectionDataset
|
||||
from supervision.detection.annotate import BoxAnnotator, MaskAnnotator
|
||||
try:
|
||||
# This will read version from pyproject.toml
|
||||
__version__ = importlib_metadata.version(__package__ or __name__)
|
||||
except importlib_metadata.PackageNotFoundError:
|
||||
__version__ = "development"
|
||||
|
||||
from supervision.annotators.core import (
|
||||
BlurAnnotator,
|
||||
BoundingBoxAnnotator,
|
||||
BoxCornerAnnotator,
|
||||
CircleAnnotator,
|
||||
EllipseAnnotator,
|
||||
LabelAnnotator,
|
||||
MaskAnnotator,
|
||||
TraceAnnotator,
|
||||
)
|
||||
from supervision.classification.core import Classifications
|
||||
from supervision.dataset.core import (
|
||||
BaseDataset,
|
||||
ClassificationDataset,
|
||||
DetectionDataset,
|
||||
)
|
||||
from supervision.detection.annotate import BoxAnnotator
|
||||
from supervision.detection.core import Detections
|
||||
from supervision.detection.line_counter import LineZone, LineZoneAnnotator
|
||||
from supervision.detection.tools.inference_slicer import InferenceSlicer
|
||||
from supervision.detection.tools.polygon_zone import PolygonZone, PolygonZoneAnnotator
|
||||
from supervision.detection.utils import (
|
||||
box_iou_batch,
|
||||
@@ -18,8 +40,10 @@ 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, 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.image import ImageSink, crop_image
|
||||
from supervision.utils.notebook import plot_image, plot_images_grid
|
||||
from supervision.utils.video import (
|
||||
VideoInfo,
|
||||
|
||||
@@ -0,0 +1,11 @@
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
import numpy as np
|
||||
|
||||
from supervision.detection.core import Detections
|
||||
|
||||
|
||||
class BaseAnnotator(ABC):
|
||||
@abstractmethod
|
||||
def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray:
|
||||
pass
|
||||
@@ -0,0 +1,707 @@
|
||||
from math import sqrt
|
||||
from typing import List, Optional, Tuple, Union
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from supervision.annotators.base import BaseAnnotator
|
||||
from supervision.annotators.utils import (
|
||||
ColorMap,
|
||||
Trace,
|
||||
resolve_color,
|
||||
resolve_color_idx,
|
||||
)
|
||||
from supervision.detection.core import Detections
|
||||
from supervision.draw.color import Color, ColorPalette
|
||||
from supervision.geometry.core import Position
|
||||
|
||||
|
||||
class BoundingBoxAnnotator(BaseAnnotator):
|
||||
"""
|
||||
A class for drawing bounding boxes on an image using provided detections.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
color: Union[Color, ColorPalette] = ColorPalette.default(),
|
||||
thickness: int = 2,
|
||||
color_map: str = "class",
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
color (Union[Color, ColorPalette]): The color or color palette to use for
|
||||
annotating detections.
|
||||
thickness (int): Thickness of the bounding box lines.
|
||||
color_map (str): Strategy for mapping colors to annotations.
|
||||
Options are `index`, `class`, or `track`.
|
||||
"""
|
||||
self.color: Union[Color, ColorPalette] = color
|
||||
self.thickness: int = thickness
|
||||
self.color_map: ColorMap = ColorMap(color_map)
|
||||
|
||||
def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray:
|
||||
"""
|
||||
Annotates the given scene with bounding boxes based on the provided detections.
|
||||
|
||||
Args:
|
||||
scene (np.ndarray): The image where bounding boxes will be drawn.
|
||||
detections (Detections): Object detections to annotate.
|
||||
|
||||
Returns:
|
||||
np.ndarray: The annotated image.
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> image = ...
|
||||
>>> detections = sv.Detections(...)
|
||||
|
||||
>>> bounding_box_annotator = sv.BoundingBoxAnnotator()
|
||||
>>> annotated_frame = bounding_box_annotator.annotate(
|
||||
... scene=image.copy(),
|
||||
... detections=detections
|
||||
... )
|
||||
```
|
||||
|
||||

|
||||
"""
|
||||
for detection_idx in range(len(detections)):
|
||||
x1, y1, x2, y2 = detections.xyxy[detection_idx].astype(int)
|
||||
idx = resolve_color_idx(
|
||||
detections=detections,
|
||||
detection_idx=detection_idx,
|
||||
color_map=self.color_map,
|
||||
)
|
||||
color = resolve_color(color=self.color, idx=idx)
|
||||
cv2.rectangle(
|
||||
img=scene,
|
||||
pt1=(x1, y1),
|
||||
pt2=(x2, y2),
|
||||
color=color.as_bgr(),
|
||||
thickness=self.thickness,
|
||||
)
|
||||
return scene
|
||||
|
||||
|
||||
class MaskAnnotator(BaseAnnotator):
|
||||
"""
|
||||
A class for drawing masks on an image using provided detections.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
color: Union[Color, ColorPalette] = ColorPalette.default(),
|
||||
opacity: float = 0.5,
|
||||
color_map: str = "class",
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
color (Union[Color, ColorPalette]): The color or color palette to use for
|
||||
annotating detections.
|
||||
opacity (float): Opacity of the overlay mask. Must be between `0` and `1`.
|
||||
color_map (str): Strategy for mapping colors to annotations.
|
||||
Options are `index`, `class`, or `track`.
|
||||
"""
|
||||
self.color: Union[Color, ColorPalette] = color
|
||||
self.opacity = opacity
|
||||
self.color_map: ColorMap = ColorMap(color_map)
|
||||
|
||||
def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray:
|
||||
"""
|
||||
Annotates the given scene with masks based on the provided detections.
|
||||
|
||||
Args:
|
||||
scene (np.ndarray): The image where masks will be drawn.
|
||||
detections (Detections): Object detections to annotate.
|
||||
|
||||
Returns:
|
||||
np.ndarray: The annotated image.
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> image = ...
|
||||
>>> detections = sv.Detections(...)
|
||||
|
||||
>>> mask_annotator = sv.MaskAnnotator()
|
||||
>>> annotated_frame = mask_annotator.annotate(
|
||||
... scene=image.copy(),
|
||||
... detections=detections
|
||||
... )
|
||||
```
|
||||
|
||||

|
||||
"""
|
||||
if detections.mask is None:
|
||||
return scene
|
||||
|
||||
for detection_idx in np.flip(np.argsort(detections.area)):
|
||||
idx = resolve_color_idx(
|
||||
detections=detections,
|
||||
detection_idx=detection_idx,
|
||||
color_map=self.color_map,
|
||||
)
|
||||
color = resolve_color(color=self.color, idx=idx)
|
||||
mask = detections.mask[detection_idx]
|
||||
colored_mask = np.zeros_like(scene, dtype=np.uint8)
|
||||
colored_mask[:] = color.as_bgr()
|
||||
|
||||
scene = np.where(
|
||||
np.expand_dims(mask, axis=-1),
|
||||
np.uint8(self.opacity * colored_mask + (1 - self.opacity) * scene),
|
||||
scene,
|
||||
)
|
||||
return scene
|
||||
|
||||
|
||||
class EllipseAnnotator(BaseAnnotator):
|
||||
"""
|
||||
A class for drawing ellipses on an image using provided detections.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
color: Union[Color, ColorPalette] = ColorPalette.default(),
|
||||
thickness: int = 2,
|
||||
start_angle: int = -45,
|
||||
end_angle: int = 235,
|
||||
color_map: str = "class",
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
color (Union[Color, ColorPalette]): The color or color palette to use for
|
||||
annotating detections.
|
||||
thickness (int): Thickness of the ellipse lines.
|
||||
start_angle (int): Starting angle of the ellipse.
|
||||
end_angle (int): Ending angle of the ellipse.
|
||||
color_map (str): Strategy for mapping colors to annotations.
|
||||
Options are `index`, `class`, or `track`.
|
||||
"""
|
||||
self.color: Union[Color, ColorPalette] = color
|
||||
self.thickness: int = thickness
|
||||
self.start_angle: int = start_angle
|
||||
self.end_angle: int = end_angle
|
||||
self.color_map: ColorMap = ColorMap(color_map)
|
||||
|
||||
def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray:
|
||||
"""
|
||||
Annotates the given scene with ellipses based on the provided detections.
|
||||
|
||||
Args:
|
||||
scene (np.ndarray): The image where ellipses will be drawn.
|
||||
detections (Detections): Object detections to annotate.
|
||||
|
||||
Returns:
|
||||
np.ndarray: The annotated image.
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> image = ...
|
||||
>>> detections = sv.Detections(...)
|
||||
|
||||
>>> ellipse_annotator = sv.EllipseAnnotator()
|
||||
>>> annotated_frame = ellipse_annotator.annotate(
|
||||
... scene=image.copy(),
|
||||
... detections=detections
|
||||
... )
|
||||
```
|
||||
|
||||

|
||||
"""
|
||||
for detection_idx in range(len(detections)):
|
||||
x1, y1, x2, y2 = detections.xyxy[detection_idx].astype(int)
|
||||
idx = resolve_color_idx(
|
||||
detections=detections,
|
||||
detection_idx=detection_idx,
|
||||
color_map=self.color_map,
|
||||
)
|
||||
color = resolve_color(color=self.color, idx=idx)
|
||||
|
||||
center = (int((x1 + x2) / 2), y2)
|
||||
width = x2 - x1
|
||||
cv2.ellipse(
|
||||
scene,
|
||||
center=center,
|
||||
axes=(int(width), int(0.35 * width)),
|
||||
angle=0.0,
|
||||
startAngle=self.start_angle,
|
||||
endAngle=self.end_angle,
|
||||
color=color.as_bgr(),
|
||||
thickness=self.thickness,
|
||||
lineType=cv2.LINE_4,
|
||||
)
|
||||
return scene
|
||||
|
||||
|
||||
class BoxCornerAnnotator(BaseAnnotator):
|
||||
"""
|
||||
A class for drawing box corners on an image using provided detections.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
color: Union[Color, ColorPalette] = ColorPalette.default(),
|
||||
thickness: int = 4,
|
||||
corner_length: int = 25,
|
||||
color_map: str = "class",
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
color (Union[Color, ColorPalette]): The color or color palette to use for
|
||||
annotating detections.
|
||||
thickness (int): Thickness of the corner lines.
|
||||
corner_length (int): Length of each corner line.
|
||||
color_map (str): Strategy for mapping colors to annotations.
|
||||
Options are `index`, `class`, or `track`.
|
||||
"""
|
||||
self.color: Union[Color, ColorPalette] = color
|
||||
self.thickness: int = thickness
|
||||
self.corner_length: int = corner_length
|
||||
self.color_map: ColorMap = ColorMap(color_map)
|
||||
|
||||
def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray:
|
||||
"""
|
||||
Annotates the given scene with box corners based on the provided detections.
|
||||
|
||||
Args:
|
||||
scene (np.ndarray): The image where box corners will be drawn.
|
||||
detections (Detections): Object detections to annotate.
|
||||
|
||||
Returns:
|
||||
np.ndarray: The annotated image.
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> image = ...
|
||||
>>> detections = sv.Detections(...)
|
||||
|
||||
>>> corner_annotator = sv.BoxCornerAnnotator()
|
||||
>>> annotated_frame = corner_annotator.annotate(
|
||||
... scene=image.copy(),
|
||||
... detections=detections
|
||||
... )
|
||||
```
|
||||
|
||||

|
||||
"""
|
||||
for detection_idx in range(len(detections)):
|
||||
x1, y1, x2, y2 = detections.xyxy[detection_idx].astype(int)
|
||||
idx = resolve_color_idx(
|
||||
detections=detections,
|
||||
detection_idx=detection_idx,
|
||||
color_map=self.color_map,
|
||||
)
|
||||
color = resolve_color(color=self.color, idx=idx)
|
||||
corners = [(x1, y1), (x2, y1), (x1, y2), (x2, y2)]
|
||||
|
||||
for x, y in corners:
|
||||
x_end = x + self.corner_length if x == x1 else x - self.corner_length
|
||||
cv2.line(
|
||||
scene, (x, y), (x_end, y), color.as_bgr(), thickness=self.thickness
|
||||
)
|
||||
|
||||
y_end = y + self.corner_length if y == y1 else y - self.corner_length
|
||||
cv2.line(
|
||||
scene, (x, y), (x, y_end), color.as_bgr(), thickness=self.thickness
|
||||
)
|
||||
return scene
|
||||
|
||||
|
||||
class CircleAnnotator(BaseAnnotator):
|
||||
"""
|
||||
A class for drawing circle on an image using provided detections.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
color: Union[Color, ColorPalette] = ColorPalette.default(),
|
||||
thickness: int = 4,
|
||||
color_map: str = "class",
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
color (Union[Color, ColorPalette]): The color or color palette to use for
|
||||
annotating detections.
|
||||
thickness (int): Thickness of the circle line.
|
||||
color_map (str): Strategy for mapping colors to annotations.
|
||||
Options are `index`, `class`, or `track`.
|
||||
"""
|
||||
|
||||
self.color: Union[Color, ColorPalette] = color
|
||||
self.thickness: int = thickness
|
||||
self.color_map: ColorMap = ColorMap(color_map)
|
||||
|
||||
def annotate(
|
||||
self,
|
||||
scene: np.ndarray,
|
||||
detections: Detections,
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
Annotates the given scene with circles based on the provided detections.
|
||||
|
||||
Args:
|
||||
scene (np.ndarray): The image where box corners will be drawn.
|
||||
detections (Detections): Object detections to annotate.
|
||||
|
||||
Returns:
|
||||
np.ndarray: The annotated image.
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> image = ...
|
||||
>>> detections = sv.Detections(...)
|
||||
|
||||
>>> circle_annotator = sv.CircleAnnotator()
|
||||
>>> annotated_frame = circle_annotator.annotate(
|
||||
... scene=image.copy(),
|
||||
... detections=detections
|
||||
... )
|
||||
```
|
||||
|
||||
|
||||

|
||||
"""
|
||||
for detection_idx in range(len(detections)):
|
||||
x1, y1, x2, y2 = detections.xyxy[detection_idx].astype(int)
|
||||
center = ((x1 + x2) // 2, (y1 + y2) // 2)
|
||||
distance = sqrt((x1 - center[0]) ** 2 + (y1 - center[1]) ** 2)
|
||||
|
||||
idx = resolve_color_idx(
|
||||
detections=detections,
|
||||
detection_idx=detection_idx,
|
||||
color_map=self.color_map,
|
||||
)
|
||||
|
||||
color = (
|
||||
self.color.by_idx(idx)
|
||||
if isinstance(self.color, ColorPalette)
|
||||
else self.color
|
||||
)
|
||||
|
||||
cv2.circle(
|
||||
img=scene,
|
||||
center=center,
|
||||
radius=int(distance),
|
||||
color=color.as_bgr(),
|
||||
thickness=self.thickness,
|
||||
)
|
||||
|
||||
return scene
|
||||
|
||||
|
||||
class LabelAnnotator:
|
||||
"""
|
||||
A class for annotating labels on an image using provided detections.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
color: Union[Color, ColorPalette] = ColorPalette.default(),
|
||||
text_color: Color = Color.black(),
|
||||
text_scale: float = 0.5,
|
||||
text_thickness: int = 1,
|
||||
text_padding: int = 10,
|
||||
text_position: Position = Position.TOP_LEFT,
|
||||
color_map: str = "class",
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
color (Union[Color, ColorPalette]): The color or color palette to use for
|
||||
annotating the text background.
|
||||
text_color (Color): The color to use for the text.
|
||||
text_scale (float): Font scale for the text.
|
||||
text_thickness (int): Thickness of the text characters.
|
||||
text_padding (int): Padding around the text within its background box.
|
||||
text_position (Position): Position of the text relative to the detection.
|
||||
Possible values are defined in the `Position` enum.
|
||||
color_map (str): Strategy for mapping colors to annotations.
|
||||
Options are `index`, `class`, or `track`.
|
||||
"""
|
||||
self.color: Union[Color, ColorPalette] = color
|
||||
self.text_color: Color = text_color
|
||||
self.text_scale: float = text_scale
|
||||
self.text_thickness: int = text_thickness
|
||||
self.text_padding: int = text_padding
|
||||
self.text_position: Position = text_position
|
||||
self.color_map: ColorMap = ColorMap(color_map)
|
||||
|
||||
@staticmethod
|
||||
def resolve_text_background_xyxy(
|
||||
detection_xyxy: Tuple[int, int, int, int],
|
||||
text_wh: Tuple[int, int],
|
||||
text_padding: int,
|
||||
position: Position,
|
||||
) -> Tuple[int, int, int, int]:
|
||||
padded_text_wh = (text_wh[0] + 2 * text_padding, text_wh[1] + 2 * text_padding)
|
||||
x1, y1, x2, y2 = detection_xyxy
|
||||
center_x = (x1 + x2) // 2
|
||||
center_y = (y1 + y2) // 2
|
||||
|
||||
if position == Position.TOP_LEFT:
|
||||
return x1, y1 - padded_text_wh[1], x1 + padded_text_wh[0], y1
|
||||
elif position == Position.TOP_RIGHT:
|
||||
return x2 - padded_text_wh[0], y1 - padded_text_wh[1], x2, y1
|
||||
elif position == Position.TOP_CENTER:
|
||||
return (
|
||||
center_x - padded_text_wh[0] // 2,
|
||||
y1 - padded_text_wh[1],
|
||||
center_x + padded_text_wh[0] // 2,
|
||||
y1,
|
||||
)
|
||||
elif position == Position.CENTER:
|
||||
return (
|
||||
center_x - padded_text_wh[0] // 2,
|
||||
center_y - padded_text_wh[1] // 2,
|
||||
center_x + padded_text_wh[0] // 2,
|
||||
center_y + padded_text_wh[1] // 2,
|
||||
)
|
||||
elif position == Position.BOTTOM_LEFT:
|
||||
return x1, y2, x1 + padded_text_wh[0], y2 + padded_text_wh[1]
|
||||
elif position == Position.BOTTOM_RIGHT:
|
||||
return x2 - padded_text_wh[0], y2, x2, y2 + padded_text_wh[1]
|
||||
elif position == Position.BOTTOM_CENTER:
|
||||
return (
|
||||
center_x - padded_text_wh[0] // 2,
|
||||
y2,
|
||||
center_x + padded_text_wh[0] // 2,
|
||||
y2 + padded_text_wh[1],
|
||||
)
|
||||
|
||||
def annotate(
|
||||
self,
|
||||
scene: np.ndarray,
|
||||
detections: Detections,
|
||||
labels: List[str] = None,
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
Annotates the given scene with labels based on the provided detections.
|
||||
|
||||
Args:
|
||||
scene (np.ndarray): The image where labels will be drawn.
|
||||
detections (Detections): Object detections to annotate.
|
||||
labels (List[str]): Optional. Custom labels for each detection.
|
||||
|
||||
Returns:
|
||||
np.ndarray: The annotated image.
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> image = ...
|
||||
>>> detections = sv.Detections(...)
|
||||
|
||||
>>> label_annotator = sv.LabelAnnotator(text_position=sv.Position.CENTER)
|
||||
>>> annotated_frame = label_annotator.annotate(
|
||||
... scene=image.copy(),
|
||||
... detections=detections
|
||||
... )
|
||||
```
|
||||
|
||||

|
||||
"""
|
||||
font = cv2.FONT_HERSHEY_SIMPLEX
|
||||
for detection_idx in range(len(detections)):
|
||||
detection_xyxy = detections.xyxy[detection_idx].astype(int)
|
||||
idx = resolve_color_idx(
|
||||
detections=detections,
|
||||
detection_idx=detection_idx,
|
||||
color_map=self.color_map,
|
||||
)
|
||||
color = resolve_color(color=self.color, idx=idx)
|
||||
text = (
|
||||
f"{detections.class_id[detection_idx]}"
|
||||
if (labels is None or len(detections) != len(labels))
|
||||
else labels[detection_idx]
|
||||
)
|
||||
text_wh = cv2.getTextSize(
|
||||
text=text,
|
||||
fontFace=font,
|
||||
fontScale=self.text_scale,
|
||||
thickness=self.text_thickness,
|
||||
)[0]
|
||||
|
||||
text_background_xyxy = self.resolve_text_background_xyxy(
|
||||
detection_xyxy=detection_xyxy,
|
||||
text_wh=text_wh,
|
||||
text_padding=self.text_padding,
|
||||
position=self.text_position,
|
||||
)
|
||||
|
||||
text_x = text_background_xyxy[0] + self.text_padding
|
||||
text_y = text_background_xyxy[1] + self.text_padding + text_wh[1]
|
||||
|
||||
cv2.rectangle(
|
||||
img=scene,
|
||||
pt1=(text_background_xyxy[0], text_background_xyxy[1]),
|
||||
pt2=(text_background_xyxy[2], text_background_xyxy[3]),
|
||||
color=color.as_bgr(),
|
||||
thickness=cv2.FILLED,
|
||||
)
|
||||
cv2.putText(
|
||||
img=scene,
|
||||
text=text,
|
||||
org=(text_x, text_y),
|
||||
fontFace=font,
|
||||
fontScale=self.text_scale,
|
||||
color=self.text_color.as_rgb(),
|
||||
thickness=self.text_thickness,
|
||||
lineType=cv2.LINE_AA,
|
||||
)
|
||||
return scene
|
||||
|
||||
|
||||
class BlurAnnotator(BaseAnnotator):
|
||||
"""
|
||||
A class for blurring regions in an image using provided detections.
|
||||
"""
|
||||
|
||||
def __init__(self, kernel_size: int = 15):
|
||||
"""
|
||||
Args:
|
||||
kernel_size (int): The size of the average pooling kernel used for blurring.
|
||||
"""
|
||||
self.kernel_size: int = kernel_size
|
||||
|
||||
def annotate(
|
||||
self,
|
||||
scene: np.ndarray,
|
||||
detections: Detections,
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
Annotates the given scene by blurring regions based on the provided detections.
|
||||
|
||||
Args:
|
||||
scene (np.ndarray): The image where blurring will be applied.
|
||||
detections (Detections): Object detections to annotate.
|
||||
|
||||
Returns:
|
||||
np.ndarray: The annotated image.
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> image = ...
|
||||
>>> detections = sv.Detections(...)
|
||||
|
||||
>>> blur_annotator = sv.BlurAnnotator()
|
||||
>>> annotated_frame = blur_annotator.annotate(
|
||||
... scene=image.copy(),
|
||||
... detections=detections
|
||||
... )
|
||||
```
|
||||
|
||||

|
||||
"""
|
||||
for detection_idx in range(len(detections)):
|
||||
x1, y1, x2, y2 = detections.xyxy[detection_idx].astype(int)
|
||||
roi = scene[y1:y2, x1:x2]
|
||||
|
||||
roi = cv2.blur(roi, (self.kernel_size, self.kernel_size))
|
||||
scene[y1:y2, x1:x2] = roi
|
||||
|
||||
return scene
|
||||
|
||||
|
||||
class TraceAnnotator:
|
||||
"""
|
||||
A class for drawing trace paths on an image based on detection coordinates.
|
||||
|
||||
!!! warning
|
||||
|
||||
This annotator utilizes the `tracker_id`. Read
|
||||
[here](https://supervision.roboflow.com/trackers/) to learn how to plug
|
||||
tracking into your inference pipeline.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
color: Union[Color, ColorPalette] = ColorPalette.default(),
|
||||
position: Optional[Position] = Position.CENTER,
|
||||
trace_length: int = 30,
|
||||
thickness: int = 2,
|
||||
color_map: str = "class",
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
color (Union[Color, ColorPalette]): The color to draw the trace, can be
|
||||
a single color or a color palette.
|
||||
position (Optional[Position]): The position of the trace.
|
||||
Defaults to `CENTER`.
|
||||
trace_length (int): The maximum length of the trace in terms of historical
|
||||
points. Defaults to `30`.
|
||||
thickness (int): The thickness of the trace lines. Defaults to `2`.
|
||||
color_map (str): Strategy for mapping colors to annotations.
|
||||
Options are `index`, `class`, or `track`.
|
||||
"""
|
||||
self.color: Union[Color, ColorPalette] = color
|
||||
self.position = position
|
||||
self.trace = Trace(max_size=trace_length)
|
||||
self.thickness = thickness
|
||||
self.color_map: ColorMap = ColorMap(color_map)
|
||||
|
||||
def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray:
|
||||
"""
|
||||
Draws trace paths on the frame based on the detection coordinates provided.
|
||||
|
||||
Args:
|
||||
scene (np.ndarray): The image on which the traces will be drawn.
|
||||
detections (Detections): The detections which include coordinates for
|
||||
which the traces will be drawn.
|
||||
|
||||
Returns:
|
||||
np.ndarray: The image with the trace paths drawn on it.
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> image = ...
|
||||
>>> detections = sv.Detections(...)
|
||||
|
||||
>>> trace_annotator = sv.TraceAnnotator()
|
||||
>>> annotated_frame = trace_annotator.annotate(
|
||||
... scene=image.copy(),
|
||||
... detections=detections
|
||||
... )
|
||||
```
|
||||
|
||||

|
||||
"""
|
||||
self.trace.put(detections)
|
||||
|
||||
for detection_idx in range(len(detections)):
|
||||
tracker_id = int(detections.tracker_id[detection_idx])
|
||||
idx = resolve_color_idx(
|
||||
detections=detections,
|
||||
detection_idx=detection_idx,
|
||||
color_map=self.color_map,
|
||||
)
|
||||
color = resolve_color(color=self.color, idx=idx)
|
||||
xy = self.trace.get(tracker_id=tracker_id)
|
||||
if len(xy) > 1:
|
||||
scene = cv2.polylines(
|
||||
scene,
|
||||
[xy.astype(np.int32)],
|
||||
False,
|
||||
color=color.as_bgr(),
|
||||
thickness=self.thickness,
|
||||
)
|
||||
return scene
|
||||
@@ -0,0 +1,89 @@
|
||||
from enum import Enum
|
||||
from typing import Optional, Union
|
||||
|
||||
import numpy as np
|
||||
|
||||
from supervision.detection.core import Detections
|
||||
from supervision.draw.color import Color, ColorPalette
|
||||
from supervision.geometry.core import Position
|
||||
|
||||
|
||||
class ColorMap(Enum):
|
||||
"""
|
||||
Enum for annotator color mapping.
|
||||
"""
|
||||
|
||||
INDEX = "index"
|
||||
CLASS = "class"
|
||||
TRACK = "track"
|
||||
|
||||
|
||||
def resolve_color_idx(
|
||||
detections: Detections, detection_idx: int, color_map: ColorMap = ColorMap.CLASS
|
||||
) -> int:
|
||||
if detection_idx >= len(detections):
|
||||
raise ValueError(
|
||||
f"Detection index {detection_idx}"
|
||||
f"is out of bounds for detections of length {len(detections)}"
|
||||
)
|
||||
|
||||
if color_map == ColorMap.INDEX:
|
||||
return detection_idx
|
||||
elif color_map == ColorMap.CLASS:
|
||||
if detections.class_id is None:
|
||||
raise ValueError(
|
||||
"Could not resolve color by class because"
|
||||
"Detections do not have class_id"
|
||||
)
|
||||
return detections.class_id[detection_idx]
|
||||
elif color_map == ColorMap.TRACK:
|
||||
if detections.tracker_id is None:
|
||||
raise ValueError(
|
||||
"Could not resolve color by track because"
|
||||
"Detections do not have tracker_id"
|
||||
)
|
||||
return detections.tracker_id[detection_idx]
|
||||
|
||||
|
||||
def resolve_color(color: Union[Color, ColorPalette], idx: int) -> Color:
|
||||
if isinstance(color, ColorPalette):
|
||||
return color.by_idx(idx)
|
||||
return color
|
||||
|
||||
|
||||
class Trace:
|
||||
def __init__(
|
||||
self,
|
||||
max_size: Optional[int] = None,
|
||||
start_frame_id: int = 0,
|
||||
anchor: Position = Position.CENTER,
|
||||
) -> None:
|
||||
self.current_frame_id = start_frame_id
|
||||
self.max_size = max_size
|
||||
self.anchor = anchor
|
||||
|
||||
self.frame_id = np.array([], dtype=int)
|
||||
self.xy = np.empty((0, 2), dtype=np.float32)
|
||||
self.tracker_id = np.array([], dtype=int)
|
||||
|
||||
def put(self, detections: Detections) -> None:
|
||||
frame_id = np.full(len(detections), self.current_frame_id, dtype=int)
|
||||
self.frame_id = np.concatenate([self.frame_id, frame_id])
|
||||
self.xy = np.concatenate(
|
||||
[self.xy, detections.get_anchor_coordinates(self.anchor)]
|
||||
)
|
||||
self.tracker_id = np.concatenate([self.tracker_id, detections.tracker_id])
|
||||
|
||||
unique_frame_id = np.unique(self.frame_id)
|
||||
|
||||
if 0 < self.max_size < len(unique_frame_id):
|
||||
max_allowed_frame_id = self.current_frame_id - self.max_size + 1
|
||||
filtering_mask = self.frame_id >= max_allowed_frame_id
|
||||
self.frame_id = self.frame_id[filtering_mask]
|
||||
self.xy = self.xy[filtering_mask]
|
||||
self.tracker_id = self.tracker_id[filtering_mask]
|
||||
|
||||
self.current_frame_id += 1
|
||||
|
||||
def get(self, tracker_id: int) -> np.ndarray:
|
||||
return self.xy[self.tracker_id == tracker_id]
|
||||
@@ -0,0 +1,140 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Optional, Tuple
|
||||
|
||||
import numpy as np
|
||||
|
||||
from supervision.utils.internal import deprecated
|
||||
|
||||
|
||||
def _validate_class_ids(class_id: Any, n: int) -> None:
|
||||
"""
|
||||
Ensure that class_id is a 1d np.ndarray with (n, ) shape.
|
||||
"""
|
||||
is_valid = isinstance(class_id, np.ndarray) and class_id.shape == (n,)
|
||||
if not is_valid:
|
||||
raise ValueError("class_id must be 1d np.ndarray with (n, ) shape")
|
||||
|
||||
|
||||
def _validate_confidence(confidence: Any, n: int) -> None:
|
||||
"""
|
||||
Ensure that confidence is a 1d np.ndarray with (n, ) shape.
|
||||
"""
|
||||
if confidence is not None:
|
||||
is_valid = isinstance(confidence, np.ndarray) and confidence.shape == (n,)
|
||||
if not is_valid:
|
||||
raise ValueError("confidence must be 1d np.ndarray with (n, ) shape")
|
||||
|
||||
|
||||
@dataclass
|
||||
class Classifications:
|
||||
class_id: np.ndarray
|
||||
confidence: Optional[np.ndarray] = None
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
"""
|
||||
Validate the classification inputs.
|
||||
"""
|
||||
n = len(self.class_id)
|
||||
|
||||
_validate_class_ids(self.class_id, n)
|
||||
_validate_confidence(self.confidence, n)
|
||||
|
||||
@classmethod
|
||||
@deprecated(
|
||||
"""
|
||||
This method is deprecated and removed in 0.16.0 release.
|
||||
Use sv.Classifications.from_ultralytics() instead as it is more generic and
|
||||
can be used for detections from any ultralytics.engine.results.Results Object
|
||||
"""
|
||||
)
|
||||
def from_yolov8(cls, yolov8_results) -> Classifications:
|
||||
"""
|
||||
Creates a Classifications instance from a
|
||||
[YOLOv8](https://github.com/ultralytics/ultralytics) inference result.
|
||||
|
||||
Args:
|
||||
yolov8_results (ultralytics.yolo.engine.results.Results):
|
||||
The output Results instance from YOLOv8
|
||||
|
||||
Returns:
|
||||
Classifications: A new Classifications object.
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> import cv2
|
||||
>>> from ultralytics import YOLO
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> image = cv2.imread(SOURCE_IMAGE_PATH)
|
||||
>>> model = YOLO('yolov8s-cls.pt')
|
||||
>>> result = model(image)[0]
|
||||
>>> classifications = sv.Classifications.from_yolov8(result)
|
||||
```
|
||||
"""
|
||||
confidence = yolov8_results.probs.data.cpu().numpy()
|
||||
return cls(class_id=np.arange(confidence.shape[0]), confidence=confidence)
|
||||
|
||||
@classmethod
|
||||
def from_ultralytics(cls, ultralytics_results) -> Classifications:
|
||||
"""
|
||||
Creates a Classifications instance from a
|
||||
(https://github.com/ultralytics/ultralytics) inference result.
|
||||
|
||||
Args:
|
||||
ultralytics_results (ultralytics.engine.results.Results):
|
||||
The output Results instance from ultralytics model
|
||||
|
||||
Returns:
|
||||
Classifications: A new Classifications object.
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> import cv2
|
||||
>>> from ultralytics import YOLO
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> image = cv2.imread(SOURCE_IMAGE_PATH)
|
||||
>>> model = YOLO('yolov8n-cls.pt')
|
||||
>>> model = YOLO('yolov8s-cls.pt')
|
||||
|
||||
>>> result = model(image)[0]
|
||||
>>> classifications = sv.Classifications.from_ultralytics(result)
|
||||
```
|
||||
"""
|
||||
confidence = ultralytics_results.probs.data.cpu().numpy()
|
||||
return cls(class_id=np.arange(confidence.shape[0]), confidence=confidence)
|
||||
|
||||
def get_top_k(self, k: int) -> Tuple[np.ndarray, np.ndarray]:
|
||||
"""
|
||||
Retrieve the top k class IDs and confidences,
|
||||
ordered in descending order by confidence.
|
||||
|
||||
Args:
|
||||
k (int): The number of top class IDs and confidences to retrieve.
|
||||
|
||||
Returns:
|
||||
Tuple[np.ndarray, np.ndarray]: A tuple containing
|
||||
the top k class IDs and confidences.
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> classifications = sv.Classifications(...)
|
||||
|
||||
>>> classifications.get_top_k(1)
|
||||
|
||||
(array([1]), array([0.9]))
|
||||
```
|
||||
"""
|
||||
if self.confidence is None:
|
||||
raise ValueError("top_k could not be calculated, confidence is None")
|
||||
|
||||
order = np.argsort(self.confidence)[::-1]
|
||||
top_k_order = order[:k]
|
||||
top_k_class_id = self.class_id[top_k_order]
|
||||
top_k_confidence = self.confidence[top_k_order]
|
||||
|
||||
return top_k_class_id, top_k_confidence
|
||||
+429
-81
@@ -1,5 +1,6 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
@@ -8,6 +9,11 @@ from typing import Dict, Iterator, List, Optional, Tuple
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from supervision.classification.core import Classifications
|
||||
from supervision.dataset.formats.coco import (
|
||||
load_coco_annotations,
|
||||
save_coco_annotations,
|
||||
)
|
||||
from supervision.dataset.formats.pascal_voc import (
|
||||
detections_to_pascal_voc,
|
||||
load_pascal_voc_annotations,
|
||||
@@ -17,9 +23,14 @@ from supervision.dataset.formats.yolo import (
|
||||
save_data_yaml,
|
||||
save_yolo_annotations,
|
||||
)
|
||||
from supervision.dataset.ultils import save_dataset_images, train_test_split
|
||||
from supervision.dataset.utils import (
|
||||
build_class_index_mapping,
|
||||
map_detections_class_id,
|
||||
merge_class_lists,
|
||||
save_dataset_images,
|
||||
train_test_split,
|
||||
)
|
||||
from supervision.detection.core import Detections
|
||||
from supervision.utils.file import list_files_with_extensions
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -43,7 +54,8 @@ class DetectionDataset(BaseDataset):
|
||||
Attributes:
|
||||
classes (List[str]): List containing dataset class names.
|
||||
images (Dict[str, np.ndarray]): Dictionary mapping image name to image.
|
||||
annotations (Dict[str, Detections]): Dictionary mapping image name to annotations.
|
||||
annotations (Dict[str, Detections]): Dictionary mapping
|
||||
image name to annotations.
|
||||
"""
|
||||
|
||||
classes: List[str]
|
||||
@@ -64,32 +76,53 @@ class DetectionDataset(BaseDataset):
|
||||
Iterate over the images and annotations in the dataset.
|
||||
|
||||
Yields:
|
||||
Iterator[Tuple[str, np.ndarray, Detections]]: An iterator that yields tuples containing the image name,
|
||||
the image data, and its corresponding annotation.
|
||||
Iterator[Tuple[str, np.ndarray, Detections]]:
|
||||
An iterator that yields tuples containing the image name,
|
||||
the image data, and its corresponding annotation.
|
||||
"""
|
||||
for image_name, image in self.images.items():
|
||||
yield image_name, image, self.annotations.get(image_name, None)
|
||||
|
||||
def __eq__(self, other):
|
||||
if not isinstance(other, DetectionDataset):
|
||||
return False
|
||||
|
||||
if set(self.classes) != set(other.classes):
|
||||
return False
|
||||
|
||||
for key in self.images:
|
||||
if not np.array_equal(self.images[key], other.images[key]):
|
||||
return False
|
||||
if not self.annotations[key] == other.annotations[key]:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def split(
|
||||
self, split_ratio=0.8, random_state=None, shuffle: bool = True
|
||||
) -> Tuple[DetectionDataset, DetectionDataset]:
|
||||
"""
|
||||
Splits the dataset into two parts (training and testing) using the provided split_ratio.
|
||||
Splits the dataset into two parts (training and testing)
|
||||
using the provided split_ratio.
|
||||
|
||||
Args:
|
||||
split_ratio (float, optional): The ratio of the training set to the entire dataset. Default is 0.8.
|
||||
random_state (int, optional): The seed for the random number generator. This is used for reproducibility. Default is None.
|
||||
shuffle (bool, optional): Whether to shuffle the data before splitting. Default is True.
|
||||
split_ratio (float, optional): The ratio of the training
|
||||
set to the entire dataset.
|
||||
random_state (int, optional): The seed for the random number generator.
|
||||
This is used for reproducibility.
|
||||
shuffle (bool, optional): Whether to shuffle the data before splitting.
|
||||
|
||||
Returns:
|
||||
Tuple[DetectionDataset, DetectionDataset]: A tuple containing the training and testing datasets.
|
||||
Tuple[DetectionDataset, DetectionDataset]: A tuple containing
|
||||
the training and testing datasets.
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> ds = sv.DetectionDataset(...)
|
||||
>>> train_ds, test_ds = ds.split(split_ratio=0.7, random_state=42, shuffle=True)
|
||||
>>> train_ds, test_ds = ds.split(split_ratio=0.7,
|
||||
... random_state=42, shuffle=True)
|
||||
>>> len(train_ds), len(test_ds)
|
||||
(700, 300)
|
||||
```
|
||||
@@ -124,37 +157,45 @@ class DetectionDataset(BaseDataset):
|
||||
approximation_percentage: float = 0.0,
|
||||
) -> None:
|
||||
"""
|
||||
Exports the dataset to PASCAL VOC format. This method saves the images and their corresponding annotations in
|
||||
PASCAL VOC format, which consists of XML files. The method allows filtering the detections based on their area
|
||||
percentage.
|
||||
Exports the dataset to PASCAL VOC format. This method saves the images
|
||||
and their corresponding annotations in PASCAL VOC format.
|
||||
|
||||
Args:
|
||||
images_directory_path (Optional[str]): The path to the directory where the images should be saved.
|
||||
images_directory_path (Optional[str]): The path to the directory
|
||||
where the images should be saved.
|
||||
If not provided, images will not be saved.
|
||||
annotations_directory_path (Optional[str]): The path to the directory where the annotations in
|
||||
PASCAL VOC format should be saved. If not provided, annotations will not be saved.
|
||||
min_image_area_percentage (float): The minimum percentage of detection area relative to
|
||||
annotations_directory_path (Optional[str]): The path to
|
||||
the directory where the annotations in
|
||||
PASCAL VOC format should be saved. If not provided,
|
||||
annotations will not be saved.
|
||||
min_image_area_percentage (float): The minimum percentage of
|
||||
detection area relative to
|
||||
the image area for a detection to be included.
|
||||
max_image_area_percentage (float): The maximum percentage of detection area relative to
|
||||
Argument is used only for segmentation datasets.
|
||||
max_image_area_percentage (float): The maximum percentage
|
||||
of detection area relative to
|
||||
the image area for a detection to be included.
|
||||
approximation_percentage (float): The percentage of polygon points to be removed from the input polygon, in the range [0, 1).
|
||||
Argument is used only for segmentation datasets.
|
||||
approximation_percentage (float): The percentage of
|
||||
polygon points to be removed from the input polygon,
|
||||
in the range [0, 1). Argument is used only for segmentation datasets.
|
||||
"""
|
||||
if images_directory_path:
|
||||
images_path = Path(images_directory_path)
|
||||
images_path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
save_dataset_images(
|
||||
images_directory_path=images_directory_path, images=self.images
|
||||
)
|
||||
if annotations_directory_path:
|
||||
annotations_path = Path(annotations_directory_path)
|
||||
annotations_path.mkdir(parents=True, exist_ok=True)
|
||||
Path(annotations_directory_path).mkdir(parents=True, exist_ok=True)
|
||||
|
||||
for image_name, image in self.images.items():
|
||||
detections = self.annotations[image_name]
|
||||
|
||||
if images_directory_path:
|
||||
cv2.imwrite(str(images_path / image_name), image)
|
||||
for image_path, image in self.images.items():
|
||||
detections = self.annotations[image_path]
|
||||
|
||||
if annotations_directory_path:
|
||||
annotation_name = Path(image_name).stem
|
||||
annotation_name = Path(image_path).stem
|
||||
annotations_path = os.path.join(
|
||||
annotations_directory_path, f"{annotation_name}.xml"
|
||||
)
|
||||
image_name = Path(image_path).name
|
||||
pascal_voc_xml = detections_to_pascal_voc(
|
||||
detections=detections,
|
||||
classes=self.classes,
|
||||
@@ -165,22 +206,29 @@ class DetectionDataset(BaseDataset):
|
||||
approximation_percentage=approximation_percentage,
|
||||
)
|
||||
|
||||
with open(annotations_path / f"{annotation_name}.xml", "w") as f:
|
||||
with open(annotations_path, "w") as f:
|
||||
f.write(pascal_voc_xml)
|
||||
|
||||
@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.
|
||||
|
||||
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.
|
||||
images_directory_path (str): Path to the directory containing the images.
|
||||
annotations_directory_path (str): 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.
|
||||
DetectionDataset: A DetectionDataset instance containing
|
||||
the loaded images and annotations.
|
||||
|
||||
Example:
|
||||
```python
|
||||
@@ -195,7 +243,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"
|
||||
... )
|
||||
@@ -204,34 +252,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
|
||||
@@ -246,13 +273,19 @@ class DetectionDataset(BaseDataset):
|
||||
Creates a Dataset instance from YOLO formatted data.
|
||||
|
||||
Args:
|
||||
images_directory_path (str): The path to the directory containing the images.
|
||||
annotations_directory_path (str): The path to the directory containing the YOLO annotation files.
|
||||
data_yaml_path (str): The path to the data YAML file containing class information.
|
||||
force_masks (bool, optional): If True, forces masks to be loaded for all annotations, regardless of whether they are present.
|
||||
images_directory_path (str): The path to the
|
||||
directory containing the images.
|
||||
annotations_directory_path (str): The path to the directory
|
||||
containing the YOLO annotation files.
|
||||
data_yaml_path (str): The path to the data
|
||||
YAML file containing class information.
|
||||
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.
|
||||
DetectionDataset: A DetectionDataset instance
|
||||
containing the loaded images and annotations.
|
||||
|
||||
Example:
|
||||
```python
|
||||
@@ -295,25 +328,32 @@ class DetectionDataset(BaseDataset):
|
||||
approximation_percentage: float = 0.0,
|
||||
) -> None:
|
||||
"""
|
||||
Exports the dataset to YOLO format. This method saves the images and their corresponding
|
||||
annotations in YOLO format, which is a simple text file that describes an object in the image. It also allows
|
||||
for the optional saving of a data.yaml file, used in YOLOv5, that contains metadata about the dataset.
|
||||
|
||||
The method allows filtering the detections based on their area percentage and offers an option for polygon approximation.
|
||||
Exports the dataset to YOLO format. This method saves the
|
||||
images and their corresponding annotations in YOLO format.
|
||||
|
||||
Args:
|
||||
images_directory_path (Optional[str]): The path to the directory where the images should be saved.
|
||||
images_directory_path (Optional[str]): The path to the
|
||||
directory where the images should be saved.
|
||||
If not provided, images will not be saved.
|
||||
annotations_directory_path (Optional[str]): The path to the directory where the annotations in
|
||||
YOLO format should be saved. If not provided, annotations will not be saved.
|
||||
data_yaml_path (Optional[str]): The path where the data.yaml file should be saved.
|
||||
annotations_directory_path (Optional[str]): The path to the
|
||||
directory where the annotations in
|
||||
YOLO format should be saved. If not provided,
|
||||
annotations will not be saved.
|
||||
data_yaml_path (Optional[str]): The path where the data.yaml
|
||||
file should be saved.
|
||||
If not provided, the file will not be saved.
|
||||
min_image_area_percentage (float): The minimum percentage of detection area relative to
|
||||
min_image_area_percentage (float): The minimum percentage of
|
||||
detection area relative to
|
||||
the image area for a detection to be included.
|
||||
max_image_area_percentage (float): The maximum percentage of detection area relative to
|
||||
Argument is used only for segmentation datasets.
|
||||
max_image_area_percentage (float): The maximum percentage
|
||||
of detection area relative to
|
||||
the image area for a detection to be included.
|
||||
approximation_percentage (float): The percentage of polygon points to be removed from the input polygon,
|
||||
in the range [0, 1). This is useful for simplifying the annotations.
|
||||
Argument is used only for segmentation datasets.
|
||||
approximation_percentage (float): The percentage of polygon points to
|
||||
be removed from the input polygon, in the range [0, 1).
|
||||
This is useful for simplifying the annotations.
|
||||
Argument is used only for segmentation datasets.
|
||||
"""
|
||||
if images_directory_path is not None:
|
||||
save_dataset_images(
|
||||
@@ -330,3 +370,311 @@ class DetectionDataset(BaseDataset):
|
||||
)
|
||||
if data_yaml_path is not None:
|
||||
save_data_yaml(data_yaml_path=data_yaml_path, classes=self.classes)
|
||||
|
||||
@classmethod
|
||||
def from_coco(
|
||||
cls,
|
||||
images_directory_path: str,
|
||||
annotations_path: str,
|
||||
force_masks: bool = False,
|
||||
) -> DetectionDataset:
|
||||
"""
|
||||
Creates a Dataset instance from COCO formatted data.
|
||||
|
||||
Args:
|
||||
images_directory_path (str): The path to the
|
||||
directory containing the images.
|
||||
annotations_path (str): The path to the json annotation files.
|
||||
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.
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> import roboflow
|
||||
>>> from roboflow import Roboflow
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> roboflow.login()
|
||||
|
||||
>>> rf = Roboflow()
|
||||
|
||||
>>> project = rf.workspace(WORKSPACE_ID).project(PROJECT_ID)
|
||||
>>> dataset = project.version(PROJECT_VERSION).download("coco")
|
||||
|
||||
>>> ds = sv.DetectionDataset.from_coco(
|
||||
... images_directory_path=f"{dataset.location}/train",
|
||||
... annotations_path=f"{dataset.location}/train/_annotations.coco.json",
|
||||
... )
|
||||
|
||||
>>> ds.classes
|
||||
['dog', 'person']
|
||||
```
|
||||
"""
|
||||
classes, images, annotations = load_coco_annotations(
|
||||
images_directory_path=images_directory_path,
|
||||
annotations_path=annotations_path,
|
||||
force_masks=force_masks,
|
||||
)
|
||||
return DetectionDataset(classes=classes, images=images, annotations=annotations)
|
||||
|
||||
def as_coco(
|
||||
self,
|
||||
images_directory_path: Optional[str] = None,
|
||||
annotations_path: Optional[str] = None,
|
||||
min_image_area_percentage: float = 0.0,
|
||||
max_image_area_percentage: float = 1.0,
|
||||
approximation_percentage: float = 0.0,
|
||||
) -> None:
|
||||
"""
|
||||
Exports the dataset to COCO format. This method saves the
|
||||
images and their corresponding annotations in COCO format.
|
||||
|
||||
Args:
|
||||
images_directory_path (Optional[str]): The path to the directory
|
||||
where the images should be saved.
|
||||
If not provided, images will not be saved.
|
||||
annotations_path (Optional[str]): The path to COCO annotation file.
|
||||
min_image_area_percentage (float): The minimum percentage of
|
||||
detection area relative to
|
||||
the image area for a detection to be included.
|
||||
Argument is used only for segmentation datasets.
|
||||
max_image_area_percentage (float): The maximum percentage of
|
||||
detection area relative to
|
||||
the image area for a detection to be included.
|
||||
Argument is used only for segmentation datasets.
|
||||
approximation_percentage (float): The percentage of polygon points
|
||||
to be removed from the input polygon,
|
||||
in the range [0, 1). This is useful for simplifying the annotations.
|
||||
Argument is used only for segmentation datasets.
|
||||
"""
|
||||
if images_directory_path is not None:
|
||||
save_dataset_images(
|
||||
images_directory_path=images_directory_path, images=self.images
|
||||
)
|
||||
if annotations_path is not None:
|
||||
save_coco_annotations(
|
||||
annotation_path=annotations_path,
|
||||
images=self.images,
|
||||
annotations=self.annotations,
|
||||
classes=self.classes,
|
||||
min_image_area_percentage=min_image_area_percentage,
|
||||
max_image_area_percentage=max_image_area_percentage,
|
||||
approximation_percentage=approximation_percentage,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def merge(cls, dataset_list: List[DetectionDataset]) -> DetectionDataset:
|
||||
"""
|
||||
Merge a list of `DetectionDataset` objects into a single
|
||||
`DetectionDataset` object.
|
||||
|
||||
This method takes a list of `DetectionDataset` objects and combines
|
||||
their respective fields (`classes`, `images`,
|
||||
`annotations`) into a single `DetectionDataset` object.
|
||||
|
||||
Args:
|
||||
dataset_list (List[DetectionDataset]): A list of `DetectionDataset`
|
||||
objects to merge.
|
||||
|
||||
Returns:
|
||||
(DetectionDataset): A single `DetectionDataset` object containing
|
||||
the merged data from the input list.
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> ds_1 = sv.DetectionDataset(...)
|
||||
>>> len(ds_1)
|
||||
100
|
||||
>>> ds_1.classes
|
||||
['dog', 'person']
|
||||
|
||||
>>> ds_2 = sv.DetectionDataset(...)
|
||||
>>> len(ds_2)
|
||||
200
|
||||
>>> ds_2.classes
|
||||
['cat']
|
||||
|
||||
>>> ds_merged = sv.DetectionDataset.merge([ds_1, ds_2])
|
||||
>>> len(ds_merged)
|
||||
300
|
||||
>>> ds_merged.classes
|
||||
['cat', 'dog', 'person']
|
||||
```
|
||||
"""
|
||||
merged_images, merged_annotations = {}, {}
|
||||
class_lists = [dataset.classes for dataset in dataset_list]
|
||||
merged_classes = merge_class_lists(class_lists=class_lists)
|
||||
|
||||
for dataset in dataset_list:
|
||||
class_index_mapping = build_class_index_mapping(
|
||||
source_classes=dataset.classes, target_classes=merged_classes
|
||||
)
|
||||
for image_name, image, detections in dataset:
|
||||
if image_name in merged_annotations:
|
||||
raise ValueError(
|
||||
f"Image name {image_name} is not unique across datasets."
|
||||
)
|
||||
|
||||
merged_images[image_name] = image
|
||||
merged_annotations[image_name] = map_detections_class_id(
|
||||
source_to_target_mapping=class_index_mapping,
|
||||
detections=detections,
|
||||
)
|
||||
|
||||
return cls(
|
||||
classes=merged_classes, images=merged_images, annotations=merged_annotations
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ClassificationDataset(BaseDataset):
|
||||
"""
|
||||
Dataclass containing information about a classification dataset.
|
||||
|
||||
Attributes:
|
||||
classes (List[str]): List containing dataset class names.
|
||||
images (Dict[str, np.ndarray]): Dictionary mapping image name to image.
|
||||
annotations (Dict[str, Detections]): Dictionary mapping
|
||||
image name to annotations.
|
||||
"""
|
||||
|
||||
classes: List[str]
|
||||
images: Dict[str, np.ndarray]
|
||||
annotations: Dict[str, Classifications]
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.images)
|
||||
|
||||
def split(
|
||||
self, split_ratio=0.8, random_state=None, shuffle: bool = True
|
||||
) -> Tuple[ClassificationDataset, ClassificationDataset]:
|
||||
"""
|
||||
Splits the dataset into two parts (training and testing)
|
||||
using the provided split_ratio.
|
||||
|
||||
Args:
|
||||
split_ratio (float, optional): The ratio of the training
|
||||
set to the entire dataset.
|
||||
random_state (int, optional): The seed for the
|
||||
random number generator.
|
||||
This is used for reproducibility.
|
||||
shuffle (bool, optional): Whether to shuffle the data before splitting.
|
||||
|
||||
Returns:
|
||||
Tuple[ClassificationDataset, ClassificationDataset]: A tuple containing
|
||||
the training and testing datasets.
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> cd = sv.ClassificationDataset(...)
|
||||
>>> train_cd,test_cd = cd.split(split_ratio=0.7,
|
||||
... random_state=42,shuffle=True)
|
||||
>>> len(train_cd), len(test_cd)
|
||||
(700, 300)
|
||||
```
|
||||
"""
|
||||
image_names = list(self.images.keys())
|
||||
train_names, test_names = train_test_split(
|
||||
data=image_names,
|
||||
train_ratio=split_ratio,
|
||||
random_state=random_state,
|
||||
shuffle=shuffle,
|
||||
)
|
||||
|
||||
train_dataset = ClassificationDataset(
|
||||
classes=self.classes,
|
||||
images={name: self.images[name] for name in train_names},
|
||||
annotations={name: self.annotations[name] for name in train_names},
|
||||
)
|
||||
test_dataset = ClassificationDataset(
|
||||
classes=self.classes,
|
||||
images={name: self.images[name] for name in test_names},
|
||||
annotations={name: self.annotations[name] for name in test_names},
|
||||
)
|
||||
return train_dataset, test_dataset
|
||||
|
||||
def as_folder_structure(self, root_directory_path: str) -> None:
|
||||
"""
|
||||
Saves the dataset as a multi-class folder structure.
|
||||
|
||||
Args:
|
||||
root_directory_path (str): The path to the directory
|
||||
where the dataset will be saved.
|
||||
"""
|
||||
os.makedirs(root_directory_path, exist_ok=True)
|
||||
|
||||
for class_name in self.classes:
|
||||
os.makedirs(os.path.join(root_directory_path, class_name), exist_ok=True)
|
||||
|
||||
for image_path in self.images:
|
||||
classification = self.annotations[image_path]
|
||||
image = self.images[image_path]
|
||||
image_name = Path(image_path).name
|
||||
class_id = (
|
||||
classification.class_id[0]
|
||||
if classification.confidence is None
|
||||
else classification.get_top_k(1)[0][0]
|
||||
)
|
||||
class_name = self.classes[class_id]
|
||||
image_path = os.path.join(root_directory_path, class_name, image_name)
|
||||
cv2.imwrite(image_path, image)
|
||||
|
||||
@classmethod
|
||||
def from_folder_structure(cls, root_directory_path: str) -> ClassificationDataset:
|
||||
"""
|
||||
Load data from a multiclass folder structure into a ClassificationDataset.
|
||||
|
||||
Args:
|
||||
root_directory_path (str): The path to the dataset directory.
|
||||
|
||||
Returns:
|
||||
ClassificationDataset: The dataset.
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> import roboflow
|
||||
>>> from roboflow import Roboflow
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> roboflow.login()
|
||||
|
||||
>>> rf = Roboflow()
|
||||
|
||||
>>> project = rf.workspace(WORKSPACE_ID).project(PROJECT_ID)
|
||||
>>> dataset = project.version(PROJECT_VERSION).download("folder")
|
||||
|
||||
>>> cd = sv.ClassificationDataset.from_folder_structure(
|
||||
... root_directory_path=f"{dataset.location}/train"
|
||||
... )
|
||||
```
|
||||
"""
|
||||
classes = os.listdir(root_directory_path)
|
||||
classes = sorted(set(classes))
|
||||
|
||||
images = {}
|
||||
annotations = {}
|
||||
|
||||
for class_name in classes:
|
||||
class_id = classes.index(class_name)
|
||||
|
||||
for image in os.listdir(os.path.join(root_directory_path, class_name)):
|
||||
image_path = str(os.path.join(root_directory_path, class_name, image))
|
||||
images[image_path] = cv2.imread(image_path)
|
||||
annotations[image_path] = Classifications(
|
||||
class_id=np.array([class_id]),
|
||||
)
|
||||
|
||||
return cls(
|
||||
classes=classes,
|
||||
images=images,
|
||||
annotations=annotations,
|
||||
)
|
||||
|
||||
@@ -0,0 +1,237 @@
|
||||
import os
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Tuple
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from supervision.dataset.utils import (
|
||||
approximate_mask_with_polygons,
|
||||
map_detections_class_id,
|
||||
)
|
||||
from supervision.detection.core import Detections
|
||||
from supervision.detection.utils import polygon_to_mask
|
||||
from supervision.utils.file import read_json_file, save_json_file
|
||||
|
||||
|
||||
def coco_categories_to_classes(coco_categories: List[dict]) -> List[str]:
|
||||
return [
|
||||
category["name"]
|
||||
for category in sorted(coco_categories, key=lambda category: category["id"])
|
||||
]
|
||||
|
||||
|
||||
def build_coco_class_index_mapping(
|
||||
coco_categories: List[dict], target_classes: List[str]
|
||||
) -> Dict[int, int]:
|
||||
source_class_to_index = {
|
||||
category["name"]: category["id"] for category in coco_categories
|
||||
}
|
||||
return {
|
||||
source_class_to_index[target_class_name]: target_class_index
|
||||
for target_class_index, target_class_name in enumerate(target_classes)
|
||||
}
|
||||
|
||||
|
||||
def classes_to_coco_categories(classes: List[str]) -> List[dict]:
|
||||
return [
|
||||
{
|
||||
"id": class_id,
|
||||
"name": class_name,
|
||||
"supercategory": "common-objects",
|
||||
}
|
||||
for class_id, class_name in enumerate(classes)
|
||||
]
|
||||
|
||||
|
||||
def group_coco_annotations_by_image_id(
|
||||
coco_annotations: List[dict],
|
||||
) -> Dict[int, List[dict]]:
|
||||
annotations = {}
|
||||
for annotation in coco_annotations:
|
||||
image_id = annotation["image_id"]
|
||||
if image_id not in annotations:
|
||||
annotations[image_id] = []
|
||||
annotations[image_id].append(annotation)
|
||||
return annotations
|
||||
|
||||
|
||||
def _polygons_to_masks(
|
||||
polygons: List[np.ndarray], resolution_wh: Tuple[int, int]
|
||||
) -> np.ndarray:
|
||||
return np.array(
|
||||
[
|
||||
polygon_to_mask(polygon=polygon, resolution_wh=resolution_wh)
|
||||
for polygon in polygons
|
||||
],
|
||||
dtype=bool,
|
||||
)
|
||||
|
||||
|
||||
def coco_annotations_to_detections(
|
||||
image_annotations: List[dict], resolution_wh: Tuple[int, int], with_masks: bool
|
||||
) -> Detections:
|
||||
if not image_annotations:
|
||||
return Detections.empty()
|
||||
|
||||
class_ids = [
|
||||
image_annotation["category_id"] for image_annotation in image_annotations
|
||||
]
|
||||
xyxy = [image_annotation["bbox"] for image_annotation in image_annotations]
|
||||
xyxy = np.asarray(xyxy)
|
||||
xyxy[:, 2:4] += xyxy[:, 0:2]
|
||||
|
||||
if with_masks:
|
||||
polygons = [
|
||||
np.reshape(
|
||||
np.asarray(image_annotation["segmentation"], dtype=np.int32), (-1, 2)
|
||||
)
|
||||
for image_annotation in image_annotations
|
||||
]
|
||||
mask = _polygons_to_masks(polygons=polygons, resolution_wh=resolution_wh)
|
||||
return Detections(
|
||||
class_id=np.asarray(class_ids, dtype=int), xyxy=xyxy, mask=mask
|
||||
)
|
||||
|
||||
return Detections(xyxy=xyxy, class_id=np.asarray(class_ids, dtype=int))
|
||||
|
||||
|
||||
def detections_to_coco_annotations(
|
||||
detections: Detections,
|
||||
image_id: int,
|
||||
annotation_id: int,
|
||||
min_image_area_percentage: float = 0.0,
|
||||
max_image_area_percentage: float = 1.0,
|
||||
approximation_percentage: float = 0.75,
|
||||
) -> Tuple[List[Dict], int]:
|
||||
coco_annotations = []
|
||||
for xyxy, mask, _, class_id, _ in detections:
|
||||
box_width, box_height = xyxy[2] - xyxy[0], xyxy[3] - xyxy[1]
|
||||
polygon = []
|
||||
if mask is not None:
|
||||
polygon = list(
|
||||
approximate_mask_with_polygons(
|
||||
mask=mask,
|
||||
min_image_area_percentage=min_image_area_percentage,
|
||||
max_image_area_percentage=max_image_area_percentage,
|
||||
approximation_percentage=approximation_percentage,
|
||||
)[0].flatten()
|
||||
)
|
||||
coco_annotation = {
|
||||
"id": annotation_id,
|
||||
"image_id": image_id,
|
||||
"category_id": int(class_id),
|
||||
"bbox": [xyxy[0], xyxy[1], box_width, box_height],
|
||||
"area": box_width * box_height,
|
||||
"segmentation": [polygon] if polygon else [],
|
||||
"iscrowd": 0,
|
||||
}
|
||||
coco_annotations.append(coco_annotation)
|
||||
annotation_id += 1
|
||||
return coco_annotations, annotation_id
|
||||
|
||||
|
||||
def load_coco_annotations(
|
||||
images_directory_path: str,
|
||||
annotations_path: str,
|
||||
force_masks: bool = False,
|
||||
) -> Tuple[List[str], Dict[str, np.ndarray], Dict[str, Detections]]:
|
||||
coco_data = read_json_file(file_path=annotations_path)
|
||||
classes = coco_categories_to_classes(coco_categories=coco_data["categories"])
|
||||
class_index_mapping = build_coco_class_index_mapping(
|
||||
coco_categories=coco_data["categories"], target_classes=classes
|
||||
)
|
||||
coco_images = coco_data["images"]
|
||||
coco_annotations_groups = group_coco_annotations_by_image_id(
|
||||
coco_annotations=coco_data["annotations"]
|
||||
)
|
||||
|
||||
images = {}
|
||||
annotations = {}
|
||||
|
||||
for coco_image in coco_images:
|
||||
image_name, image_width, image_height = (
|
||||
coco_image["file_name"],
|
||||
coco_image["width"],
|
||||
coco_image["height"],
|
||||
)
|
||||
image_annotations = coco_annotations_groups.get(coco_image["id"], [])
|
||||
image_path = os.path.join(images_directory_path, image_name)
|
||||
|
||||
image = cv2.imread(image_path)
|
||||
annotation = coco_annotations_to_detections(
|
||||
image_annotations=image_annotations,
|
||||
resolution_wh=(image_width, image_height),
|
||||
with_masks=force_masks,
|
||||
)
|
||||
annotation = map_detections_class_id(
|
||||
source_to_target_mapping=class_index_mapping,
|
||||
detections=annotation,
|
||||
)
|
||||
|
||||
images[image_path] = image
|
||||
annotations[image_path] = annotation
|
||||
|
||||
return classes, images, annotations
|
||||
|
||||
|
||||
def save_coco_annotations(
|
||||
annotation_path: str,
|
||||
images: Dict[str, np.ndarray],
|
||||
annotations: Dict[str, Detections],
|
||||
classes: List[str],
|
||||
min_image_area_percentage: float = 0.0,
|
||||
max_image_area_percentage: float = 1.0,
|
||||
approximation_percentage: float = 0.75,
|
||||
) -> None:
|
||||
Path(annotation_path).parent.mkdir(parents=True, exist_ok=True)
|
||||
info = {}
|
||||
licenses = [
|
||||
{
|
||||
"id": 1,
|
||||
"url": "https://creativecommons.org/licenses/by/4.0/",
|
||||
"name": "CC BY 4.0",
|
||||
}
|
||||
]
|
||||
|
||||
coco_annotations = []
|
||||
coco_images = []
|
||||
coco_categories = classes_to_coco_categories(classes=classes)
|
||||
|
||||
image_id, annotation_id = 1, 1
|
||||
for image_path, image in images.items():
|
||||
image_height, image_width, _ = image.shape
|
||||
image_name = f"{Path(image_path).stem}{Path(image_path).suffix}"
|
||||
coco_image = {
|
||||
"id": image_id,
|
||||
"license": 1,
|
||||
"file_name": image_name,
|
||||
"height": image_height,
|
||||
"width": image_width,
|
||||
"date_captured": datetime.now().strftime("%m/%d/%Y,%H:%M:%S"),
|
||||
}
|
||||
|
||||
coco_images.append(coco_image)
|
||||
detections = annotations[image_path]
|
||||
|
||||
coco_annotation, annotation_id = detections_to_coco_annotations(
|
||||
detections=detections,
|
||||
image_id=image_id,
|
||||
annotation_id=annotation_id,
|
||||
min_image_area_percentage=min_image_area_percentage,
|
||||
max_image_area_percentage=max_image_area_percentage,
|
||||
approximation_percentage=approximation_percentage,
|
||||
)
|
||||
|
||||
coco_annotations.extend(coco_annotation)
|
||||
image_id += 1
|
||||
|
||||
annotation_dict = {
|
||||
"info": info,
|
||||
"licenses": licenses,
|
||||
"categories": coco_categories,
|
||||
"images": coco_images,
|
||||
"annotations": coco_annotations,
|
||||
}
|
||||
save_json_file(annotation_dict, file_path=annotation_path)
|
||||
@@ -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.ultils import approximate_mask_with_polygons
|
||||
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(
|
||||
@@ -17,6 +21,9 @@ def object_to_pascal_voc(
|
||||
object_name = SubElement(root, "name")
|
||||
object_name.text = name
|
||||
|
||||
# https://github.com/roboflow/supervision/issues/144
|
||||
xyxy += 1
|
||||
|
||||
bndbox = SubElement(root, "bndbox")
|
||||
xmin = SubElement(bndbox, "xmin")
|
||||
xmin.text = str(int(xyxy[0]))
|
||||
@@ -28,6 +35,8 @@ def object_to_pascal_voc(
|
||||
ymax.text = str(int(xyxy[3]))
|
||||
|
||||
if polygon is not None:
|
||||
# https://github.com/roboflow/supervision/issues/144
|
||||
polygon += 1
|
||||
object_polygon = SubElement(root, "polygon")
|
||||
for index, point in enumerate(polygon, start=1):
|
||||
x_coordinate, y_coordinate = point
|
||||
@@ -52,13 +61,19 @@ def detections_to_pascal_voc(
|
||||
Converts Detections object to Pascal VOC XML format.
|
||||
|
||||
Args:
|
||||
detections (Detections): A Detections object containing bounding boxes, class ids, and other relevant information.
|
||||
classes (List[str]): A list of class names corresponding to the class ids in the Detections object.
|
||||
detections (Detections): A Detections object containing bounding boxes,
|
||||
class ids, and other relevant information.
|
||||
classes (List[str]): A list of class names corresponding to the
|
||||
class ids in the Detections object.
|
||||
filename (str): The name of the image file associated with the detections.
|
||||
image_shape (Tuple[int, int, int]): The shape of the image file associated with the detections.
|
||||
min_image_area_percentage (float): Minimum detection area relative to area of image associated with it.
|
||||
max_image_area_percentage (float): Maximum detection area relative to area of image associated with it.
|
||||
approximation_percentage (float): The percentage of polygon points to be removed from the input polygon, in the range [0, 1).
|
||||
image_shape (Tuple[int, int, int]): The shape of the image
|
||||
file associated with the detections.
|
||||
min_image_area_percentage (float): Minimum detection area
|
||||
relative to area of image associated with it.
|
||||
max_image_area_percentage (float): Maximum detection area
|
||||
relative to area of image associated with it.
|
||||
approximation_percentage (float): The percentage of
|
||||
polygon points to be removed from the input polygon, in the range [0, 1).
|
||||
Returns:
|
||||
str: An XML string in Pascal VOC format representing the detections.
|
||||
"""
|
||||
@@ -120,24 +135,101 @@ 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 XML annotations and returns the image name,
|
||||
a Detections instance, and a list of class names.
|
||||
|
||||
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_path = str(image_path)
|
||||
image = cv2.imread(image_path)
|
||||
|
||||
annotation_path = os.path.join(annotations_directory_path, f"{image_name}.xml")
|
||||
if not os.path.exists(annotation_path):
|
||||
images[image_path] = image
|
||||
annotations[image_path] = 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] = image
|
||||
annotations[image_path] = 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:
|
||||
<annotation>
|
||||
...
|
||||
<object>
|
||||
<name>dog</name>
|
||||
<bndbox>
|
||||
<xmin>48</xmin>
|
||||
<ymin>240</ymin>
|
||||
<xmax>195</xmax>
|
||||
<ymax>371</ymax>
|
||||
</bndbox>
|
||||
<polygon>
|
||||
<x1>48</x1>
|
||||
<y1>240</y1>
|
||||
<x2>195</x2>
|
||||
<y2>240</y2>
|
||||
<x3>195</x3>
|
||||
<y3>371</y3>
|
||||
<x4>48</x4>
|
||||
<y4>371</y4>
|
||||
</polygon>
|
||||
</object>
|
||||
</annotation>
|
||||
|
||||
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 +242,43 @@ 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 = parse_polygon_points(polygon)
|
||||
# https://github.com/roboflow/supervision/issues/144
|
||||
polygon -= 1
|
||||
|
||||
mask_from_polygon = polygon_to_mask(
|
||||
polygon=polygon,
|
||||
resolution_wh=resolution_wh,
|
||||
)
|
||||
masks.append(mask_from_polygon)
|
||||
|
||||
xyxy = np.array(xyxy) if len(xyxy) > 0 else np.empty((0, 4))
|
||||
|
||||
# https://github.com/roboflow/supervision/issues/144
|
||||
xyxy -= 1
|
||||
|
||||
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) -> np.ndarray:
|
||||
coordinates = [int(coord.text) for coord in polygon.findall(".//*")]
|
||||
return np.array(
|
||||
[(coordinates[i], coordinates[i + 1]) for i in range(0, len(coordinates), 2)]
|
||||
)
|
||||
|
||||
@@ -1,18 +1,19 @@
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Optional, Tuple, Union
|
||||
from typing import Dict, List, Optional, Tuple
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import yaml
|
||||
|
||||
from supervision.dataset.ultils import approximate_mask_with_polygons
|
||||
from supervision.dataset.utils import approximate_mask_with_polygons
|
||||
from supervision.detection.core import Detections
|
||||
from supervision.detection.utils import polygon_to_mask, polygon_to_xyxy
|
||||
from supervision.utils.file import (
|
||||
list_files_with_extensions,
|
||||
read_txt_file,
|
||||
read_yaml_file,
|
||||
save_text_file,
|
||||
save_yaml_file,
|
||||
)
|
||||
|
||||
|
||||
@@ -51,37 +52,15 @@ def _polygons_to_masks(
|
||||
)
|
||||
|
||||
|
||||
def _pair_image_with_annotation(
|
||||
image_paths: List[Union[str, Path]], annotation_paths: List[Union[str, Path]]
|
||||
) -> List[Tuple[Union[str, Path], Union[str, Path]]]:
|
||||
image_dict = {p.stem: p for p in image_paths}
|
||||
return [
|
||||
(image_dict[annotation_path.stem], annotation_path)
|
||||
for annotation_path in annotation_paths
|
||||
if annotation_path.stem in image_dict
|
||||
]
|
||||
|
||||
|
||||
def _with_mask(lines: List[str]) -> bool:
|
||||
return any([len(line.split()) > 5 for line in lines])
|
||||
|
||||
|
||||
def _extract_class_names(file_path: str) -> List[str]:
|
||||
names = []
|
||||
|
||||
with open(file_path, "r") as file:
|
||||
lines = file.readlines()
|
||||
for line in lines:
|
||||
if line.strip().startswith("names:"):
|
||||
start_index = lines.index(line) + 1
|
||||
break
|
||||
|
||||
for line in lines[start_index:]:
|
||||
if line.strip().startswith("-"):
|
||||
names.append(line.strip().replace("-", "").strip())
|
||||
else:
|
||||
break
|
||||
|
||||
data = read_yaml_file(file_path=file_path)
|
||||
names = data["names"]
|
||||
if isinstance(names, dict):
|
||||
names = [names[key] for key in sorted(names.keys())]
|
||||
return names
|
||||
|
||||
|
||||
@@ -133,31 +112,43 @@ def load_yolo_annotations(
|
||||
force_masks: bool = False,
|
||||
) -> Tuple[List[str], Dict[str, np.ndarray], Dict[str, Detections]]:
|
||||
"""
|
||||
Loads YOLO annotations and returns class names, images, and their corresponding detections.
|
||||
Loads YOLO annotations and returns class names, images,
|
||||
and their corresponding detections.
|
||||
|
||||
Args:
|
||||
images_directory_path (str): The path to the directory containing the images.
|
||||
annotations_directory_path (str): The path to the directory containing the YOLO annotation files.
|
||||
data_yaml_path (str): The path to the data YAML file containing class information.
|
||||
force_masks (bool, optional): If True, forces masks to be loaded for all annotations, regardless of whether they are present.
|
||||
annotations_directory_path (str): The path to the directory
|
||||
containing the YOLO annotation files.
|
||||
data_yaml_path (str): The path to the data
|
||||
YAML file containing class information.
|
||||
force_masks (bool, optional): If True, forces masks to be loaded
|
||||
for all annotations, regardless of whether they are present.
|
||||
|
||||
Returns:
|
||||
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.
|
||||
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.
|
||||
"""
|
||||
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"]
|
||||
)
|
||||
path_pairs = _pair_image_with_annotation(
|
||||
image_paths=image_paths, annotation_paths=annotation_paths
|
||||
)
|
||||
|
||||
classes = _extract_class_names(file_path=data_yaml_path)
|
||||
images = {}
|
||||
annotations = {}
|
||||
for image_path, annotation_path in path_pairs:
|
||||
image = cv2.imread(str(image_path))
|
||||
|
||||
for image_path in image_paths:
|
||||
image_stem = Path(image_path).stem
|
||||
image_path = str(image_path)
|
||||
image = cv2.imread(image_path)
|
||||
|
||||
annotation_path = os.path.join(annotations_directory_path, f"{image_stem}.txt")
|
||||
if not os.path.exists(annotation_path):
|
||||
images[image_path] = image
|
||||
annotations[image_path] = Detections.empty()
|
||||
continue
|
||||
|
||||
lines = read_txt_file(str(annotation_path))
|
||||
h, w, _ = image.shape
|
||||
resolution_wh = (w, h)
|
||||
@@ -168,8 +159,8 @@ def load_yolo_annotations(
|
||||
lines=lines, resolution_wh=resolution_wh, with_masks=with_masks
|
||||
)
|
||||
|
||||
images[image_path.name] = image
|
||||
annotations[image_path.name] = annotation
|
||||
images[image_path] = image
|
||||
annotations[image_path] = annotation
|
||||
return classes, images, annotations
|
||||
|
||||
|
||||
@@ -237,8 +228,9 @@ def save_yolo_annotations(
|
||||
approximation_percentage: float = 0.75,
|
||||
) -> None:
|
||||
Path(annotations_directory_path).mkdir(parents=True, exist_ok=True)
|
||||
for image_name, image in images.items():
|
||||
detections = annotations[image_name]
|
||||
for image_path, image in images.items():
|
||||
detections = annotations[image_path]
|
||||
image_name = Path(image_path).name
|
||||
yolo_annotations_name = _image_name_to_annotation_name(image_name=image_name)
|
||||
yolo_annotations_path = os.path.join(
|
||||
annotations_directory_path, yolo_annotations_name
|
||||
@@ -256,5 +248,4 @@ def save_yolo_annotations(
|
||||
def save_data_yaml(data_yaml_path: str, classes: List[str]) -> None:
|
||||
data = {"nc": len(classes), "names": classes}
|
||||
Path(data_yaml_path).parent.mkdir(parents=True, exist_ok=True)
|
||||
with open(data_yaml_path, "w") as outfile:
|
||||
yaml.dump(data, outfile, default_flow_style=False)
|
||||
save_yaml_file(data=data, file_path=data_yaml_path)
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
import copy
|
||||
import os
|
||||
import random
|
||||
from pathlib import Path
|
||||
@@ -6,6 +7,7 @@ from typing import Dict, List, Optional, Tuple, TypeVar
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from supervision.detection.core import Detections
|
||||
from supervision.detection.utils import (
|
||||
approximate_polygon,
|
||||
filter_polygons_by_area,
|
||||
@@ -43,12 +45,60 @@ def approximate_mask_with_polygons(
|
||||
]
|
||||
|
||||
|
||||
def merge_class_lists(class_lists: List[List[str]]) -> List[str]:
|
||||
unique_classes = set()
|
||||
|
||||
for class_list in class_lists:
|
||||
for class_name in class_list:
|
||||
unique_classes.add(class_name.lower())
|
||||
|
||||
return sorted(list(unique_classes))
|
||||
|
||||
|
||||
def build_class_index_mapping(
|
||||
source_classes: List[str], target_classes: List[str]
|
||||
) -> Dict[int, int]:
|
||||
index_mapping = {}
|
||||
|
||||
for i, class_name in enumerate(source_classes):
|
||||
if class_name not in target_classes:
|
||||
raise ValueError(
|
||||
f"Class {class_name} not found in target classes. "
|
||||
"source_classes must be a subset of target_classes."
|
||||
)
|
||||
corresponding_index = target_classes.index(class_name)
|
||||
index_mapping[i] = corresponding_index
|
||||
|
||||
return index_mapping
|
||||
|
||||
|
||||
def map_detections_class_id(
|
||||
source_to_target_mapping: Dict[int, int], detections: Detections
|
||||
) -> Detections:
|
||||
if detections.class_id is None:
|
||||
raise ValueError("Detections must have class_id attribute.")
|
||||
if set(np.unique(detections.class_id)) - set(source_to_target_mapping.keys()):
|
||||
raise ValueError(
|
||||
"Detections class_id must be a subset of source_to_target_mapping keys."
|
||||
)
|
||||
|
||||
detections_copy = copy.deepcopy(detections)
|
||||
|
||||
if len(detections) > 0:
|
||||
detections_copy.class_id = np.vectorize(source_to_target_mapping.get)(
|
||||
detections_copy.class_id
|
||||
)
|
||||
|
||||
return detections_copy
|
||||
|
||||
|
||||
def save_dataset_images(
|
||||
images_directory_path: str, images: Dict[str, np.ndarray]
|
||||
) -> None:
|
||||
Path(images_directory_path).mkdir(parents=True, exist_ok=True)
|
||||
|
||||
for image_name, image in images.items():
|
||||
for image_path, image in images.items():
|
||||
image_name = Path(image_path).name
|
||||
target_image_path = os.path.join(images_directory_path, image_name)
|
||||
cv2.imwrite(target_image_path, image)
|
||||
|
||||
@@ -12,12 +12,15 @@ class BoxAnnotator:
|
||||
A class for drawing bounding boxes on an image using detections provided.
|
||||
|
||||
Attributes:
|
||||
color (Union[Color, ColorPalette]): The color to draw the bounding box, can be a single color or a color palette
|
||||
color (Union[Color, ColorPalette]): The color to draw the bounding box,
|
||||
can be a single color or a color palette
|
||||
thickness (int): The thickness of the bounding box lines, default is 2
|
||||
text_color (Color): The color of the text on the bounding box, default is white
|
||||
text_scale (float): The scale of the text on the bounding box, default is 0.5
|
||||
text_thickness (int): The thickness of the text on the bounding box, default is 1
|
||||
text_padding (int): The padding around the text on the bounding box, default is 5
|
||||
text_thickness (int): The thickness of the text on the bounding box,
|
||||
default is 1
|
||||
text_padding (int): The padding around the text on the bounding box,
|
||||
default is 5
|
||||
|
||||
"""
|
||||
|
||||
@@ -49,8 +52,11 @@ class BoxAnnotator:
|
||||
|
||||
Args:
|
||||
scene (np.ndarray): The image on which the bounding boxes will be drawn
|
||||
detections (Detections): The detections for which the bounding boxes will be drawn
|
||||
labels (Optional[List[str]]): An optional list of labels corresponding to each detection. If `labels` are not provided, corresponding `class_id` will be used as label.
|
||||
detections (Detections): The detections for which the
|
||||
bounding boxes will be drawn
|
||||
labels (Optional[List[str]]): An optional list of labels
|
||||
corresponding to each detection. If `labels` are not provided,
|
||||
corresponding `class_id` will be used as label.
|
||||
skip_label (bool): Is set to `True`, skips bounding box label annotation.
|
||||
Returns:
|
||||
np.ndarray: The image with the bounding boxes drawn on it
|
||||
@@ -138,58 +144,3 @@ class BoxAnnotator:
|
||||
lineType=cv2.LINE_AA,
|
||||
)
|
||||
return scene
|
||||
|
||||
|
||||
class MaskAnnotator:
|
||||
"""
|
||||
A class for overlaying masks on an image using detections provided.
|
||||
|
||||
Attributes:
|
||||
color (Union[Color, ColorPalette]): The color to fill the mask, can be a single color or a color palette
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
color: Union[Color, ColorPalette] = ColorPalette.default(),
|
||||
):
|
||||
self.color: Union[Color, ColorPalette] = color
|
||||
|
||||
def annotate(
|
||||
self, scene: np.ndarray, detections: Detections, opacity: float = 0.5
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
Overlays the masks on the given image based on the provided detections, with a specified opacity.
|
||||
|
||||
Args:
|
||||
scene (np.ndarray): The image on which the masks will be overlaid
|
||||
detections (Detections): The detections for which the masks will be overlaid
|
||||
opacity (float): The opacity of the masks, between 0 and 1, default is 0.5
|
||||
|
||||
Returns:
|
||||
np.ndarray: The image with the masks overlaid
|
||||
"""
|
||||
if detections.mask is None:
|
||||
return scene
|
||||
|
||||
for i in np.flip(np.argsort(detections.area)):
|
||||
class_id = (
|
||||
detections.class_id[i] if detections.class_id is not None else None
|
||||
)
|
||||
idx = class_id if class_id is not None else i
|
||||
color = (
|
||||
self.color.by_idx(idx)
|
||||
if isinstance(self.color, ColorPalette)
|
||||
else self.color
|
||||
)
|
||||
|
||||
mask = detections.mask[i]
|
||||
colored_mask = np.zeros_like(scene, dtype=np.uint8)
|
||||
colored_mask[:] = color.as_bgr()
|
||||
|
||||
scene = np.where(
|
||||
np.expand_dims(mask, axis=-1),
|
||||
np.uint8(opacity * colored_mask + (1 - opacity) * scene),
|
||||
scene,
|
||||
)
|
||||
|
||||
return scene
|
||||
|
||||
+323
-87
@@ -3,12 +3,12 @@ from __future__ import annotations
|
||||
from dataclasses import astuple, dataclass
|
||||
from typing import Any, Iterator, List, Optional, Tuple, Union
|
||||
|
||||
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
|
||||
@@ -26,7 +26,14 @@ 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_inference_callback(callback) -> None:
|
||||
tmp_img = np.zeros((256, 256, 3), dtype=np.uint8)
|
||||
res = callback(tmp_img)
|
||||
if not isinstance(res, Detections):
|
||||
raise ValueError("Callback function must return sv.Detection type")
|
||||
|
||||
|
||||
def _validate_class_id(class_id: Any, n: int) -> None:
|
||||
@@ -58,15 +65,20 @@ 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.
|
||||
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.
|
||||
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, 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.
|
||||
"""
|
||||
|
||||
xyxy: np.ndarray
|
||||
mask: np.Optional[np.ndarray] = None
|
||||
mask: Optional[np.ndarray] = None
|
||||
confidence: Optional[np.ndarray] = None
|
||||
class_id: Optional[np.ndarray] = None
|
||||
tracker_id: Optional[np.ndarray] = None
|
||||
@@ -97,7 +109,8 @@ class Detections:
|
||||
]
|
||||
]:
|
||||
"""
|
||||
Iterates over the Detections object and yield a tuple of `(xyxy, mask, confidence, class_id, tracker_id)` for each detection.
|
||||
Iterates over the Detections object and yield a tuple of
|
||||
`(xyxy, mask, confidence, class_id, tracker_id)` for each detection.
|
||||
"""
|
||||
for i in range(len(self.xyxy)):
|
||||
yield (
|
||||
@@ -142,10 +155,12 @@ class Detections:
|
||||
@classmethod
|
||||
def from_yolov5(cls, yolov5_results) -> Detections:
|
||||
"""
|
||||
Creates a Detections instance from a [YOLOv5](https://github.com/ultralytics/yolov5) inference result.
|
||||
Creates a Detections instance from a
|
||||
[YOLOv5](https://github.com/ultralytics/yolov5) inference result.
|
||||
|
||||
Args:
|
||||
yolov5_results (yolov5.models.common.Detections): The output Detections instance from YOLOv5
|
||||
yolov5_results (yolov5.models.common.Detections):
|
||||
The output Detections instance from YOLOv5
|
||||
|
||||
Returns:
|
||||
Detections: A new Detections object.
|
||||
@@ -163,6 +178,7 @@ class Detections:
|
||||
```
|
||||
"""
|
||||
yolov5_detections_predictions = yolov5_results.pred[0].cpu().cpu().numpy()
|
||||
|
||||
return cls(
|
||||
xyxy=yolov5_detections_predictions[:, :4],
|
||||
confidence=yolov5_detections_predictions[:, 4],
|
||||
@@ -170,12 +186,21 @@ class Detections:
|
||||
)
|
||||
|
||||
@classmethod
|
||||
@deprecated(
|
||||
"""
|
||||
This method is deprecated and removed in 0.15.0 release.
|
||||
Use sv.Detections.from_ultralytics() instead as it is more generic and
|
||||
can be used for detections from any ultralytics.engine.results.Results Object
|
||||
"""
|
||||
)
|
||||
def from_yolov8(cls, yolov8_results) -> Detections:
|
||||
"""
|
||||
Creates a Detections instance from a [YOLOv8](https://github.com/ultralytics/ultralytics) inference result.
|
||||
Creates a Detections instance from a
|
||||
[YOLOv8](https://github.com/ultralytics/ultralytics) inference result.
|
||||
|
||||
Args:
|
||||
yolov8_results (ultralytics.yolo.engine.results.Results): The output Results instance from YOLOv8
|
||||
yolov8_results (ultralytics.yolo.engine.results.Results):
|
||||
The output Results instance from YOLOv8
|
||||
|
||||
Returns:
|
||||
Detections: A new Detections object.
|
||||
@@ -192,20 +217,69 @@ class Detections:
|
||||
>>> detections = sv.Detections.from_yolov8(result)
|
||||
```
|
||||
"""
|
||||
|
||||
return cls(
|
||||
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
|
||||
def from_yolo_nas(cls, yolo_nas_results) -> Detections:
|
||||
"""
|
||||
Creates a Detections instance from a [YOLO-NAS](https://github.com/Deci-AI/super-gradients/blob/master/YOLONAS.md) inference result.
|
||||
Creates a Detections instance from a
|
||||
[YOLO-NAS](https://github.com/Deci-AI/super-gradients/blob/master/YOLONAS.md)
|
||||
inference result.
|
||||
|
||||
Args:
|
||||
yolo_nas_results (super_gradients.training.models.prediction_results.ImageDetectionPrediction): The output Results instance from YOLO-NAS
|
||||
yolo_nas_results (ImageDetectionPrediction):
|
||||
The output Results instance from YOLO-NAS
|
||||
ImageDetectionPrediction is coming from
|
||||
'super_gradients.training.models.prediction_results'
|
||||
|
||||
Returns:
|
||||
Detections: A new Detections object.
|
||||
@@ -222,20 +296,95 @@ class Detections:
|
||||
>>> detections = sv.Detections.from_yolo_nas(result)
|
||||
```
|
||||
"""
|
||||
if np.asarray(yolo_nas_results.bboxes_xyxy).shape[0] == 0:
|
||||
return cls.empty()
|
||||
|
||||
return cls(
|
||||
xyxy=yolo_nas_results.prediction.bboxes_xyxy,
|
||||
confidence=yolo_nas_results.prediction.confidence,
|
||||
class_id=yolo_nas_results.prediction.labels.astype(int),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_deepsparse(cls, deepsparse_results) -> Detections:
|
||||
"""
|
||||
Creates a Detections instance from a
|
||||
[DeepSparse](https://github.com/neuralmagic/deepsparse)
|
||||
inference result.
|
||||
|
||||
Args:
|
||||
deepsparse_results (deepsparse.yolo.schemas.YOLOOutput):
|
||||
The output Results instance from DeepSparse.
|
||||
|
||||
Returns:
|
||||
Detections: A new Detections object.
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> from deepsparse import Pipeline
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> yolo_pipeline = Pipeline.create(
|
||||
... task="yolo",
|
||||
... model_path = "zoo:cv/detection/yolov5-l/pytorch/" \
|
||||
... "ultralytics/coco/pruned80_quant-none"
|
||||
>>> pipeline_outputs = yolo_pipeline(SOURCE_IMAGE_PATH,
|
||||
... iou_thres=0.6, conf_thres=0.001)
|
||||
>>> detections = sv.Detections.from_deepsparse(result)
|
||||
```
|
||||
"""
|
||||
if np.asarray(deepsparse_results.boxes[0]).shape[0] == 0:
|
||||
return cls.empty()
|
||||
|
||||
return cls(
|
||||
xyxy=np.array(deepsparse_results.boxes[0]),
|
||||
confidence=np.array(deepsparse_results.scores[0]),
|
||||
class_id=np.array(deepsparse_results.labels[0]).astype(float).astype(int),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_mmdetection(cls, mmdet_results) -> Detections:
|
||||
"""
|
||||
Creates a Detections instance from
|
||||
a [mmdetection](https://github.com/open-mmlab/mmdetection) inference result.
|
||||
Also supported for [mmyolo](https://github.com/open-mmlab/mmyolo)
|
||||
|
||||
Args:
|
||||
mmdet_results (mmdet.structures.DetDataSample):
|
||||
The output Results instance from MMDetection.
|
||||
|
||||
Returns:
|
||||
Detections: A new Detections object.
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> import cv2
|
||||
>>> import supervision as sv
|
||||
>>> from mmdet.apis import DetInferencer
|
||||
|
||||
>>> inferencer = DetInferencer(model_name, checkpoint, device)
|
||||
>>> mmdet_result = inferencer(SOURCE_IMAGE_PATH, out_dir='./output',
|
||||
... return_datasample=True)["predictions"][0]
|
||||
>>> detections = sv.Detections.from_mmdet(mmdet_result)
|
||||
```
|
||||
"""
|
||||
|
||||
return cls(
|
||||
xyxy=mmdet_results.pred_instances.bboxes.cpu().numpy(),
|
||||
confidence=mmdet_results.pred_instances.scores.cpu().numpy(),
|
||||
class_id=mmdet_results.pred_instances.labels.cpu().numpy().astype(int),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_transformers(cls, transformers_results: dict) -> Detections:
|
||||
"""
|
||||
Creates a Detections instance from object detection [transformer](https://github.com/huggingface/transformers) inference result.
|
||||
Creates a Detections instance from object detection
|
||||
[transformer](https://github.com/huggingface/transformers) inference result.
|
||||
|
||||
Returns:
|
||||
Detections: A new Detections object.
|
||||
"""
|
||||
|
||||
return cls(
|
||||
xyxy=transformers_results["boxes"].cpu().numpy(),
|
||||
confidence=transformers_results["scores"].cpu().numpy(),
|
||||
@@ -245,13 +394,16 @@ class Detections:
|
||||
@classmethod
|
||||
def from_detectron2(cls, detectron2_results) -> Detections:
|
||||
"""
|
||||
Create a Detections object from the [Detectron2](https://github.com/facebookresearch/detectron2) inference result.
|
||||
Create a Detections object from the
|
||||
[Detectron2](https://github.com/facebookresearch/detectron2) inference result.
|
||||
|
||||
Args:
|
||||
detectron2_results: The output of a Detectron2 model containing instances with prediction data.
|
||||
detectron2_results: The output of a
|
||||
Detectron2 model containing instances with prediction data.
|
||||
|
||||
Returns:
|
||||
(Detections): A Detections object containing the bounding boxes, class IDs, and confidences of the predictions.
|
||||
(Detections): A Detections object containing the bounding boxes,
|
||||
class IDs, and confidences of the predictions.
|
||||
|
||||
Example:
|
||||
```python
|
||||
@@ -266,10 +418,10 @@ class Detections:
|
||||
>>> cfg.MODEL.WEIGHTS = "path/to/model_weights.pth"
|
||||
>>> predictor = DefaultPredictor(cfg)
|
||||
>>> result = predictor(image)
|
||||
|
||||
>>> detections = sv.Detections.from_detectron2(result)
|
||||
```
|
||||
"""
|
||||
|
||||
return cls(
|
||||
xyxy=detectron2_results["instances"].pred_boxes.tensor.cpu().numpy(),
|
||||
confidence=detectron2_results["instances"].scores.cpu().numpy(),
|
||||
@@ -280,16 +432,18 @@ class Detections:
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_roboflow(cls, roboflow_result: dict, class_list: List[str]) -> Detections:
|
||||
def from_roboflow(cls, roboflow_result: dict) -> Detections:
|
||||
"""
|
||||
Create a Detections object from the [Roboflow](https://roboflow.com/) API inference result.
|
||||
Create a Detections object from the [Roboflow](https://roboflow.com/)
|
||||
API inference result.
|
||||
|
||||
Args:
|
||||
roboflow_result (dict): The result from the Roboflow API containing predictions.
|
||||
class_list (List[str]): A list of class names corresponding to the class IDs in the API result.
|
||||
roboflow_result (dict): The result from the
|
||||
Roboflow API containing predictions.
|
||||
|
||||
Returns:
|
||||
(Detections): A Detections object containing the bounding boxes, class IDs, and confidences of the predictions.
|
||||
(Detections): A Detections object containing the bounding boxes, class IDs,
|
||||
and confidences of the predictions.
|
||||
|
||||
Example:
|
||||
```python
|
||||
@@ -302,44 +456,37 @@ class Detections:
|
||||
... "y": 0.5,
|
||||
... "width": 0.2,
|
||||
... "height": 0.3,
|
||||
... "class_id": 0,
|
||||
... "class": "person",
|
||||
... "confidence": 0.9
|
||||
... },
|
||||
... # ... more predictions ...
|
||||
... ]
|
||||
... }
|
||||
>>> class_list = ["person", "car", "dog"]
|
||||
|
||||
>>> detections = sv.Detections.from_roboflow(roboflow_result, class_list)
|
||||
>>> detections = sv.Detections.from_roboflow(roboflow_result)
|
||||
```
|
||||
"""
|
||||
xyxy = []
|
||||
confidence = []
|
||||
class_id = []
|
||||
xyxy, confidence, class_id, masks = process_roboflow_result(
|
||||
roboflow_result=roboflow_result
|
||||
)
|
||||
|
||||
for prediction in roboflow_result["predictions"]:
|
||||
x = prediction["x"]
|
||||
y = prediction["y"]
|
||||
width = prediction["width"]
|
||||
height = prediction["height"]
|
||||
x_min = x - width / 2
|
||||
y_min = y - height / 2
|
||||
x_max = x_min + width
|
||||
y_max = y_min + height
|
||||
xyxy.append([x_min, y_min, x_max, y_max])
|
||||
class_id.append(class_list.index(prediction["class"]))
|
||||
confidence.append(prediction["confidence"])
|
||||
if np.asarray(xyxy).shape[0] == 0:
|
||||
return cls.empty()
|
||||
|
||||
return Detections(
|
||||
xyxy=np.array(xyxy),
|
||||
confidence=np.array(confidence),
|
||||
class_id=np.array(class_id).astype(int),
|
||||
return cls(
|
||||
xyxy=xyxy,
|
||||
confidence=confidence,
|
||||
class_id=class_id,
|
||||
mask=masks,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_sam(cls, sam_result: List[dict]) -> Detections:
|
||||
"""
|
||||
Creates a Detections instance from [Segment Anything Model](https://github.com/facebookresearch/segment-anything) inference result.
|
||||
Creates a Detections instance from
|
||||
[Segment Anything Model](https://github.com/facebookresearch/segment-anything)
|
||||
inference result.
|
||||
|
||||
Args:
|
||||
sam_result (List[dict]): The output Results instance from SAM
|
||||
@@ -350,14 +497,19 @@ class Detections:
|
||||
Example:
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
>>> from segment_anything import sam_model_registry, SamAutomaticMaskGenerator
|
||||
>>> from segment_anything import (
|
||||
... sam_model_registry,
|
||||
... SamAutomaticMaskGenerator
|
||||
... )
|
||||
|
||||
>>> sam = sam_model_registry[MODEL_TYPE](checkpoint=CHECKPOINT_PATH).to(device=DEVICE)
|
||||
>>> sam_model_reg = sam_model_registry[MODEL_TYPE]
|
||||
>>> sam = sam_model_reg(checkpoint=CHECKPOINT_PATH).to(device=DEVICE)
|
||||
>>> mask_generator = SamAutomaticMaskGenerator(sam)
|
||||
>>> sam_result = mask_generator.generate(IMAGE)
|
||||
>>> detections = sv.Detections.from_sam(sam_result=sam_result)
|
||||
```
|
||||
"""
|
||||
|
||||
sorted_generated_masks = sorted(
|
||||
sam_result, key=lambda x: x["area"], reverse=True
|
||||
)
|
||||
@@ -365,26 +517,59 @@ class Detections:
|
||||
xywh = np.array([mask["bbox"] for mask in sorted_generated_masks])
|
||||
mask = np.array([mask["segmentation"] for mask in sorted_generated_masks])
|
||||
|
||||
return Detections(xyxy=xywh_to_xyxy(boxes_xywh=xywh), mask=mask)
|
||||
if np.asarray(xywh).shape[0] == 0:
|
||||
return cls.empty()
|
||||
|
||||
xyxy = xywh_to_xyxy(boxes_xywh=xywh)
|
||||
return cls(xyxy=xyxy, mask=mask)
|
||||
|
||||
@classmethod
|
||||
@deprecated(
|
||||
"Dataset loading and saving is going to be executed by supervision.dataset.core.Dataset"
|
||||
)
|
||||
def from_coco_annotations(cls, coco_annotation: dict) -> Detections:
|
||||
xyxy, class_id = [], []
|
||||
def from_paddledet(cls, paddledet_result) -> Detections:
|
||||
"""
|
||||
Creates a Detections instance from
|
||||
[PaddleDetection](https://github.com/PaddlePaddle/PaddleDetection)
|
||||
inference result.
|
||||
|
||||
for annotation in coco_annotation:
|
||||
x_min, y_min, width, height = annotation["bbox"]
|
||||
xyxy.append([x_min, y_min, x_min + width, y_min + height])
|
||||
class_id.append(annotation["category_id"])
|
||||
Args:
|
||||
paddledet_result (List[dict]): The output Results instance from PaddleDet
|
||||
|
||||
return cls(xyxy=np.array(xyxy), class_id=np.array(class_id))
|
||||
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)
|
||||
```
|
||||
"""
|
||||
|
||||
if np.asarray(paddledet_result["bbox"][:, 2:6]).shape[0] == 0:
|
||||
return cls.empty()
|
||||
|
||||
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:
|
||||
"""
|
||||
Create an empty Detections object with no bounding boxes, confidences, or class IDs.
|
||||
Create an empty Detections object with no bounding boxes,
|
||||
confidences, or class IDs.
|
||||
|
||||
Returns:
|
||||
(Detections): An empty Detections object.
|
||||
@@ -407,15 +592,18 @@ class Detections:
|
||||
"""
|
||||
Merge a list of Detections objects into a single Detections object.
|
||||
|
||||
This method takes a list of Detections objects and combines their respective fields (`xyxy`, `mask`,
|
||||
`confidence`, `class_id`, and `tracker_id`) into a single Detections object. If all elements in a field are not
|
||||
`None`, the corresponding field will be stacked. Otherwise, the field will be set to `None`.
|
||||
This method takes a list of Detections objects and combines their
|
||||
respective fields (`xyxy`, `mask`, `confidence`, `class_id`, and `tracker_id`)
|
||||
into a single Detections object. If all elements in a field are not
|
||||
`None`, the corresponding field will be stacked.
|
||||
Otherwise, the field will be set to `None`.
|
||||
|
||||
Args:
|
||||
detections_list (List[Detections]): A list of Detections objects to merge.
|
||||
|
||||
Returns:
|
||||
(Detections): A single Detections object containing the merged data from the input list.
|
||||
(Detections): A single Detections object containing
|
||||
the merged data from the input list.
|
||||
|
||||
Example:
|
||||
```python
|
||||
@@ -435,13 +623,14 @@ class Detections:
|
||||
list(field) for field in zip(*detections_tuples_list)
|
||||
]
|
||||
|
||||
all_not_none = lambda l: all(x is not None for x in l)
|
||||
def __all_not_none(item_list: List[Any]):
|
||||
return all(x is not None for x in item_list)
|
||||
|
||||
xyxy = np.vstack(xyxy)
|
||||
mask = np.vstack(mask) if all_not_none(mask) else None
|
||||
confidence = np.hstack(confidence) if all_not_none(confidence) else None
|
||||
class_id = np.hstack(class_id) if all_not_none(class_id) else None
|
||||
tracker_id = np.hstack(tracker_id) if all_not_none(tracker_id) else None
|
||||
mask = np.vstack(mask) if __all_not_none(mask) else None
|
||||
confidence = np.hstack(confidence) if __all_not_none(confidence) else None
|
||||
class_id = np.hstack(class_id) if __all_not_none(class_id) else None
|
||||
tracker_id = np.hstack(tracker_id) if __all_not_none(tracker_id) else None
|
||||
|
||||
return cls(
|
||||
xyxy=xyxy,
|
||||
@@ -453,13 +642,23 @@ class Detections:
|
||||
|
||||
def get_anchor_coordinates(self, anchor: Position) -> np.ndarray:
|
||||
"""
|
||||
Returns the bounding box coordinates for a specific anchor.
|
||||
Calculates and returns the coordinates of a specific anchor point
|
||||
within the bounding boxes defined by the `xyxy` attribute. The anchor
|
||||
point can be any of the predefined positions in the `Position` enum,
|
||||
such as `CENTER`, `CENTER_LEFT`, `BOTTOM_RIGHT`, etc.
|
||||
|
||||
Args:
|
||||
anchor (Position): Position of bounding box anchor for which to return the coordinates.
|
||||
anchor (Position): An enum specifying the position of the anchor point
|
||||
within the bounding box. Supported positions are defined in the
|
||||
`Position` enum.
|
||||
|
||||
Returns:
|
||||
np.ndarray: An array of shape `(n, 2)` containing the bounding box anchor coordinates in format `[x, y]`.
|
||||
np.ndarray: An array of shape `(n, 2)`, where `n` is the number of bounding
|
||||
boxes. Each row contains the `[x, y]` coordinates of the specified
|
||||
anchor point for the corresponding bounding box.
|
||||
|
||||
Raises:
|
||||
ValueError: If the provided `anchor` is not supported.
|
||||
"""
|
||||
if anchor == Position.CENTER:
|
||||
return np.array(
|
||||
@@ -468,10 +667,36 @@ class Detections:
|
||||
(self.xyxy[:, 1] + self.xyxy[:, 3]) / 2,
|
||||
]
|
||||
).transpose()
|
||||
elif anchor == Position.CENTER_LEFT:
|
||||
return np.array(
|
||||
[
|
||||
self.xyxy[:, 0],
|
||||
(self.xyxy[:, 1] + self.xyxy[:, 3]) / 2,
|
||||
]
|
||||
).transpose()
|
||||
elif anchor == Position.CENTER_RIGHT:
|
||||
return np.array(
|
||||
[
|
||||
self.xyxy[:, 2],
|
||||
(self.xyxy[:, 1] + self.xyxy[:, 3]) / 2,
|
||||
]
|
||||
).transpose()
|
||||
elif anchor == Position.BOTTOM_CENTER:
|
||||
return np.array(
|
||||
[(self.xyxy[:, 0] + self.xyxy[:, 2]) / 2, self.xyxy[:, 3]]
|
||||
).transpose()
|
||||
elif anchor == Position.BOTTOM_LEFT:
|
||||
return np.array([self.xyxy[:, 0], self.xyxy[:, 3]]).transpose()
|
||||
elif anchor == Position.BOTTOM_RIGHT:
|
||||
return np.array([self.xyxy[:, 2], self.xyxy[:, 3]]).transpose()
|
||||
elif anchor == Position.TOP_CENTER:
|
||||
return np.array(
|
||||
[(self.xyxy[:, 0] + self.xyxy[:, 2]) / 2, self.xyxy[:, 1]]
|
||||
).transpose()
|
||||
elif anchor == Position.TOP_LEFT:
|
||||
return np.array([self.xyxy[:, 0], self.xyxy[:, 1]]).transpose()
|
||||
elif anchor == Position.TOP_RIGHT:
|
||||
return np.array([self.xyxy[:, 2], self.xyxy[:, 1]]).transpose()
|
||||
|
||||
raise ValueError(f"{anchor} is not supported.")
|
||||
|
||||
@@ -482,7 +707,8 @@ class Detections:
|
||||
Get a subset of the Detections object.
|
||||
|
||||
Args:
|
||||
index (Union[int, slice, List[int], np.ndarray]): The index or indices of the subset of the Detections
|
||||
index (Union[int, slice, List[int], np.ndarray]):
|
||||
The index or indices of the subset of the Detections
|
||||
|
||||
Returns:
|
||||
(Detections): A subset of the Detections object.
|
||||
@@ -517,11 +743,14 @@ class Detections:
|
||||
@property
|
||||
def area(self) -> np.ndarray:
|
||||
"""
|
||||
Calculate the area of each detection in the set of object detections. If masks field is defined property
|
||||
returns are of each mask. If only box is given property return area of each box.
|
||||
Calculate the area of each detection in the set of object detections.
|
||||
If masks field is defined property returns are of each mask.
|
||||
If only box is given property return area of each box.
|
||||
|
||||
Returns:
|
||||
np.ndarray: An array of floats containing the area of each detection in the format of `(area_1, area_2, ..., area_n)`, where n is the number of detections.
|
||||
np.ndarray: An array of floats containing the area of each detection
|
||||
in the format of `(area_1, area_2, ..., area_n)`,
|
||||
where n is the number of detections.
|
||||
"""
|
||||
if self.mask is not None:
|
||||
return np.array([np.sum(mask) for mask in self.mask])
|
||||
@@ -534,7 +763,9 @@ class Detections:
|
||||
Calculate the area of each bounding box in the set of object detections.
|
||||
|
||||
Returns:
|
||||
np.ndarray: An array of floats containing the area of each bounding box in the format of `(area_1, area_2, ..., area_n)`, where n is the number of detections.
|
||||
np.ndarray: An array of floats containing the area of each bounding
|
||||
box in the format of `(area_1, area_2, ..., area_n)`,
|
||||
where n is the number of detections.
|
||||
"""
|
||||
return (self.xyxy[:, 3] - self.xyxy[:, 1]) * (self.xyxy[:, 2] - self.xyxy[:, 0])
|
||||
|
||||
@@ -545,21 +776,26 @@ class Detections:
|
||||
Perform non-maximum suppression on the current set of object detections.
|
||||
|
||||
Args:
|
||||
threshold (float, optional): The intersection-over-union threshold to use for non-maximum suppression. Defaults to 0.5.
|
||||
class_agnostic (bool, optional): Whether to perform class-agnostic non-maximum suppression. If True, the class_id of each detection will be ignored. Defaults to False.
|
||||
threshold (float, optional): The intersection-over-union threshold
|
||||
to use for non-maximum suppression. Defaults to 0.5.
|
||||
class_agnostic (bool, optional): Whether to perform class-agnostic
|
||||
non-maximum suppression. If True, the class_id of each detection
|
||||
will be ignored. Defaults to False.
|
||||
|
||||
Returns:
|
||||
Detections: A new Detections object containing the subset of detections after non-maximum suppression.
|
||||
Detections: A new Detections object containing the subset of detections
|
||||
after non-maximum suppression.
|
||||
|
||||
Raises:
|
||||
AssertionError: If `confidence` is None and class_agnostic is False. If `class_id` is None and class_agnostic is False.
|
||||
AssertionError: If `confidence` is None and class_agnostic is False.
|
||||
If `class_id` is None and class_agnostic is False.
|
||||
"""
|
||||
if len(self) == 0:
|
||||
return self
|
||||
|
||||
assert (
|
||||
self.confidence is not None
|
||||
), f"Detections confidence must be given for NMS to be executed."
|
||||
), "Detections confidence must be given for NMS to be executed."
|
||||
|
||||
if class_agnostic:
|
||||
predictions = np.hstack((self.xyxy, self.confidence.reshape(-1, 1)))
|
||||
@@ -569,8 +805,8 @@ class Detections:
|
||||
return self[indices]
|
||||
|
||||
assert self.class_id is not None, (
|
||||
f"Detections class_id must be given for NMS to be executed. If you intended to perform class agnostic "
|
||||
f"NMS set class_agnostic=True."
|
||||
"Detections class_id must be given for NMS to be executed. If you intended"
|
||||
" to perform class agnostic NMS set class_agnostic=True."
|
||||
)
|
||||
|
||||
predictions = np.hstack(
|
||||
|
||||
@@ -113,7 +113,8 @@ class LineZoneAnnotator:
|
||||
|
||||
Attributes:
|
||||
frame (np.ndarray): The image on which the line will be drawn.
|
||||
line_counter (LineCounter): The line counter that will be used to draw the line.
|
||||
line_counter (LineCounter): The line counter
|
||||
that will be used to draw the line.
|
||||
|
||||
Returns:
|
||||
np.ndarray: The image with the line drawn on it.
|
||||
|
||||
@@ -0,0 +1,181 @@
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
from typing import Callable, Optional, Tuple
|
||||
|
||||
import numpy as np
|
||||
|
||||
from supervision.detection.core import Detections, validate_inference_callback
|
||||
from supervision.detection.utils import move_boxes
|
||||
from supervision.utils.image import crop_image
|
||||
|
||||
|
||||
def move_detections(detections: Detections, offset: np.array) -> Detections:
|
||||
"""
|
||||
Args:
|
||||
detections (sv.Detections): Detections object to be moved.
|
||||
offset (np.array): An array of shape `(2,)` containing offset values in format
|
||||
is `[dx, dy]`.
|
||||
Returns:
|
||||
(sv.Detections) repositioned Detections object.
|
||||
"""
|
||||
detections.xyxy = move_boxes(xyxy=detections.xyxy, offset=offset)
|
||||
return detections
|
||||
|
||||
|
||||
class InferenceSlicer:
|
||||
"""
|
||||
InferenceSlicer performs slicing-based inference for small target detection. This
|
||||
method, often referred to as
|
||||
[Slicing Adaptive Inference (SAHI)](https://ieeexplore.ieee.org/document/9897990),
|
||||
involves dividing a larger image into smaller slices, performing inference on each
|
||||
slice, and then merging the detections.
|
||||
|
||||
Args:
|
||||
slice_wh (Tuple[int, int]): Dimensions of each slice in the format
|
||||
`(width, height)`.
|
||||
overlap_ratio_wh (Tuple[float, float]): Overlap ratio between consecutive
|
||||
slices in the format `(width_ratio, height_ratio)`.
|
||||
iou_threshold (Optional[float]): Intersection over Union (IoU) threshold
|
||||
used for non-max suppression.
|
||||
callback (Callable): A function that performs inference on a given image
|
||||
slice and returns detections.
|
||||
thread_workers (int): Number of threads for parallel execution.
|
||||
|
||||
Note:
|
||||
The class ensures that slices do not exceed the boundaries of the original
|
||||
image. As a result, the final slices in the row and column dimensions might be
|
||||
smaller than the specified slice dimensions if the image's width or height is
|
||||
not a multiple of the slice's width or height minus the overlap.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
callback: Callable[[np.ndarray], Detections],
|
||||
slice_wh: Tuple[int, int] = (320, 320),
|
||||
overlap_ratio_wh: Tuple[float, float] = (0.2, 0.2),
|
||||
iou_threshold: Optional[float] = 0.5,
|
||||
thread_workers: int = 1,
|
||||
):
|
||||
self.slice_wh = slice_wh
|
||||
self.overlap_ratio_wh = overlap_ratio_wh
|
||||
self.iou_threshold = iou_threshold
|
||||
self.callback = callback
|
||||
self.thread_workers = thread_workers
|
||||
validate_inference_callback(callback=callback)
|
||||
|
||||
def __call__(self, image: np.ndarray) -> Detections:
|
||||
"""
|
||||
Performs slicing-based inference on the provided image using the specified
|
||||
callback.
|
||||
|
||||
Args:
|
||||
image (np.ndarray): The input image on which inference needs to be
|
||||
performed. The image should be in the format
|
||||
`(height, width, channels)`.
|
||||
|
||||
Returns:
|
||||
Detections: A collection of detections for the entire image after merging
|
||||
results from all slices and applying NMS.
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> import cv2
|
||||
>>> import supervision as sv
|
||||
>>> from ultralytics import YOLO
|
||||
|
||||
>>> image = cv2.imread(SOURCE_IMAGE_PATH)
|
||||
>>> model = YOLO(...)
|
||||
|
||||
>>> def callback(image_slice: np.ndarray) -> sv.Detections:
|
||||
... result = model(image_slice)[0]
|
||||
... return sv.Detections.from_ultralytics(result)
|
||||
|
||||
>>> slicer = sv.InferenceSlicer(callback = callback)
|
||||
|
||||
>>> detections = slicer(image)
|
||||
```
|
||||
"""
|
||||
detections_list = []
|
||||
resolution_wh = (image.shape[1], image.shape[0])
|
||||
offsets = self._generate_offset(
|
||||
resolution_wh=resolution_wh,
|
||||
slice_wh=self.slice_wh,
|
||||
overlap_ratio_wh=self.overlap_ratio_wh,
|
||||
)
|
||||
|
||||
with ThreadPoolExecutor(max_workers=self.thread_workers) as executor:
|
||||
futures = [
|
||||
executor.submit(self._run_callback, image, offset) for offset in offsets
|
||||
]
|
||||
for future in as_completed(futures):
|
||||
detections_list.append(future.result())
|
||||
|
||||
return Detections.merge(detections_list=detections_list).with_nms(
|
||||
threshold=self.iou_threshold
|
||||
)
|
||||
|
||||
def _run_callback(self, image, offset) -> Detections:
|
||||
"""
|
||||
Run the provided callback on a slice of an image.
|
||||
|
||||
Args:
|
||||
image (np.ndarray): The input image on which inference needs to run
|
||||
offset (np.ndarray): An array of shape `(4,)` containing coordinates
|
||||
for the slice.
|
||||
|
||||
Returns:
|
||||
Detections: A collection of detections for the slice.
|
||||
"""
|
||||
image_slice = crop_image(image=image, xyxy=offset)
|
||||
detections = self.callback(image_slice)
|
||||
detections = move_detections(detections=detections, offset=offset[:2])
|
||||
|
||||
return detections
|
||||
|
||||
@staticmethod
|
||||
def _generate_offset(
|
||||
resolution_wh: Tuple[int, int],
|
||||
slice_wh: Tuple[int, int],
|
||||
overlap_ratio_wh: Tuple[float, float],
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
Generate offset coordinates for slicing an image based on the given resolution,
|
||||
slice dimensions, and overlap ratios.
|
||||
|
||||
Args:
|
||||
resolution_wh (Tuple[int, int]): A tuple representing the width and height
|
||||
of the image to be sliced.
|
||||
slice_wh (Tuple[int, int]): A tuple representing the desired width and
|
||||
height of each slice.
|
||||
overlap_ratio_wh (Tuple[float, float]): A tuple representing the desired
|
||||
overlap ratio for width and height between consecutive slices. Each
|
||||
value should be in the range [0, 1), where 0 means no overlap and a
|
||||
value close to 1 means high overlap.
|
||||
|
||||
Returns:
|
||||
np.ndarray: An array of shape `(n, 4)` containing coordinates for each
|
||||
slice in the format `[xmin, ymin, xmax, ymax]`.
|
||||
|
||||
Note:
|
||||
The function ensures that slices do not exceed the boundaries of the
|
||||
original image. As a result, the final slices in the row and column
|
||||
dimensions might be smaller than the specified slice dimensions if the
|
||||
image's width or height is not a multiple of the slice's width or
|
||||
height minus the overlap.
|
||||
"""
|
||||
slice_width, slice_height = slice_wh
|
||||
image_width, image_height = resolution_wh
|
||||
overlap_ratio_width, overlap_ratio_height = overlap_ratio_wh
|
||||
|
||||
width_stride = slice_width - int(overlap_ratio_width * slice_width)
|
||||
height_stride = slice_height - int(overlap_ratio_height * slice_height)
|
||||
|
||||
ws = np.arange(0, image_width, width_stride)
|
||||
hs = np.arange(0, image_height, height_stride)
|
||||
|
||||
xmin, ymin = np.meshgrid(ws, hs)
|
||||
xmax = np.clip(xmin + slice_width, 0, image_width)
|
||||
ymax = np.clip(ymin + slice_height, 0, image_height)
|
||||
|
||||
offsets = np.stack([xmin, ymin, xmax, ymax], axis=-1).reshape(-1, 4)
|
||||
|
||||
return offsets
|
||||
@@ -17,9 +17,11 @@ class PolygonZone:
|
||||
A class for defining a polygon-shaped zone within a frame for detecting objects.
|
||||
|
||||
Attributes:
|
||||
polygon (np.ndarray): A polygon represented by a numpy array of shape `(N, 2)`, containing the `x`, `y` coordinates of the points.
|
||||
polygon (np.ndarray): A polygon represented by a numpy array of shape
|
||||
`(N, 2)`, containing the `x`, `y` coordinates of the points.
|
||||
frame_resolution_wh (Tuple[int, int]): The frame resolution (width, height)
|
||||
triggering_position (Position): The position within the bounding box that triggers the zone (default: Position.BOTTOM_CENTER)
|
||||
triggering_position (Position): The position within the bounding
|
||||
box that triggers the zone (default: Position.BOTTOM_CENTER)
|
||||
current_count (int): The current count of detected objects within the zone
|
||||
mask (np.ndarray): The 2D bool mask for the polygon zone
|
||||
"""
|
||||
@@ -45,10 +47,12 @@ class PolygonZone:
|
||||
Determines if the detections are within the polygon zone.
|
||||
|
||||
Parameters:
|
||||
detections (Detections): The detections to be checked against the polygon zone
|
||||
detections (Detections): The detections
|
||||
to be checked against the polygon zone
|
||||
|
||||
Returns:
|
||||
np.ndarray: A boolean numpy array indicating if each detection is within the polygon zone
|
||||
np.ndarray: A boolean numpy array indicating
|
||||
if each detection is within the polygon zone
|
||||
"""
|
||||
|
||||
clipped_xyxy = clip_boxes(
|
||||
@@ -59,13 +63,14 @@ class PolygonZone:
|
||||
clipped_detections.get_anchor_coordinates(anchor=self.triggering_position)
|
||||
).astype(int)
|
||||
is_in_zone = self.mask[clipped_anchors[:, 1], clipped_anchors[:, 0]]
|
||||
self.current_count = np.sum(is_in_zone)
|
||||
self.current_count = int(np.sum(is_in_zone))
|
||||
return is_in_zone.astype(bool)
|
||||
|
||||
|
||||
class PolygonZoneAnnotator:
|
||||
"""
|
||||
A class for annotating a polygon-shaped zone within a frame with a count of detected objects.
|
||||
A class for annotating a polygon-shaped zone within a
|
||||
frame with a count of detected objects.
|
||||
|
||||
Attributes:
|
||||
zone (PolygonZone): The polygon zone to be annotated
|
||||
@@ -75,7 +80,8 @@ class PolygonZoneAnnotator:
|
||||
text_scale (float): The scale of the text on the polygon, default is 0.5
|
||||
text_thickness (int): The thickness of the text on the polygon, default is 1
|
||||
text_padding (int): The padding around the text on the polygon, default is 10
|
||||
font (int): The font type for the text on the polygon, default is cv2.FONT_HERSHEY_SIMPLEX
|
||||
font (int): The font type for the text on the polygon,
|
||||
default is cv2.FONT_HERSHEY_SIMPLEX
|
||||
center (Tuple[int, int]): The center of the polygon for text placement
|
||||
"""
|
||||
|
||||
@@ -105,7 +111,8 @@ class PolygonZoneAnnotator:
|
||||
|
||||
Parameters:
|
||||
scene (np.ndarray): The image on which the polygon zone will be annotated
|
||||
label (Optional[str]): An optional label for the count of detected objects within the polygon zone (default: None)
|
||||
label (Optional[str]): An optional label for the count of detected objects
|
||||
within the polygon zone (default: None)
|
||||
|
||||
Returns:
|
||||
np.ndarray: The image with the polygon zone and count of detected objects
|
||||
|
||||
+137
-40
@@ -10,29 +10,37 @@ def polygon_to_mask(polygon: np.ndarray, resolution_wh: Tuple[int, int]) -> np.n
|
||||
"""Generate a mask from a polygon.
|
||||
|
||||
Args:
|
||||
polygon (np.ndarray): The polygon for which the mask should be generated, given as a list of vertices.
|
||||
polygon (np.ndarray): The polygon for which the mask should be generated,
|
||||
given as a list of vertices.
|
||||
resolution_wh (Tuple[int, int]): The width and height of the desired resolution.
|
||||
|
||||
Returns:
|
||||
np.ndarray: The generated 2D mask, where the polygon is marked with `1`'s and the rest is filled with `0`'s.
|
||||
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
|
||||
|
||||
|
||||
def box_iou_batch(boxes_true: np.ndarray, boxes_detection: np.ndarray) -> np.ndarray:
|
||||
"""
|
||||
Compute Intersection over Union (IoU) of two sets of bounding boxes - `boxes_true` and `boxes_detection`. Both sets
|
||||
of boxes are expected to be in `(x_min, y_min, x_max, y_max)` format.
|
||||
Compute Intersection over Union (IoU) of two sets of bounding boxes -
|
||||
`boxes_true` and `boxes_detection`. Both sets
|
||||
of boxes are expected to be in `(x_min, y_min, x_max, y_max)` format.
|
||||
|
||||
Args:
|
||||
boxes_true (np.ndarray): 2D `np.ndarray` representing ground-truth boxes. `shape = (N, 4)` where `N` is number of true objects.
|
||||
boxes_detection (np.ndarray): 2D `np.ndarray` representing detection boxes. `shape = (M, 4)` where `M` is number of detected objects.
|
||||
boxes_true (np.ndarray): 2D `np.ndarray` representing ground-truth boxes.
|
||||
`shape = (N, 4)` where `N` is number of true objects.
|
||||
boxes_detection (np.ndarray): 2D `np.ndarray` representing detection boxes.
|
||||
`shape = (M, 4)` where `M` is number of detected objects.
|
||||
|
||||
Returns:
|
||||
np.ndarray: Pairwise IoU of boxes from `boxes_true` and `boxes_detection`. `shape = (N, M)` where `N` is number of true objects and `M` is number of detected objects.
|
||||
np.ndarray: Pairwise IoU of boxes from `boxes_true` and `boxes_detection`.
|
||||
`shape = (N, M)` where `N` is number of true objects and
|
||||
`M` is number of detected objects.
|
||||
"""
|
||||
|
||||
def box_area(box):
|
||||
@@ -55,17 +63,22 @@ def non_max_suppression(
|
||||
Perform Non-Maximum Suppression (NMS) on object detection predictions.
|
||||
|
||||
Args:
|
||||
predictions (np.ndarray): An array of object detection predictions in the format of `(x_min, y_min, x_max, y_max, score)` or `(x_min, y_min, x_max, y_max, score, class)`.
|
||||
iou_threshold (float, optional): The intersection-over-union threshold to use for non-maximum suppression.
|
||||
predictions (np.ndarray): An array of object detection predictions in
|
||||
the format of `(x_min, y_min, x_max, y_max, score)`
|
||||
or `(x_min, y_min, x_max, y_max, score, class)`.
|
||||
iou_threshold (float, optional): The intersection-over-union threshold
|
||||
to use for non-maximum suppression.
|
||||
|
||||
Returns:
|
||||
np.ndarray: A boolean array indicating which predictions to keep after non-maximum suppression.
|
||||
np.ndarray: A boolean array indicating which predictions to keep after n
|
||||
on-maximum suppression.
|
||||
|
||||
Raises:
|
||||
AssertionError: If `iou_threshold` is not within the closed range from `0` to `1`.
|
||||
AssertionError: If `iou_threshold` is not within the
|
||||
closed range from `0` to `1`.
|
||||
"""
|
||||
assert 0 <= iou_threshold <= 1, (
|
||||
f"Value of `iou_threshold` must be in the closed range from 0 to 1, "
|
||||
"Value of `iou_threshold` must be in the closed range from 0 to 1, "
|
||||
f"{iou_threshold} given."
|
||||
)
|
||||
rows, columns = predictions.shape
|
||||
@@ -89,7 +102,8 @@ def non_max_suppression(
|
||||
if not keep[index]:
|
||||
continue
|
||||
|
||||
# drop detections with iou > iou_threshold and same category as current detections
|
||||
# drop detections with iou > iou_threshold and
|
||||
# same category as current detections
|
||||
condition = (iou > iou_threshold) & (categories == category)
|
||||
keep = keep & ~condition
|
||||
|
||||
@@ -103,14 +117,16 @@ def clip_boxes(
|
||||
Clips bounding boxes coordinates to fit within the frame resolution.
|
||||
|
||||
Args:
|
||||
boxes_xyxy (np.ndarray): A numpy array of shape `(N, 4)` where each row corresponds to a bounding box in
|
||||
boxes_xyxy (np.ndarray): A numpy array of shape `(N, 4)` where each
|
||||
row corresponds to a bounding box in
|
||||
the format `(x_min, y_min, x_max, y_max)`.
|
||||
frame_resolution_wh (Tuple[int, int]): A tuple of the form `(width, height)` representing the resolution of the
|
||||
frame.
|
||||
frame_resolution_wh (Tuple[int, int]): A tuple of the form `(width, height)`
|
||||
representing the resolution of the frame.
|
||||
|
||||
Returns:
|
||||
np.ndarray: A numpy array of shape `(N, 4)` where each row corresponds to a bounding box with coordinates
|
||||
clipped to fit within the frame resolution.
|
||||
np.ndarray: A numpy array of shape `(N, 4)` where each row
|
||||
corresponds to a bounding box with coordinates clipped to fit
|
||||
within the frame resolution.
|
||||
"""
|
||||
result = np.copy(boxes_xyxy)
|
||||
width, height = frame_resolution_wh
|
||||
@@ -131,10 +147,12 @@ def mask_to_xyxy(masks: np.ndarray) -> np.ndarray:
|
||||
Converts a 3D `np.array` of 2D bool masks into a 2D `np.array` of bounding boxes.
|
||||
|
||||
Parameters:
|
||||
masks (np.ndarray): A 3D `np.array` of shape `(N, W, H)` containing 2D bool masks
|
||||
masks (np.ndarray): A 3D `np.array` of shape `(N, W, H)`
|
||||
containing 2D bool masks
|
||||
|
||||
Returns:
|
||||
np.ndarray: A 2D `np.array` of shape `(N, 4)` containing the bounding boxes `(x_min, y_min, x_max, y_max)` for each mask
|
||||
np.ndarray: A 2D `np.array` of shape `(N, 4)` containing the bounding boxes
|
||||
`(x_min, y_min, x_max, y_max)` for each mask
|
||||
"""
|
||||
n = masks.shape[0]
|
||||
bboxes = np.zeros((n, 4), dtype=int)
|
||||
@@ -155,12 +173,14 @@ def mask_to_polygons(mask: np.ndarray) -> List[np.ndarray]:
|
||||
Converts a binary mask to a list of polygons.
|
||||
|
||||
Parameters:
|
||||
mask (np.ndarray): A binary mask represented as a 2D NumPy array of shape `(H, W)`,
|
||||
where H and W are the height and width of the mask, respectively.
|
||||
mask (np.ndarray): A binary mask represented as a 2D NumPy array of
|
||||
shape `(H, W)`, where H and W are the height and width of
|
||||
the mask, respectively.
|
||||
|
||||
Returns:
|
||||
List[np.ndarray]: A list of polygons, where each polygon is represented by a NumPy array of shape `(N, 2)`,
|
||||
containing the `x`, `y` coordinates of the points. Polygons with fewer points than `MIN_POLYGON_POINT_COUNT = 3`
|
||||
List[np.ndarray]: A list of polygons, where each polygon is represented by a
|
||||
NumPy array of shape `(N, 2)`, containing the `x`, `y` coordinates
|
||||
of the points. Polygons with fewer points than `MIN_POLYGON_POINT_COUNT = 3`
|
||||
are excluded from the output.
|
||||
"""
|
||||
|
||||
@@ -183,15 +203,21 @@ def filter_polygons_by_area(
|
||||
Filters a list of polygons based on their area.
|
||||
|
||||
Parameters:
|
||||
polygons (List[np.ndarray]): A list of polygons, where each polygon is represented by a NumPy array of shape `(N, 2)`,
|
||||
polygons (List[np.ndarray]): A list of polygons, where each polygon is
|
||||
represented by a NumPy array of shape `(N, 2)`,
|
||||
containing the `x`, `y` coordinates of the points.
|
||||
min_area (Optional[float]): The minimum area threshold. Only polygons with an area greater than or equal to this value
|
||||
will be included in the output. If set to None, no minimum area constraint will be applied.
|
||||
max_area (Optional[float]): The maximum area threshold. Only polygons with an area less than or equal to this value
|
||||
will be included in the output. If set to None, no maximum area constraint will be applied.
|
||||
min_area (Optional[float]): The minimum area threshold.
|
||||
Only polygons with an area greater than or equal to this value
|
||||
will be included in the output. If set to None,
|
||||
no minimum area constraint will be applied.
|
||||
max_area (Optional[float]): The maximum area threshold.
|
||||
Only polygons with an area less than or equal to this value
|
||||
will be included in the output. If set to None,
|
||||
no maximum area constraint will be applied.
|
||||
|
||||
Returns:
|
||||
List[np.ndarray]: A new list of polygons containing only those with areas within the specified thresholds.
|
||||
List[np.ndarray]: A new list of polygons containing only those with
|
||||
areas within the specified thresholds.
|
||||
"""
|
||||
if min_area is None and max_area is None:
|
||||
return polygons
|
||||
@@ -213,7 +239,8 @@ def polygon_to_xyxy(polygon: np.ndarray) -> np.ndarray:
|
||||
containing the `x`, `y` coordinates of the points.
|
||||
|
||||
Returns:
|
||||
np.ndarray: A 1D NumPy array containing the bounding box `(x_min, y_min, x_max, y_max)` of the input polygon.
|
||||
np.ndarray: A 1D NumPy array containing the bounding box
|
||||
`(x_min, y_min, x_max, y_max)` of the input polygon.
|
||||
"""
|
||||
x_min, y_min = np.min(polygon, axis=0)
|
||||
x_max, y_max = np.max(polygon, axis=0)
|
||||
@@ -226,16 +253,23 @@ def approximate_polygon(
|
||||
"""
|
||||
Approximates a given polygon by reducing a certain percentage of points.
|
||||
|
||||
This function uses the Ramer-Douglas-Peucker algorithm to simplify the input polygon by reducing the number of points
|
||||
while preserving the general shape.
|
||||
This function uses the Ramer-Douglas-Peucker algorithm to simplify the input
|
||||
polygon by reducing the number of points
|
||||
while preserving the general shape.
|
||||
|
||||
Parameters:
|
||||
polygon (np.ndarray): A 2D NumPy array of shape `(N, 2)` containing the `x`, `y` coordinates of the input polygon's points.
|
||||
percentage (float): The percentage of points to be removed from the input polygon, in the range `[0, 1)`.
|
||||
epsilon_step (float): Approximation accuracy step. Epsilon is the maximum distance between the original curve and its approximation.
|
||||
polygon (np.ndarray): A 2D NumPy array of shape `(N, 2)` containing
|
||||
the `x`, `y` coordinates of the input polygon's points.
|
||||
percentage (float): The percentage of points to be removed from the
|
||||
input polygon, in the range `[0, 1)`.
|
||||
epsilon_step (float): Approximation accuracy step.
|
||||
Epsilon is the maximum distance between the original curve
|
||||
and its approximation.
|
||||
|
||||
Returns:
|
||||
np.ndarray: A new 2D NumPy array of shape `(M, 2)`, where `M <= N * (1 - percentage)`, containing the `x`, `y` coordinates of the
|
||||
np.ndarray: A new 2D NumPy array of shape `(M, 2)`,
|
||||
where `M <= N * (1 - percentage)`, containing
|
||||
the `x`, `y` coordinates of the
|
||||
approximated polygon's points.
|
||||
"""
|
||||
|
||||
@@ -260,7 +294,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 +322,70 @@ 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)
|
||||
|
||||
|
||||
def process_roboflow_result(
|
||||
roboflow_result: dict,
|
||||
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, Optional[np.ndarray]]:
|
||||
if not roboflow_result["predictions"]:
|
||||
return np.empty((0, 4)), np.empty(0), np.empty(0), None
|
||||
|
||||
xyxy = []
|
||||
confidence = []
|
||||
class_id = []
|
||||
masks = []
|
||||
|
||||
image_width = int(roboflow_result["image"]["width"])
|
||||
image_height = int(roboflow_result["image"]["height"])
|
||||
|
||||
for prediction in roboflow_result["predictions"]:
|
||||
x = prediction["x"]
|
||||
y = prediction["y"]
|
||||
width = prediction["width"]
|
||||
height = prediction["height"]
|
||||
x_min = x - width / 2
|
||||
y_min = y - height / 2
|
||||
x_max = x_min + width
|
||||
y_max = y_min + height
|
||||
|
||||
if "points" not in prediction:
|
||||
xyxy.append([x_min, y_min, x_max, y_max])
|
||||
class_id.append(prediction["class_id"])
|
||||
confidence.append(prediction["confidence"])
|
||||
elif len(prediction["points"]) >= 3:
|
||||
polygon = np.array(
|
||||
[[point["x"], point["y"]] for point in prediction["points"]], dtype=int
|
||||
)
|
||||
mask = polygon_to_mask(polygon, resolution_wh=(image_width, image_height))
|
||||
xyxy.append([x_min, y_min, x_max, y_max])
|
||||
class_id.append(prediction["class_id"])
|
||||
confidence.append(prediction["confidence"])
|
||||
masks.append(mask)
|
||||
|
||||
xyxy = np.array(xyxy) if len(xyxy) > 0 else np.empty((0, 4))
|
||||
confidence = np.array(confidence) if len(confidence) > 0 else np.empty(0)
|
||||
class_id = np.array(class_id).astype(int) if len(class_id) > 0 else np.empty(0)
|
||||
masks = np.array(masks, dtype=bool) if len(masks) > 0 else None
|
||||
|
||||
return xyxy, confidence, class_id, masks
|
||||
|
||||
|
||||
def move_boxes(xyxy: np.ndarray, offset: np.ndarray) -> np.ndarray:
|
||||
"""
|
||||
Args:
|
||||
xyxy (np.ndarray): An array of shape `(n, 4)` containing the bounding boxes
|
||||
coordinates in format `[x1, y1, x2, y2]`
|
||||
offset (np.array): An array of shape `(2,)` containing offset values in format
|
||||
is `[dx, dy]`.
|
||||
|
||||
Returns:
|
||||
(np.ndarray) repositioned bounding boxes
|
||||
"""
|
||||
return xyxy + np.hstack([offset, offset])
|
||||
|
||||
@@ -42,7 +42,8 @@ class Color:
|
||||
"""
|
||||
Creates a Color instance from a color hex string
|
||||
|
||||
:param color_hex: str : The color hex string in the format of "fff", "ffffff", "#fff", or "#ffffff"
|
||||
:param color_hex: str : The color hex string in the format
|
||||
of "fff", "ffffff", "#fff", or "#ffffff"
|
||||
:return: Color : A Color instance representing the color
|
||||
|
||||
Example:
|
||||
@@ -105,7 +106,8 @@ class ColorPalette:
|
||||
"""
|
||||
Creates a ColorPalette instance from a list of color hex strings
|
||||
|
||||
:param color_hex_list: List[str] : A list of color hex strings in the format of "fff", "ffffff", "#fff", or "#ffffff"
|
||||
:param color_hex_list: List[str] : A list of color hex strings in the
|
||||
format of "fff", "ffffff", "#fff", or "#ffffff"
|
||||
:return: ColorPalette : A ColorPalette instance representing the color palette
|
||||
|
||||
Example:
|
||||
|
||||
@@ -115,15 +115,19 @@ def draw_text(
|
||||
Draw text with background on a scene.
|
||||
|
||||
Parameters:
|
||||
scene (np.ndarray): A 2-dimensional numpy ndarray representing an image or scene.
|
||||
scene (np.ndarray): A 2-dimensional numpy ndarray representing an image or scene
|
||||
text (str): The text to be drawn.
|
||||
text_anchor (Point): The anchor point for the text, represented as a Point object with x and y attributes.
|
||||
text_anchor (Point): The anchor point for the text, represented as a
|
||||
Point object with x and y attributes.
|
||||
text_color (Color, optional): The color of the text. Defaults to black.
|
||||
text_scale (float, optional): The scale of the text. Defaults to 0.5.
|
||||
text_thickness (int, optional): The thickness of the text. Defaults to 1.
|
||||
text_padding (int, optional): The amount of padding to add around the text when drawing a rectangle in the background. Defaults to 10.
|
||||
text_font (int, optional): The font to use for the text. Defaults to cv2.FONT_HERSHEY_SIMPLEX.
|
||||
background_color (Color, optional): The color of the background rectangle, if one is to be drawn. Defaults to None.
|
||||
text_padding (int, optional): The amount of padding to add around the text
|
||||
when drawing a rectangle in the background. Defaults to 10.
|
||||
text_font (int, optional): The font to use for the text.
|
||||
Defaults to cv2.FONT_HERSHEY_SIMPLEX.
|
||||
background_color (Color, optional): The color of the background rectangle,
|
||||
if one is to be drawn. Defaults to None.
|
||||
|
||||
Returns:
|
||||
np.ndarray: The input scene with the text drawn on it.
|
||||
@@ -132,7 +136,7 @@ def draw_text(
|
||||
```python
|
||||
>>> scene = np.zeros((100, 100, 3), dtype=np.uint8)
|
||||
>>> text_anchor = Point(x=50, y=50)
|
||||
>>> scene = draw_text(scene=scene, text="Hello, world!", text_anchor=text_anchor)
|
||||
>>> scene = draw_text(scene=scene, text="Hello, world!",text_anchor=text_anchor)
|
||||
```
|
||||
"""
|
||||
text_width, text_height = cv2.getTextSize(
|
||||
|
||||
@@ -6,8 +6,19 @@ from typing import Tuple
|
||||
|
||||
|
||||
class Position(Enum):
|
||||
"""
|
||||
Enum representing the position of an anchor point.
|
||||
"""
|
||||
|
||||
CENTER = "CENTER"
|
||||
CENTER_LEFT = "CENTER_LEFT"
|
||||
CENTER_RIGHT = "CENTER_RIGHT"
|
||||
TOP_CENTER = "TOP_CENTER"
|
||||
TOP_LEFT = "TOP_LEFT"
|
||||
TOP_RIGHT = "TOP_RIGHT"
|
||||
BOTTOM_LEFT = "BOTTOM_LEFT"
|
||||
BOTTOM_CENTER = "BOTTOM_CENTER"
|
||||
BOTTOM_RIGHT = "BOTTOM_RIGHT"
|
||||
|
||||
@classmethod
|
||||
def list(cls):
|
||||
|
||||
@@ -7,13 +7,18 @@ def get_polygon_center(polygon: np.ndarray) -> Point:
|
||||
"""
|
||||
Calculate the center of a polygon.
|
||||
|
||||
This function takes in a polygon as a 2-dimensional numpy ndarray and returns the center of the polygon as a Point object. The center is calculated as the mean of the polygon's vertices along each axis, and is rounded down to the nearest integer.
|
||||
This function takes in a polygon as a 2-dimensional numpy ndarray and
|
||||
returns the center of the polygon as a Point object.
|
||||
The center is calculated as the mean of the polygon's vertices along each axis,
|
||||
and is rounded down to the nearest integer.
|
||||
|
||||
Parameters:
|
||||
polygon (np.ndarray): A 2-dimensional numpy ndarray representing the vertices of the polygon.
|
||||
polygon (np.ndarray): A 2-dimensional numpy ndarray representing the
|
||||
vertices of the polygon.
|
||||
|
||||
Returns:
|
||||
Point: The center of the polygon, represented as a Point object with x and y attributes.
|
||||
Point: The center of the polygon, represented as a
|
||||
Point object with x and y attributes.
|
||||
|
||||
Examples:
|
||||
```python
|
||||
|
||||
@@ -0,0 +1,843 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Callable, List, Optional, Tuple
|
||||
|
||||
import matplotlib
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
|
||||
from supervision.dataset.core import DetectionDataset
|
||||
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"
|
||||
f"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."
|
||||
f"Got {type(predictions[0])} and {type(targets[0])} instead."
|
||||
)
|
||||
if predictions[0].shape[1] != 6:
|
||||
raise ValueError(
|
||||
f"Predictions must have shape (N, 6)."
|
||||
f"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:
|
||||
"""
|
||||
Confusion matrix for object detection tasks.
|
||||
|
||||
Attributes:
|
||||
matrix (np.ndarray): An 2D `np.ndarray` of shape
|
||||
`(len(classes) + 1, len(classes) + 1)`
|
||||
containing the number of `TP`, `FP`, `FN` and `TN` for each class.
|
||||
classes (List[str]): Model class names.
|
||||
conf_threshold (float): Detection confidence threshold between `0` and `1`.
|
||||
Detections with lower confidence will be excluded from the matrix.
|
||||
iou_threshold (float): Detection IoU threshold between `0` and `1`.
|
||||
Detections with lower IoU will be classified as `FP`.
|
||||
"""
|
||||
|
||||
matrix: np.ndarray
|
||||
classes: List[str]
|
||||
conf_threshold: float
|
||||
iou_threshold: float
|
||||
|
||||
@classmethod
|
||||
def from_detections(
|
||||
cls,
|
||||
predictions: List[Detections],
|
||||
targets: List[Detections],
|
||||
classes: List[str],
|
||||
conf_threshold: float = 0.3,
|
||||
iou_threshold: float = 0.5,
|
||||
) -> ConfusionMatrix:
|
||||
"""
|
||||
Calculate confusion matrix 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.
|
||||
classes (List[str]): Model class names.
|
||||
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`.
|
||||
|
||||
Returns:
|
||||
ConfusionMatrix: New instance of ConfusionMatrix.
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> targets = [
|
||||
... sv.Detections(...),
|
||||
... sv.Detections(...)
|
||||
... ]
|
||||
|
||||
>>> predictions = [
|
||||
... sv.Detections(...),
|
||||
... sv.Detections(...)
|
||||
... ]
|
||||
|
||||
>>> confusion_matrix = sv.ConfusionMatrix.from_detections(
|
||||
... predictions=predictions,
|
||||
... targets=target,
|
||||
... classes=['person', ...]
|
||||
... )
|
||||
|
||||
>>> confusion_matrix.matrix
|
||||
array([
|
||||
[0., 0., 0., 0.],
|
||||
[0., 1., 0., 1.],
|
||||
[0., 1., 1., 0.],
|
||||
[1., 1., 0., 0.]
|
||||
])
|
||||
```
|
||||
"""
|
||||
|
||||
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,
|
||||
classes=classes,
|
||||
conf_threshold=conf_threshold,
|
||||
iou_threshold=iou_threshold,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_tensors(
|
||||
cls,
|
||||
predictions: List[np.ndarray],
|
||||
targets: List[np.ndarray],
|
||||
classes: List[str],
|
||||
conf_threshold: float = 0.3,
|
||||
iou_threshold: float = 0.5,
|
||||
) -> ConfusionMatrix:
|
||||
"""
|
||||
Calculate confusion matrix based on predicted and ground-truth detections.
|
||||
|
||||
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.
|
||||
classes (List[str]): Model class names.
|
||||
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`.
|
||||
|
||||
Returns:
|
||||
ConfusionMatrix: New instance of ConfusionMatrix.
|
||||
|
||||
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]])
|
||||
... ]
|
||||
|
||||
>>> confusion_matrix = sv.ConfusionMatrix.from_tensors(
|
||||
... predictions=predictions,
|
||||
... targets=targets,
|
||||
... classes=['person', ...]
|
||||
... )
|
||||
|
||||
>>> confusion_matrix.matrix
|
||||
array([
|
||||
[0., 0., 0., 0.],
|
||||
[0., 1., 0., 1.],
|
||||
[0., 1., 1., 0.],
|
||||
[1., 1., 0., 0.]
|
||||
])
|
||||
```
|
||||
"""
|
||||
validate_input_tensors(predictions, targets)
|
||||
|
||||
num_classes = len(classes)
|
||||
matrix = np.zeros((num_classes + 1, num_classes + 1))
|
||||
for true_batch, detection_batch in zip(targets, predictions):
|
||||
matrix += cls.evaluate_detection_batch(
|
||||
predictions=detection_batch,
|
||||
targets=true_batch,
|
||||
num_classes=num_classes,
|
||||
conf_threshold=conf_threshold,
|
||||
iou_threshold=iou_threshold,
|
||||
)
|
||||
return cls(
|
||||
matrix=matrix,
|
||||
classes=classes,
|
||||
conf_threshold=conf_threshold,
|
||||
iou_threshold=iou_threshold,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def evaluate_detection_batch(
|
||||
predictions: np.ndarray,
|
||||
targets: np.ndarray,
|
||||
num_classes: int,
|
||||
conf_threshold: float,
|
||||
iou_threshold: float,
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
Calculate confusion matrix for a batch of detections for a single image.
|
||||
|
||||
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.
|
||||
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`.
|
||||
|
||||
Returns:
|
||||
np.ndarray: Confusion matrix based on a single image.
|
||||
"""
|
||||
result_matrix = np.zeros((num_classes + 1, num_classes + 1))
|
||||
|
||||
conf_idx = 5
|
||||
confidence = predictions[:, conf_idx]
|
||||
detection_batch_filtered = predictions[confidence > conf_threshold]
|
||||
|
||||
class_id_idx = 4
|
||||
true_classes = np.array(targets[:, class_id_idx], dtype=np.int16)
|
||||
detection_classes = np.array(
|
||||
detection_batch_filtered[:, class_id_idx], dtype=np.int16
|
||||
)
|
||||
true_boxes = targets[:, :class_id_idx]
|
||||
detection_boxes = detection_batch_filtered[:, :class_id_idx]
|
||||
|
||||
iou_batch = box_iou_batch(
|
||||
boxes_true=true_boxes, boxes_detection=detection_boxes
|
||||
)
|
||||
matched_idx = np.asarray(iou_batch > iou_threshold).nonzero()
|
||||
|
||||
if matched_idx[0].shape[0]:
|
||||
matches = np.stack(
|
||||
(matched_idx[0], matched_idx[1], iou_batch[matched_idx]), axis=1
|
||||
)
|
||||
matches = ConfusionMatrix._drop_extra_matches(matches=matches)
|
||||
else:
|
||||
matches = np.zeros((0, 3))
|
||||
|
||||
matched_true_idx, matched_detection_idx, _ = matches.transpose().astype(
|
||||
np.int16
|
||||
)
|
||||
|
||||
for i, true_class_value in enumerate(true_classes):
|
||||
j = matched_true_idx == i
|
||||
if matches.shape[0] > 0 and sum(j) == 1:
|
||||
result_matrix[
|
||||
true_class_value, detection_classes[matched_detection_idx[j]]
|
||||
] += 1 # TP
|
||||
else:
|
||||
result_matrix[true_class_value, num_classes] += 1 # FN
|
||||
|
||||
for i, detection_class_value in enumerate(detection_classes):
|
||||
if not any(matched_detection_idx == i):
|
||||
result_matrix[num_classes, detection_class_value] += 1 # FP
|
||||
|
||||
return result_matrix
|
||||
|
||||
@staticmethod
|
||||
def _drop_extra_matches(matches: np.ndarray) -> np.ndarray:
|
||||
"""
|
||||
Deduplicate matches. If there are multiple matches for the same true or
|
||||
predicted box, only the one with the highest IoU is kept.
|
||||
"""
|
||||
if matches.shape[0] > 0:
|
||||
matches = matches[matches[:, 2].argsort()[::-1]]
|
||||
matches = matches[np.unique(matches[:, 1], return_index=True)[1]]
|
||||
matches = matches[matches[:, 2].argsort()[::-1]]
|
||||
matches = matches[np.unique(matches[:, 0], return_index=True)[1]]
|
||||
return matches
|
||||
|
||||
@classmethod
|
||||
def benchmark(
|
||||
cls,
|
||||
dataset: DetectionDataset,
|
||||
callback: Callable[[np.ndarray], Detections],
|
||||
conf_threshold: float = 0.3,
|
||||
iou_threshold: float = 0.5,
|
||||
) -> ConfusionMatrix:
|
||||
"""
|
||||
Calculate confusion matrix 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.
|
||||
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`.
|
||||
|
||||
Returns:
|
||||
ConfusionMatrix: New instance of ConfusionMatrix.
|
||||
|
||||
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_ultralytics(result)
|
||||
|
||||
>>> confusion_matrix = sv.ConfusionMatrix.benchmark(
|
||||
... dataset = dataset,
|
||||
... callback = callback
|
||||
... )
|
||||
|
||||
>>> confusion_matrix.matrix
|
||||
array([
|
||||
[0., 0., 0., 0.],
|
||||
[0., 1., 0., 1.],
|
||||
[0., 1., 1., 0.],
|
||||
[1., 1., 0., 0.]
|
||||
])
|
||||
```
|
||||
"""
|
||||
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,
|
||||
classes=dataset.classes,
|
||||
conf_threshold=conf_threshold,
|
||||
iou_threshold=iou_threshold,
|
||||
)
|
||||
|
||||
def plot(
|
||||
self,
|
||||
save_path: Optional[str] = None,
|
||||
title: Optional[str] = None,
|
||||
classes: Optional[List[str]] = None,
|
||||
normalize: bool = False,
|
||||
fig_size: Tuple[int, int] = (12, 10),
|
||||
) -> matplotlib.figure.Figure:
|
||||
"""
|
||||
Create confusion matrix plot and save it at selected location.
|
||||
|
||||
Args:
|
||||
save_path (Optional[str]): Path to save the plot. If not provided,
|
||||
plot will be displayed.
|
||||
title (Optional[str]): Title of the plot.
|
||||
classes (Optional[List[str]]): List of classes to be displayed on the plot.
|
||||
If not provided, all classes will be displayed.
|
||||
normalize (bool): If True, normalize the confusion matrix.
|
||||
fig_size (Tuple[int, int]): Size of the plot.
|
||||
|
||||
Returns:
|
||||
matplotlib.figure.Figure: Confusion matrix plot.
|
||||
"""
|
||||
|
||||
array = self.matrix.copy()
|
||||
|
||||
if normalize:
|
||||
eps = 1e-8
|
||||
array = array / (array.sum(0).reshape(1, -1) + eps)
|
||||
|
||||
array[array < 0.005] = np.nan
|
||||
|
||||
fig, ax = plt.subplots(figsize=fig_size, tight_layout=True, facecolor="white")
|
||||
|
||||
class_names = classes if classes is not None else self.classes
|
||||
use_labels_for_ticks = class_names is not None and (0 < len(class_names) < 99)
|
||||
if use_labels_for_ticks:
|
||||
x_tick_labels = class_names + ["FN"]
|
||||
y_tick_labels = class_names + ["FP"]
|
||||
num_ticks = len(x_tick_labels)
|
||||
else:
|
||||
x_tick_labels = None
|
||||
y_tick_labels = None
|
||||
num_ticks = len(array)
|
||||
im = ax.imshow(array, cmap="Blues")
|
||||
|
||||
cbar = ax.figure.colorbar(im, ax=ax)
|
||||
cbar.mappable.set_clim(vmin=0, vmax=np.nanmax(array))
|
||||
|
||||
if x_tick_labels is None:
|
||||
tick_interval = 2
|
||||
else:
|
||||
tick_interval = 1
|
||||
ax.set_xticks(np.arange(0, num_ticks, tick_interval), labels=x_tick_labels)
|
||||
ax.set_yticks(np.arange(0, num_ticks, tick_interval), labels=y_tick_labels)
|
||||
|
||||
plt.setp(ax.get_xticklabels(), rotation=90, ha="right", rotation_mode="default")
|
||||
|
||||
labelsize = 10 if num_ticks < 50 else 8
|
||||
ax.tick_params(axis="both", which="both", labelsize=labelsize)
|
||||
|
||||
if num_ticks < 30:
|
||||
for i in range(array.shape[0]):
|
||||
for j in range(array.shape[1]):
|
||||
n_preds = array[i, j]
|
||||
if not np.isnan(n_preds):
|
||||
ax.text(
|
||||
j,
|
||||
i,
|
||||
f"{n_preds:.2f}" if normalize else f"{n_preds:.0f}",
|
||||
ha="center",
|
||||
va="center",
|
||||
color="black"
|
||||
if n_preds < 0.5 * np.nanmax(array)
|
||||
else "white",
|
||||
)
|
||||
|
||||
if title:
|
||||
ax.set_title(title, fontsize=20)
|
||||
|
||||
ax.set_xlabel("Predicted")
|
||||
ax.set_ylabel("True")
|
||||
ax.set_facecolor("white")
|
||||
if save_path:
|
||||
fig.savefig(
|
||||
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_ultralytics(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
|
||||
@@ -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
|
||||
@@ -0,0 +1,475 @@
|
||||
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 = []
|
||||
multi_covariance = []
|
||||
for i, st in enumerate(stracks):
|
||||
multi_mean.append(st.mean.copy())
|
||||
multi_covariance.append(st.covariance)
|
||||
if st.state != TrackState.Tracked:
|
||||
multi_mean[i][7] = 0
|
||||
|
||||
multi_mean, multi_covariance = STrack.shared_kalman.multi_predict(
|
||||
np.asarray(multi_mean), np.asarray(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_ultralytics(results)
|
||||
... detections = byte_tracker.update_with_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
|
||||
@@ -0,0 +1,205 @@
|
||||
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
|
||||
@@ -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
|
||||
@@ -1,16 +1,33 @@
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import List, Optional, Union
|
||||
|
||||
import numpy as np
|
||||
import yaml
|
||||
|
||||
|
||||
class NumpyJsonEncoder(json.JSONEncoder):
|
||||
def default(self, obj):
|
||||
if isinstance(obj, np.integer):
|
||||
return int(obj)
|
||||
if isinstance(obj, np.floating):
|
||||
return float(obj)
|
||||
if isinstance(obj, np.ndarray):
|
||||
return obj.tolist()
|
||||
return super(NumpyJsonEncoder, self).default(obj)
|
||||
|
||||
|
||||
def list_files_with_extensions(
|
||||
directory: Union[str, Path], extensions: Optional[List[str]] = None
|
||||
) -> List[Path]:
|
||||
"""
|
||||
List files in a directory with specified extensions or all files if no extensions are provided.
|
||||
List files in a directory with specified extensions or
|
||||
all files if no extensions are provided.
|
||||
|
||||
Args:
|
||||
directory (Union[str, Path]): The directory path as a string or Path object.
|
||||
extensions (Optional[List[str]]): A list of file extensions to filter. Default is None, which lists all files.
|
||||
extensions (Optional[List[str]]): A list of file extensions to filter.
|
||||
Default is None, which lists all files.
|
||||
|
||||
Returns:
|
||||
(List[Path]): A list of Path objects for the matching files.
|
||||
@@ -23,9 +40,11 @@ def list_files_with_extensions(
|
||||
>>> files = sv.list_files_with_extensions(directory='my_directory')
|
||||
|
||||
>>> # List only files with '.txt' and '.md' extensions
|
||||
>>> files = sv.list_files_with_extensions(directory='my_directory', extensions=['txt', 'md'])
|
||||
>>> files = sv.list_files_with_extensions(
|
||||
... directory='my_directory', extensions=['txt', 'md'])
|
||||
```
|
||||
"""
|
||||
|
||||
directory = Path(directory)
|
||||
files_with_extensions = []
|
||||
|
||||
@@ -66,3 +85,60 @@ def save_text_file(lines: List[str], file_path: str):
|
||||
with open(file_path, "w") as file:
|
||||
for line in lines:
|
||||
file.write(line + "\n")
|
||||
|
||||
|
||||
def read_json_file(file_path: str) -> dict:
|
||||
"""
|
||||
Read a json file and return a dict.
|
||||
|
||||
Args:
|
||||
file_path (str): The path to the json file.
|
||||
|
||||
Returns:
|
||||
dict: A dict of annotations information
|
||||
"""
|
||||
with open(file_path, "r") as file:
|
||||
data = json.load(file)
|
||||
return data
|
||||
|
||||
|
||||
def save_json_file(data: dict, file_path: str, indent: int = 3) -> None:
|
||||
"""
|
||||
Write a dict to a json file.
|
||||
|
||||
Args:
|
||||
indent:
|
||||
data (dict): dict with unique keys and value as pair.
|
||||
file_path (str): The path to the json file.
|
||||
"""
|
||||
with open(file_path, "w") as fp:
|
||||
json.dump(data, fp, cls=NumpyJsonEncoder, indent=indent)
|
||||
|
||||
|
||||
def read_yaml_file(file_path: str) -> dict:
|
||||
"""
|
||||
Read a yaml file and return a dict.
|
||||
|
||||
Args:
|
||||
file_path (str): The path to the yaml file.
|
||||
|
||||
Returns:
|
||||
dict: A dict of content information
|
||||
"""
|
||||
with open(file_path, "r") as file:
|
||||
data = yaml.safe_load(file)
|
||||
return data
|
||||
|
||||
|
||||
def save_yaml_file(data: dict, file_path: str) -> None:
|
||||
"""
|
||||
Save a dict to a json file.
|
||||
|
||||
Args:
|
||||
indent:
|
||||
data (dict): dict with unique keys and value as pair.
|
||||
file_path (str): The path to the json file.
|
||||
"""
|
||||
|
||||
with open(file_path, "w") as outfile:
|
||||
yaml.dump(data, outfile, sort_keys=False, default_flow_style=None)
|
||||
|
||||
@@ -6,13 +6,14 @@ import cv2
|
||||
import numpy as np
|
||||
|
||||
|
||||
def crop(image: np.ndarray, xyxy: np.ndarray) -> np.ndarray:
|
||||
def crop_image(image: np.ndarray, xyxy: np.ndarray) -> np.ndarray:
|
||||
"""
|
||||
Crops the given image based on the given bounding box.
|
||||
|
||||
Args:
|
||||
image (np.ndarray): The image to be cropped, represented as a numpy array.
|
||||
xyxy (np.ndarray): A numpy array containing the bounding box coordinates in the format (x1, y1, x2, y2).
|
||||
xyxy (np.ndarray): A numpy array containing the bounding box coordinates
|
||||
in the format (x1, y1, x2, y2).
|
||||
|
||||
Returns:
|
||||
(np.ndarray): The cropped image as a numpy array.
|
||||
@@ -24,7 +25,7 @@ def crop(image: np.ndarray, xyxy: np.ndarray) -> np.ndarray:
|
||||
>>> detection = sv.Detections(...)
|
||||
>>> with sv.ImageSink(target_dir_path='target/directory/path') as sink:
|
||||
... for xyxy in detection.xyxy:
|
||||
... cropped_image = sv.crop(image=image, xyxy=xyxy)
|
||||
... cropped_image = sv.crop_image(image=image, xyxy=xyxy)
|
||||
... sink.save_image(image=image)
|
||||
```
|
||||
"""
|
||||
@@ -47,18 +48,23 @@ class ImageSink:
|
||||
|
||||
Args:
|
||||
target_dir_path (str): The target directory where images will be saved.
|
||||
overwrite (bool, optional): Whether to overwrite the existing directory. Defaults to False.
|
||||
image_name_pattern (str, optional): The image file name pattern. Defaults to "image_{:05d}.png".
|
||||
overwrite (bool, optional): Whether to overwrite the existing directory.
|
||||
Defaults to False.
|
||||
image_name_pattern (str, optional): The image file name pattern.
|
||||
Defaults to "image_{:05d}.png".
|
||||
|
||||
Examples:
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> with sv.ImageSink(target_dir_path='target/directory/path', overwrite=True) as sink:
|
||||
... for image in sv.get_video_frames_generator(source_path='source_video.mp4', stride=2):
|
||||
>>> with sv.ImageSink(target_dir_path='target/directory/path',
|
||||
... overwrite=True) as sink:
|
||||
... for image in sv.get_video_frames_generator(
|
||||
... source_path='source_video.mp4', stride=2):
|
||||
... sink.save_image(image=image)
|
||||
```
|
||||
"""
|
||||
|
||||
self.target_dir_path = target_dir_path
|
||||
self.overwrite = overwrite
|
||||
self.image_name_pattern = image_name_pattern
|
||||
@@ -80,7 +86,9 @@ class ImageSink:
|
||||
|
||||
Args:
|
||||
image (np.ndarray): The image to be saved.
|
||||
image_name (str, optional): The name to use for the saved image. If not provided, a name will be generated using the `image_name_pattern`.
|
||||
image_name (str, optional): The name to use for the saved image.
|
||||
If not provided, a name will be
|
||||
generated using the `image_name_pattern`.
|
||||
"""
|
||||
if image_name is None:
|
||||
image_name = self.image_name_pattern.format(self.image_count)
|
||||
|
||||
@@ -50,9 +50,12 @@ def plot_images_grid(
|
||||
|
||||
Args:
|
||||
images (List[np.ndarray]): A list of images as numpy arrays.
|
||||
grid_size (Tuple[int, int]): A tuple specifying the number of rows and columns for the grid.
|
||||
titles (Optional[List[str]]): A list of titles for each image. Defaults to None.
|
||||
size (Tuple[int, int]): A tuple specifying the width and height of the entire plot in inches.
|
||||
grid_size (Tuple[int, int]): A tuple specifying the number
|
||||
of rows and columns for the grid.
|
||||
titles (Optional[List[str]]): A list of titles for each image.
|
||||
Defaults to None.
|
||||
size (Tuple[int, int]): A tuple specifying the width and
|
||||
height of the entire plot in inches.
|
||||
cmap (str): the colormap to use for single channel images.
|
||||
|
||||
Raises:
|
||||
@@ -78,7 +81,8 @@ def plot_images_grid(
|
||||
|
||||
if len(images) > nrows * ncols:
|
||||
raise ValueError(
|
||||
"The number of images exceeds the grid size. Please increase the grid size or reduce the number of images."
|
||||
"The number of images exceeds the grid size. Please increase the grid size"
|
||||
" or reduce the number of images."
|
||||
)
|
||||
|
||||
fig, axes = plt.subplots(nrows=nrows, ncols=ncols, figsize=size)
|
||||
|
||||
+64
-29
@@ -10,13 +10,15 @@ import numpy as np
|
||||
@dataclass
|
||||
class VideoInfo:
|
||||
"""
|
||||
A class to store video information, including width, height, fps and total number of frames.
|
||||
A class to store video information, including width, height, fps and
|
||||
total number of frames.
|
||||
|
||||
Attributes:
|
||||
width (int): width of the video in pixels
|
||||
height (int): height of the video in pixels
|
||||
fps (int): frames per second of the video
|
||||
total_frames (int, optional): total number of frames in the video, default is None
|
||||
total_frames (int, optional): total number of frames in the video,
|
||||
default is None
|
||||
|
||||
Examples:
|
||||
```python
|
||||
@@ -61,7 +63,9 @@ class VideoSink:
|
||||
|
||||
Attributes:
|
||||
target_path (str): The path to the output file where the video will be saved.
|
||||
video_info (VideoInfo): Information about the video resolution, fps, and total frame count.
|
||||
video_info (VideoInfo): Information about the video resolution, fps,
|
||||
and total frame count.
|
||||
codec (str): FOURCC code for video format
|
||||
|
||||
Example:
|
||||
```python
|
||||
@@ -69,19 +73,27 @@ class VideoSink:
|
||||
|
||||
>>> video_info = sv.VideoInfo.from_video_path(video_path='source_video.mp4')
|
||||
|
||||
>>> with sv.VideoSink(target_path='target_video.mp4', video_info=video_info) as sink:
|
||||
... for frame in get_video_frames_generator(source_path='source_video.mp4', stride=2):
|
||||
>>> with sv.VideoSink(target_path='target_video.mp4',
|
||||
... video_info=video_info,
|
||||
codec='H264') as sink:
|
||||
... for frame in get_video_frames_generator(source_path='source_video.mp4',
|
||||
... stride=2):
|
||||
... sink.write_frame(frame=frame)
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(self, target_path: str, video_info: VideoInfo):
|
||||
def __init__(self, target_path: str, video_info: VideoInfo, codec: str = "mp4v"):
|
||||
self.target_path = target_path
|
||||
self.video_info = video_info
|
||||
self.__fourcc = cv2.VideoWriter_fourcc(*"mp4v")
|
||||
self.__codec = codec
|
||||
self.__writer = None
|
||||
|
||||
def __enter__(self):
|
||||
try:
|
||||
self.__fourcc = cv2.VideoWriter_fourcc(*self.__codec)
|
||||
except TypeError as e:
|
||||
print(str(e) + ". Defaulting to mp4v...")
|
||||
self.__fourcc = cv2.VideoWriter_fourcc(*"mp4v")
|
||||
self.__writer = cv2.VideoWriter(
|
||||
self.target_path,
|
||||
self.__fourcc,
|
||||
@@ -97,39 +109,58 @@ class VideoSink:
|
||||
self.__writer.release()
|
||||
|
||||
|
||||
def _validate_and_setup_video(source_path: str, start: int, end: Optional[int]):
|
||||
video = cv2.VideoCapture(source_path)
|
||||
if not video.isOpened():
|
||||
raise Exception(f"Could not open video at {source_path}")
|
||||
total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
|
||||
if end is not None and end > total_frames:
|
||||
raise Exception("Requested frames are outbound")
|
||||
start = max(start, 0)
|
||||
end = min(end, total_frames) if end is not None else total_frames
|
||||
video.set(cv2.CAP_PROP_POS_FRAMES, start)
|
||||
return video, start, end
|
||||
|
||||
|
||||
def get_video_frames_generator(
|
||||
source_path: str, stride: int = 1
|
||||
source_path: str, stride: int = 1, start: int = 0, end: Optional[int] = None
|
||||
) -> Generator[np.ndarray, None, None]:
|
||||
"""
|
||||
Get a generator that yields the frames of the video.
|
||||
|
||||
Args:
|
||||
source_path (str): The path of the video file.
|
||||
stride (int): The number of frames to skip before returning the next one.
|
||||
stride (int): Indicates the interval at which frames are returned,
|
||||
skipping stride - 1 frames between each.
|
||||
start (int): Indicates the starting position from which
|
||||
video should generate frames
|
||||
end (Optional[int]): Indicates the ending position at which video
|
||||
should stop generating frames. If None, video will be read to the end.
|
||||
|
||||
Returns:
|
||||
(Generator[np.ndarray, None, None]): A generator that yields the frames of the video.
|
||||
(Generator[np.ndarray, None, None]): A generator that yields the
|
||||
frames of the video.
|
||||
|
||||
Examples:
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> for frame in sv.get_video_frames_generator(source_path='source_video.mp4', stride=2):
|
||||
>>> for frame in sv.get_video_frames_generator(source_path='source_video.mp4'):
|
||||
... ...
|
||||
```
|
||||
"""
|
||||
video = cv2.VideoCapture(source_path)
|
||||
if not video.isOpened():
|
||||
raise Exception(f"Could not open video at {source_path}")
|
||||
|
||||
frame_count = 0
|
||||
success, frame = video.read()
|
||||
while success:
|
||||
if frame_count % stride == 0:
|
||||
yield frame
|
||||
video, start, end = _validate_and_setup_video(source_path, start, end)
|
||||
frame_position = start
|
||||
while True:
|
||||
success, frame = video.read()
|
||||
frame_count += 1
|
||||
|
||||
if not success or frame_position >= end:
|
||||
break
|
||||
yield frame
|
||||
for _ in range(stride - 1):
|
||||
success = video.grab()
|
||||
if not success:
|
||||
break
|
||||
frame_position += stride
|
||||
video.release()
|
||||
|
||||
|
||||
@@ -139,24 +170,28 @@ def process_video(
|
||||
callback: Callable[[np.ndarray, int], np.ndarray],
|
||||
) -> None:
|
||||
"""
|
||||
Process a video file by applying a callback function on each frame and saving the result to a target video file.
|
||||
Process a video file by applying a callback function on each frame
|
||||
and saving the result to a target video file.
|
||||
|
||||
Args:
|
||||
source_path (str): The path to the source video file.
|
||||
target_path (str): The path to the target video file.
|
||||
callback (Callable[[np.ndarray, int], np.ndarray]): A function that takes in a numpy ndarray representation of a video frame and an int index of the frame and returns a processed numpy ndarray representation of the frame.
|
||||
callback (Callable[[np.ndarray, int], np.ndarray]): A function that takes in
|
||||
a numpy ndarray representation of a video frame and an
|
||||
int index of the frame and returns a processed numpy ndarray
|
||||
representation of the frame.
|
||||
|
||||
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
|
||||
... )
|
||||
```
|
||||
"""
|
||||
|
||||
@@ -0,0 +1,94 @@
|
||||
from contextlib import ExitStack as DoesNotRaise
|
||||
from test.utils import mock_detections
|
||||
from typing import Optional
|
||||
|
||||
import pytest
|
||||
|
||||
from supervision.annotators.utils import ColorMap, resolve_color_idx
|
||||
from supervision.detection.core import Detections
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"detections, detection_idx, color_map, expected_result, exception",
|
||||
[
|
||||
(
|
||||
mock_detections(
|
||||
xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]],
|
||||
class_id=[5, 3],
|
||||
tracker_id=[2, 6],
|
||||
),
|
||||
0,
|
||||
ColorMap.INDEX,
|
||||
0,
|
||||
DoesNotRaise(),
|
||||
), # multiple detections; index mapping
|
||||
(
|
||||
mock_detections(
|
||||
xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]],
|
||||
class_id=[5, 3],
|
||||
tracker_id=[2, 6],
|
||||
),
|
||||
0,
|
||||
ColorMap.CLASS,
|
||||
5,
|
||||
DoesNotRaise(),
|
||||
), # multiple detections; class mapping
|
||||
(
|
||||
mock_detections(
|
||||
xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]],
|
||||
class_id=[5, 3],
|
||||
tracker_id=[2, 6],
|
||||
),
|
||||
0,
|
||||
ColorMap.TRACK,
|
||||
2,
|
||||
DoesNotRaise(),
|
||||
), # multiple detections; track mapping
|
||||
(
|
||||
Detections.empty(),
|
||||
0,
|
||||
ColorMap.INDEX,
|
||||
None,
|
||||
pytest.raises(ValueError),
|
||||
), # no detections; index mapping; out of bounds
|
||||
(
|
||||
mock_detections(
|
||||
xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]],
|
||||
class_id=[5, 3],
|
||||
tracker_id=[2, 6],
|
||||
),
|
||||
2,
|
||||
ColorMap.INDEX,
|
||||
None,
|
||||
pytest.raises(ValueError),
|
||||
), # multiple detections; index mapping; out of bounds
|
||||
(
|
||||
mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]),
|
||||
0,
|
||||
ColorMap.CLASS,
|
||||
None,
|
||||
pytest.raises(ValueError),
|
||||
), # multiple detections; class mapping; no class_id
|
||||
(
|
||||
mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]),
|
||||
0,
|
||||
ColorMap.TRACK,
|
||||
None,
|
||||
pytest.raises(ValueError),
|
||||
), # multiple detections; class mapping; no track_id
|
||||
],
|
||||
)
|
||||
def test_resolve_color_idx(
|
||||
detections: Detections,
|
||||
detection_idx: int,
|
||||
color_map: ColorMap,
|
||||
expected_result: Optional[int],
|
||||
exception: Exception,
|
||||
) -> None:
|
||||
with exception:
|
||||
result = resolve_color_idx(
|
||||
detections=detections,
|
||||
detection_idx=detection_idx,
|
||||
color_map=color_map,
|
||||
)
|
||||
assert result == expected_result
|
||||
@@ -0,0 +1,63 @@
|
||||
from contextlib import ExitStack as DoesNotRaise
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from supervision.classification.core import Classifications
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"class_id, confidence, k, expected_result, exception",
|
||||
[
|
||||
(
|
||||
np.array([0, 1, 2, 3, 4]),
|
||||
np.array([0.1, 0.2, 0.9, 0.4, 0.5]),
|
||||
5,
|
||||
(np.array([2, 4, 3, 1, 0]), np.array([0.9, 0.5, 0.4, 0.2, 0.1])),
|
||||
DoesNotRaise(),
|
||||
), # class_id with 5 numbers and 5 confidences
|
||||
(
|
||||
np.array([5, 1, 2, 3, 4]),
|
||||
np.array([0.1, 0.2, 0.9, 0.4, 0.5]),
|
||||
1,
|
||||
(np.array([2]), np.array([0.9])),
|
||||
DoesNotRaise(),
|
||||
), # class_id with 5 numbers and 5 confidences, retrieve where k = 1
|
||||
(
|
||||
np.array([4, 1, 2, 3, 6, 5]),
|
||||
np.array([0.8, 0.2, 0.9, 0.4, 0.5, 0.1]),
|
||||
2,
|
||||
(np.array([2, 4]), np.array([0.9, 0.8])),
|
||||
DoesNotRaise(),
|
||||
), # class_id with 5 numbers and 5 confidences, retrieve where k = 3
|
||||
(
|
||||
np.array([0, 1, 2, 3, 4]),
|
||||
np.array([]),
|
||||
5,
|
||||
None,
|
||||
pytest.raises(ValueError),
|
||||
), # class_id with 5 numbers and 0 confidences
|
||||
(
|
||||
[0, 1, 2, 3, 4],
|
||||
[0.1, 0.2, 0.3, 0.4],
|
||||
5,
|
||||
None,
|
||||
pytest.raises(ValueError),
|
||||
), # class_id with 5 numbers and 4 confidences
|
||||
],
|
||||
)
|
||||
def test_top_k(
|
||||
class_id: np.ndarray,
|
||||
confidence: Optional[np.ndarray],
|
||||
k: int,
|
||||
expected_result: Optional[Tuple[np.ndarray, np.ndarray]],
|
||||
exception: Exception,
|
||||
) -> None:
|
||||
with exception:
|
||||
result = Classifications(
|
||||
class_id=np.array(class_id), confidence=np.array(confidence)
|
||||
).get_top_k(k)
|
||||
|
||||
assert np.array_equal(result[0], expected_result[0])
|
||||
assert np.array_equal(result[1], expected_result[1])
|
||||
@@ -0,0 +1,303 @@
|
||||
from contextlib import ExitStack as DoesNotRaise
|
||||
from typing import Dict, List, Tuple
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from supervision import Detections
|
||||
from supervision.dataset.formats.coco import (
|
||||
build_coco_class_index_mapping,
|
||||
classes_to_coco_categories,
|
||||
coco_annotations_to_detections,
|
||||
coco_categories_to_classes,
|
||||
group_coco_annotations_by_image_id,
|
||||
)
|
||||
|
||||
|
||||
def mock_cock_coco_annotation(
|
||||
annotation_id: int = 0,
|
||||
image_id: int = 0,
|
||||
category_id: int = 0,
|
||||
bbox: Tuple[float, float, float, float] = (0.0, 0.0, 0.0, 0.0),
|
||||
area: float = 0.0,
|
||||
) -> dict:
|
||||
return {
|
||||
"id": annotation_id,
|
||||
"image_id": image_id,
|
||||
"category_id": category_id,
|
||||
"bbox": list(bbox),
|
||||
"area": area,
|
||||
"iscrowd": 0,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"coco_categories, expected_result, exception",
|
||||
[
|
||||
([], [], DoesNotRaise()), # empty coco categories
|
||||
(
|
||||
[{"id": 0, "name": "fashion-assistant", "supercategory": "none"}],
|
||||
["fashion-assistant"],
|
||||
DoesNotRaise(),
|
||||
), # single coco category with supercategory == "none"
|
||||
(
|
||||
[
|
||||
{"id": 0, "name": "fashion-assistant", "supercategory": "none"},
|
||||
{"id": 1, "name": "baseball cap", "supercategory": "fashion-assistant"},
|
||||
],
|
||||
["fashion-assistant", "baseball cap"],
|
||||
DoesNotRaise(),
|
||||
), # two coco categories; one with supercategory == "none" and
|
||||
# one with supercategory != "none"
|
||||
(
|
||||
[
|
||||
{"id": 0, "name": "fashion-assistant", "supercategory": "none"},
|
||||
{"id": 1, "name": "baseball cap", "supercategory": "fashion-assistant"},
|
||||
{"id": 2, "name": "hoodie", "supercategory": "fashion-assistant"},
|
||||
],
|
||||
["fashion-assistant", "baseball cap", "hoodie"],
|
||||
DoesNotRaise(),
|
||||
), # three coco categories; one with supercategory == "none" and
|
||||
# two with supercategory != "none"
|
||||
(
|
||||
[
|
||||
{"id": 0, "name": "fashion-assistant", "supercategory": "none"},
|
||||
{"id": 2, "name": "hoodie", "supercategory": "fashion-assistant"},
|
||||
{"id": 1, "name": "baseball cap", "supercategory": "fashion-assistant"},
|
||||
],
|
||||
["fashion-assistant", "baseball cap", "hoodie"],
|
||||
DoesNotRaise(),
|
||||
), # three coco categories; one with supercategory == "none" and
|
||||
# two with supercategory != "none" (different order)
|
||||
],
|
||||
)
|
||||
def test_coco_categories_to_classes(
|
||||
coco_categories: List[dict], expected_result: List[str], exception: Exception
|
||||
) -> None:
|
||||
with exception:
|
||||
result = coco_categories_to_classes(coco_categories=coco_categories)
|
||||
assert result == expected_result
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"classes, exception",
|
||||
[
|
||||
([], DoesNotRaise()), # empty classes
|
||||
(["baseball cap"], DoesNotRaise()), # single class
|
||||
(["baseball cap", "hoodie"], DoesNotRaise()), # two classes
|
||||
],
|
||||
)
|
||||
def test_classes_to_coco_categories_and_back_to_classes(
|
||||
classes: List[str], exception: Exception
|
||||
) -> None:
|
||||
with exception:
|
||||
coco_categories = classes_to_coco_categories(classes=classes)
|
||||
result = coco_categories_to_classes(coco_categories=coco_categories)
|
||||
assert result == classes
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"coco_annotations, expected_result, exception",
|
||||
[
|
||||
([], {}, DoesNotRaise()), # empty coco annotations
|
||||
(
|
||||
[mock_cock_coco_annotation(annotation_id=0, image_id=0, category_id=0)],
|
||||
{
|
||||
0: [
|
||||
mock_cock_coco_annotation(
|
||||
annotation_id=0, image_id=0, category_id=0
|
||||
)
|
||||
]
|
||||
},
|
||||
DoesNotRaise(),
|
||||
), # single coco annotation
|
||||
(
|
||||
[
|
||||
mock_cock_coco_annotation(annotation_id=0, image_id=0, category_id=0),
|
||||
mock_cock_coco_annotation(annotation_id=1, image_id=1, category_id=0),
|
||||
],
|
||||
{
|
||||
0: [
|
||||
mock_cock_coco_annotation(
|
||||
annotation_id=0, image_id=0, category_id=0
|
||||
)
|
||||
],
|
||||
1: [
|
||||
mock_cock_coco_annotation(
|
||||
annotation_id=1, image_id=1, category_id=0
|
||||
)
|
||||
],
|
||||
},
|
||||
DoesNotRaise(),
|
||||
), # two coco annotations
|
||||
(
|
||||
[
|
||||
mock_cock_coco_annotation(annotation_id=0, image_id=0, category_id=0),
|
||||
mock_cock_coco_annotation(annotation_id=1, image_id=1, category_id=1),
|
||||
mock_cock_coco_annotation(annotation_id=2, image_id=1, category_id=2),
|
||||
mock_cock_coco_annotation(annotation_id=3, image_id=2, category_id=3),
|
||||
mock_cock_coco_annotation(annotation_id=4, image_id=3, category_id=1),
|
||||
mock_cock_coco_annotation(annotation_id=5, image_id=3, category_id=2),
|
||||
mock_cock_coco_annotation(annotation_id=5, image_id=3, category_id=3),
|
||||
],
|
||||
{
|
||||
0: [
|
||||
mock_cock_coco_annotation(
|
||||
annotation_id=0, image_id=0, category_id=0
|
||||
),
|
||||
],
|
||||
1: [
|
||||
mock_cock_coco_annotation(
|
||||
annotation_id=1, image_id=1, category_id=1
|
||||
),
|
||||
mock_cock_coco_annotation(
|
||||
annotation_id=2, image_id=1, category_id=2
|
||||
),
|
||||
],
|
||||
2: [
|
||||
mock_cock_coco_annotation(
|
||||
annotation_id=3, image_id=2, category_id=3
|
||||
),
|
||||
],
|
||||
3: [
|
||||
mock_cock_coco_annotation(
|
||||
annotation_id=4, image_id=3, category_id=1
|
||||
),
|
||||
mock_cock_coco_annotation(
|
||||
annotation_id=5, image_id=3, category_id=2
|
||||
),
|
||||
mock_cock_coco_annotation(
|
||||
annotation_id=5, image_id=3, category_id=3
|
||||
),
|
||||
],
|
||||
},
|
||||
DoesNotRaise(),
|
||||
), # two coco annotations
|
||||
],
|
||||
)
|
||||
def test_group_coco_annotations_by_image_id(
|
||||
coco_annotations: List[dict], expected_result: dict, exception: Exception
|
||||
) -> None:
|
||||
with exception:
|
||||
result = group_coco_annotations_by_image_id(coco_annotations=coco_annotations)
|
||||
assert result == expected_result
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"image_annotations, resolution_wh, with_masks, expected_result, exception",
|
||||
[
|
||||
(
|
||||
[],
|
||||
(1000, 1000),
|
||||
False,
|
||||
Detections.empty(),
|
||||
DoesNotRaise(),
|
||||
), # empty image annotations
|
||||
(
|
||||
[
|
||||
mock_cock_coco_annotation(
|
||||
category_id=0, bbox=(0, 0, 100, 100), area=100 * 100
|
||||
)
|
||||
],
|
||||
(1000, 1000),
|
||||
False,
|
||||
Detections(
|
||||
xyxy=np.array([[0, 0, 100, 100]], dtype=np.float32),
|
||||
class_id=np.array([0], dtype=int),
|
||||
),
|
||||
DoesNotRaise(),
|
||||
), # single image annotations
|
||||
(
|
||||
[
|
||||
mock_cock_coco_annotation(
|
||||
category_id=0, bbox=(0, 0, 100, 100), area=100 * 100
|
||||
),
|
||||
mock_cock_coco_annotation(
|
||||
category_id=0, bbox=(100, 100, 100, 100), area=100 * 100
|
||||
),
|
||||
],
|
||||
(1000, 1000),
|
||||
False,
|
||||
Detections(
|
||||
xyxy=np.array(
|
||||
[[0, 0, 100, 100], [100, 100, 200, 200]], dtype=np.float32
|
||||
),
|
||||
class_id=np.array([0, 0], dtype=int),
|
||||
),
|
||||
DoesNotRaise(),
|
||||
), # two image annotations
|
||||
],
|
||||
)
|
||||
def test_coco_annotations_to_detections(
|
||||
image_annotations: List[dict],
|
||||
resolution_wh: Tuple[int, int],
|
||||
with_masks: bool,
|
||||
expected_result: Detections,
|
||||
exception: Exception,
|
||||
) -> None:
|
||||
with exception:
|
||||
result = coco_annotations_to_detections(
|
||||
image_annotations=image_annotations,
|
||||
resolution_wh=resolution_wh,
|
||||
with_masks=with_masks,
|
||||
)
|
||||
assert result == expected_result
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"coco_categories, target_classes, expected_result, exception",
|
||||
[
|
||||
([], [], {}, DoesNotRaise()), # empty coco categories
|
||||
(
|
||||
[{"id": 0, "name": "fashion-assistant", "supercategory": "none"}],
|
||||
["fashion-assistant"],
|
||||
{0: 0},
|
||||
DoesNotRaise(),
|
||||
), # single coco category starting from 0
|
||||
(
|
||||
[{"id": 1, "name": "fashion-assistant", "supercategory": "none"}],
|
||||
["fashion-assistant"],
|
||||
{1: 0},
|
||||
DoesNotRaise(),
|
||||
), # single coco category starting from 1
|
||||
(
|
||||
[
|
||||
{"id": 0, "name": "fashion-assistant", "supercategory": "none"},
|
||||
{"id": 2, "name": "hoodie", "supercategory": "fashion-assistant"},
|
||||
{"id": 1, "name": "baseball cap", "supercategory": "fashion-assistant"},
|
||||
],
|
||||
["fashion-assistant", "baseball cap", "hoodie"],
|
||||
{0: 0, 1: 1, 2: 2},
|
||||
DoesNotRaise(),
|
||||
), # three coco categories
|
||||
(
|
||||
[
|
||||
{"id": 2, "name": "hoodie", "supercategory": "fashion-assistant"},
|
||||
{"id": 1, "name": "baseball cap", "supercategory": "fashion-assistant"},
|
||||
],
|
||||
["baseball cap", "hoodie"],
|
||||
{2: 1, 1: 0},
|
||||
DoesNotRaise(),
|
||||
), # two coco categories
|
||||
(
|
||||
[
|
||||
{"id": 3, "name": "hoodie", "supercategory": "fashion-assistant"},
|
||||
{"id": 1, "name": "baseball cap", "supercategory": "fashion-assistant"},
|
||||
],
|
||||
["baseball cap", "hoodie"],
|
||||
{3: 1, 1: 0},
|
||||
DoesNotRaise(),
|
||||
), # two coco categories with missing category
|
||||
],
|
||||
)
|
||||
def test_build_coco_class_index_mapping(
|
||||
coco_categories: List[dict],
|
||||
target_classes: List[str],
|
||||
expected_result: Dict[int, int],
|
||||
exception: Exception,
|
||||
) -> None:
|
||||
with exception:
|
||||
result = build_coco_class_index_mapping(
|
||||
coco_categories=coco_categories, target_classes=target_classes
|
||||
)
|
||||
assert result == expected_result
|
||||
@@ -0,0 +1,164 @@
|
||||
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",
|
||||
[
|
||||
(
|
||||
np.array([0, 0, 10, 10]),
|
||||
"test",
|
||||
None,
|
||||
ET.fromstring(
|
||||
"""<object><name>test</name><bndbox><xmin>1</xmin><ymin>1</ymin>
|
||||
<xmax>11</xmax><ymax>11</ymax></bndbox></object>"""
|
||||
),
|
||||
DoesNotRaise(),
|
||||
),
|
||||
(
|
||||
np.array([0, 0, 10, 10]),
|
||||
"test",
|
||||
np.array([[0, 0], [10, 0], [10, 10], [0, 10]]),
|
||||
ET.fromstring(
|
||||
"""<object><name>test</name><bndbox><xmin>1</xmin><ymin>1</ymin>
|
||||
<xmax>11</xmax><ymax>11</ymax>
|
||||
</bndbox><polygon><x1>1</x1><y1>1</y1><x2>11</x2>
|
||||
<y2>1</y2><x3>11</x3><y3>11</y3><x4>1</x4><y4>11</y4>
|
||||
</polygon></object>"""
|
||||
),
|
||||
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(
|
||||
"""<polygon><x1>0</x1><y1>0</y1><x2>10</x2><y2>0</y2><x3>10</x3>
|
||||
<y3>10</y3><x4>0</x4><y4>10</y4></polygon>"""
|
||||
),
|
||||
np.array([[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 np.array_equal(result, expected_result)
|
||||
|
||||
|
||||
ONE_CLASS_N_BBOX = """<annotation><object><name>test</name><bndbox><xmin>1</xmin>
|
||||
<ymin>1</ymin><xmax>11</xmax><ymax>11</ymax>
|
||||
</bndbox></object><object><name>test</name><bndbox><xmin>11</xmin><ymin>11</ymin>
|
||||
<xmax>21</xmax><ymax>21</ymax></bndbox></object></annotation>"""
|
||||
|
||||
|
||||
ONE_CLASS_ONE_BBOX = """<annotation><object><name>test</name><bndbox>
|
||||
<xmin>1</xmin><ymin>1</ymin><xmax>11</xmax><ymax>11</ymax></bndbox></object>
|
||||
</annotation>"""
|
||||
|
||||
|
||||
N_CLASS_N_BBOX = """<annotation><object><name>test</name><bndbox><xmin>1</xmin>
|
||||
<ymin>1</ymin><xmax>11</xmax><ymax>11</ymax>
|
||||
</bndbox></object><object><name>test</name><bndbox>
|
||||
<xmin>21</xmin><ymin>31</ymin><xmax>31</xmax><ymax>41</ymax></bndbox>
|
||||
</object><object><name>test2</name><bndbox><xmin>
|
||||
11</xmin><ymin>11</ymin><xmax>21</xmax><ymax>
|
||||
21</ymax></bndbox></object></annotation>"""
|
||||
|
||||
NO_DETECTIONS = """<annotation></annotation>"""
|
||||
|
||||
|
||||
@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
|
||||
+129
-154
@@ -1,12 +1,16 @@
|
||||
from contextlib import ExitStack as DoesNotRaise
|
||||
from typing import List, Tuple, Optional
|
||||
from typing import List, Optional, Tuple
|
||||
|
||||
import pytest
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from supervision.dataset.formats.yolo import (
|
||||
_image_name_to_annotation_name,
|
||||
_with_mask,
|
||||
object_to_yolo,
|
||||
yolo_annotations_to_detections,
|
||||
)
|
||||
from supervision.detection.core import Detections
|
||||
from supervision.dataset.formats.yolo import yolo_annotations_to_detections, _with_mask, _image_name_to_annotation_name, \
|
||||
object_to_yolo
|
||||
|
||||
|
||||
def _mock_simple_mask(resolution_wh: Tuple[int, int], box: List[int]) -> np.array:
|
||||
@@ -17,223 +21,188 @@ def _mock_simple_mask(resolution_wh: Tuple[int, int], box: List[int]) -> np.arra
|
||||
|
||||
|
||||
# The result of _mock_simple_mask is a little different from the result produced by cv2.
|
||||
def _arrays_almost_equal(arr1: np.ndarray, arr2: np.ndarray, threshold: float = 0.99) -> bool:
|
||||
def _arrays_almost_equal(
|
||||
arr1: np.ndarray, arr2: np.ndarray, threshold: float = 0.99
|
||||
) -> bool:
|
||||
equal_elements = np.equal(arr1, arr2)
|
||||
proportion_equal = np.mean(equal_elements)
|
||||
return proportion_equal >= threshold
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
'lines, expected_result, exception',
|
||||
"lines, expected_result, exception",
|
||||
[
|
||||
([], False, DoesNotRaise()), # empty yolo annotation file
|
||||
(
|
||||
[],
|
||||
["0 0.5 0.5 0.2 0.2"],
|
||||
False,
|
||||
DoesNotRaise()
|
||||
), # empty yolo annotation file
|
||||
(
|
||||
[
|
||||
'0 0.5 0.5 0.2 0.2'
|
||||
],
|
||||
False,
|
||||
DoesNotRaise()
|
||||
DoesNotRaise(),
|
||||
), # yolo annotation file with single line with box
|
||||
(
|
||||
[
|
||||
'0 0.50 0.50 0.20 0.20',
|
||||
'1 0.11 0.47 0.22 0.30'
|
||||
],
|
||||
False,
|
||||
DoesNotRaise()
|
||||
), # yolo annotation file with two lines with box
|
||||
(
|
||||
[
|
||||
'0 0.5 0.5 0.2 0.2'
|
||||
],
|
||||
["0 0.50 0.50 0.20 0.20", "1 0.11 0.47 0.22 0.30"],
|
||||
False,
|
||||
DoesNotRaise()
|
||||
),
|
||||
DoesNotRaise(),
|
||||
), # yolo annotation file with two lines with box
|
||||
(["0 0.5 0.5 0.2 0.2"], False, DoesNotRaise()),
|
||||
(
|
||||
[
|
||||
'0 0.4 0.4 0.6 0.4 0.6 0.6 0.4 0.6'
|
||||
],
|
||||
["0 0.4 0.4 0.6 0.4 0.6 0.6 0.4 0.6"],
|
||||
True,
|
||||
DoesNotRaise()
|
||||
DoesNotRaise(),
|
||||
), # yolo annotation file with single line with polygon
|
||||
(
|
||||
[
|
||||
'0 0.4 0.4 0.6 0.4 0.6 0.6 0.4 0.6',
|
||||
'1 0.11 0.47 0.22 0.30'
|
||||
],
|
||||
["0 0.4 0.4 0.6 0.4 0.6 0.6 0.4 0.6", "1 0.11 0.47 0.22 0.30"],
|
||||
True,
|
||||
DoesNotRaise()
|
||||
DoesNotRaise(),
|
||||
), # yolo annotation file with two lines - one box and one polygon
|
||||
]
|
||||
],
|
||||
)
|
||||
def test_with_mask(lines: List[str], expected_result: Optional[bool], exception: Exception) -> None:
|
||||
def test_with_mask(
|
||||
lines: List[str], expected_result: Optional[bool], exception: Exception
|
||||
) -> None:
|
||||
with exception:
|
||||
result = _with_mask(lines=lines)
|
||||
assert result == expected_result
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
'lines, resolution_wh, with_masks, expected_result, exception',
|
||||
"lines, resolution_wh, with_masks, expected_result, exception",
|
||||
[
|
||||
(
|
||||
[],
|
||||
(1000, 1000),
|
||||
False,
|
||||
Detections.empty(),
|
||||
DoesNotRaise()
|
||||
DoesNotRaise(),
|
||||
), # empty yolo annotation file
|
||||
(
|
||||
[
|
||||
'0 0.5 0.5 0.2 0.2'
|
||||
],
|
||||
["0 0.5 0.5 0.2 0.2"],
|
||||
(1000, 1000),
|
||||
False,
|
||||
Detections(
|
||||
xyxy=np.array([
|
||||
[400, 400, 600, 600]
|
||||
], dtype=np.float32),
|
||||
class_id=np.array([0], dtype=int)
|
||||
xyxy=np.array([[400, 400, 600, 600]], dtype=np.float32),
|
||||
class_id=np.array([0], dtype=int),
|
||||
),
|
||||
DoesNotRaise()
|
||||
DoesNotRaise(),
|
||||
), # yolo annotation file with single line with box
|
||||
(
|
||||
[
|
||||
'0 0.50 0.50 0.20 0.20',
|
||||
'1 0.11 0.47 0.22 0.30'
|
||||
],
|
||||
["0 0.50 0.50 0.20 0.20", "1 0.11 0.47 0.22 0.30"],
|
||||
(1000, 1000),
|
||||
False,
|
||||
Detections(
|
||||
xyxy=np.array([
|
||||
[400, 400, 600, 600],
|
||||
[ 0, 320, 220, 620]
|
||||
], dtype=np.float32),
|
||||
class_id=np.array([0, 1], dtype=int)
|
||||
xyxy=np.array(
|
||||
[[400, 400, 600, 600], [0, 320, 220, 620]], dtype=np.float32
|
||||
),
|
||||
class_id=np.array([0, 1], dtype=int),
|
||||
),
|
||||
DoesNotRaise()
|
||||
DoesNotRaise(),
|
||||
), # yolo annotation file with two lines with box
|
||||
(
|
||||
[
|
||||
'0 0.5 0.5 0.2 0.2'
|
||||
],
|
||||
["0 0.5 0.5 0.2 0.2"],
|
||||
(1000, 1000),
|
||||
True,
|
||||
Detections(
|
||||
xyxy=np.array([
|
||||
[400, 400, 600, 600]
|
||||
], dtype=np.float32),
|
||||
xyxy=np.array([[400, 400, 600, 600]], dtype=np.float32),
|
||||
class_id=np.array([0], dtype=int),
|
||||
mask=np.array([
|
||||
_mock_simple_mask(resolution_wh=(1000, 1000), box=[400, 400, 600, 600])
|
||||
], dtype=bool)
|
||||
mask=np.array(
|
||||
[
|
||||
_mock_simple_mask(
|
||||
resolution_wh=(1000, 1000), box=[400, 400, 600, 600]
|
||||
)
|
||||
],
|
||||
dtype=bool,
|
||||
),
|
||||
),
|
||||
DoesNotRaise()
|
||||
DoesNotRaise(),
|
||||
), # yolo annotation file with single line with box in with_masks mode
|
||||
(
|
||||
[
|
||||
'0 0.4 0.4 0.6 0.4 0.6 0.6 0.4 0.6'
|
||||
],
|
||||
["0 0.4 0.4 0.6 0.4 0.6 0.6 0.4 0.6"],
|
||||
(1000, 1000),
|
||||
True,
|
||||
Detections(
|
||||
xyxy=np.array([
|
||||
[400, 400, 600, 600]
|
||||
], dtype=np.float32),
|
||||
xyxy=np.array([[400, 400, 600, 600]], dtype=np.float32),
|
||||
class_id=np.array([0], dtype=int),
|
||||
mask=np.array([
|
||||
_mock_simple_mask(resolution_wh=(1000, 1000), box=[400, 400, 600, 600])
|
||||
], dtype=bool)
|
||||
mask=np.array(
|
||||
[
|
||||
_mock_simple_mask(
|
||||
resolution_wh=(1000, 1000), box=[400, 400, 600, 600]
|
||||
)
|
||||
],
|
||||
dtype=bool,
|
||||
),
|
||||
),
|
||||
DoesNotRaise()
|
||||
DoesNotRaise(),
|
||||
), # yolo annotation file with single line with polygon
|
||||
(
|
||||
[
|
||||
'0 0.4 0.4 0.6 0.4 0.6 0.6 0.4 0.6',
|
||||
'1 0.11 0.47 0.22 0.30'
|
||||
],
|
||||
["0 0.4 0.4 0.6 0.4 0.6 0.6 0.4 0.6", "1 0.11 0.47 0.22 0.30"],
|
||||
(1000, 1000),
|
||||
True,
|
||||
Detections(
|
||||
xyxy=np.array([
|
||||
[400, 400, 600, 600],
|
||||
[ 0, 320, 220, 620]
|
||||
], dtype=np.float32),
|
||||
xyxy=np.array(
|
||||
[[400, 400, 600, 600], [0, 320, 220, 620]], dtype=np.float32
|
||||
),
|
||||
class_id=np.array([0, 1], dtype=int),
|
||||
mask=np.array([
|
||||
_mock_simple_mask(resolution_wh=(1000, 1000), box=[400, 400, 600, 600]),
|
||||
_mock_simple_mask(resolution_wh=(1000, 1000), box=[ 0, 320, 220, 620])
|
||||
], dtype=bool)
|
||||
mask=np.array(
|
||||
[
|
||||
_mock_simple_mask(
|
||||
resolution_wh=(1000, 1000), box=[400, 400, 600, 600]
|
||||
),
|
||||
_mock_simple_mask(
|
||||
resolution_wh=(1000, 1000), box=[0, 320, 220, 620]
|
||||
),
|
||||
],
|
||||
dtype=bool,
|
||||
),
|
||||
),
|
||||
DoesNotRaise()
|
||||
), # yolo annotation file with two lines - one box and one polygon in with_masks mode
|
||||
DoesNotRaise(),
|
||||
), # yolo annotation file with two lines -
|
||||
# one box and one polygon in with_masks mode
|
||||
(
|
||||
[
|
||||
'0 0.4 0.4 0.6 0.4 0.6 0.6 0.4 0.6',
|
||||
'1 0.11 0.47 0.22 0.30'
|
||||
],
|
||||
["0 0.4 0.4 0.6 0.4 0.6 0.6 0.4 0.6", "1 0.11 0.47 0.22 0.30"],
|
||||
(1000, 1000),
|
||||
False,
|
||||
Detections(
|
||||
xyxy=np.array([
|
||||
[400, 400, 600, 600],
|
||||
[ 0, 320, 220, 620]
|
||||
], dtype=np.float32),
|
||||
class_id=np.array([0, 1], dtype=int)
|
||||
xyxy=np.array(
|
||||
[[400, 400, 600, 600], [0, 320, 220, 620]], dtype=np.float32
|
||||
),
|
||||
class_id=np.array([0, 1], dtype=int),
|
||||
),
|
||||
DoesNotRaise()
|
||||
DoesNotRaise(),
|
||||
), # yolo annotation file with two lines - one box and one polygon
|
||||
]
|
||||
],
|
||||
)
|
||||
def test_yolo_annotations_to_detections(
|
||||
lines: List[str],
|
||||
resolution_wh: Tuple[int, int],
|
||||
with_masks: bool,
|
||||
expected_result: Optional[Detections],
|
||||
exception: Exception
|
||||
exception: Exception,
|
||||
) -> None:
|
||||
with exception:
|
||||
result = yolo_annotations_to_detections(
|
||||
lines=lines,
|
||||
resolution_wh=resolution_wh,
|
||||
with_masks=with_masks)
|
||||
lines=lines, resolution_wh=resolution_wh, with_masks=with_masks
|
||||
)
|
||||
assert np.array_equal(result.xyxy, expected_result.xyxy)
|
||||
assert np.array_equal(result.class_id, expected_result.class_id)
|
||||
assert (result.mask is None and expected_result.mask is None) or _arrays_almost_equal(result.mask, expected_result.mask)
|
||||
assert (
|
||||
result.mask is None and expected_result.mask is None
|
||||
) or _arrays_almost_equal(result.mask, expected_result.mask)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
'image_name, expected_result, exception',
|
||||
"image_name, expected_result, exception",
|
||||
[
|
||||
("image.png", "image.txt", DoesNotRaise()), # simple png image
|
||||
("image.jpeg", "image.txt", DoesNotRaise()), # simple jpeg image
|
||||
("image.jpg", "image.txt", DoesNotRaise()), # simple jpg image
|
||||
(
|
||||
'image.png',
|
||||
'image.txt',
|
||||
DoesNotRaise()
|
||||
), # simple png image
|
||||
(
|
||||
'image.jpeg',
|
||||
'image.txt',
|
||||
DoesNotRaise()
|
||||
), # simple jpeg image
|
||||
(
|
||||
'image.jpg',
|
||||
'image.txt',
|
||||
DoesNotRaise()
|
||||
), # simple jpg image
|
||||
(
|
||||
'image.000.jpg',
|
||||
'image.000.txt',
|
||||
DoesNotRaise()
|
||||
"image.000.jpg",
|
||||
"image.000.txt",
|
||||
DoesNotRaise(),
|
||||
), # jpg image with multiple dots in name
|
||||
]
|
||||
],
|
||||
)
|
||||
def test_image_name_to_annotation_name(
|
||||
image_name: str,
|
||||
expected_result: Optional[str],
|
||||
exception: Exception
|
||||
image_name: str, expected_result: Optional[str], exception: Exception
|
||||
) -> None:
|
||||
with exception:
|
||||
result = _image_name_to_annotation_name(image_name=image_name)
|
||||
@@ -241,64 +210,70 @@ def test_image_name_to_annotation_name(
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
'xyxy, class_id, image_shape, polygon, expected_result, exception',
|
||||
"xyxy, class_id, image_shape, polygon, expected_result, exception",
|
||||
[
|
||||
(
|
||||
np.array([100, 100, 200, 200], dtype=np.float32),
|
||||
1,
|
||||
(1000, 1000, 3),
|
||||
None,
|
||||
'1 0.15000 0.15000 0.10000 0.10000',
|
||||
DoesNotRaise()
|
||||
"1 0.15000 0.15000 0.10000 0.10000",
|
||||
DoesNotRaise(),
|
||||
), # square bounding box on square image
|
||||
(
|
||||
np.array([100, 100, 200, 200], dtype=np.float32),
|
||||
1,
|
||||
(800, 1000, 3),
|
||||
None,
|
||||
'1 0.15000 0.18750 0.10000 0.12500',
|
||||
DoesNotRaise()
|
||||
"1 0.15000 0.18750 0.10000 0.12500",
|
||||
DoesNotRaise(),
|
||||
), # square bounding box on horizontal image
|
||||
(
|
||||
np.array([100, 100, 200, 200], dtype=np.float32),
|
||||
1,
|
||||
(1000, 800, 3),
|
||||
None,
|
||||
'1 0.18750 0.15000 0.12500 0.10000',
|
||||
DoesNotRaise()
|
||||
"1 0.18750 0.15000 0.12500 0.10000",
|
||||
DoesNotRaise(),
|
||||
), # square bounding box on vertical image
|
||||
(
|
||||
np.array([100, 200, 200, 400], dtype=np.float32),
|
||||
1,
|
||||
(1000, 1000, 3),
|
||||
None,
|
||||
'1 0.15000 0.30000 0.10000 0.20000',
|
||||
DoesNotRaise()
|
||||
"1 0.15000 0.30000 0.10000 0.20000",
|
||||
DoesNotRaise(),
|
||||
), # horizontal bounding box on square image
|
||||
(
|
||||
np.array([200, 100, 400, 200], dtype=np.float32),
|
||||
1,
|
||||
(1000, 1000, 3),
|
||||
None,
|
||||
'1 0.30000 0.15000 0.20000 0.10000',
|
||||
DoesNotRaise()
|
||||
"1 0.30000 0.15000 0.20000 0.10000",
|
||||
DoesNotRaise(),
|
||||
), # vertical bounding box on square image
|
||||
(
|
||||
np.array([100, 100, 200, 200], dtype=np.float32),
|
||||
1,
|
||||
(1000, 1000, 3),
|
||||
np.array([
|
||||
[100, 100],
|
||||
[200, 100],
|
||||
[200, 200],
|
||||
[100, 100]
|
||||
], dtype=np.float32),
|
||||
'1 0.10000 0.10000 0.20000 0.10000 0.20000 0.20000 0.10000 0.10000',
|
||||
DoesNotRaise()
|
||||
np.array(
|
||||
[[100, 100], [200, 100], [200, 200], [100, 100]], dtype=np.float32
|
||||
),
|
||||
"1 0.10000 0.10000 0.20000 0.10000 0.20000 0.20000 0.10000 0.10000",
|
||||
DoesNotRaise(),
|
||||
), # square mask on square image
|
||||
]
|
||||
],
|
||||
)
|
||||
def test_object_to_yolo(xyxy: np.ndarray, class_id: int, image_shape: Tuple[int, int, int], polygon: Optional[np.ndarray], expected_result: Optional[str], exception: Exception) -> None:
|
||||
def test_object_to_yolo(
|
||||
xyxy: np.ndarray,
|
||||
class_id: int,
|
||||
image_shape: Tuple[int, int, int],
|
||||
polygon: Optional[np.ndarray],
|
||||
expected_result: Optional[str],
|
||||
exception: Exception,
|
||||
) -> None:
|
||||
with exception:
|
||||
result = object_to_yolo(xyxy=xyxy, class_id=class_id, image_shape=image_shape, polygon=polygon)
|
||||
result = object_to_yolo(
|
||||
xyxy=xyxy, class_id=class_id, image_shape=image_shape, polygon=polygon
|
||||
)
|
||||
assert result == expected_result
|
||||
|
||||
@@ -0,0 +1,195 @@
|
||||
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 import DetectionDataset
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"dataset_list, expected_result, exception",
|
||||
[
|
||||
(
|
||||
[],
|
||||
DetectionDataset(classes=[], images={}, annotations={}),
|
||||
DoesNotRaise(),
|
||||
), # empty dataset list
|
||||
(
|
||||
[DetectionDataset(classes=[], images={}, annotations={})],
|
||||
DetectionDataset(classes=[], images={}, annotations={}),
|
||||
DoesNotRaise(),
|
||||
), # single empty dataset
|
||||
(
|
||||
[
|
||||
DetectionDataset(classes=["dog", "person"], images={}, annotations={}),
|
||||
DetectionDataset(classes=["dog", "person"], images={}, annotations={}),
|
||||
],
|
||||
DetectionDataset(classes=["dog", "person"], images={}, annotations={}),
|
||||
DoesNotRaise(),
|
||||
), # two datasets; no images and annotations, the same classes
|
||||
(
|
||||
[
|
||||
DetectionDataset(classes=["dog", "person"], images={}, annotations={}),
|
||||
DetectionDataset(classes=["cat"], images={}, annotations={}),
|
||||
],
|
||||
DetectionDataset(
|
||||
classes=["cat", "dog", "person"], images={}, annotations={}
|
||||
),
|
||||
DoesNotRaise(),
|
||||
), # two datasets; no images and annotations, different classes
|
||||
(
|
||||
[
|
||||
DetectionDataset(
|
||||
classes=["dog", "person"],
|
||||
images={
|
||||
"image-1.png": np.zeros((100, 100, 3), dtype=np.uint8),
|
||||
"image-2.png": np.zeros((100, 100, 3), dtype=np.uint8),
|
||||
},
|
||||
annotations={
|
||||
"image-1.png": mock_detections(
|
||||
xyxy=[[0, 0, 10, 10]], class_id=[0]
|
||||
),
|
||||
"image-2.png": mock_detections(
|
||||
xyxy=[[0, 0, 10, 10]], class_id=[1]
|
||||
),
|
||||
},
|
||||
),
|
||||
DetectionDataset(classes=[], images={}, annotations={}),
|
||||
],
|
||||
DetectionDataset(
|
||||
classes=["dog", "person"],
|
||||
images={
|
||||
"image-1.png": np.zeros((100, 100, 3), dtype=np.uint8),
|
||||
"image-2.png": np.zeros((100, 100, 3), dtype=np.uint8),
|
||||
},
|
||||
annotations={
|
||||
"image-1.png": mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[0]),
|
||||
"image-2.png": mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[1]),
|
||||
},
|
||||
),
|
||||
DoesNotRaise(),
|
||||
), # two datasets; images and annotations, the same classes
|
||||
(
|
||||
[
|
||||
DetectionDataset(
|
||||
classes=["dog", "person"],
|
||||
images={
|
||||
"image-1.png": np.zeros((100, 100, 3), dtype=np.uint8),
|
||||
"image-2.png": np.zeros((100, 100, 3), dtype=np.uint8),
|
||||
},
|
||||
annotations={
|
||||
"image-1.png": mock_detections(
|
||||
xyxy=[[0, 0, 10, 10]], class_id=[0]
|
||||
),
|
||||
"image-2.png": mock_detections(
|
||||
xyxy=[[0, 0, 10, 10]], class_id=[1]
|
||||
),
|
||||
},
|
||||
),
|
||||
DetectionDataset(classes=["cat"], images={}, annotations={}),
|
||||
],
|
||||
DetectionDataset(
|
||||
classes=["cat", "dog", "person"],
|
||||
images={
|
||||
"image-1.png": np.zeros((100, 100, 3), dtype=np.uint8),
|
||||
"image-2.png": np.zeros((100, 100, 3), dtype=np.uint8),
|
||||
},
|
||||
annotations={
|
||||
"image-1.png": mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[1]),
|
||||
"image-2.png": mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[2]),
|
||||
},
|
||||
),
|
||||
DoesNotRaise(),
|
||||
), # two datasets; images and annotations, different classes
|
||||
(
|
||||
[
|
||||
DetectionDataset(
|
||||
classes=["dog", "person"],
|
||||
images={
|
||||
"image-1.png": np.zeros((100, 100, 3), dtype=np.uint8),
|
||||
"image-2.png": np.zeros((100, 100, 3), dtype=np.uint8),
|
||||
},
|
||||
annotations={
|
||||
"image-1.png": mock_detections(
|
||||
xyxy=[[0, 0, 10, 10]], class_id=[0]
|
||||
),
|
||||
"image-2.png": mock_detections(
|
||||
xyxy=[[0, 0, 10, 10]], class_id=[1]
|
||||
),
|
||||
},
|
||||
),
|
||||
DetectionDataset(
|
||||
classes=["cat"],
|
||||
images={
|
||||
"image-3.png": np.zeros((100, 100, 3), dtype=np.uint8),
|
||||
},
|
||||
annotations={
|
||||
"image-3.png": mock_detections(
|
||||
xyxy=[[0, 0, 10, 10]], class_id=[0]
|
||||
),
|
||||
},
|
||||
),
|
||||
],
|
||||
DetectionDataset(
|
||||
classes=["cat", "dog", "person"],
|
||||
images={
|
||||
"image-1.png": np.zeros((100, 100, 3), dtype=np.uint8),
|
||||
"image-2.png": np.zeros((100, 100, 3), dtype=np.uint8),
|
||||
"image-3.png": np.zeros((100, 100, 3), dtype=np.uint8),
|
||||
},
|
||||
annotations={
|
||||
"image-1.png": mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[1]),
|
||||
"image-2.png": mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[2]),
|
||||
"image-3.png": mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[0]),
|
||||
},
|
||||
),
|
||||
DoesNotRaise(),
|
||||
), # two datasets; images and annotations, different classes
|
||||
(
|
||||
[
|
||||
DetectionDataset(
|
||||
classes=["dog", "person"],
|
||||
images={
|
||||
"image-1.png": np.zeros((100, 100, 3), dtype=np.uint8),
|
||||
"image-2.png": np.zeros((100, 100, 3), dtype=np.uint8),
|
||||
},
|
||||
annotations={
|
||||
"image-1.png": mock_detections(
|
||||
xyxy=[[0, 0, 10, 10]], class_id=[0]
|
||||
),
|
||||
"image-2.png": mock_detections(
|
||||
xyxy=[[0, 0, 10, 10]], class_id=[1]
|
||||
),
|
||||
},
|
||||
),
|
||||
DetectionDataset(
|
||||
classes=["dog", "person"],
|
||||
images={
|
||||
"image-2.png": np.zeros((100, 100, 3), dtype=np.uint8),
|
||||
"image-3.png": np.zeros((100, 100, 3), dtype=np.uint8),
|
||||
},
|
||||
annotations={
|
||||
"image-2.png": mock_detections(
|
||||
xyxy=[[0, 0, 10, 10]], class_id=[0]
|
||||
),
|
||||
"image-3.png": mock_detections(
|
||||
xyxy=[[0, 0, 10, 10]], class_id=[1]
|
||||
),
|
||||
},
|
||||
),
|
||||
],
|
||||
None,
|
||||
pytest.raises(ValueError),
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_dataset_merge(
|
||||
dataset_list: List[DetectionDataset],
|
||||
expected_result: Optional[DetectionDataset],
|
||||
exception: Exception,
|
||||
) -> None:
|
||||
with exception:
|
||||
result = DetectionDataset.merge(dataset_list=dataset_list)
|
||||
assert result == expected_result
|
||||
+171
-22
@@ -1,31 +1,31 @@
|
||||
from typing import List, TypeVar, Optional, Tuple
|
||||
from contextlib import ExitStack as DoesNotRaise
|
||||
from test.utils import mock_detections
|
||||
from typing import Dict, List, Optional, Tuple, TypeVar
|
||||
|
||||
import pytest
|
||||
|
||||
from supervision.dataset.ultils import train_test_split
|
||||
from supervision import Detections
|
||||
from supervision.dataset.utils import (
|
||||
build_class_index_mapping,
|
||||
map_detections_class_id,
|
||||
merge_class_lists,
|
||||
train_test_split,
|
||||
)
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
'data, train_ratio, random_state, shuffle, expected_result, exception',
|
||||
"data, train_ratio, random_state, shuffle, expected_result, exception",
|
||||
[
|
||||
(
|
||||
[],
|
||||
0.5,
|
||||
None,
|
||||
False,
|
||||
([], []),
|
||||
DoesNotRaise()
|
||||
), # empty data
|
||||
([], 0.5, None, False, ([], []), DoesNotRaise()), # empty data
|
||||
(
|
||||
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
|
||||
0.5,
|
||||
None,
|
||||
False,
|
||||
([0, 1, 2, 3, 4], [5, 6, 7, 8, 9]),
|
||||
DoesNotRaise()
|
||||
DoesNotRaise(),
|
||||
), # data with 10 numbers and 50% train split
|
||||
(
|
||||
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
|
||||
@@ -33,7 +33,7 @@ T = TypeVar("T")
|
||||
None,
|
||||
False,
|
||||
([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], []),
|
||||
DoesNotRaise()
|
||||
DoesNotRaise(),
|
||||
), # data with 10 numbers and 100% train split
|
||||
(
|
||||
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
|
||||
@@ -41,15 +41,15 @@ T = TypeVar("T")
|
||||
None,
|
||||
False,
|
||||
([], [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]),
|
||||
DoesNotRaise()
|
||||
DoesNotRaise(),
|
||||
), # data with 10 numbers and 0% train split
|
||||
(
|
||||
['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j'],
|
||||
["a", "b", "c", "d", "e", "f", "g", "h", "i", "j"],
|
||||
0.5,
|
||||
None,
|
||||
False,
|
||||
(['a', 'b', 'c', 'd', 'e'], ['f', 'g', 'h', 'i', 'j']),
|
||||
DoesNotRaise()
|
||||
(["a", "b", "c", "d", "e"], ["f", "g", "h", "i", "j"]),
|
||||
DoesNotRaise(),
|
||||
), # data with 10 chars and 50% train split
|
||||
(
|
||||
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
|
||||
@@ -57,7 +57,7 @@ T = TypeVar("T")
|
||||
23,
|
||||
True,
|
||||
([7, 8, 5, 6, 3], [2, 9, 0, 1, 4]),
|
||||
DoesNotRaise()
|
||||
DoesNotRaise(),
|
||||
), # data with 10 numbers and 50% train split with 23 random seed
|
||||
(
|
||||
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
|
||||
@@ -65,9 +65,9 @@ T = TypeVar("T")
|
||||
32,
|
||||
True,
|
||||
([4, 6, 0, 8, 9], [5, 7, 2, 3, 1]),
|
||||
DoesNotRaise()
|
||||
DoesNotRaise(),
|
||||
), # data with 10 numbers and 50% train split with 23 random seed
|
||||
]
|
||||
],
|
||||
)
|
||||
def test_train_test_split(
|
||||
data: List[T],
|
||||
@@ -75,8 +75,157 @@ def test_train_test_split(
|
||||
random_state: int,
|
||||
shuffle: bool,
|
||||
expected_result: Optional[Tuple[List[T], List[T]]],
|
||||
exception: Exception
|
||||
exception: Exception,
|
||||
) -> None:
|
||||
with exception:
|
||||
result = train_test_split(data=data, train_ratio=train_ratio, random_state=random_state, shuffle=shuffle)
|
||||
result = train_test_split(
|
||||
data=data,
|
||||
train_ratio=train_ratio,
|
||||
random_state=random_state,
|
||||
shuffle=shuffle,
|
||||
)
|
||||
assert result == expected_result
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"class_lists, expected_result, exception",
|
||||
[
|
||||
([], [], DoesNotRaise()), # empty class lists
|
||||
(
|
||||
[["dog", "person"]],
|
||||
["dog", "person"],
|
||||
DoesNotRaise(),
|
||||
), # single class list; already alphabetically sorted
|
||||
(
|
||||
[["person", "dog"]],
|
||||
["dog", "person"],
|
||||
DoesNotRaise(),
|
||||
), # single class list; not alphabetically sorted
|
||||
(
|
||||
[["dog", "person"], ["dog", "person"]],
|
||||
["dog", "person"],
|
||||
DoesNotRaise(),
|
||||
), # two class lists; the same classes; already alphabetically sorted
|
||||
(
|
||||
[["dog", "person"], ["cat"]],
|
||||
["cat", "dog", "person"],
|
||||
DoesNotRaise(),
|
||||
), # two class lists; different classes; already alphabetically sorted
|
||||
],
|
||||
)
|
||||
def test_merge_class_maps(
|
||||
class_lists: List[List[str]], expected_result: List[str], exception: Exception
|
||||
) -> None:
|
||||
with exception:
|
||||
result = merge_class_lists(class_lists=class_lists)
|
||||
assert result == expected_result
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"source_classes, target_classes, expected_result, exception",
|
||||
[
|
||||
([], [], {}, DoesNotRaise()), # empty class lists
|
||||
([], ["dog", "person"], {}, DoesNotRaise()), # empty source class list
|
||||
(
|
||||
["dog", "person"],
|
||||
[],
|
||||
None,
|
||||
pytest.raises(ValueError),
|
||||
), # empty target class list
|
||||
(
|
||||
["dog", "person"],
|
||||
["dog", "person"],
|
||||
{0: 0, 1: 1},
|
||||
DoesNotRaise(),
|
||||
), # same class lists
|
||||
(
|
||||
["dog", "person"],
|
||||
["person", "dog"],
|
||||
{0: 1, 1: 0},
|
||||
DoesNotRaise(),
|
||||
), # same class lists but not alphabetically sorted
|
||||
(
|
||||
["dog", "person"],
|
||||
["cat", "dog", "person"],
|
||||
{0: 1, 1: 2},
|
||||
DoesNotRaise(),
|
||||
), # source class list is a subset of target class list
|
||||
(
|
||||
["dog", "person"],
|
||||
["cat", "dog"],
|
||||
None,
|
||||
pytest.raises(ValueError),
|
||||
), # source class list is not a subset of target class list
|
||||
],
|
||||
)
|
||||
def test_build_class_index_mapping(
|
||||
source_classes: List[str],
|
||||
target_classes: List[str],
|
||||
expected_result: Optional[Dict[int, int]],
|
||||
exception: Exception,
|
||||
) -> None:
|
||||
with exception:
|
||||
result = build_class_index_mapping(
|
||||
source_classes=source_classes, target_classes=target_classes
|
||||
)
|
||||
assert result == expected_result
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"source_to_target_mapping, detections, expected_result, exception",
|
||||
[
|
||||
(
|
||||
{},
|
||||
mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[0]),
|
||||
None,
|
||||
pytest.raises(ValueError),
|
||||
), # empty mapping
|
||||
(
|
||||
{0: 1},
|
||||
mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[0]),
|
||||
mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[1]),
|
||||
DoesNotRaise(),
|
||||
), # single mapping
|
||||
(
|
||||
{0: 1, 1: 2},
|
||||
Detections.empty(),
|
||||
Detections.empty(),
|
||||
DoesNotRaise(),
|
||||
), # empty detections
|
||||
(
|
||||
{0: 1, 1: 2},
|
||||
mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[0]),
|
||||
mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[1]),
|
||||
DoesNotRaise(),
|
||||
), # multiple mappings
|
||||
(
|
||||
{0: 1, 1: 2},
|
||||
mock_detections(xyxy=[[0, 0, 10, 10], [0, 0, 10, 10]], class_id=[0, 1]),
|
||||
mock_detections(xyxy=[[0, 0, 10, 10], [0, 0, 10, 10]], class_id=[1, 2]),
|
||||
DoesNotRaise(),
|
||||
), # multiple mappings
|
||||
(
|
||||
{0: 1, 1: 2},
|
||||
mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[2]),
|
||||
None,
|
||||
pytest.raises(ValueError),
|
||||
), # class_id not in mapping
|
||||
(
|
||||
{0: 1, 1: 2},
|
||||
mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[0], confidence=[0.5]),
|
||||
mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[1], confidence=[0.5]),
|
||||
DoesNotRaise(),
|
||||
), # confidence is not None
|
||||
],
|
||||
)
|
||||
def test_map_detections_class_id(
|
||||
source_to_target_mapping: Dict[int, int],
|
||||
detections: Detections,
|
||||
expected_result: Optional[Detections],
|
||||
exception: Exception,
|
||||
) -> None:
|
||||
with exception:
|
||||
result = map_detections_class_id(
|
||||
source_to_target_mapping=source_to_target_mapping, detections=detections
|
||||
)
|
||||
assert result == expected_result
|
||||
|
||||
+179
-195
@@ -1,203 +1,137 @@
|
||||
from contextlib import ExitStack as DoesNotRaise
|
||||
|
||||
import pytest
|
||||
|
||||
from supervision import Detections
|
||||
|
||||
from typing import Optional, Union, List
|
||||
from test.utils import mock_detections
|
||||
from typing import List, Optional, Union
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from supervision.detection.core import Detections
|
||||
from supervision.geometry.core import Position
|
||||
|
||||
PREDICTIONS = np.array([
|
||||
[ 2254, 906, 2447, 1353, 0.90538, 0],
|
||||
[ 2049, 1133, 2226, 1371, 0.59002, 56],
|
||||
[ 727, 1224, 838, 1601, 0.51119, 39],
|
||||
[ 808, 1214, 910, 1564, 0.45287, 39],
|
||||
[ 6, 52, 1131, 2133, 0.45057, 72],
|
||||
[ 299, 1225, 512, 1663, 0.45029, 39],
|
||||
[ 529, 874, 645, 945, 0.31101, 39],
|
||||
[ 8, 47, 1935, 2135, 0.28192, 72],
|
||||
[ 2265, 813, 2328, 901, 0.2714, 62]
|
||||
], dtype=np.float32)
|
||||
PREDICTIONS = np.array(
|
||||
[
|
||||
[2254, 906, 2447, 1353, 0.90538, 0],
|
||||
[2049, 1133, 2226, 1371, 0.59002, 56],
|
||||
[727, 1224, 838, 1601, 0.51119, 39],
|
||||
[808, 1214, 910, 1564, 0.45287, 39],
|
||||
[6, 52, 1131, 2133, 0.45057, 72],
|
||||
[299, 1225, 512, 1663, 0.45029, 39],
|
||||
[529, 874, 645, 945, 0.31101, 39],
|
||||
[8, 47, 1935, 2135, 0.28192, 72],
|
||||
[2265, 813, 2328, 901, 0.2714, 62],
|
||||
],
|
||||
dtype=np.float32,
|
||||
)
|
||||
|
||||
DETECTIONS = Detections(
|
||||
xyxy=PREDICTIONS[:, :4],
|
||||
confidence=PREDICTIONS[:, 4],
|
||||
class_id=PREDICTIONS[:, 5].astype(int)
|
||||
class_id=PREDICTIONS[:, 5].astype(int),
|
||||
)
|
||||
|
||||
|
||||
def mock_detections(xyxy, confidence = None, class_id = None, tracker_id = None) -> Detections:
|
||||
return Detections(
|
||||
xyxy = np.array(xyxy, dtype=np.float32),
|
||||
confidence = confidence if confidence is None else np.array(confidence, dtype=np.float32),
|
||||
class_id = class_id if class_id is None else np.array(class_id, dtype=int),
|
||||
tracker_id = tracker_id if tracker_id is None else np.array(tracker_id, dtype=int),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
'detections, index, expected_result, exception',
|
||||
"detections, index, expected_result, exception",
|
||||
[
|
||||
(
|
||||
DETECTIONS,
|
||||
DETECTIONS.class_id == 0,
|
||||
Detections(
|
||||
xyxy=np.array([
|
||||
[ 2254, 906, 2447, 1353]
|
||||
], dtype=np.float32),
|
||||
confidence=np.array([
|
||||
0.90538
|
||||
], dtype=np.float32),
|
||||
class_id=np.array([
|
||||
0
|
||||
], dtype=int)
|
||||
mock_detections(
|
||||
xyxy=[[2254, 906, 2447, 1353]], confidence=[0.90538], class_id=[0]
|
||||
),
|
||||
DoesNotRaise()
|
||||
DoesNotRaise(),
|
||||
), # take only detections with class_id = 0
|
||||
(
|
||||
DETECTIONS,
|
||||
DETECTIONS.confidence > 0.5,
|
||||
Detections(
|
||||
xyxy=np.array([
|
||||
[ 2254, 906, 2447, 1353],
|
||||
[ 2049, 1133, 2226, 1371],
|
||||
[ 727, 1224, 838, 1601]
|
||||
], dtype=np.float32),
|
||||
confidence=np.array([
|
||||
0.90538,
|
||||
0.59002,
|
||||
0.51119
|
||||
], dtype=np.float32),
|
||||
class_id=np.array([
|
||||
0,
|
||||
56,
|
||||
39
|
||||
], dtype=int)
|
||||
mock_detections(
|
||||
xyxy=[
|
||||
[2254, 906, 2447, 1353],
|
||||
[2049, 1133, 2226, 1371],
|
||||
[727, 1224, 838, 1601],
|
||||
],
|
||||
confidence=[0.90538, 0.59002, 0.51119],
|
||||
class_id=[0, 56, 39],
|
||||
),
|
||||
DoesNotRaise()
|
||||
DoesNotRaise(),
|
||||
), # take only detections with confidence > 0.5
|
||||
(
|
||||
DETECTIONS,
|
||||
np.array([True, True, True, True, True, True, True, True, True], dtype=bool),
|
||||
np.array(
|
||||
[True, True, True, True, True, True, True, True, True], dtype=bool
|
||||
),
|
||||
DETECTIONS,
|
||||
DoesNotRaise()
|
||||
DoesNotRaise(),
|
||||
), # take all detections
|
||||
(
|
||||
DETECTIONS,
|
||||
np.array([False, False, False, False, False, False, False, False, False], dtype=bool),
|
||||
np.array(
|
||||
[False, False, False, False, False, False, False, False, False],
|
||||
dtype=bool,
|
||||
),
|
||||
Detections(
|
||||
xyxy=np.empty((0, 4), dtype=np.float32),
|
||||
confidence=np.array([], dtype=np.float32),
|
||||
class_id=np.array([], dtype=int)
|
||||
class_id=np.array([], dtype=int),
|
||||
),
|
||||
DoesNotRaise()
|
||||
DoesNotRaise(),
|
||||
), # take no detections
|
||||
(
|
||||
DETECTIONS,
|
||||
[0, 2],
|
||||
Detections(
|
||||
xyxy=np.array([
|
||||
[ 2254, 906, 2447, 1353],
|
||||
[ 727, 1224, 838, 1601]
|
||||
], dtype=np.float32),
|
||||
confidence=np.array([
|
||||
0.90538,
|
||||
0.51119
|
||||
], dtype=np.float32),
|
||||
class_id=np.array([
|
||||
0,
|
||||
39
|
||||
], dtype=int)
|
||||
mock_detections(
|
||||
xyxy=[[2254, 906, 2447, 1353], [727, 1224, 838, 1601]],
|
||||
confidence=[0.90538, 0.51119],
|
||||
class_id=[0, 39],
|
||||
),
|
||||
DoesNotRaise()
|
||||
DoesNotRaise(),
|
||||
), # take only first and third detection using List[int] index
|
||||
(
|
||||
DETECTIONS,
|
||||
np.array([0, 2]),
|
||||
Detections(
|
||||
xyxy=np.array([
|
||||
[ 2254, 906, 2447, 1353],
|
||||
[ 727, 1224, 838, 1601]
|
||||
], dtype=np.float32),
|
||||
confidence=np.array([
|
||||
0.90538,
|
||||
0.51119
|
||||
], dtype=np.float32),
|
||||
class_id=np.array([
|
||||
0,
|
||||
39
|
||||
], dtype=int)
|
||||
mock_detections(
|
||||
xyxy=[[2254, 906, 2447, 1353], [727, 1224, 838, 1601]],
|
||||
confidence=[0.90538, 0.51119],
|
||||
class_id=[0, 39],
|
||||
),
|
||||
DoesNotRaise()
|
||||
DoesNotRaise(),
|
||||
), # take only first and third detection using np.ndarray index
|
||||
(
|
||||
DETECTIONS,
|
||||
0,
|
||||
Detections(
|
||||
xyxy=np.array([
|
||||
[ 2254, 906, 2447, 1353]
|
||||
], dtype=np.float32),
|
||||
confidence=np.array([
|
||||
0.90538
|
||||
], dtype=np.float32),
|
||||
class_id=np.array([
|
||||
0
|
||||
], dtype=int)
|
||||
mock_detections(
|
||||
xyxy=[[2254, 906, 2447, 1353]], confidence=[0.90538], class_id=[0]
|
||||
),
|
||||
DoesNotRaise()
|
||||
DoesNotRaise(),
|
||||
), # take only first detection by index
|
||||
(
|
||||
DETECTIONS,
|
||||
slice(1, 3),
|
||||
Detections(
|
||||
xyxy=np.array([
|
||||
[ 2049, 1133, 2226, 1371],
|
||||
[ 727, 1224, 838, 1601]
|
||||
], dtype=np.float32),
|
||||
confidence=np.array([
|
||||
0.59002,
|
||||
0.51119
|
||||
], dtype=np.float32),
|
||||
class_id=np.array([
|
||||
56,
|
||||
39
|
||||
], dtype=int)
|
||||
mock_detections(
|
||||
xyxy=[[2049, 1133, 2226, 1371], [727, 1224, 838, 1601]],
|
||||
confidence=[0.59002, 0.51119],
|
||||
class_id=[56, 39],
|
||||
),
|
||||
DoesNotRaise()
|
||||
DoesNotRaise(),
|
||||
), # take only first detection by index slice (1, 3)
|
||||
(DETECTIONS, 10, None, pytest.raises(IndexError)), # index out of range
|
||||
(DETECTIONS, [0, 2, 10], None, pytest.raises(IndexError)), # index out of range
|
||||
(DETECTIONS, np.array([0, 2, 10]), None, pytest.raises(IndexError)),
|
||||
(
|
||||
DETECTIONS,
|
||||
10,
|
||||
np.array(
|
||||
[True, True, True, True, True, True, True, True, True, True, True]
|
||||
),
|
||||
None,
|
||||
pytest.raises(IndexError)
|
||||
), # index out of range
|
||||
(
|
||||
DETECTIONS,
|
||||
[0, 2, 10],
|
||||
None,
|
||||
pytest.raises(IndexError)
|
||||
), # index out of range
|
||||
(
|
||||
DETECTIONS,
|
||||
np.array([0, 2, 10]),
|
||||
None,
|
||||
pytest.raises(IndexError)
|
||||
pytest.raises(IndexError),
|
||||
),
|
||||
(
|
||||
DETECTIONS,
|
||||
np.array([True, True, True, True, True, True, True, True, True, True, True]),
|
||||
None,
|
||||
pytest.raises(IndexError)
|
||||
)
|
||||
]
|
||||
],
|
||||
)
|
||||
def test_getitem(
|
||||
detections: Detections,
|
||||
index: Union[int, slice, List[int], np.ndarray],
|
||||
expected_result: Optional[Detections],
|
||||
exception: Exception
|
||||
detections: Detections,
|
||||
index: Union[int, slice, List[int], np.ndarray],
|
||||
expected_result: Optional[Detections],
|
||||
exception: Exception,
|
||||
) -> None:
|
||||
with exception:
|
||||
result = detections[index]
|
||||
@@ -205,87 +139,137 @@ def test_getitem(
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
'detections_list, expected_result, exception',
|
||||
"detections_list, expected_result, exception",
|
||||
[
|
||||
([], Detections.empty(), DoesNotRaise()), # empty detections list
|
||||
(
|
||||
[],
|
||||
[Detections.empty()],
|
||||
Detections.empty(),
|
||||
DoesNotRaise()
|
||||
), # empty detections list
|
||||
(
|
||||
[
|
||||
Detections.empty()
|
||||
],
|
||||
Detections.empty(),
|
||||
DoesNotRaise()
|
||||
DoesNotRaise(),
|
||||
), # single empty detections
|
||||
(
|
||||
[
|
||||
mock_detections(xyxy=[[10, 10, 20, 20]])
|
||||
],
|
||||
[mock_detections(xyxy=[[10, 10, 20, 20]])],
|
||||
mock_detections(xyxy=[[10, 10, 20, 20]]),
|
||||
DoesNotRaise()
|
||||
DoesNotRaise(),
|
||||
), # single detection with xyxy field
|
||||
(
|
||||
[
|
||||
mock_detections(xyxy=[[10, 10, 20, 20]]),
|
||||
Detections.empty()
|
||||
],
|
||||
[mock_detections(xyxy=[[10, 10, 20, 20]]), Detections.empty()],
|
||||
mock_detections(xyxy=[[10, 10, 20, 20]]),
|
||||
DoesNotRaise()
|
||||
DoesNotRaise(),
|
||||
), # single detection with xyxy field + empty detection
|
||||
(
|
||||
[
|
||||
mock_detections(xyxy=[[10, 10, 20, 20]]),
|
||||
mock_detections(xyxy=[[20, 20, 30, 30]])
|
||||
mock_detections(xyxy=[[20, 20, 30, 30]]),
|
||||
],
|
||||
mock_detections(
|
||||
xyxy=[
|
||||
[10, 10, 20, 20],
|
||||
[20, 20, 30, 30]
|
||||
]),
|
||||
DoesNotRaise()
|
||||
mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]),
|
||||
DoesNotRaise(),
|
||||
), # two detections with xyxy field
|
||||
(
|
||||
[
|
||||
mock_detections(
|
||||
xyxy=[[10, 10, 20, 20]],
|
||||
class_id=[0]),
|
||||
mock_detections(
|
||||
xyxy=[[20, 20, 30, 30]])
|
||||
mock_detections(xyxy=[[10, 10, 20, 20]], class_id=[0]),
|
||||
mock_detections(xyxy=[[20, 20, 30, 30]]),
|
||||
],
|
||||
mock_detections(
|
||||
xyxy=[
|
||||
[10, 10, 20, 20],
|
||||
[20, 20, 30, 30]
|
||||
]),
|
||||
DoesNotRaise()
|
||||
mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]),
|
||||
DoesNotRaise(),
|
||||
), # detection with xyxy, class_id fields + detection with xyxy field
|
||||
(
|
||||
(
|
||||
[
|
||||
mock_detections(
|
||||
xyxy=[[10, 10, 20, 20]],
|
||||
class_id=[0]),
|
||||
mock_detections(
|
||||
xyxy=[[20, 20, 30, 30]],
|
||||
class_id=[1]),
|
||||
mock_detections(xyxy=[[10, 10, 20, 20]], class_id=[0]),
|
||||
mock_detections(xyxy=[[20, 20, 30, 30]], class_id=[1]),
|
||||
],
|
||||
mock_detections(
|
||||
xyxy=[
|
||||
[10, 10, 20, 20],
|
||||
[20, 20, 30, 30]
|
||||
],
|
||||
class_id=[0, 1]
|
||||
),
|
||||
DoesNotRaise()
|
||||
mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]], class_id=[0, 1]),
|
||||
DoesNotRaise(),
|
||||
), # two detections with xyxy, class_id fields
|
||||
]
|
||||
],
|
||||
)
|
||||
def test_merge(
|
||||
detections_list: List[Detections],
|
||||
expected_result: Optional[Detections],
|
||||
exception: Exception
|
||||
detections_list: List[Detections],
|
||||
expected_result: Optional[Detections],
|
||||
exception: Exception,
|
||||
) -> None:
|
||||
with exception:
|
||||
result = Detections.merge(detections_list=detections_list)
|
||||
assert result == expected_result
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"detections, anchor, expected_result, exception",
|
||||
[
|
||||
(
|
||||
Detections.empty(),
|
||||
Position.CENTER,
|
||||
np.empty((0, 2), dtype=np.float32),
|
||||
DoesNotRaise(),
|
||||
), # empty detections
|
||||
(
|
||||
mock_detections(xyxy=[[10, 10, 20, 20]]),
|
||||
Position.CENTER,
|
||||
np.array([[15, 15]], dtype=np.float32),
|
||||
DoesNotRaise(),
|
||||
), # single detection; center anchor
|
||||
(
|
||||
mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]),
|
||||
Position.CENTER,
|
||||
np.array([[15, 15], [25, 25]], dtype=np.float32),
|
||||
DoesNotRaise(),
|
||||
), # two detections; center anchor
|
||||
(
|
||||
mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]),
|
||||
Position.CENTER_LEFT,
|
||||
np.array([[10, 15], [20, 25]], dtype=np.float32),
|
||||
DoesNotRaise(),
|
||||
), # two detections; center left anchor
|
||||
(
|
||||
mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]),
|
||||
Position.CENTER_RIGHT,
|
||||
np.array([[20, 15], [30, 25]], dtype=np.float32),
|
||||
DoesNotRaise(),
|
||||
), # two detections; center right anchor
|
||||
(
|
||||
mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]),
|
||||
Position.TOP_CENTER,
|
||||
np.array([[15, 10], [25, 20]], dtype=np.float32),
|
||||
DoesNotRaise(),
|
||||
), # two detections; top center anchor
|
||||
(
|
||||
mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]),
|
||||
Position.TOP_LEFT,
|
||||
np.array([[10, 10], [20, 20]], dtype=np.float32),
|
||||
DoesNotRaise(),
|
||||
), # two detections; top left anchor
|
||||
(
|
||||
mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]),
|
||||
Position.TOP_RIGHT,
|
||||
np.array([[20, 10], [30, 20]], dtype=np.float32),
|
||||
DoesNotRaise(),
|
||||
), # two detections; top right anchor
|
||||
(
|
||||
mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]),
|
||||
Position.BOTTOM_CENTER,
|
||||
np.array([[15, 20], [25, 30]], dtype=np.float32),
|
||||
DoesNotRaise(),
|
||||
), # two detections; bottom center anchor
|
||||
(
|
||||
mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]),
|
||||
Position.BOTTOM_LEFT,
|
||||
np.array([[10, 20], [20, 30]], dtype=np.float32),
|
||||
DoesNotRaise(),
|
||||
), # two detections; bottom left anchor
|
||||
(
|
||||
mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]),
|
||||
Position.BOTTOM_RIGHT,
|
||||
np.array([[20, 20], [30, 30]], dtype=np.float32),
|
||||
DoesNotRaise(),
|
||||
), # two detections; bottom right anchor
|
||||
],
|
||||
)
|
||||
def test_get_anchor_coordinates(
|
||||
detections: Detections,
|
||||
anchor: Position,
|
||||
expected_result: np.ndarray,
|
||||
exception: Exception,
|
||||
) -> None:
|
||||
result = detections.get_anchor_coordinates(anchor)
|
||||
with exception:
|
||||
assert np.array_equal(result, expected_result)
|
||||
|
||||
+351
-139
@@ -1,11 +1,19 @@
|
||||
from contextlib import ExitStack as DoesNotRaise
|
||||
from typing import Optional, Tuple, List
|
||||
|
||||
import pytest
|
||||
from typing import List, Optional, Tuple
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from supervision.detection.utils import non_max_suppression, clip_boxes, filter_polygons_by_area
|
||||
from supervision.detection.utils import (
|
||||
clip_boxes,
|
||||
filter_polygons_by_area,
|
||||
move_boxes,
|
||||
non_max_suppression,
|
||||
process_roboflow_result,
|
||||
)
|
||||
|
||||
TEST_MASK = np.zeros((1, 1000, 1000), dtype=bool)
|
||||
TEST_MASK[:, 300:351, 200:251] = True
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
@@ -15,116 +23,101 @@ from supervision.detection.utils import non_max_suppression, clip_boxes, filter_
|
||||
np.empty(shape=(0, 5)),
|
||||
0.5,
|
||||
np.array([]),
|
||||
DoesNotRaise()
|
||||
DoesNotRaise(),
|
||||
), # single box with no category
|
||||
(
|
||||
np.array([
|
||||
[10.0, 10.0, 40.0, 40.0, 0.8]
|
||||
]),
|
||||
np.array([[10.0, 10.0, 40.0, 40.0, 0.8]]),
|
||||
0.5,
|
||||
np.array([
|
||||
True
|
||||
]),
|
||||
DoesNotRaise()
|
||||
np.array([True]),
|
||||
DoesNotRaise(),
|
||||
), # single box with no category
|
||||
(
|
||||
np.array([
|
||||
[10.0, 10.0, 40.0, 40.0, 0.8, 0]
|
||||
]),
|
||||
np.array([[10.0, 10.0, 40.0, 40.0, 0.8, 0]]),
|
||||
0.5,
|
||||
np.array([
|
||||
True
|
||||
]),
|
||||
DoesNotRaise()
|
||||
np.array([True]),
|
||||
DoesNotRaise(),
|
||||
), # single box with category
|
||||
(
|
||||
np.array([
|
||||
[10.0, 10.0, 40.0, 40.0, 0.8],
|
||||
[15.0, 15.0, 40.0, 40.0, 0.9],
|
||||
]),
|
||||
np.array(
|
||||
[
|
||||
[10.0, 10.0, 40.0, 40.0, 0.8],
|
||||
[15.0, 15.0, 40.0, 40.0, 0.9],
|
||||
]
|
||||
),
|
||||
0.5,
|
||||
np.array([
|
||||
False,
|
||||
True
|
||||
]),
|
||||
DoesNotRaise()
|
||||
np.array([False, True]),
|
||||
DoesNotRaise(),
|
||||
), # two boxes with no category
|
||||
(
|
||||
np.array([
|
||||
[10.0, 10.0, 40.0, 40.0, 0.8, 0],
|
||||
[15.0, 15.0, 40.0, 40.0, 0.9, 1],
|
||||
]),
|
||||
np.array(
|
||||
[
|
||||
[10.0, 10.0, 40.0, 40.0, 0.8, 0],
|
||||
[15.0, 15.0, 40.0, 40.0, 0.9, 1],
|
||||
]
|
||||
),
|
||||
0.5,
|
||||
np.array([
|
||||
True,
|
||||
True
|
||||
]),
|
||||
DoesNotRaise()
|
||||
np.array([True, True]),
|
||||
DoesNotRaise(),
|
||||
), # two boxes with different category
|
||||
(
|
||||
np.array([
|
||||
[10.0, 10.0, 40.0, 40.0, 0.8, 0],
|
||||
[15.0, 15.0, 40.0, 40.0, 0.9, 0],
|
||||
]),
|
||||
np.array(
|
||||
[
|
||||
[10.0, 10.0, 40.0, 40.0, 0.8, 0],
|
||||
[15.0, 15.0, 40.0, 40.0, 0.9, 0],
|
||||
]
|
||||
),
|
||||
0.5,
|
||||
np.array([
|
||||
False,
|
||||
True
|
||||
]),
|
||||
DoesNotRaise()
|
||||
np.array([False, True]),
|
||||
DoesNotRaise(),
|
||||
), # two boxes with same category
|
||||
(
|
||||
np.array([
|
||||
[0.0, 0.0, 30.0, 40.0, 0.8],
|
||||
[5.0, 5.0, 35.0, 45.0, 0.9],
|
||||
[10.0, 10.0, 40.0, 50.0, 0.85],
|
||||
]),
|
||||
np.array(
|
||||
[
|
||||
[0.0, 0.0, 30.0, 40.0, 0.8],
|
||||
[5.0, 5.0, 35.0, 45.0, 0.9],
|
||||
[10.0, 10.0, 40.0, 50.0, 0.85],
|
||||
]
|
||||
),
|
||||
0.5,
|
||||
np.array([
|
||||
False,
|
||||
True,
|
||||
False
|
||||
]),
|
||||
DoesNotRaise()
|
||||
np.array([False, True, False]),
|
||||
DoesNotRaise(),
|
||||
), # three boxes with no category
|
||||
(
|
||||
np.array([
|
||||
[0.0, 0.0, 30.0, 40.0, 0.8, 0],
|
||||
[5.0, 5.0, 35.0, 45.0, 0.9, 1],
|
||||
[10.0, 10.0, 40.0, 50.0, 0.85, 2],
|
||||
]),
|
||||
np.array(
|
||||
[
|
||||
[0.0, 0.0, 30.0, 40.0, 0.8, 0],
|
||||
[5.0, 5.0, 35.0, 45.0, 0.9, 1],
|
||||
[10.0, 10.0, 40.0, 50.0, 0.85, 2],
|
||||
]
|
||||
),
|
||||
0.5,
|
||||
np.array([
|
||||
True,
|
||||
True,
|
||||
True
|
||||
]),
|
||||
DoesNotRaise()
|
||||
np.array([True, True, True]),
|
||||
DoesNotRaise(),
|
||||
), # three boxes with same category
|
||||
(
|
||||
np.array([
|
||||
[0.0, 0.0, 30.0, 40.0, 0.8, 0],
|
||||
[5.0, 5.0, 35.0, 45.0, 0.9, 0],
|
||||
[10.0, 10.0, 40.0, 50.0, 0.85, 1],
|
||||
]),
|
||||
np.array(
|
||||
[
|
||||
[0.0, 0.0, 30.0, 40.0, 0.8, 0],
|
||||
[5.0, 5.0, 35.0, 45.0, 0.9, 0],
|
||||
[10.0, 10.0, 40.0, 50.0, 0.85, 1],
|
||||
]
|
||||
),
|
||||
0.5,
|
||||
np.array([
|
||||
False,
|
||||
True,
|
||||
True
|
||||
]),
|
||||
DoesNotRaise()
|
||||
np.array([False, True, True]),
|
||||
DoesNotRaise(),
|
||||
), # three boxes with different category
|
||||
]
|
||||
],
|
||||
)
|
||||
def test_non_max_suppression(
|
||||
predictions: np.ndarray,
|
||||
iou_threshold: float,
|
||||
expected_result: Optional[np.ndarray],
|
||||
exception: Exception
|
||||
predictions: np.ndarray,
|
||||
iou_threshold: float,
|
||||
expected_result: Optional[np.ndarray],
|
||||
exception: Exception,
|
||||
) -> None:
|
||||
with exception:
|
||||
result = non_max_suppression(predictions=predictions, iou_threshold=iou_threshold)
|
||||
result = non_max_suppression(
|
||||
predictions=predictions, iou_threshold=iou_threshold
|
||||
)
|
||||
assert np.array_equal(result, expected_result)
|
||||
|
||||
|
||||
@@ -137,53 +130,37 @@ def test_non_max_suppression(
|
||||
np.empty(shape=(0, 4)),
|
||||
),
|
||||
(
|
||||
np.array([
|
||||
[1.0, 1.0, 1279.0, 719.0]
|
||||
]),
|
||||
np.array([[1.0, 1.0, 1279.0, 719.0]]),
|
||||
(1280, 720),
|
||||
np.array([
|
||||
[1.0, 1.0, 1279.0, 719.0]
|
||||
]),
|
||||
np.array([[1.0, 1.0, 1279.0, 719.0]]),
|
||||
),
|
||||
(
|
||||
np.array([
|
||||
[-1.0, 1.0, 1279.0, 719.0]
|
||||
]),
|
||||
np.array([[-1.0, 1.0, 1279.0, 719.0]]),
|
||||
(1280, 720),
|
||||
np.array([
|
||||
[0.0, 1.0, 1279.0, 719.0]
|
||||
]),
|
||||
np.array([[0.0, 1.0, 1279.0, 719.0]]),
|
||||
),
|
||||
(
|
||||
np.array([
|
||||
[1.0, -1.0, 1279.0, 719.0]
|
||||
]),
|
||||
np.array([[1.0, -1.0, 1279.0, 719.0]]),
|
||||
(1280, 720),
|
||||
np.array([
|
||||
[1.0, 0.0, 1279.0, 719.0]
|
||||
]),
|
||||
np.array([[1.0, 0.0, 1279.0, 719.0]]),
|
||||
),
|
||||
(
|
||||
np.array([
|
||||
[1.0, 1.0, 1281.0, 719.0]
|
||||
]),
|
||||
np.array([[1.0, 1.0, 1281.0, 719.0]]),
|
||||
(1280, 720),
|
||||
np.array([
|
||||
[1.0, 1.0, 1280.0, 719.0]
|
||||
]),
|
||||
np.array([[1.0, 1.0, 1280.0, 719.0]]),
|
||||
),
|
||||
(
|
||||
np.array([
|
||||
[1.0, 1.0, 1279.0, 721.0]
|
||||
]),
|
||||
np.array([[1.0, 1.0, 1279.0, 721.0]]),
|
||||
(1280, 720),
|
||||
np.array([
|
||||
[1.0, 1.0, 1279.0, 720.0]
|
||||
]),
|
||||
np.array([[1.0, 1.0, 1279.0, 720.0]]),
|
||||
),
|
||||
]
|
||||
],
|
||||
)
|
||||
def test_clip_boxes(boxes_xyxy: np.ndarray, frame_resolution_wh: Tuple[int, int], expected_result: np.ndarray) -> None:
|
||||
def test_clip_boxes(
|
||||
boxes_xyxy: np.ndarray,
|
||||
frame_resolution_wh: Tuple[int, int],
|
||||
expected_result: np.ndarray,
|
||||
) -> None:
|
||||
result = clip_boxes(boxes_xyxy=boxes_xyxy, frame_resolution_wh=frame_resolution_wh)
|
||||
assert np.array_equal(result, expected_result)
|
||||
|
||||
@@ -196,83 +173,318 @@ def test_clip_boxes(boxes_xyxy: np.ndarray, frame_resolution_wh: Tuple[int, int]
|
||||
None,
|
||||
None,
|
||||
[np.array([[0, 0], [0, 10], [10, 10], [10, 0]])],
|
||||
DoesNotRaise()
|
||||
DoesNotRaise(),
|
||||
), # single polygon without area constraints
|
||||
(
|
||||
[np.array([[0, 0], [0, 10], [10, 10], [10, 0]])],
|
||||
50,
|
||||
None,
|
||||
[np.array([[0, 0], [0, 10], [10, 10], [10, 0]])],
|
||||
DoesNotRaise()
|
||||
DoesNotRaise(),
|
||||
), # single polygon with min_area constraint
|
||||
(
|
||||
[np.array([[0, 0], [0, 10], [10, 10], [10, 0]])],
|
||||
None,
|
||||
50,
|
||||
[],
|
||||
DoesNotRaise()
|
||||
DoesNotRaise(),
|
||||
), # single polygon with max_area constraint
|
||||
(
|
||||
[
|
||||
np.array([[0, 0], [0, 10], [10, 10], [10, 0]]),
|
||||
np.array([[0, 0], [0, 20], [20, 20], [20, 0]])
|
||||
np.array([[0, 0], [0, 20], [20, 20], [20, 0]]),
|
||||
],
|
||||
200,
|
||||
None,
|
||||
[np.array([[0, 0], [0, 20], [20, 20], [20, 0]])],
|
||||
DoesNotRaise()
|
||||
DoesNotRaise(),
|
||||
), # two polygons with min_area constraint
|
||||
(
|
||||
[
|
||||
np.array([[0, 0], [0, 10], [10, 10], [10, 0]]),
|
||||
np.array([[0, 0], [0, 20], [20, 20], [20, 0]])
|
||||
np.array([[0, 0], [0, 20], [20, 20], [20, 0]]),
|
||||
],
|
||||
None,
|
||||
200,
|
||||
[np.array([[0, 0], [0, 10], [10, 10], [10, 0]])],
|
||||
DoesNotRaise()
|
||||
DoesNotRaise(),
|
||||
), # two polygons with max_area constraint
|
||||
(
|
||||
[
|
||||
np.array([[0, 0], [0, 10], [10, 10], [10, 0]]),
|
||||
np.array([[0, 0], [0, 20], [20, 20], [20, 0]])
|
||||
np.array([[0, 0], [0, 20], [20, 20], [20, 0]]),
|
||||
],
|
||||
200,
|
||||
200,
|
||||
[],
|
||||
DoesNotRaise()
|
||||
DoesNotRaise(),
|
||||
), # two polygons with both area constraints
|
||||
(
|
||||
[
|
||||
np.array([[0, 0], [0, 10], [10, 10], [10, 0]]),
|
||||
np.array([[0, 0], [0, 20], [20, 20], [20, 0]])
|
||||
np.array([[0, 0], [0, 20], [20, 20], [20, 0]]),
|
||||
],
|
||||
100,
|
||||
100,
|
||||
[np.array([[0, 0], [0, 10], [10, 10], [10, 0]])],
|
||||
DoesNotRaise()
|
||||
), # two polygons with min_area and max_area equal to the area of the first polygon
|
||||
DoesNotRaise(),
|
||||
), # two polygons with min_area and
|
||||
# max_area equal to the area of the first polygon
|
||||
(
|
||||
[
|
||||
np.array([[0, 0], [0, 10], [10, 10], [10, 0]]),
|
||||
np.array([[0, 0], [0, 20], [20, 20], [20, 0]])
|
||||
np.array([[0, 0], [0, 20], [20, 20], [20, 0]]),
|
||||
],
|
||||
400,
|
||||
400,
|
||||
[np.array([[0, 0], [0, 20], [20, 20], [20, 0]])],
|
||||
DoesNotRaise()
|
||||
), # two polygons with min_area and max_area equal to the area of the second polygon
|
||||
]
|
||||
DoesNotRaise(),
|
||||
), # two polygons with min_area and
|
||||
# max_area equal to the area of the second polygon
|
||||
],
|
||||
)
|
||||
def test_filter_polygons_by_area(
|
||||
polygons: List[np.ndarray],
|
||||
min_area: Optional[float],
|
||||
max_area: Optional[float],
|
||||
expected_result: List[np.ndarray],
|
||||
exception: Exception
|
||||
polygons: List[np.ndarray],
|
||||
min_area: Optional[float],
|
||||
max_area: Optional[float],
|
||||
expected_result: List[np.ndarray],
|
||||
exception: Exception,
|
||||
) -> None:
|
||||
with exception:
|
||||
result = filter_polygons_by_area(polygons=polygons, min_area=min_area, max_area=max_area)
|
||||
result = filter_polygons_by_area(
|
||||
polygons=polygons, min_area=min_area, max_area=max_area
|
||||
)
|
||||
assert len(result) == len(expected_result)
|
||||
for result_polygon, expected_result_polygon in zip(result, expected_result):
|
||||
assert np.array_equal(result_polygon, expected_result_polygon)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"roboflow_result, expected_result, exception",
|
||||
[
|
||||
(
|
||||
{"predictions": [], "image": {"width": 1000, "height": 1000}},
|
||||
(np.empty((0, 4)), np.empty(0), np.empty(0), None),
|
||||
DoesNotRaise(),
|
||||
), # empty result
|
||||
(
|
||||
{
|
||||
"predictions": [
|
||||
{
|
||||
"x": 200.0,
|
||||
"y": 300.0,
|
||||
"width": 50.0,
|
||||
"height": 50.0,
|
||||
"confidence": 0.9,
|
||||
"class_id": 0,
|
||||
"class": "person",
|
||||
}
|
||||
],
|
||||
"image": {"width": 1000, "height": 1000},
|
||||
},
|
||||
(
|
||||
np.array([[175.0, 275.0, 225.0, 325.0]]),
|
||||
np.array([0.9]),
|
||||
np.array([0]),
|
||||
None,
|
||||
),
|
||||
DoesNotRaise(),
|
||||
), # single correct object detection result
|
||||
(
|
||||
{
|
||||
"predictions": [
|
||||
{
|
||||
"x": 200.0,
|
||||
"y": 300.0,
|
||||
"width": 50.0,
|
||||
"height": 50.0,
|
||||
"confidence": 0.9,
|
||||
"class_id": 0,
|
||||
"class": "person",
|
||||
},
|
||||
{
|
||||
"x": 500.0,
|
||||
"y": 500.0,
|
||||
"width": 100.0,
|
||||
"height": 100.0,
|
||||
"confidence": 0.8,
|
||||
"class_id": 7,
|
||||
"class": "truck",
|
||||
},
|
||||
],
|
||||
"image": {"width": 1000, "height": 1000},
|
||||
},
|
||||
(
|
||||
np.array([[175.0, 275.0, 225.0, 325.0], [450.0, 450.0, 550.0, 550.0]]),
|
||||
np.array([0.9, 0.8]),
|
||||
np.array([0, 7]),
|
||||
None,
|
||||
),
|
||||
DoesNotRaise(),
|
||||
), # two correct object detection result
|
||||
(
|
||||
{
|
||||
"predictions": [
|
||||
{
|
||||
"x": 200.0,
|
||||
"y": 300.0,
|
||||
"width": 50.0,
|
||||
"height": 50.0,
|
||||
"confidence": 0.9,
|
||||
"class_id": 0,
|
||||
"class": "person",
|
||||
"points": [],
|
||||
}
|
||||
],
|
||||
"image": {"width": 1000, "height": 1000},
|
||||
},
|
||||
(np.empty((0, 4)), np.empty(0), np.empty(0), None),
|
||||
DoesNotRaise(),
|
||||
), # single incorrect instance segmentation result with no points
|
||||
(
|
||||
{
|
||||
"predictions": [
|
||||
{
|
||||
"x": 200.0,
|
||||
"y": 300.0,
|
||||
"width": 50.0,
|
||||
"height": 50.0,
|
||||
"confidence": 0.9,
|
||||
"class_id": 0,
|
||||
"class": "person",
|
||||
"points": [{"x": 200.0, "y": 300.0}, {"x": 250.0, "y": 300.0}],
|
||||
}
|
||||
],
|
||||
"image": {"width": 1000, "height": 1000},
|
||||
},
|
||||
(np.empty((0, 4)), np.empty(0), np.empty(0), None),
|
||||
DoesNotRaise(),
|
||||
), # single incorrect instance segmentation result with no enough points
|
||||
(
|
||||
{
|
||||
"predictions": [
|
||||
{
|
||||
"x": 200.0,
|
||||
"y": 300.0,
|
||||
"width": 50.0,
|
||||
"height": 50.0,
|
||||
"confidence": 0.9,
|
||||
"class_id": 0,
|
||||
"class": "person",
|
||||
"points": [
|
||||
{"x": 200.0, "y": 300.0},
|
||||
{"x": 250.0, "y": 300.0},
|
||||
{"x": 250.0, "y": 350.0},
|
||||
{"x": 200.0, "y": 350.0},
|
||||
],
|
||||
}
|
||||
],
|
||||
"image": {"width": 1000, "height": 1000},
|
||||
},
|
||||
(
|
||||
np.array([[175.0, 275.0, 225.0, 325.0]]),
|
||||
np.array([0.9]),
|
||||
np.array([0]),
|
||||
TEST_MASK,
|
||||
),
|
||||
DoesNotRaise(),
|
||||
), # single incorrect instance segmentation result with no enough points
|
||||
(
|
||||
{
|
||||
"predictions": [
|
||||
{
|
||||
"x": 200.0,
|
||||
"y": 300.0,
|
||||
"width": 50.0,
|
||||
"height": 50.0,
|
||||
"confidence": 0.9,
|
||||
"class_id": 0,
|
||||
"class": "person",
|
||||
"points": [
|
||||
{"x": 200.0, "y": 300.0},
|
||||
{"x": 250.0, "y": 300.0},
|
||||
{"x": 250.0, "y": 350.0},
|
||||
{"x": 200.0, "y": 350.0},
|
||||
],
|
||||
},
|
||||
{
|
||||
"x": 500.0,
|
||||
"y": 500.0,
|
||||
"width": 100.0,
|
||||
"height": 100.0,
|
||||
"confidence": 0.8,
|
||||
"class_id": 7,
|
||||
"class": "truck",
|
||||
"points": [],
|
||||
},
|
||||
],
|
||||
"image": {"width": 1000, "height": 1000},
|
||||
},
|
||||
(
|
||||
np.array([[175.0, 275.0, 225.0, 325.0]]),
|
||||
np.array([0.9]),
|
||||
np.array([0]),
|
||||
TEST_MASK,
|
||||
),
|
||||
DoesNotRaise(),
|
||||
), # two instance segmentation results - one correct, one incorrect
|
||||
],
|
||||
)
|
||||
def test_process_roboflow_result(
|
||||
roboflow_result: dict,
|
||||
expected_result: Tuple[np.ndarray, np.ndarray, np.ndarray, Optional[np.ndarray]],
|
||||
exception: Exception,
|
||||
) -> None:
|
||||
with exception:
|
||||
result = process_roboflow_result(roboflow_result=roboflow_result)
|
||||
assert np.array_equal(result[0], expected_result[0])
|
||||
assert np.array_equal(result[1], expected_result[1])
|
||||
assert np.array_equal(result[2], expected_result[2])
|
||||
assert (result[3] is None and expected_result[3] is None) or (
|
||||
np.array_equal(result[3], expected_result[3])
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"xyxy, offset, expected_result, exception",
|
||||
[
|
||||
(
|
||||
np.empty(shape=(0, 4)),
|
||||
np.array([0, 0]),
|
||||
np.empty(shape=(0, 4)),
|
||||
DoesNotRaise(),
|
||||
), # empty xyxy array
|
||||
(
|
||||
np.array([[0, 0, 10, 10]]),
|
||||
np.array([0, 0]),
|
||||
np.array([[0, 0, 10, 10]]),
|
||||
DoesNotRaise(),
|
||||
), # single box with zero offset
|
||||
(
|
||||
np.array([[0, 0, 10, 10]]),
|
||||
np.array([10, 10]),
|
||||
np.array([[10, 10, 20, 20]]),
|
||||
DoesNotRaise(),
|
||||
), # single box with non-zero offset
|
||||
(
|
||||
np.array([[0, 0, 10, 10], [0, 0, 10, 10]]),
|
||||
np.array([10, 10]),
|
||||
np.array([[10, 10, 20, 20], [10, 10, 20, 20]]),
|
||||
DoesNotRaise(),
|
||||
), # two boxes with non-zero offset
|
||||
(
|
||||
np.array([[0, 0, 10, 10], [0, 0, 10, 10]]),
|
||||
np.array([-10, -10]),
|
||||
np.array([[-10, -10, 0, 0], [-10, -10, 0, 0]]),
|
||||
DoesNotRaise(),
|
||||
), # two boxes with negative offset
|
||||
],
|
||||
)
|
||||
def test_move_boxes(
|
||||
xyxy: np.ndarray,
|
||||
offset: np.ndarray,
|
||||
expected_result: np.ndarray,
|
||||
exception: Exception,
|
||||
) -> None:
|
||||
result = move_boxes(xyxy=xyxy, offset=offset)
|
||||
assert np.array_equal(result, expected_result)
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
import pytest
|
||||
|
||||
from supervision.geometry.core import Vector, Point
|
||||
from supervision.geometry.core import Point, Vector
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
|
||||
@@ -0,0 +1,459 @@
|
||||
from contextlib import ExitStack as DoesNotRaise
|
||||
from test.utils import assert_almost_equal, mock_detections
|
||||
from typing import Optional, Union
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from supervision.detection.core import Detections
|
||||
from supervision.metrics.detection import (
|
||||
ConfusionMatrix,
|
||||
MeanAveragePrecision,
|
||||
detections_to_tensor,
|
||||
)
|
||||
|
||||
CLASSES = np.arange(80)
|
||||
NUM_CLASSES = len(CLASSES)
|
||||
|
||||
PREDICTIONS = np.array(
|
||||
[
|
||||
[2254, 906, 2447, 1353, 0.90538, 0],
|
||||
[2049, 1133, 2226, 1371, 0.59002, 56],
|
||||
[727, 1224, 838, 1601, 0.51119, 39],
|
||||
[808, 1214, 910, 1564, 0.45287, 39],
|
||||
[6, 52, 1131, 2133, 0.45057, 72],
|
||||
[299, 1225, 512, 1663, 0.45029, 39],
|
||||
[529, 874, 645, 945, 0.31101, 39],
|
||||
[8, 47, 1935, 2135, 0.28192, 72],
|
||||
[2265, 813, 2328, 901, 0.2714, 62],
|
||||
],
|
||||
dtype=np.float32,
|
||||
)
|
||||
|
||||
TARGET_TENSORS = [
|
||||
np.array(
|
||||
[
|
||||
[2254, 906, 2447, 1353, 0],
|
||||
[2049, 1133, 2226, 1371, 56],
|
||||
[727, 1224, 838, 1601, 39],
|
||||
[808, 1214, 910, 1564, 39],
|
||||
[6, 52, 1131, 2133, 72],
|
||||
[299, 1225, 512, 1663, 39],
|
||||
[529, 874, 645, 945, 39],
|
||||
[8, 47, 1935, 2135, 72],
|
||||
[2265, 813, 2328, 901, 62],
|
||||
]
|
||||
)
|
||||
]
|
||||
|
||||
DETECTIONS = Detections(
|
||||
xyxy=PREDICTIONS[:, :4],
|
||||
confidence=PREDICTIONS[:, 4],
|
||||
class_id=PREDICTIONS[:, 5].astype(int),
|
||||
)
|
||||
CERTAIN_DETECTIONS = Detections(
|
||||
xyxy=PREDICTIONS[:, :4],
|
||||
confidence=np.ones(len(PREDICTIONS)),
|
||||
class_id=PREDICTIONS[:, 5].astype(int),
|
||||
)
|
||||
|
||||
DETECTION_TENSORS = [
|
||||
np.concatenate(
|
||||
[
|
||||
det.xyxy,
|
||||
np.expand_dims(det.class_id, 1),
|
||||
np.expand_dims(det.confidence, 1),
|
||||
],
|
||||
axis=1,
|
||||
)
|
||||
for det in [DETECTIONS]
|
||||
]
|
||||
CERTAIN_DETECTION_TENSORS = [
|
||||
np.concatenate(
|
||||
[
|
||||
det.xyxy,
|
||||
np.expand_dims(det.class_id, 1),
|
||||
np.ones((len(det), 1)),
|
||||
],
|
||||
axis=1,
|
||||
)
|
||||
for det in [DETECTIONS]
|
||||
]
|
||||
|
||||
IDEAL_MATCHES = np.stack(
|
||||
[
|
||||
np.arange(len(PREDICTIONS)),
|
||||
np.arange(len(PREDICTIONS)),
|
||||
np.ones(len(PREDICTIONS)),
|
||||
],
|
||||
axis=1,
|
||||
)
|
||||
|
||||
|
||||
def create_empty_conf_matrix(num_classes: int, do_add_dummy_class: bool = True):
|
||||
if do_add_dummy_class:
|
||||
num_classes += 1
|
||||
return np.zeros((num_classes, num_classes))
|
||||
|
||||
|
||||
def update_ideal_conf_matrix(conf_matrix: np.ndarray, class_ids: np.ndarray):
|
||||
for class_id, count in zip(*np.unique(class_ids, return_counts=True)):
|
||||
class_id = int(class_id)
|
||||
conf_matrix[class_id, class_id] += count
|
||||
return conf_matrix
|
||||
|
||||
|
||||
def worsen_ideal_conf_matrix(
|
||||
conf_matrix: np.ndarray, class_ids: Union[np.ndarray, list]
|
||||
):
|
||||
for class_id in class_ids:
|
||||
class_id = int(class_id)
|
||||
conf_matrix[class_id, class_id] -= 1
|
||||
conf_matrix[class_id, 80] += 1
|
||||
return conf_matrix
|
||||
|
||||
|
||||
IDEAL_CONF_MATRIX = create_empty_conf_matrix(NUM_CLASSES)
|
||||
IDEAL_CONF_MATRIX = update_ideal_conf_matrix(IDEAL_CONF_MATRIX, PREDICTIONS[:, 5])
|
||||
|
||||
GOOD_CONF_MATRIX = worsen_ideal_conf_matrix(IDEAL_CONF_MATRIX.copy(), [62, 72])
|
||||
|
||||
BAD_CONF_MATRIX = worsen_ideal_conf_matrix(
|
||||
IDEAL_CONF_MATRIX.copy(), [62, 72, 72, 39, 39, 39, 39, 56]
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"detections, with_confidence, expected_result, exception",
|
||||
[
|
||||
(
|
||||
Detections.empty(),
|
||||
False,
|
||||
np.empty((0, 5), dtype=np.float32),
|
||||
DoesNotRaise(),
|
||||
), # empty detections; no confidence
|
||||
(
|
||||
Detections.empty(),
|
||||
True,
|
||||
np.empty((0, 6), dtype=np.float32),
|
||||
DoesNotRaise(),
|
||||
), # empty detections; with confidence
|
||||
(
|
||||
mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[0], confidence=[0.5]),
|
||||
False,
|
||||
np.array([[0, 0, 10, 10, 0]], dtype=np.float32),
|
||||
DoesNotRaise(),
|
||||
), # single detection; no confidence
|
||||
(
|
||||
mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[0], confidence=[0.5]),
|
||||
True,
|
||||
np.array([[0, 0, 10, 10, 0, 0.5]], dtype=np.float32),
|
||||
DoesNotRaise(),
|
||||
), # single detection; with confidence
|
||||
(
|
||||
mock_detections(
|
||||
xyxy=[[0, 0, 10, 10], [0, 0, 20, 20]],
|
||||
class_id=[0, 1],
|
||||
confidence=[0.5, 0.2],
|
||||
),
|
||||
False,
|
||||
np.array([[0, 0, 10, 10, 0], [0, 0, 20, 20, 1]], dtype=np.float32),
|
||||
DoesNotRaise(),
|
||||
), # multiple detections; no confidence
|
||||
(
|
||||
mock_detections(
|
||||
xyxy=[[0, 0, 10, 10], [0, 0, 20, 20]],
|
||||
class_id=[0, 1],
|
||||
confidence=[0.5, 0.2],
|
||||
),
|
||||
True,
|
||||
np.array(
|
||||
[[0, 0, 10, 10, 0, 0.5], [0, 0, 20, 20, 1, 0.2]], dtype=np.float32
|
||||
),
|
||||
DoesNotRaise(),
|
||||
), # multiple detections; with confidence
|
||||
],
|
||||
)
|
||||
def test_detections_to_tensor(
|
||||
detections: Detections,
|
||||
with_confidence: bool,
|
||||
expected_result: Optional[np.ndarray],
|
||||
exception: Exception,
|
||||
):
|
||||
with exception:
|
||||
result = detections_to_tensor(
|
||||
detections=detections, with_confidence=with_confidence
|
||||
)
|
||||
assert np.array_equal(result, expected_result)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"predictions, targets, classes, conf_threshold, iou_threshold, expected_result,"
|
||||
" exception",
|
||||
[
|
||||
(
|
||||
DETECTION_TENSORS,
|
||||
TARGET_TENSORS,
|
||||
CLASSES,
|
||||
0.2,
|
||||
0.5,
|
||||
IDEAL_CONF_MATRIX,
|
||||
DoesNotRaise(),
|
||||
),
|
||||
(
|
||||
[],
|
||||
[],
|
||||
CLASSES,
|
||||
0.2,
|
||||
0.5,
|
||||
create_empty_conf_matrix(NUM_CLASSES),
|
||||
DoesNotRaise(),
|
||||
),
|
||||
(
|
||||
DETECTION_TENSORS,
|
||||
TARGET_TENSORS,
|
||||
CLASSES,
|
||||
0.3,
|
||||
0.5,
|
||||
GOOD_CONF_MATRIX,
|
||||
DoesNotRaise(),
|
||||
),
|
||||
(
|
||||
DETECTION_TENSORS,
|
||||
TARGET_TENSORS,
|
||||
CLASSES,
|
||||
0.6,
|
||||
0.5,
|
||||
BAD_CONF_MATRIX,
|
||||
DoesNotRaise(),
|
||||
),
|
||||
(
|
||||
[
|
||||
np.array(
|
||||
[
|
||||
[0.0, 0.0, 3.0, 3.0, 0, 0.9], # correct detection of [0]
|
||||
[
|
||||
0.1,
|
||||
0.1,
|
||||
3.0,
|
||||
3.0,
|
||||
0,
|
||||
0.9,
|
||||
], # additional detection of [0] - FP
|
||||
[
|
||||
6.0,
|
||||
1.0,
|
||||
8.0,
|
||||
3.0,
|
||||
1,
|
||||
0.8,
|
||||
], # correct detection with incorrect class
|
||||
[1.0, 6.0, 2.0, 7.0, 1, 0.8], # incorrect detection - FP
|
||||
[
|
||||
1.0,
|
||||
2.0,
|
||||
2.0,
|
||||
4.0,
|
||||
1,
|
||||
0.8,
|
||||
], # incorrect detection with low IoU - FP
|
||||
]
|
||||
)
|
||||
],
|
||||
[
|
||||
np.array(
|
||||
[
|
||||
[0.0, 0.0, 3.0, 3.0, 0], # [0] detected
|
||||
[2.0, 2.0, 5.0, 5.0, 1], # [1] undetected - FN
|
||||
[
|
||||
6.0,
|
||||
1.0,
|
||||
8.0,
|
||||
3.0,
|
||||
2,
|
||||
], # [2] correct detection with incorrect class
|
||||
]
|
||||
)
|
||||
],
|
||||
CLASSES[:3],
|
||||
0.6,
|
||||
0.5,
|
||||
np.array([[1, 0, 0, 0], [0, 0, 0, 1], [0, 1, 0, 0], [1, 2, 0, 0]]),
|
||||
DoesNotRaise(),
|
||||
),
|
||||
(
|
||||
[
|
||||
np.array(
|
||||
[
|
||||
[0.0, 0.0, 3.0, 3.0, 0, 0.9], # correct detection of [0]
|
||||
[
|
||||
0.1,
|
||||
0.1,
|
||||
3.0,
|
||||
3.0,
|
||||
0,
|
||||
0.9,
|
||||
], # additional detection of [0] - FP
|
||||
[
|
||||
6.0,
|
||||
1.0,
|
||||
8.0,
|
||||
3.0,
|
||||
1,
|
||||
0.8,
|
||||
], # correct detection with incorrect class
|
||||
[1.0, 6.0, 2.0, 7.0, 1, 0.8], # incorrect detection - FP
|
||||
[
|
||||
1.0,
|
||||
2.0,
|
||||
2.0,
|
||||
4.0,
|
||||
1,
|
||||
0.8,
|
||||
], # incorrect detection with low IoU - FP
|
||||
]
|
||||
)
|
||||
],
|
||||
[
|
||||
np.array(
|
||||
[
|
||||
[0.0, 0.0, 3.0, 3.0, 0], # [0] detected
|
||||
[2.0, 2.0, 5.0, 5.0, 1], # [1] undetected - FN
|
||||
[
|
||||
6.0,
|
||||
1.0,
|
||||
8.0,
|
||||
3.0,
|
||||
2,
|
||||
], # [2] correct detection with incorrect class
|
||||
]
|
||||
)
|
||||
],
|
||||
CLASSES[:3],
|
||||
0.6,
|
||||
1.0,
|
||||
np.array([[0, 0, 0, 1], [0, 0, 0, 1], [0, 0, 0, 1], [2, 3, 0, 0]]),
|
||||
DoesNotRaise(),
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_from_tensors(
|
||||
predictions,
|
||||
targets,
|
||||
classes,
|
||||
conf_threshold,
|
||||
iou_threshold,
|
||||
expected_result: Optional[np.ndarray],
|
||||
exception: Exception,
|
||||
):
|
||||
with exception:
|
||||
result = ConfusionMatrix.from_tensors(
|
||||
predictions=predictions,
|
||||
targets=targets,
|
||||
classes=classes,
|
||||
conf_threshold=conf_threshold,
|
||||
iou_threshold=iou_threshold,
|
||||
)
|
||||
|
||||
assert result.matrix.diagonal().sum() == expected_result.diagonal().sum()
|
||||
assert np.array_equal(result.matrix, expected_result)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"predictions, targets, num_classes, conf_threshold, iou_threshold, expected_result,"
|
||||
" exception",
|
||||
[
|
||||
(
|
||||
DETECTION_TENSORS[0],
|
||||
CERTAIN_DETECTION_TENSORS[0],
|
||||
NUM_CLASSES,
|
||||
0.2,
|
||||
0.5,
|
||||
IDEAL_CONF_MATRIX,
|
||||
DoesNotRaise(),
|
||||
)
|
||||
],
|
||||
)
|
||||
def test_evaluate_detection_batch(
|
||||
predictions,
|
||||
targets,
|
||||
num_classes,
|
||||
conf_threshold,
|
||||
iou_threshold,
|
||||
expected_result: Optional[np.ndarray],
|
||||
exception: Exception,
|
||||
):
|
||||
with exception:
|
||||
result = ConfusionMatrix.evaluate_detection_batch(
|
||||
predictions=predictions,
|
||||
targets=targets,
|
||||
num_classes=num_classes,
|
||||
conf_threshold=conf_threshold,
|
||||
iou_threshold=iou_threshold,
|
||||
)
|
||||
|
||||
assert result.diagonal().sum() == result.sum()
|
||||
assert np.array_equal(result, expected_result)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"matches, expected_result, exception",
|
||||
[
|
||||
(
|
||||
IDEAL_MATCHES,
|
||||
IDEAL_MATCHES,
|
||||
DoesNotRaise(),
|
||||
)
|
||||
],
|
||||
)
|
||||
def test_drop_extra_matches(
|
||||
matches,
|
||||
expected_result: Optional[np.ndarray],
|
||||
exception: Exception,
|
||||
):
|
||||
with exception:
|
||||
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)
|
||||
@@ -0,0 +1,27 @@
|
||||
from typing import List
|
||||
|
||||
import numpy as np
|
||||
|
||||
from supervision.detection.core import Detections
|
||||
|
||||
|
||||
def mock_detections(
|
||||
xyxy: List[List[float]],
|
||||
confidence: List[float] = None,
|
||||
class_id: List[int] = None,
|
||||
tracker_id: List[int] = None,
|
||||
) -> Detections:
|
||||
return Detections(
|
||||
xyxy=np.array(xyxy, dtype=np.float32),
|
||||
confidence=confidence
|
||||
if confidence is None
|
||||
else np.array(confidence, dtype=np.float32),
|
||||
class_id=class_id if class_id is None else np.array(class_id, dtype=int),
|
||||
tracker_id=tracker_id
|
||||
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}."
|
||||
Reference in New Issue
Block a user