Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
93c1b947fc | ||
|
|
66d1f4c67d | ||
|
|
70a13fbf1e | ||
|
|
1db8e12e1b | ||
|
|
8853321058 | ||
|
|
0a5e830159 | ||
|
|
fb041f9b6c | ||
|
|
28f83b18f9 | ||
|
|
feb779bae9 | ||
|
|
521dc57b76 | ||
|
|
a0775544df | ||
|
|
14a0e299a1 | ||
|
|
1e783b623b | ||
|
|
c7a8654390 | ||
|
|
a7f8aee2c7 | ||
|
|
5ed03491f9 | ||
|
|
5d979df2f9 | ||
|
|
5006f8c2f7 | ||
|
|
8570ce68e2 | ||
|
|
3d3bbb4095 | ||
|
|
a1e3543897 | ||
|
|
4b82679b00 | ||
|
|
5d4abe6e0d | ||
|
|
0405c74cf3 | ||
|
|
13fe6e0942 | ||
|
|
03c92f7652 | ||
|
|
084926ded3 | ||
|
|
382ed8e609 | ||
|
|
2ed2a793cd | ||
|
|
e0fb73f18a | ||
|
|
de896189b8 | ||
|
|
fbfd078795 | ||
|
|
e31aad242a | ||
|
|
93c5c39f8f | ||
|
|
810148550b | ||
|
|
172877a3cf | ||
|
|
a91fc1e4f3 | ||
|
|
f4fd595783 | ||
|
|
e690b92688 | ||
|
|
c8f771ec69 | ||
|
|
6b09e1d67c | ||
|
|
12195705f5 | ||
|
|
9f7871b4be | ||
|
|
44890dc712 | ||
|
|
57ae7467b9 | ||
|
|
1e5e25c54a | ||
|
|
ba7c4f4404 | ||
|
|
f3b52c08a7 | ||
|
|
1a36842407 | ||
|
|
a8b12914c8 | ||
|
|
d3db15fcd0 | ||
|
|
c40734fb8d | ||
|
|
55a9cef270 | ||
|
|
4b9b70fbb2 | ||
|
|
e61a1074ac | ||
|
|
e6c7f970f0 | ||
|
|
f06dcf73a5 | ||
|
|
0bb25fdd91 | ||
|
|
9501a0c72e | ||
|
|
d9b9143f07 | ||
|
|
b7e6bcc15f | ||
|
|
08fbcbd8ea | ||
|
|
e5114a233c | ||
|
|
dea7f83e50 | ||
|
|
034d3eaaac | ||
|
|
7397d8a240 | ||
|
|
a023259445 | ||
|
|
09a06e10a1 | ||
|
|
8d26a9291b | ||
|
|
28d219c86d | ||
|
|
89ef841955 | ||
|
|
b164ef1111 | ||
|
|
2f39a2c3c6 | ||
|
|
d8f990231c | ||
|
|
74d19dd069 | ||
|
|
552c8dbede | ||
|
|
e421af494f | ||
|
|
5b3fc12aea | ||
|
|
854fb18b34 | ||
|
|
185ae462a4 | ||
|
|
3676f07168 | ||
|
|
099258baa6 | ||
|
|
0bb3839d29 | ||
|
|
f66444bc91 | ||
|
|
6de21048d2 | ||
|
|
db9d6c623b | ||
|
|
fb7ecd73d9 | ||
|
|
2bc19d69af | ||
|
|
a377b54a60 | ||
|
|
db1998c155 | ||
|
|
ef6369b93d | ||
|
|
a21a17ba29 | ||
|
|
c973647cb8 | ||
|
|
71bdb01547 | ||
|
|
92a5a931a8 | ||
|
|
0ecd88d4d1 | ||
|
|
4d12853b95 | ||
|
|
9333460ee9 | ||
|
|
d5cebd4fd2 | ||
|
|
61579af25a | ||
|
|
bcd3e1be17 | ||
|
|
9ef697cd10 | ||
|
|
19bf417810 | ||
|
|
3df33f0f9d | ||
|
|
b8b31f6f00 | ||
|
|
fc11f4d7af | ||
|
|
ec517124f8 | ||
|
|
f00dcede1a | ||
|
|
6fb90f6b75 | ||
|
|
23e63501eb | ||
|
|
0df56b1d66 | ||
|
|
ec75a86a83 | ||
|
|
4599548eeb | ||
|
|
bb513f880a | ||
|
|
fbce6e45d4 | ||
|
|
a00e2d22ba | ||
|
|
04c5f18da0 | ||
|
|
8b933f640e | ||
|
|
f01ca33511 | ||
|
|
2f88bb2b8e | ||
|
|
b2b251a424 | ||
|
|
eec3615250 | ||
|
|
272a9cb35e | ||
|
|
0d0328ef22 | ||
|
|
d821fca2af | ||
|
|
ff9a63a68b | ||
|
|
a6222a065c | ||
|
|
1489f82e0f | ||
|
|
12722df65b | ||
|
|
900952acae | ||
|
|
95308126c7 | ||
|
|
e7de6859e9 | ||
|
|
feeea38e6e | ||
|
|
6ea07a1c2a | ||
|
|
fbb770a486 | ||
|
|
09f663139c | ||
|
|
9efd04925c | ||
|
|
14ce0799ce | ||
|
|
a81f24e08d | ||
|
|
257e73fb37 | ||
|
|
dcfb122816 | ||
|
|
79ea168fd4 | ||
|
|
d35830ddeb | ||
|
|
c5ff58910d | ||
|
|
15d9b75225 | ||
|
|
49318ffe31 | ||
|
|
b52a0297f1 | ||
|
|
eceebb33e3 | ||
|
|
fa7821b1c4 | ||
|
|
a56a16937f | ||
|
|
0c95a6126d | ||
|
|
c14642accf | ||
|
|
8c0e2fb362 | ||
|
|
e5fb43aea6 | ||
|
|
32b70c5d37 | ||
|
|
b9e98a5e20 | ||
|
|
246b8a720c | ||
|
|
a43970fb65 | ||
|
|
ff5709def2 | ||
|
|
2805f41675 | ||
|
|
4eaa897f2c | ||
|
|
f96e7766f3 | ||
|
|
07cb0d8ccf | ||
|
|
3725beb2db | ||
|
|
d5f8b94d65 | ||
|
|
6f7f728044 | ||
|
|
aad3c4fce4 | ||
|
|
52724ca673 | ||
|
|
e3fc238c70 | ||
|
|
cfe29cd05f | ||
|
|
0f76f731ee | ||
|
|
c239291541 | ||
|
|
a446ffd021 | ||
|
|
c2084c4d0f | ||
|
|
e3cb417129 | ||
|
|
27de665940 | ||
|
|
93b21a8d2d | ||
|
|
a4bdfc2062 | ||
|
|
fd1a07cb6b | ||
|
|
edfb814fc4 | ||
|
|
3d5437b960 | ||
|
|
68c2504d61 | ||
|
|
6984d72f4a | ||
|
|
6f5b477ac0 | ||
|
|
b19511968d | ||
|
|
165c77d830 | ||
|
|
558581943b | ||
|
|
8acbc422c9 | ||
|
|
9995771604 | ||
|
|
bcc1bc1ff2 | ||
|
|
cce69d893e | ||
|
|
3e2fe835db | ||
|
|
47f1c44e20 | ||
|
|
420543534d | ||
|
|
1ea9345845 | ||
|
|
6ca09ccbfc | ||
|
|
7f78f93407 | ||
|
|
05ceca70ef | ||
|
|
e8f8c38f39 | ||
|
|
e50c761b83 | ||
|
|
0739e24b25 | ||
|
|
560760993a | ||
|
|
859f0258bc | ||
|
|
2f644c52e3 | ||
|
|
a82c37c7fa | ||
|
|
fc2710c9c9 | ||
|
|
29161a485e | ||
|
|
523d218dd4 | ||
|
|
0200fd307b | ||
|
|
6a199c5d7e | ||
|
|
d19c44c991 | ||
|
|
2d5de66a9b | ||
|
|
13b46dda67 | ||
|
|
a3f299ceb0 | ||
|
|
9a42372ff6 | ||
|
|
47b9ab8f96 | ||
|
|
b476acc1f3 | ||
|
|
3cee266f71 | ||
|
|
111d4c6c91 | ||
|
|
a77a2ba8aa | ||
|
|
fb9024bc75 | ||
|
|
ceecb8bbd7 | ||
|
|
c28c93bd4f | ||
|
|
46bd267079 | ||
|
|
d025791805 | ||
|
|
7e4fc8a485 | ||
|
|
ca482967c4 | ||
|
|
57fa637eea | ||
|
|
59c5ab5252 | ||
|
|
c27a28235e | ||
|
|
cfbb6681cc | ||
|
|
24b1db9e25 | ||
|
|
8c92072ecd | ||
|
|
38d442c176 | ||
|
|
2f09e51458 | ||
|
|
56eb8e1698 | ||
|
|
5bfd51fdec | ||
|
|
a5d4612402 | ||
|
|
941b7e30aa | ||
|
|
f4f2626b44 | ||
|
|
8f198e80cd | ||
|
|
612ac6145a | ||
|
|
26d70219bf | ||
|
|
daf13442f3 | ||
|
|
2b5d430382 | ||
|
|
b6ff9242fc | ||
|
|
d2f169bcfa | ||
|
|
33e10ff722 | ||
|
|
883a6bed3c | ||
|
|
f7bd7011fe | ||
|
|
6193128774 | ||
|
|
bdc10f83fe | ||
|
|
dc1249969f | ||
|
|
34353aac54 | ||
|
|
6ad32406af | ||
|
|
315c4295ff | ||
|
|
4dc001d9bf | ||
|
|
9d60b2009f | ||
|
|
c98698d442 | ||
|
|
3bda0bfc11 | ||
|
|
0e19c581e4 | ||
|
|
8fb843ec8e | ||
|
|
ba2b60ffb4 | ||
|
|
a6b985ff4b | ||
|
|
d59467b4f4 | ||
|
|
c0486b7f94 | ||
|
|
f24c6f38a0 | ||
|
|
a6c995c9a9 | ||
|
|
3c3a0792f3 | ||
|
|
1d13397503 | ||
|
|
f0e88b1982 | ||
|
|
9bc2907c49 | ||
|
|
e7dfc04fd3 | ||
|
|
a0d0d45890 | ||
|
|
4a365fb907 | ||
|
|
a60cfe08e2 | ||
|
|
c1ebb81d47 | ||
|
|
2ee9e08446 | ||
|
|
530e1d01e1 | ||
|
|
4fea0031ba | ||
|
|
490a341d5d | ||
|
|
2c3fa27929 | ||
|
|
404be16126 | ||
|
|
72d677ad89 | ||
|
|
d6ae86e97b | ||
|
|
ed5c26527a | ||
|
|
d23cc008f8 | ||
|
|
9d54d91791 | ||
|
|
80feb2b12f | ||
|
|
51c7ebdf60 | ||
|
|
2653f5a090 | ||
|
|
fb30bab0fd | ||
|
|
a09b6fac6f | ||
|
|
4e2eb0afc2 | ||
|
|
1a021d5792 | ||
|
|
909ae2a8ff | ||
|
|
7d0488efc0 | ||
|
|
2723735b01 | ||
|
|
7757d86abc | ||
|
|
f40dd38b1a | ||
|
|
ad7e643c83 | ||
|
|
eae4040b17 | ||
|
|
9afca17854 | ||
|
|
19296a73f6 | ||
|
|
0721bc289b | ||
|
|
db1b4737fe | ||
|
|
6fbca8333e | ||
|
|
ad2220bc1d | ||
|
|
0115ef8b7d | ||
|
|
d4541e1469 | ||
|
|
c9c2dad316 | ||
|
|
e59e2822d8 | ||
|
|
99854a5b66 | ||
|
|
251185ba67 | ||
|
|
8b4884ec9e | ||
|
|
07c36c6223 | ||
|
|
bf43d6566b | ||
|
|
b59d17fdae | ||
|
|
b81c5e8475 | ||
|
|
7eb918282c | ||
|
|
36a73f7d49 | ||
|
|
a2195aa378 | ||
|
|
2bd4348e5d | ||
|
|
6b7861d4d1 | ||
|
|
8a2f5d4726 | ||
|
|
dcfa916c75 | ||
|
|
54a0422b58 | ||
|
|
089916e26f | ||
|
|
1bbe03ca27 | ||
|
|
5840889818 | ||
|
|
c3855c9828 | ||
|
|
40b537c0be | ||
|
|
1d43042c55 | ||
|
|
996b0a3af3 | ||
|
|
f2716252d5 | ||
|
|
b916dc547c | ||
|
|
e1be566861 | ||
|
|
32f773548e | ||
|
|
414d3e9afb | ||
|
|
6dbf3f61e3 | ||
|
|
ab55e81087 | ||
|
|
7b42d8381b | ||
|
|
dfac30cba7 | ||
|
|
c26eb1a70f | ||
|
|
f0fc76c8c7 | ||
|
|
68f07e18f3 | ||
|
|
7ee84c48e0 | ||
|
|
179d00bead | ||
|
|
2626263ff3 | ||
|
|
a40bf78bb2 | ||
|
|
41ca70aab6 | ||
|
|
3fbce7b6bb | ||
|
|
b761a5c32b | ||
|
|
12a455eafa | ||
|
|
6c9f8302a5 | ||
|
|
da904ad2da | ||
|
|
8d432c78b1 | ||
|
|
0cd2c9d146 | ||
|
|
aad5e2a212 | ||
|
|
ab775d9f76 | ||
|
|
1779014fd9 | ||
|
|
84d8d98619 | ||
|
|
879ae9416a | ||
|
|
9024396f6c | ||
|
|
de5349ad6d | ||
|
|
d391376912 | ||
|
|
eb84d41c61 | ||
|
|
0252cf210c | ||
|
|
b3b745508b | ||
|
|
511f0ab9b2 | ||
|
|
3ee72e38c3 | ||
|
|
5acfb9531e | ||
|
|
2c769e2f72 | ||
|
|
ce05a0c544 | ||
|
|
1e8ee47db2 | ||
|
|
eefa05cf01 | ||
|
|
d8679d9c46 | ||
|
|
25ceb46ed4 | ||
|
|
a615dfccb6 | ||
|
|
1776bf1670 | ||
|
|
d983177c3a | ||
|
|
7b259145bb | ||
|
|
24d1840ea1 | ||
|
|
f8f3647a98 | ||
|
|
559ef90d83 | ||
|
|
0e2eec08c8 | ||
|
|
53f345e916 | ||
|
|
6c40935266 | ||
|
|
145b5fe56c | ||
|
|
aed9221307 | ||
|
|
f08404eb7d | ||
|
|
aab861c2db | ||
|
|
9c0d647fdf | ||
|
|
5ef934e780 | ||
|
|
2d740bdcb6 | ||
|
|
d2d50fbe46 | ||
|
|
26bafec8f7 | ||
|
|
8014e88944 | ||
|
|
c3b77d05c0 | ||
|
|
2eb0c7c242 | ||
|
|
41a08b87be | ||
|
|
2bcd81e5fd | ||
|
|
19396771f1 | ||
|
|
f5324897a3 | ||
|
|
cd8a2be04b | ||
|
|
8901192b92 | ||
|
|
a9044e4fed | ||
|
|
85e4c0ec65 | ||
|
|
93398108ee | ||
|
|
f41adca384 | ||
|
|
d64bb3879c | ||
|
|
452e1f1ea4 | ||
|
|
45e5f49c5a | ||
|
|
4c36b39f94 | ||
|
|
e13c7dc22c | ||
|
|
713719a582 | ||
|
|
051fe5178b | ||
|
|
1ebccebf4c | ||
|
|
f56d7a1733 | ||
|
|
82ab6c57cb | ||
|
|
cccacab621 | ||
|
|
59b273ba87 | ||
|
|
9a1c11218c | ||
|
|
d21c98a4d2 | ||
|
|
816aecb51b | ||
|
|
3c8995758d | ||
|
|
14b57345a7 | ||
|
|
d292cbf80e | ||
|
|
8c857dd9a6 | ||
|
|
acfabc247f | ||
|
|
61a63c83d9 | ||
|
|
528c960587 | ||
|
|
6c6171b5d5 | ||
|
|
b8efa25e93 | ||
|
|
17c8e41fb3 | ||
|
|
0c3594432c | ||
|
|
a8c44cfaff | ||
|
|
3d0c3d9182 | ||
|
|
2364abf481 | ||
|
|
7bd92b096f | ||
|
|
39477feaf6 | ||
|
|
e4e66f8a9a | ||
|
|
3b61ee674e | ||
|
|
01f7eb5be6 | ||
|
|
4cd1fbcef9 | ||
|
|
b6a55694f2 | ||
|
|
8c58ebe30d | ||
|
|
77034700dd | ||
|
|
1ebbe3a8d1 | ||
|
|
617bd30c5f | ||
|
|
2e8720f156 | ||
|
|
c47b0ce990 | ||
|
|
45563dfddf | ||
|
|
008711aac4 | ||
|
|
649e27323d | ||
|
|
36a301e7fc | ||
|
|
aac1d8862a | ||
|
|
d59ec0f223 | ||
|
|
3618ac35ec | ||
|
|
a9f821c28a | ||
|
|
7bb94ce966 | ||
|
|
fb49be8197 | ||
|
|
d7b6af84e6 | ||
|
|
d3b845fea5 | ||
|
|
330eae4411 | ||
|
|
6b5dacd8e3 | ||
|
|
88a43cc795 | ||
|
|
a941583bb0 | ||
|
|
2727152465 | ||
|
|
4d25437459 | ||
|
|
3ade7f060a | ||
|
|
18a790e9ea | ||
|
|
478fbd5751 | ||
|
|
7823e0de89 | ||
|
|
ab0882dd55 | ||
|
|
ec6864ff04 | ||
|
|
81786ff9d3 | ||
|
|
95c243c058 | ||
|
|
204669b08c | ||
|
|
97c407101a | ||
|
|
bee3252110 | ||
|
|
d7e52bee26 | ||
|
|
b1598733c2 | ||
|
|
5c539c3bda | ||
|
|
5d38120743 | ||
|
|
aaf5df2f6a | ||
|
|
447ef41fc4 | ||
|
|
edd944fff8 | ||
|
|
a75efbbf73 | ||
|
|
e9c8a80823 | ||
|
|
0b92417948 | ||
|
|
aebbd01830 | ||
|
|
c3ebccabe8 | ||
|
|
ff52352b00 | ||
|
|
f7d21acd3a | ||
|
|
92d56eda15 | ||
|
|
d37a5f90a4 | ||
|
|
869204d42d | ||
|
|
dc1ce02860 | ||
|
|
34a5b34c88 | ||
|
|
f72d06b05e | ||
|
|
4a068ff8d7 | ||
|
|
044a8fac06 | ||
|
|
a88089ced7 | ||
|
|
29b0c514cf | ||
|
|
7d5a0bce45 | ||
|
|
6913ecf54d | ||
|
|
7f114cba4c | ||
|
|
cc5e72dd5b | ||
|
|
177ab6ad4d | ||
|
|
7e79fa3d07 | ||
|
|
347b75935b | ||
|
|
079008cd17 | ||
|
|
e30d4a9abb | ||
|
|
aab3ed2ba9 | ||
|
|
ba30244add | ||
|
|
55e8f16029 | ||
|
|
33e69b2da3 | ||
|
|
9fe3170ead | ||
|
|
ffda94b1e5 | ||
|
|
7fd5d80a3b | ||
|
|
b67d9b475f | ||
|
|
d859842fa3 | ||
|
|
131c6ffefb | ||
|
|
71fd3178cc | ||
|
|
71919f1d6f | ||
|
|
4732b6f1d0 | ||
|
|
c024b62020 | ||
|
|
d4b1375db7 | ||
|
|
71200a5942 | ||
|
|
66f6a86bbf | ||
|
|
a70588284a | ||
|
|
67160c96b0 | ||
|
|
7331577e69 | ||
|
|
5f00d8b595 | ||
|
|
88209f233d | ||
|
|
d51067164a | ||
|
|
0dd8d7dca8 | ||
|
|
7c39d9dc8d | ||
|
|
d2d69f986e | ||
|
|
df0b2df2f6 | ||
|
|
31844c165b | ||
|
|
c5a2f87bb9 | ||
|
|
6ef9ba7390 | ||
|
|
a532670611 | ||
|
|
ba2d30d110 | ||
|
|
619da9575a | ||
|
|
25bdb560eb | ||
|
|
9b447296bf | ||
|
|
4a16d66bfb | ||
|
|
8436f52f8a | ||
|
|
059583e2ba | ||
|
|
4bd593fe3c | ||
|
|
399342513a | ||
|
|
c8c2ff5c45 | ||
|
|
18b3c17ed4 | ||
|
|
cab06f2c2b | ||
|
|
2b53216a0b | ||
|
|
e253f86e8f | ||
|
|
df4c546c10 | ||
|
|
25fc1ab6fa | ||
|
|
a45341ce63 | ||
|
|
e93ef041d4 | ||
|
|
df4a6f857c | ||
|
|
cb220479eb | ||
|
|
8bb8ce0316 | ||
|
|
8255bcb1ca | ||
|
|
8329d155e4 | ||
|
|
a41eea7c45 | ||
|
|
d531c59f85 | ||
|
|
15b1473c2f | ||
|
|
26c2a7ef9c | ||
|
|
361e444f53 | ||
|
|
07756328db | ||
|
|
f98ded49a2 | ||
|
|
bcb31f8809 | ||
|
|
b1d5d4faf0 | ||
|
|
7924a13dc5 | ||
|
|
755f292ed1 | ||
|
|
6166d92458 | ||
|
|
166a8da9a0 | ||
|
|
2c9a1f69c5 | ||
|
|
a2ab113f21 | ||
|
|
eaff97f902 | ||
|
|
5f0dcc2cb2 | ||
|
|
9e906c78f5 | ||
|
|
2417b36006 | ||
|
|
1a7309e7c9 | ||
|
|
ad0e6c3b51 | ||
|
|
d309d74f37 | ||
|
|
27b61ab4a3 | ||
|
|
bddb72fb28 | ||
|
|
823dc56057 | ||
|
|
6f4704625b | ||
|
|
9f222736e1 | ||
|
|
57b12e6e00 | ||
|
|
c78ae33e43 |
@@ -28,7 +28,7 @@ repos:
|
||||
|
||||
|
||||
- repo: https://github.com/PyCQA/bandit
|
||||
rev: '1.7.8'
|
||||
rev: '1.7.9'
|
||||
hooks:
|
||||
- id: bandit
|
||||
args: ["-c", "pyproject.toml"]
|
||||
@@ -45,7 +45,7 @@ repos:
|
||||
|
||||
|
||||
- repo: https://github.com/astral-sh/ruff-pre-commit
|
||||
rev: v0.4.1
|
||||
rev: v0.5.1
|
||||
hooks:
|
||||
- id: ruff
|
||||
args: [--fix, --exit-non-zero-on-fix]
|
||||
|
||||
+122
-35
@@ -4,6 +4,25 @@ Thank you for your interest in contributing to Supervision!
|
||||
|
||||
We are actively improving this library to reduce the amount of work you need to do to solve common computer vision problems.
|
||||
|
||||
## Code of Conduct
|
||||
|
||||
Please read and adhere to our [Code of Conduct](CODE_OF_CONDUCT.md). This document outlines the expected behavior for all participants in our project.
|
||||
|
||||
## Table of Contents
|
||||
|
||||
- [Contribution Guidelines](#contribution-guidelines)
|
||||
- [Contributing Features](#contributing-features-)
|
||||
- [How to Contribute Changes](#how-to-contribute-changes)
|
||||
- [Installation for Contributors](#installation-for-contributors)
|
||||
- [Code Style and Quality](#-code-style-and-quality)
|
||||
- [Pre-commit tool](#pre-commit-tool)
|
||||
- [Docstrings](#docstrings)
|
||||
- [Type checking](#type-checking)
|
||||
- [Documentation](#-documentation)
|
||||
- [Cookbooks](#-cookbooks)
|
||||
- [Tests](#-tests)
|
||||
- [License](#-license)
|
||||
|
||||
## Contribution Guidelines
|
||||
|
||||
We welcome contributions to:
|
||||
@@ -32,49 +51,49 @@ First, fork this repository to your own GitHub account. Click "fork" in the top
|
||||
|
||||
Then, run `git clone` to download the project code to your computer.
|
||||
|
||||
You should also set up `roboflow/supervision` as an "upstream" remote (that is, tell git that the reference Supervision repository was the source of your fork of it):
|
||||
|
||||
```bash
|
||||
git remote add upstream https://github.com/roboflow/supervision.git
|
||||
git fetch upstream
|
||||
```
|
||||
|
||||
Move to a new branch using the `git checkout` command:
|
||||
|
||||
```bash
|
||||
git checkout -b <your_branch_name>
|
||||
git checkout -b <scope>/<your_branch_name> upstream/develop
|
||||
```
|
||||
|
||||
The name you choose for your branch should describe the change you want to make (i.e. `line-counter-docs`).
|
||||
The name you choose for your branch should describe the change you want to make and start with an appropriate prefix:
|
||||
|
||||
- `feat/`: for new features (e.g., `feat/line-counter`)
|
||||
- `fix/`: for bug fixes (e.g., `fix/memory-leak`)
|
||||
- `docs/`: for documentation changes (e.g., `docs/update-readme`)
|
||||
- `chore/`: for routine tasks, maintenance, or tooling changes (e.g., `chore/update-dependencies`)
|
||||
- `test/`: for adding or modifying tests (e.g., `test/add-unit-tests`)
|
||||
- `refactor/`: for code refactoring (e.g., `refactor/simplify-algorithm`)
|
||||
|
||||
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
|
||||
git add -A
|
||||
git commit -m "feat: add line counter functionality"
|
||||
git push -u origin <your_branch_name>
|
||||
```
|
||||
|
||||
## 🎨 Code quality
|
||||
Use conventional commit messages to clearly describe your changes. The format is:
|
||||
|
||||
### Pre-commit tool
|
||||
<type>[optional scope]: <description>
|
||||
|
||||
This project uses 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).
|
||||
Common types include:
|
||||
- feat: A new feature
|
||||
- fix: A bug fix
|
||||
- docs: Documentation only changes
|
||||
- style: Changes that do not affect the meaning of the code (white-space, formatting, etc)
|
||||
- refactor: A code change that neither fixes a bug nor adds a feature
|
||||
- perf: A code change that improves performance
|
||||
- test: Adding missing tests or correcting existing tests
|
||||
- chore: Changes to the build process or auxiliary tools and libraries
|
||||
|
||||
Then, go back to your fork of the `supervision` repository, click "Pull Requests", and click "New Pull Request".
|
||||
|
||||
@@ -104,7 +123,75 @@ All pull requests will be reviewed by the maintainers of the project. We will pr
|
||||
|
||||
PRs must pass all tests and linting requirements before they can be merged.
|
||||
|
||||
## 📝 documentation
|
||||
## Installation for Contributors
|
||||
|
||||
Before starting your work on the project, set up your development environment:
|
||||
|
||||
1. Clone your fork of the project:
|
||||
```bash
|
||||
git clone https://github.com/YOUR_USERNAME/supervision.git
|
||||
cd supervision
|
||||
```
|
||||
Replace `YOUR_USERNAME` with your GitHub username.
|
||||
|
||||
2. Create and activate a virtual environment:
|
||||
```bash
|
||||
python3 -m venv .venv
|
||||
source .venv/bin/activate
|
||||
```
|
||||
|
||||
3. Install Poetry:
|
||||
|
||||
Using pip:
|
||||
```bash
|
||||
pip install -U pip setuptools
|
||||
pip install poetry
|
||||
```
|
||||
|
||||
Or using pipx (recommended for global installation):
|
||||
```bash
|
||||
pipx install poetry
|
||||
```
|
||||
|
||||
4. Install project dependencies:
|
||||
```bash
|
||||
poetry install
|
||||
```
|
||||
|
||||
5. Run pytest to verify the setup:
|
||||
```bash
|
||||
poetry run pytest
|
||||
```
|
||||
|
||||
## 🎨 Code Style and Quality
|
||||
|
||||
### Pre-commit tool
|
||||
|
||||
This project uses 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 executing `pre-commit install`. Every time you do a `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).
|
||||
|
||||
## 📝 Documentation
|
||||
|
||||
The `supervision` documentation is stored in a folder called `docs`. The project documentation is built using `mkdocs`.
|
||||
|
||||
@@ -112,13 +199,13 @@ To run the documentation, install the project requirements with `poetry install
|
||||
|
||||
You can learn more about mkdocs on the [mkdocs website](https://www.mkdocs.org/).
|
||||
|
||||
## 🧑🍳 cookbooks
|
||||
## 🧑🍳 Cookbooks
|
||||
|
||||
We are always looking for new examples and cookbooks to add to the `supervision`
|
||||
documentation. If you have a use case that you think would be helpful to others, please
|
||||
submit a PR with your example. Here are some guidelines for submitting a new example:
|
||||
|
||||
- Create a new notebook in the [`docs/nodebooks`](https://github.com/roboflow/supervision/tree/develop/docs/notebooks) folder.
|
||||
- Create a new notebook in the [`docs/notebooks`](https://github.com/roboflow/supervision/tree/develop/docs/notebooks) folder.
|
||||
- Add a link to the new notebook in [`docs/theme/cookbooks.html`](https://github.com/roboflow/supervision/blob/develop/docs/theme/cookbooks.html). Make sure to add the path to the new notebook, as well as a title, labels, author and supervision version.
|
||||
- Use the [Count Objects Crossing the Line](https://supervision.roboflow.com/develop/notebooks/count-objects-crossing-the-line/) example as a template for your new example.
|
||||
- Freeze the version of `supervision` you are using.
|
||||
@@ -126,10 +213,10 @@ submit a PR with your example. Here are some guidelines for submitting a new exa
|
||||
- Notebook should be self-contained. If you rely on external data ( videos, images, etc.) or libraries, include download and installation commands in the notebook.
|
||||
- Annotate the code with appropriate comments, including links to the documentation describing each of the tools you have used.
|
||||
|
||||
## 🧪 tests
|
||||
## 🧪 Tests
|
||||
|
||||
[`pytests`](https://docs.pytest.org/en/7.1.x/) is used to run our tests.
|
||||
|
||||
## 📄 license
|
||||
## 📄 License
|
||||
|
||||
By contributing, you agree that your contributions will be licensed under an [MIT license](https://github.com/roboflow/supervision/blob/develop/LICENSE.md).
|
||||
|
||||
@@ -28,7 +28,7 @@
|
||||
|
||||
**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! 🤝
|
||||
|
||||
[](https://github.com/orgs/roboflow/projects/10)
|
||||
[](https://github.com/orgs/roboflow/projects)
|
||||
|
||||
## 💻 install
|
||||
|
||||
@@ -86,7 +86,7 @@ len(detections)
|
||||
|
||||
### annotators
|
||||
|
||||
Supervision offers a wide range of highly customizable [annotators](https://supervision.roboflow.com/latest/annotators/), allowing you to compose the perfect visualization for your use case.
|
||||
Supervision offers a wide range of highly customizable [annotators](https://supervision.roboflow.com/latest/detection/annotators/), allowing you to compose the perfect visualization for your use case.
|
||||
|
||||
```python
|
||||
import cv2
|
||||
@@ -95,8 +95,8 @@ import supervision as sv
|
||||
image = cv2.imread(...)
|
||||
detections = sv.Detections(...)
|
||||
|
||||
bounding_box_annotator = sv.BoundingBoxAnnotator()
|
||||
annotated_frame = bounding_box_annotator.annotate(
|
||||
box_annotator = sv.BoxAnnotator()
|
||||
annotated_frame = box_annotator.annotate(
|
||||
scene=image.copy(),
|
||||
detections=detections
|
||||
)
|
||||
@@ -106,7 +106,7 @@ https://github.com/roboflow/supervision/assets/26109316/691e219c-0565-4403-9218-
|
||||
|
||||
### datasets
|
||||
|
||||
Supervision provides a set of [utils](https://supervision.roboflow.com/latest/datasets/) that allow you to load, split, merge, and save datasets in one of the supported formats.
|
||||
Supervision provides a set of [utils](https://supervision.roboflow.com/latest/datasets/core/) that allow you to load, split, merge, and save datasets in one of the supported formats.
|
||||
|
||||
```python
|
||||
import supervision as sv
|
||||
@@ -216,7 +216,7 @@ len(dataset)
|
||||
|
||||
## 🎬 tutorials
|
||||
|
||||
Want to learn how to use Supervision? Explore our [how-to guides](https://supervision.roboflow.com/develop/how_to/detect_and_annotate/), [end-to-end examples](https://github.com/roboflow/supervision/tree/develop/examples), and [cookbooks](https://supervision.roboflow.com/develop/cookbooks/)!
|
||||
Want to learn how to use Supervision? Explore our [how-to guides](https://supervision.roboflow.com/develop/how_to/detect_and_annotate/), [end-to-end examples](https://github.com/roboflow/supervision/tree/develop/examples), [cheatsheet](https://roboflow.github.io/cheatsheet-supervision/), and [cookbooks](https://supervision.roboflow.com/develop/cookbooks/)!
|
||||
|
||||
<br/>
|
||||
|
||||
|
||||
Vendored
+1478
-1342
File diff suppressed because one or more lines are too long
+230
-6
@@ -1,3 +1,232 @@
|
||||
### 0.22.0 <small>Jul 12, 2024</small>
|
||||
|
||||
- Added [#1326](https://github.com/roboflow/supervision/pull/1326): [`sv.DetectionsDataset`](https://supervision.roboflow.com/latest/datasets/core/#supervision.dataset.core.DetectionDataset) and [`sv.ClassificationDataset`](https://supervision.roboflow.com/latest/datasets/core/#supervision.dataset.core.ClassificationDataset) allowing to load the images into memory only when necessary (lazy loading).
|
||||
|
||||
!!! failure "Deprecated"
|
||||
|
||||
Constructing `DetectionDataset` with parameter `images` as `Dict[str, np.ndarray]` is deprecated and will be removed in `supervision-0.26.0`. Please pass a list of paths `List[str]` instead.
|
||||
|
||||
!!! failure "Deprecated"
|
||||
|
||||
The `DetectionDataset.images` property is deprecated and will be removed in `supervision-0.26.0`. Please loop over images with `for path, image, annotation in dataset:`, as that does not require loading all images into memory.
|
||||
|
||||
```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_train = sv.DetectionDataset.from_coco(
|
||||
images_directory_path=f"{dataset.location}/train",
|
||||
annotations_path=f"{dataset.location}/train/_annotations.coco.json",
|
||||
)
|
||||
|
||||
path, image, annotation = ds_train[0]
|
||||
# loads image on demand
|
||||
|
||||
for path, image, annotation in ds_train:
|
||||
# loads image on demand
|
||||
```
|
||||
|
||||
- Added [#1296](https://github.com/roboflow/supervision/pull/1296): [`sv.Detections.from_lmm`](https://supervision.roboflow.com/latest/detection/core/#supervision.detection.core.Detections.from_lmm) now supports parsing results from the [Florence 2](https://huggingface.co/microsoft/Florence-2-large) model, extending the capability to handle outputs from this Large Multimodal Model (LMM). This includes detailed object detection, OCR with region proposals, segmentation, and more. Find out more in our [Colab notebook](https://colab.research.google.com/github/roboflow-ai/notebooks/blob/main/notebooks/how-to-finetune-florence-2-on-detection-dataset.ipynb).
|
||||
|
||||
- Added [#1232](https://github.com/roboflow/supervision/pull/1232) to support keypoint detection with Mediapipe. Both [legacy](https://colab.research.google.com/github/googlesamples/mediapipe/blob/main/examples/pose_landmarker/python/%5BMediaPipe_Python_Tasks%5D_Pose_Landmarker.ipynb) and [modern](https://ai.google.dev/edge/mediapipe/solutions/vision/pose_landmarker/python) pipelines are supported. See [`sv.KeyPoints.from_mediapipe`](https://supervision.roboflow.com/latest/keypoint/core/#supervision.keypoint.core.KeyPoints.from_mediapipe) for more.
|
||||
|
||||
- Added [#1316](https://github.com/roboflow/supervision/pull/1316): [`sv.KeyPoints.from_mediapipe`](https://supervision.roboflow.com/latest/keypoint/core/#supervision.keypoint.core.KeyPoints.from_mediapipe) extended to support FaceMesh from Mediapipe. This enhancement allows for processing both face landmarks from `FaceLandmarker`, and legacy results from `FaceMesh`.
|
||||
|
||||
- Added [#1310](https://github.com/roboflow/supervision/pull/1310): [`sv.KeyPoints.from_detectron2`](https://supervision.roboflow.com/latest/keypoint/core/#supervision.keypoint.core.KeyPoints.from_detectron2) is a new `KeyPoints` method, adding support for extracting keypoints from the popular [Detectron 2](https://github.com/facebookresearch/detectron2) platform.
|
||||
|
||||
- Added [#1300](https://github.com/roboflow/supervision/pull/1300): [`sv.Detections.from_detectron2`](https://supervision.roboflow.com/latest/detection/core/#supervision.detection.core.Detections.from_detectron2) now supports segmentation models detectron2. The resulting masks can be used with [`sv.MaskAnnotator`](https://supervision.roboflow.com/latest/annotators/#supervision.annotators.core.MaskAnnotator) for displaying annotations.
|
||||
|
||||
```python
|
||||
import supervision as sv
|
||||
from detectron2 import model_zoo
|
||||
from detectron2.engine import DefaultPredictor
|
||||
from detectron2.config import get_cfg
|
||||
import cv2
|
||||
|
||||
image = cv2.imread(<SOURCE_IMAGE_PATH>)
|
||||
cfg = get_cfg()
|
||||
cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
|
||||
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")
|
||||
predictor = DefaultPredictor(cfg)
|
||||
|
||||
result = predictor(image)
|
||||
detections = sv.Detections.from_detectron2(result)
|
||||
|
||||
mask_annotator = sv.MaskAnnotator()
|
||||
annotated_frame = mask_annotator.annotate(scene=image.copy(), detections=detections)
|
||||
```
|
||||
|
||||
- Added [#1277](https://github.com/roboflow/supervision/pull/1277): if you provide a font that supports symbols of a language, [`sv.RichLabelAnnotator`](https://supervision.roboflow.com/latest/detection/annotators/#supervision.annotators.core.LabelAnnotator.annotate) will draw them on your images.
|
||||
- Various other annotators have been revised to ensure proper in-place functionality when used with `numpy` arrays. Additionally, we fixed a bug where `sv.ColorAnnotator` was filling boxes with solid color when used in-place.
|
||||
|
||||
```python
|
||||
import cv2
|
||||
import supervision as sv
|
||||
import
|
||||
|
||||
image = cv2.imread(<SOURCE_IMAGE_PATH>)
|
||||
|
||||
model = get_model(model_id="yolov8n-640")
|
||||
results = model.infer(image)[0]
|
||||
detections = sv.Detections.from_inference(results)
|
||||
|
||||
rich_label_annotator = sv.RichLabelAnnotator(font_path=<TTF_FONT_PATH>)
|
||||
annotated_image = rich_label_annotator.annotate(scene=image.copy(), detections=detections)
|
||||
```
|
||||
|
||||
- Added [#1227](https://github.com/roboflow/supervision/pull/1227): Added support for loading Oriented Bounding Boxes dataset in YOLO format.
|
||||
|
||||
```python
|
||||
import supervision as sv
|
||||
|
||||
train_ds = sv.DetectionDataset.from_yolo(
|
||||
images_directory_path="/content/dataset/train/images",
|
||||
annotations_directory_path="/content/dataset/train/labels",
|
||||
data_yaml_path="/content/dataset/data.yaml",
|
||||
is_obb=True
|
||||
)
|
||||
|
||||
_, image, detections in train_ds[0]
|
||||
|
||||
obb_annotator = OrientedBoxAnnotator()
|
||||
annotated_image = obb_annotator.annotate(scene=image.copy(), detections=detections)
|
||||
```
|
||||
|
||||
- Fixed [#1312](https://github.com/roboflow/supervision/pull/1312): Fixed [`CropAnnotator`](https://supervision.roboflow.com/latest/detection/annotators/#supervision.annotators.core.TraceAnnotator.annotate).
|
||||
|
||||
!!! failure "Removed"
|
||||
|
||||
`BoxAnnotator` was removed, however `BoundingBoxAnnotator` has been renamed to `BoxAnnotator`. Use a combination of [`BoxAnnotator`](https://supervision.roboflow.com/latest/detection/annotators/#supervision.annotators.core.BoxAnnotator) and [`LabelAnnotator`](https://supervision.roboflow.com/latest/detection/annotators/#supervision.annotators.core.LabelAnnotator) to simulate old `BoundingBox` behavior.
|
||||
|
||||
!!! failure "Deprecated"
|
||||
|
||||
The name `BoundingBoxAnnotator` has been deprecated and will be removed in `supervision-0.26.0`. It has been renamed to [`BoxAnnotator`](https://supervision.roboflow.com/latest/detection/annotators/#supervision.annotators.core.BoxAnnotator).
|
||||
|
||||
- Added [#975](https://github.com/roboflow/supervision/pull/975) 📝 New Cookbooks: serialize detections into [json](https://github.com/roboflow/supervision/blob/de896189b83a1f9434c0a37dd9192ee00d2a1283/docs/notebooks/serialise-detections-to-json.ipynb) and [csv](https://github.com/roboflow/supervision/blob/de896189b83a1f9434c0a37dd9192ee00d2a1283/docs/notebooks/serialise-detections-to-csv.ipynb).
|
||||
|
||||
- Added [#1290](https://github.com/roboflow/supervision/pull/1290): Mostly an internal change, our file utility function now support both `str` and `pathlib` paths.
|
||||
|
||||
- Added [#1340](https://github.com/roboflow/supervision/pull/1340): Two new methods for converting between bounding box formats - [`xywh_to_xyxy`](https://supervision.roboflow.com/latest/detection/utils/#supervision.detection.utils.xywh_to_xyxy) and [`xcycwh_to_xyxy`](https://supervision.roboflow.com/latest/detection/utils/#supervision.detection.utils.xcycwh_to_xyxy)
|
||||
|
||||
!!! failure "Removed"
|
||||
|
||||
`from_roboflow` method has been removed due to deprecation. Use [from_inference](https://supervision.roboflow.com/latest/detection/core/#supervision.detection.core.Detections.from_inference) instead.
|
||||
|
||||
!!! failure "Removed"
|
||||
|
||||
`Color.white()` has been removed due to deprecation. Use `color.WHITE` instead.
|
||||
|
||||
!!! failure "Removed"
|
||||
|
||||
`Color.black()` has been removed due to deprecation. Use `color.BLACK` instead.
|
||||
|
||||
!!! failure "Removed"
|
||||
|
||||
`Color.red()` has been removed due to deprecation. Use `color.RED` instead.
|
||||
|
||||
!!! failure "Removed"
|
||||
|
||||
`Color.green()` has been removed due to deprecation. Use `color.GREEN` instead.
|
||||
|
||||
!!! failure "Removed"
|
||||
|
||||
`Color.blue()` has been removed due to deprecation. Use `color.BLUE` instead.
|
||||
|
||||
!!! failure "Removed"
|
||||
|
||||
`ColorPalette.default()` has been removed due to deprecation. Use [ColorPalette.DEFAULT](https://supervision.roboflow.com/latest/utils/draw/#supervision.draw.color.ColorPalette.DEFAULT) instead.
|
||||
|
||||
!!! failure "Removed"
|
||||
|
||||
`FPSMonitor.__call__` has been removed due to deprecation. Use the attribute [FPSMonitor.fps](https://supervision.roboflow.com/latest/utils/video/#supervision.utils.video.FPSMonitor.fps) instead.
|
||||
|
||||
### 0.21.0 <small>Jun 5, 2024</small>
|
||||
|
||||
- Added [#500](https://github.com/roboflow/supervision/pull/500): [`sv.Detections.with_nmm`](https://supervision.roboflow.com/latest/detection/core/#supervision.detection.core.Detections.with_nmm) to perform non-maximum merging on the current set of object detections.
|
||||
|
||||
- Added [#1221](https://github.com/roboflow/supervision/pull/1221): [`sv.Detections.from_lmm`](https://supervision.roboflow.com/latest/detection/core/#supervision.detection.core.Detections.from_lmm) allowing to parse Large Multimodal Model (LMM) text result into [`sv.Detections`](https://supervision.roboflow.com/latest/detection/core/) object. For now `from_lmm` supports only [PaliGemma](https://colab.research.google.com/github/roboflow-ai/notebooks/blob/main/notebooks/how-to-finetune-paligemma-on-detection-dataset.ipynb) result parsing.
|
||||
|
||||
```python
|
||||
import supervision as sv
|
||||
|
||||
paligemma_result = "<loc0256><loc0256><loc0768><loc0768> cat"
|
||||
detections = sv.Detections.from_lmm(
|
||||
sv.LMM.PALIGEMMA,
|
||||
paligemma_result,
|
||||
resolution_wh=(1000, 1000),
|
||||
classes=['cat', 'dog']
|
||||
)
|
||||
detections.xyxy
|
||||
# array([[250., 250., 750., 750.]])
|
||||
|
||||
detections.class_id
|
||||
# array([0])
|
||||
```
|
||||
|
||||
- Added [#1236](https://github.com/roboflow/supervision/pull/1236): [`sv.VertexLabelAnnotator`](https://supervision.roboflow.com/latest/keypoint/annotators/#supervision.keypoint.annotators.EdgeAnnotator.annotate) allowing to annotate every vertex of a keypoint skeleton with custom text and color.
|
||||
|
||||
```python
|
||||
import supervision as sv
|
||||
|
||||
image = ...
|
||||
key_points = sv.KeyPoints(...)
|
||||
|
||||
edge_annotator = sv.EdgeAnnotator(
|
||||
color=sv.Color.GREEN,
|
||||
thickness=5
|
||||
)
|
||||
annotated_frame = edge_annotator.annotate(
|
||||
scene=image.copy(),
|
||||
key_points=key_points
|
||||
)
|
||||
```
|
||||
|
||||
- Added [#1147](https://github.com/roboflow/supervision/pull/1147): [`sv.KeyPoints.from_inference`](https://supervision.roboflow.com/develop/keypoint/core/#supervision.keypoint.core.KeyPoints.from_inference) allowing to create [`sv.KeyPoints`](https://supervision.roboflow.com/develop/keypoint/core/#supervision.keypoint.core.KeyPoints) from [Inference](https://github.com/roboflow/inference) result.
|
||||
|
||||
- Added [#1138](https://github.com/roboflow/supervision/pull/1138): [`sv.KeyPoints.from_yolo_nas`](https://supervision.roboflow.com/develop/keypoint/core/#supervision.keypoint.core.KeyPoints.from_yolo_nas) allowing to create [`sv.KeyPoints`](https://supervision.roboflow.com/develop/keypoint/core/#supervision.keypoint.core.KeyPoints) from [YOLO-NAS](https://github.com/Deci-AI/super-gradients/blob/master/YOLONAS.md) result.
|
||||
|
||||
- Added [#1163](https://github.com/roboflow/supervision/pull/1163): [`sv.mask_to_rle`](https://supervision.roboflow.com/develop/datasets/utils/#supervision.dataset.utils.rle_to_mask) and [`sv.rle_to_mask`](https://supervision.roboflow.com/develop/datasets/utils/#supervision.dataset.utils.rle_to_mask) allowing for easy conversion between mask and rle formats.
|
||||
|
||||
- Changed [#1236](https://github.com/roboflow/supervision/pull/1236): [`sv.InferenceSlicer`](https://supervision.roboflow.com/develop/detection/tools/inference_slicer/) allowing to select overlap filtering strategy (`NONE`, `NON_MAX_SUPPRESSION` and `NON_MAX_MERGE`).
|
||||
|
||||
- Changed [#1178](https://github.com/roboflow/supervision/pull/1178): [`sv.InferenceSlicer`](https://supervision.roboflow.com/develop/detection/tools/inference_slicer/) adding instance segmentation model support.
|
||||
|
||||
```python
|
||||
import cv2
|
||||
import numpy as np
|
||||
import supervision as sv
|
||||
from inference import get_model
|
||||
|
||||
model = get_model(model_id="yolov8x-seg-640")
|
||||
image = cv2.imread(<SOURCE_IMAGE_PATH>)
|
||||
|
||||
def callback(image_slice: np.ndarray) -> sv.Detections:
|
||||
results = model.infer(image_slice)[0]
|
||||
return sv.Detections.from_inference(results)
|
||||
|
||||
slicer = sv.InferenceSlicer(callback = callback)
|
||||
detections = slicer(image)
|
||||
|
||||
mask_annotator = sv.MaskAnnotator()
|
||||
label_annotator = sv.LabelAnnotator()
|
||||
|
||||
annotated_image = mask_annotator.annotate(
|
||||
scene=image, detections=detections)
|
||||
annotated_image = label_annotator.annotate(
|
||||
scene=annotated_image, detections=detections)
|
||||
```
|
||||
|
||||
- Changed [#1228](https://github.com/roboflow/supervision/pull/1228): [`sv.LineZone`](https://supervision.roboflow.com/develop/detection/tools/line_zone/) making it 10-20 times faster, depending on the use case.
|
||||
|
||||
- Changed [#1163](https://github.com/roboflow/supervision/pull/1163): [`sv.DetectionDataset.from_coco`](https://supervision.roboflow.com/develop/datasets/core/#supervision.dataset.core.DetectionDataset.from_coco) and [`sv.DetectionDataset.as_coco`](https://supervision.roboflow.com/develop/datasets/core/#supervision.dataset.core.DetectionDataset.as_coco) adding support for run-length encoding (RLE) mask format.
|
||||
|
||||
### 0.20.0 <small>April 24, 2024</small>
|
||||
|
||||
- Added [#1128](https://github.com/roboflow/supervision/pull/1128): [`sv.KeyPoints`](/0.20.0/keypoint/core/#supervision.keypoint.core.KeyPoints) to provide initial support for pose estimation and broader keypoint detection models.
|
||||
@@ -84,8 +313,6 @@ with csv_sink:
|
||||
|
||||
- Added [#819](https://github.com/roboflow/supervision/pull/819): [`sv.JSONSink`](/0.19.0/detection/tools/save_detections/#supervision.detection.tools.csv_sink.JSONSink) allowing for the straightforward saving of image, video, or stream inference results in a `.json` file.
|
||||
|
||||
```python
|
||||
|
||||
```python
|
||||
import supervision as sv
|
||||
from ultralytics import YOLO
|
||||
@@ -205,7 +432,6 @@ ColorPalette(colors=[Color(r=68, g=1, b=84), Color(r=59, g=82, b=139), ...])
|
||||
|
||||
`sv.ColorPalette.default()` is deprecated and will be removed in `supervision-0.22.0`. Use `sv.ColorPalette.DEFAULT` instead.
|
||||
|
||||
|
||||
- Changed [#769](https://github.com/roboflow/supervision/pull/769): [`sv.ColorPalette.DEFAULT`](/0.18.0/draw/color/#colorpalette) value, giving users a more extensive set of annotation colors.
|
||||
|
||||
- Changed [#677](https://github.com/roboflow/supervision/pull/677): `sv.Detections.from_roboflow` to [`sv.Detections.from_inference`](/0.18.0/detection/core/#supervision.detection.core.Detections.from_inference) streamlining its functionality to be compatible with both the both [inference](https://github.com/roboflow/inference) pip package and the Robloflow [hosted API](https://docs.roboflow.com/deploy/hosted-api).
|
||||
@@ -214,7 +440,6 @@ ColorPalette(colors=[Color(r=68, g=1, b=84), Color(r=59, g=82, b=139), ...])
|
||||
|
||||
`Detections.from_roboflow()` is deprecated and will be removed in `supervision-0.22.0`. Use `Detections.from_inference` instead.
|
||||
|
||||
|
||||
- Fixed [#735](https://github.com/roboflow/supervision/pull/735): [`sv.LineZone`](/0.18.0/detection/tools/line_zone/#linezone) functionality to accurately update the counter when an object crosses a line from any direction, including from the side. This enhancement enables more precise tracking and analytics, such as calculating individual in/out counts for each lane on the road.
|
||||
|
||||
### 0.17.0 <small>December 06, 2023</small>
|
||||
@@ -307,13 +532,12 @@ ColorPalette(colors=[Color(r=68, g=1, b=84), Color(r=59, g=82, b=139), ...])
|
||||
|
||||
- Fixed [#477](https://github.com/roboflow/supervision/pull/477): Poetry env definition allowing proper local installation.
|
||||
|
||||
- Fixed [#430](https://github.com/roboflow/supervision/pull/430): [`sv.ByteTrack`](/0.16.0/trackers/#supervision.tracker.byte_tracker.core.ByteTrack) to return `np.array([], dtype=int)` when `svDetections` is empty.
|
||||
- Fixed [#430](https://github.com/roboflow/supervision/pull/430): [`sv.ByteTrack`](/0.16.0/trackers/#supervision.tracker.byte_tracker.core.ByteTrack) to return `np.array([], dtype=int)` when `svDetections` is empty.
|
||||
|
||||
!!! failure "Deprecated"
|
||||
|
||||
`sv.Detections.from_yolov8` and `sv.Classifications.from_yolov8` as those are now replaced by [`sv.Detections.from_ultralytics`](/0.16.0/detection/core/#supervision.detection.core.Detections.from_ultralytics) and [`sv.Classifications.from_ultralytics`](/0.16.0/classification/core/#supervision.classification.core.Classifications.from_ultralytics).
|
||||
|
||||
|
||||
### 0.15.0 <small>October 5, 2023</small>
|
||||
|
||||
- Added [#170](https://github.com/roboflow/supervision/pull/170): [`sv.BoundingBoxAnnotator`](/0.15.0/annotators/#supervision.annotators.core.BoundingBoxAnnotator) allowing to annotate images and videos with bounding boxes.
|
||||
|
||||
@@ -1,7 +1,6 @@
|
||||
---
|
||||
template: cookbooks.html
|
||||
comments: true
|
||||
status: new
|
||||
hide:
|
||||
- navigation
|
||||
- toc
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
---
|
||||
comments: true
|
||||
status: new
|
||||
---
|
||||
|
||||
# Datasets
|
||||
@@ -0,0 +1,17 @@
|
||||
---
|
||||
comments: true
|
||||
---
|
||||
|
||||
# Datasets Utils
|
||||
|
||||
<div class="md-typeset">
|
||||
<h2><a href="#supervision.dataset.utils.rle_to_mask">rle_to_mask</a></h2>
|
||||
</div>
|
||||
|
||||
:::supervision.dataset.utils.rle_to_mask
|
||||
|
||||
<div class="md-typeset">
|
||||
<h2><a href="#supervision.dataset.utils.mask_to_rle">mask_to_rle</a></h2>
|
||||
</div>
|
||||
|
||||
:::supervision.dataset.utils.mask_to_rle
|
||||
+20
-12
@@ -3,17 +3,25 @@ comments: true
|
||||
status: deprecated
|
||||
---
|
||||
|
||||
These features are phased out due to better alternatives or potential issues in future versions. Deprecated functionalities are supported for **three subsequent releases**, providing time for users to transition to updated methods.
|
||||
# Deprecated
|
||||
|
||||
These features are phased out due to better alternatives or potential issues in future versions. Deprecated functionalities are supported for **five subsequent releases**, providing time for users to transition to updated methods.
|
||||
|
||||
- [`Detections.from_froboflow`](detection/core.md/#supervision.detection.core.Detections.from_roboflow) is deprecated and will be removed in `supervision-0.22.0`. Use [`Detections.from_inference`](detection/core.md/#supervision.detection.core.Detections.from_inference) instead.
|
||||
- The method `Color.white()` is deprecated and will be removed in `supervision-0.22.0`. Use the constant `Color.WHITE` instead.
|
||||
- The method `Color.black()` is deprecated and will be removed in `supervision-0.22.0`. Use the constant `Color.BLACK` instead.
|
||||
- The method `Color.red()` is deprecated and will be removed in `supervision-0.22.0`. Use the constant `Color.RED` instead.
|
||||
- The method `Color.green()` is deprecated and will be removed in `supervision-0.22.0`. Use the constant `Color.GREEN` instead.
|
||||
- The method `Color.blue()` is deprecated and will be removed in `supervision-0.22.0`. Use the constant `Color.BLUE` instead.
|
||||
- The method [`ColorPalette.default()`](draw/color.md/#supervision.draw.color.ColorPalette.default) is deprecated and will be removed in `supervision-0.22.0`. Use the constant [`ColorPalette.DEFAULT`](draw/color.md/#supervision.draw.color.ColorPalette.DEFAULT) instead.
|
||||
- `BoxAnnotator` is deprecated and will be removed in `supervision-0.22.0`. Use [`BoundingBoxAnnotator`](detection/annotators.md/#supervision.annotators.core.BoundingBoxAnnotator) and [`LabelAnnotator`](detection/annotators.md/#supervision.annotators.core.LabelAnnotator) instead.
|
||||
- The method [`FPSMonitor.__call__`](utils/video.md/#supervision.utils.video.FPSMonitor.__call__) is deprecated and will be removed in `supervision-0.22.0`. Use the attribute [`FPSMonitor.fps`](utils/video.md/#supervision.utils.video.FPSMonitor.fps) instead.
|
||||
- The `track_buffer`, `track_thresh`, and `match_thresh` parameters in [`ByterTrack`](trackers.md/#supervision.tracker.byte_tracker.core.ByteTrack) are deprecated and will be removed in `supervision-0.23.0`. Use `lost_track_buffer,` `track_activation_threshold`, and `minimum_matching_threshold` instead.
|
||||
- The `triggering_position ` parameter in [`sv.PolygonZone`](detection/tools/polygon_zone.md/#supervision.detection.tools.polygon_zone.PolygonZone) is deprecated and will be removed in `supervision-0.23.0`. Use `triggering_anchors ` instead.
|
||||
- The `frame_resolution_wh ` parameter in [`sv.PolygonZone`](detection/tools/polygon_zone.md/#supervision.detection.tools.polygon_zone.PolygonZone) is deprecated and will be removed in `supervision-0.24.0`.
|
||||
- The `triggering_position ` parameter in [`sv.PolygonZone`](detection/tools/polygon_zone.md/#supervision.detection.tools.polygon_zone.PolygonZone) will be removed in `supervision-0.23.0`. Use `triggering_anchors ` instead.
|
||||
- The `frame_resolution_wh ` parameter in [`sv.PolygonZone`](detection/tools/polygon_zone.md/#supervision.detection.tools.polygon_zone.PolygonZone) will be removed in `supervision-0.24.0`.
|
||||
- Constructing `DetectionDataset` and `ClassificationDataset` with parameter `images` as `Dict[str, np.ndarray]` will be removed in `supervision-0.26.0`. Please pass a list of paths `List[str]` instead.
|
||||
- The `DetectionDataset.images` property will be removed in `supervision-0.26.0`. Please loop over images with `for path, image, annotation in dataset:`, as that does not require loading all images into memory.
|
||||
- `BoundingBoxAnnotator` has been renamed to `BoxAnnotator` after the old implementation of `BoxAnnotator` has been removed. `BoundingBoxAnnotator` will be removed in `supervision-0.26.0`.
|
||||
|
||||
# Removed
|
||||
|
||||
- [`Detections.from_froboflow`](detection/core.md/#supervision.detection.core.Detections.from_roboflow) is removed as of `supervision-0.22.0`. Use [`Detections.from_inference`](detection/core.md/#supervision.detection.core.Detections.from_inference) instead.
|
||||
- The method `Color.white()` was removed as of `supervision-0.22.0`. Use the constant `Color.WHITE` instead.
|
||||
- The method `Color.black()` was removed as of `supervision-0.22.0`. Use the constant `Color.BLACK` instead.
|
||||
- The method `Color.red()` was removed as of `supervision-0.22.0`. Use the constant `Color.RED` instead.
|
||||
- The method `Color.green()` was removed as of `supervision-0.22.0`. Use the constant `Color.GREEN` instead.
|
||||
- The method `Color.blue()` was removed as of `supervision-0.22.0`. Use the constant `Color.BLUE` instead.
|
||||
- The method [`ColorPalette.default()`](draw/color.md/#supervision.draw.color.ColorPalette.default) was removed as of `supervision-0.22.0`. Use the constant [`ColorPalette.DEFAULT`](draw/color.md/#supervision.draw.color.ColorPalette.DEFAULT) instead.
|
||||
- `BoxAnnotator` was removed as of `supervision-0.22.0`, however `BoundingBoxAnnotator` was immediately renamed to `BoxAnnotator`. Use [`BoxAnnotator`](detection/annotators.md/#supervision.annotators.core.BoxAnnotator) and [`LabelAnnotator`](detection/annotators.md/#supervision.annotators.core.LabelAnnotator) instead of the old `BoxAnnotator`.
|
||||
- The method [`FPSMonitor.__call__`](utils/video.md/#supervision.utils.video.FPSMonitor.__call__) was removed as of `supervision-0.22.0`. Use the attribute [`FPSMonitor.fps`](utils/video.md/#supervision.utils.video.FPSMonitor.fps) instead.
|
||||
|
||||
@@ -5,7 +5,7 @@ status: new
|
||||
|
||||
# Annotators
|
||||
|
||||
=== "BoundingBox"
|
||||
=== "Box"
|
||||
|
||||
```python
|
||||
import supervision as sv
|
||||
@@ -13,8 +13,8 @@ status: new
|
||||
image = ...
|
||||
detections = sv.Detections(...)
|
||||
|
||||
bounding_box_annotator = sv.BoundingBoxAnnotator()
|
||||
annotated_frame = bounding_box_annotator.annotate(
|
||||
box_annotator = sv.BoxAnnotator()
|
||||
annotated_frame = box_annotator.annotate(
|
||||
scene=image.copy(),
|
||||
detections=detections
|
||||
)
|
||||
@@ -285,6 +285,37 @@ status: new
|
||||
|
||||
</div>
|
||||
|
||||
=== "RichLabel"
|
||||
|
||||
```python
|
||||
import supervision as sv
|
||||
|
||||
image = ...
|
||||
detections = sv.Detections(...)
|
||||
|
||||
labels = [
|
||||
f"{class_name} {confidence:.2f}"
|
||||
for class_name, confidence
|
||||
in zip(detections['class_name'], detections.confidence)
|
||||
]
|
||||
|
||||
rich_label_annotator = sv.RichLabelAnnotator(
|
||||
font_path="<TTF_FONT_PATH>",
|
||||
text_position=sv.Position.CENTER
|
||||
)
|
||||
annotated_frame = label_annotator.annotate(
|
||||
scene=image.copy(),
|
||||
detections=detections,
|
||||
labels=labels
|
||||
)
|
||||
```
|
||||
|
||||
<div class="result" markdown>
|
||||
|
||||
{ align=center width="800" }
|
||||
|
||||
</div>
|
||||
|
||||
=== "Crop"
|
||||
|
||||
```python
|
||||
@@ -403,10 +434,10 @@ status: new
|
||||
</div>
|
||||
|
||||
<div class="md-typeset">
|
||||
<h2><a href="#supervision.annotators.core.BoundingBoxAnnotator">BoundingBoxAnnotator</a></h2>
|
||||
<h2><a href="#supervision.annotators.core.BoxAnnotator">BoxAnnotator</a></h2>
|
||||
</div>
|
||||
|
||||
:::supervision.annotators.core.BoundingBoxAnnotator
|
||||
:::supervision.annotators.core.BoxAnnotator
|
||||
|
||||
<div class="md-typeset">
|
||||
<h2><a href="#supervision.annotators.core.RoundBoxAnnotator">RoundBoxAnnotator</a></h2>
|
||||
@@ -492,6 +523,12 @@ status: new
|
||||
|
||||
:::supervision.annotators.core.LabelAnnotator
|
||||
|
||||
<div class="md-typeset">
|
||||
<h2><a href="#supervision.annotators.core.RichLabelAnnotator">RichLabelAnnotator</a></h2>
|
||||
</div>
|
||||
|
||||
:::supervision.annotators.core.RichLabelAnnotator
|
||||
|
||||
<div class="md-typeset">
|
||||
<h2><a href="#supervision.annotators.core.BlurAnnotator">BlurAnnotator</a></h2>
|
||||
</div>
|
||||
|
||||
@@ -0,0 +1,29 @@
|
||||
---
|
||||
comments: true
|
||||
---
|
||||
|
||||
# Double Detection Filter
|
||||
|
||||
<div class="md-typeset">
|
||||
<h2><a href="#supervision.detection.overlap_filter.OverlapFilter">OverlapFilter</a></h2>
|
||||
</div>
|
||||
|
||||
:::supervision.detection.overlap_filter.OverlapFilter
|
||||
|
||||
<div class="md-typeset">
|
||||
<h2><a href="#supervision.detection.overlap_filter.box_non_max_suppression">box_non_max_suppression</a></h2>
|
||||
</div>
|
||||
|
||||
:::supervision.detection.overlap_filter.box_non_max_suppression
|
||||
|
||||
<div class="md-typeset">
|
||||
<h2><a href="#supervision.detection.overlap_filter.mask_non_max_suppression">mask_non_max_suppression</a></h2>
|
||||
</div>
|
||||
|
||||
:::supervision.detection.overlap_filter.mask_non_max_suppression
|
||||
|
||||
<div class="md-typeset">
|
||||
<h2><a href="#supervision.detection.overlap_filter.box_non_max_merge">box_non_max_merge</a></h2>
|
||||
</div>
|
||||
|
||||
:::supervision.detection.overlap_filter.box_non_max_merge
|
||||
@@ -1,6 +1,5 @@
|
||||
---
|
||||
comments: true
|
||||
status: new
|
||||
---
|
||||
|
||||
# Save Detections
|
||||
|
||||
+42
-12
@@ -17,18 +17,6 @@ status: new
|
||||
|
||||
:::supervision.detection.utils.mask_iou_batch
|
||||
|
||||
<div class="md-typeset">
|
||||
<h2><a href="#supervision.detection.utils.box_non_max_suppression">box_non_max_suppression</a></h2>
|
||||
</div>
|
||||
|
||||
:::supervision.detection.utils.box_non_max_suppression
|
||||
|
||||
<div class="md-typeset">
|
||||
<h2><a href="#supervision.detection.utils.mask_non_max_suppression">mask_non_max_suppression</a></h2>
|
||||
</div>
|
||||
|
||||
:::supervision.detection.utils.mask_non_max_suppression
|
||||
|
||||
<div class="md-typeset">
|
||||
<h2><a href="#supervision.detection.utils.polygon_to_mask">polygon_to_mask</a></h2>
|
||||
</div>
|
||||
@@ -65,8 +53,50 @@ status: new
|
||||
|
||||
:::supervision.detection.utils.move_boxes
|
||||
|
||||
<div class="md-typeset">
|
||||
<h2><a href="#supervision.detection.utils.move_masks">move_masks</a></h2>
|
||||
</div>
|
||||
|
||||
:::supervision.detection.utils.move_masks
|
||||
|
||||
<div class="md-typeset">
|
||||
<h2><a href="#supervision.detection.utils.scale_boxes">scale_boxes</a></h2>
|
||||
</div>
|
||||
|
||||
:::supervision.detection.utils.scale_boxes
|
||||
|
||||
<div class="md-typeset">
|
||||
<h2><a href="#supervision.detection.utils.clip_boxes">clip_boxes</a></h2>
|
||||
</div>
|
||||
|
||||
:::supervision.detection.utils.clip_boxes
|
||||
|
||||
<div class="md-typeset">
|
||||
<h2><a href="#supervision.detection.utils.pad_boxes">pad_boxes</a></h2>
|
||||
</div>
|
||||
|
||||
:::supervision.detection.utils.pad_boxes
|
||||
|
||||
<div class="md-typeset">
|
||||
<h2><a href="#supervision.detection.utils.contains_holes">contains_holes</a></h2>
|
||||
</div>
|
||||
|
||||
:::supervision.detection.utils.xywh_to_xyxy
|
||||
|
||||
<div class="md-typeset">
|
||||
<h2><a href="#supervision.detection.utils.xywh_to_xyxy">xywh_to_xyxy</a></h2>
|
||||
</div>
|
||||
|
||||
:::supervision.detection.utils.xcycwh_to_xyxy
|
||||
|
||||
<div class="md-typeset">
|
||||
<h2><a href="#supervision.detection.utils.xcycwh_to_xyxy">xcycwh_to_xyxy</a></h2>
|
||||
</div>
|
||||
|
||||
:::supervision.detection.utils.contains_holes
|
||||
|
||||
<div class="md-typeset">
|
||||
<h2><a href="#supervision.detection.utils.contains_multiple_segments">contains_multiple_segments</a></h2>
|
||||
</div>
|
||||
|
||||
:::supervision.detection.utils.contains_multiple_segments
|
||||
|
||||
@@ -1,6 +1,5 @@
|
||||
---
|
||||
comments: true
|
||||
status: new
|
||||
---
|
||||
|
||||
# Detect and Annotate
|
||||
@@ -136,7 +135,7 @@ You can load predictions from other computer vision frameworks and libraries usi
|
||||
|
||||
## Annotate Image with Detections
|
||||
|
||||
Finally, we can annotate the image with the predictions. Since we are working with an object detection model, we will use the [`sv.BoundingBoxAnnotator`](/latest/annotators/#supervision.annotators.core.BoundingBoxAnnotator) and [`sv.LabelAnnotator`](/latest/annotators/#supervision.annotators.core.LabelAnnotator) classes.
|
||||
Finally, we can annotate the image with the predictions. Since we are working with an object detection model, we will use the [`sv.BoxAnnotator`](/latest/annotators/#supervision.annotators.core.BoxAnnotator) and [`sv.LabelAnnotator`](/latest/annotators/#supervision.annotators.core.LabelAnnotator) classes.
|
||||
|
||||
=== "Inference"
|
||||
|
||||
@@ -150,10 +149,10 @@ Finally, we can annotate the image with the predictions. Since we are working wi
|
||||
results = model.infer(image)[0]
|
||||
detections = sv.Detections.from_inference(results)
|
||||
|
||||
bounding_box_annotator = sv.BoundingBoxAnnotator()
|
||||
box_annotator = sv.BoxAnnotator()
|
||||
label_annotator = sv.LabelAnnotator()
|
||||
|
||||
annotated_image = bounding_box_annotator.annotate(
|
||||
annotated_image = box_annotator.annotate(
|
||||
scene=image, detections=detections)
|
||||
annotated_image = label_annotator.annotate(
|
||||
scene=annotated_image, detections=detections)
|
||||
@@ -171,10 +170,10 @@ Finally, we can annotate the image with the predictions. Since we are working wi
|
||||
results = model(image)[0]
|
||||
detections = sv.Detections.from_ultralytics(results)
|
||||
|
||||
bounding_box_annotator = sv.BoundingBoxAnnotator()
|
||||
box_annotator = sv.BoxAnnotator()
|
||||
label_annotator = sv.LabelAnnotator()
|
||||
|
||||
annotated_image = bounding_box_annotator.annotate(
|
||||
annotated_image = box_annotator.annotate(
|
||||
scene=image, detections=detections)
|
||||
annotated_image = label_annotator.annotate(
|
||||
scene=annotated_image, detections=detections)
|
||||
@@ -205,10 +204,10 @@ Finally, we can annotate the image with the predictions. Since we are working wi
|
||||
transformers_results=results,
|
||||
id2label=model.config.id2label)
|
||||
|
||||
bounding_box_annotator = sv.BoundingBoxAnnotator()
|
||||
box_annotator = sv.BoxAnnotator()
|
||||
label_annotator = sv.LabelAnnotator()
|
||||
|
||||
annotated_image = bounding_box_annotator.annotate(
|
||||
annotated_image = box_annotator.annotate(
|
||||
scene=image, detections=detections)
|
||||
annotated_image = label_annotator.annotate(
|
||||
scene=annotated_image, detections=detections)
|
||||
@@ -234,7 +233,7 @@ override this behavior by passing a list of custom `labels` to the `annotate` me
|
||||
results = model.infer(image)[0]
|
||||
detections = sv.Detections.from_inference(results)
|
||||
|
||||
bounding_box_annotator = sv.BoundingBoxAnnotator()
|
||||
box_annotator = sv.BoxAnnotator()
|
||||
label_annotator = sv.LabelAnnotator()
|
||||
|
||||
labels = [
|
||||
@@ -243,7 +242,7 @@ override this behavior by passing a list of custom `labels` to the `annotate` me
|
||||
in zip(detections['class_name'], detections.confidence)
|
||||
]
|
||||
|
||||
annotated_image = bounding_box_annotator.annotate(
|
||||
annotated_image = box_annotator.annotate(
|
||||
scene=image, detections=detections)
|
||||
annotated_image = label_annotator.annotate(
|
||||
scene=annotated_image, detections=detections, labels=labels)
|
||||
@@ -261,7 +260,7 @@ override this behavior by passing a list of custom `labels` to the `annotate` me
|
||||
results = model(image)[0]
|
||||
detections = sv.Detections.from_ultralytics(results)
|
||||
|
||||
bounding_box_annotator = sv.BoundingBoxAnnotator()
|
||||
box_annotator = sv.BoxAnnotator()
|
||||
label_annotator = sv.LabelAnnotator()
|
||||
|
||||
labels = [
|
||||
@@ -270,7 +269,7 @@ override this behavior by passing a list of custom `labels` to the `annotate` me
|
||||
in zip(detections['class_name'], detections.confidence)
|
||||
]
|
||||
|
||||
annotated_image = bounding_box_annotator.annotate(
|
||||
annotated_image = box_annotator.annotate(
|
||||
scene=image, detections=detections)
|
||||
annotated_image = label_annotator.annotate(
|
||||
scene=annotated_image, detections=detections, labels=labels)
|
||||
@@ -301,7 +300,7 @@ override this behavior by passing a list of custom `labels` to the `annotate` me
|
||||
transformers_results=results,
|
||||
id2label=model.config.id2label)
|
||||
|
||||
bounding_box_annotator = sv.BoundingBoxAnnotator()
|
||||
box_annotator = sv.BoxAnnotator()
|
||||
label_annotator = sv.LabelAnnotator()
|
||||
|
||||
labels = [
|
||||
@@ -310,7 +309,7 @@ override this behavior by passing a list of custom `labels` to the `annotate` me
|
||||
in zip(detections['class_name'], detections.confidence)
|
||||
]
|
||||
|
||||
annotated_image = bounding_box_annotator.annotate(
|
||||
annotated_image = box_annotator.annotate(
|
||||
scene=image, detections=detections)
|
||||
annotated_image = label_annotator.annotate(
|
||||
scene=annotated_image, detections=detections, labels=labels)
|
||||
@@ -323,7 +322,7 @@ override this behavior by passing a list of custom `labels` to the `annotate` me
|
||||
If you are running the segmentation model
|
||||
[`sv.MaskAnnotator`](/latest/annotators/#supervision.annotators.core.MaskAnnotator)
|
||||
is a drop-in replacement for
|
||||
[`sv.BoundingBoxAnnotator`](/latest/annotators/#supervision.annotators.core.BoundingBoxAnnotator)
|
||||
[`sv.BoxAnnotator`](/latest/annotators/#supervision.annotators.core.BoxAnnotator)
|
||||
that will allow you to draw masks instead of boxes.
|
||||
|
||||
=== "Inference"
|
||||
|
||||
@@ -1,12 +1,11 @@
|
||||
---
|
||||
comments: true
|
||||
status: new
|
||||
---
|
||||
|
||||
# Detect Small Objects
|
||||
|
||||
This guide shows how to detect small objects
|
||||
with the [Inference](https://github.com/roboflow/inference),
|
||||
with the [Inference](https://github.com/roboflow/inference),
|
||||
[Ultralytics](https://github.com/ultralytics/ultralytics) or
|
||||
[Transformers](https://github.com/huggingface/transformers) packages using
|
||||
[`InferenceSlicer`](/latest/detection/tools/inference_slicer/#supervision.detection.tools.inference_slicer.InferenceSlicer).
|
||||
@@ -32,10 +31,10 @@ size relative to the image resolution.
|
||||
results = model.infer(image)[0]
|
||||
detections = sv.Detections.from_inference(results)
|
||||
|
||||
bounding_box_annotator = sv.BoundingBoxAnnotator()
|
||||
box_annotator = sv.BoxAnnotator()
|
||||
label_annotator = sv.LabelAnnotator()
|
||||
|
||||
annotated_image = bounding_box_annotator.annotate(
|
||||
annotated_image = box_annotator.annotate(
|
||||
scene=image, detections=detections)
|
||||
annotated_image = label_annotator.annotate(
|
||||
scene=annotated_image, detections=detections)
|
||||
@@ -53,10 +52,10 @@ size relative to the image resolution.
|
||||
results = model(image)[0]
|
||||
detections = sv.Detections.from_ultralytics(results)
|
||||
|
||||
bounding_box_annotator = sv.BoundingBoxAnnotator()
|
||||
box_annotator = sv.BoxAnnotator()
|
||||
label_annotator = sv.LabelAnnotator()
|
||||
|
||||
annotated_image = bounding_box_annotator.annotate(
|
||||
annotated_image = box_annotator.annotate(
|
||||
scene=image, detections=detections)
|
||||
annotated_image = label_annotator.annotate(
|
||||
scene=annotated_image, detections=detections)
|
||||
@@ -68,10 +67,10 @@ size relative to the image resolution.
|
||||
import torch
|
||||
import supervision as sv
|
||||
from PIL import Image
|
||||
from transformers import DetrImageProcessor, DetrForObjectDetection
|
||||
from transformers import DetrImageProcessor, DetrForSegmentation
|
||||
|
||||
processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
|
||||
model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")
|
||||
model = DetrForSegmentation.from_pretrained("facebook/detr-resnet-50")
|
||||
|
||||
image = Image.open(<SOURCE_IMAGE_PATH>)
|
||||
inputs = processor(images=image, return_tensors="pt")
|
||||
@@ -79,13 +78,13 @@ size relative to the image resolution.
|
||||
with torch.no_grad():
|
||||
outputs = model(**inputs)
|
||||
|
||||
width, height = image.size
|
||||
target_size = torch.tensor([[height, width]])
|
||||
width, height = image_slice.size
|
||||
target_size = torch.tensor([[width, height]])
|
||||
results = processor.post_process_object_detection(
|
||||
outputs=outputs, target_sizes=target_size)[0]
|
||||
detections = sv.Detections.from_transformers(results)
|
||||
|
||||
bounding_box_annotator = sv.BoundingBoxAnnotator()
|
||||
box_annotator = sv.BoxAnnotator()
|
||||
label_annotator = sv.LabelAnnotator()
|
||||
|
||||
labels = [
|
||||
@@ -94,7 +93,7 @@ size relative to the image resolution.
|
||||
in detections.class_id
|
||||
]
|
||||
|
||||
annotated_image = bounding_box_annotator.annotate(
|
||||
annotated_image = box_annotator.annotate(
|
||||
scene=image, detections=detections)
|
||||
annotated_image = label_annotator.annotate(
|
||||
scene=annotated_image, detections=detections, labels=labels)
|
||||
@@ -120,10 +119,10 @@ is less effective for ultra-high-resolution images (4K and above).
|
||||
results = model.infer(image)[0]
|
||||
detections = sv.Detections.from_inference(results)
|
||||
|
||||
bounding_box_annotator = sv.BoundingBoxAnnotator()
|
||||
box_annotator = sv.BoxAnnotator()
|
||||
label_annotator = sv.LabelAnnotator()
|
||||
|
||||
annotated_image = bounding_box_annotator.annotate(
|
||||
annotated_image = box_annotator.annotate(
|
||||
scene=image, detections=detections)
|
||||
annotated_image = label_annotator.annotate(
|
||||
scene=annotated_image, detections=detections)
|
||||
@@ -141,10 +140,10 @@ is less effective for ultra-high-resolution images (4K and above).
|
||||
results = model(image, imgsz=1280)[0]
|
||||
detections = sv.Detections.from_ultralytics(results)
|
||||
|
||||
bounding_box_annotator = sv.BoundingBoxAnnotator()
|
||||
box_annotator = sv.BoxAnnotator()
|
||||
label_annotator = sv.LabelAnnotator()
|
||||
|
||||
annotated_image = bounding_box_annotator.annotate(
|
||||
annotated_image = box_annotator.annotate(
|
||||
scene=image, detections=detections)
|
||||
annotated_image = label_annotator.annotate(
|
||||
scene=annotated_image, detections=detections)
|
||||
@@ -175,15 +174,15 @@ objects within each, and aggregating the results.
|
||||
|
||||
def callback(image_slice: np.ndarray) -> sv.Detections:
|
||||
results = model.infer(image_slice)[0]
|
||||
detections = sv.Detections.from_inference(results)
|
||||
return sv.Detections.from_inference(results)
|
||||
|
||||
slicer = sv.InferenceSlicer(callback = callback)
|
||||
detections = slicer(image)
|
||||
|
||||
bounding_box_annotator = sv.BoundingBoxAnnotator()
|
||||
box_annotator = sv.BoxAnnotator()
|
||||
label_annotator = sv.LabelAnnotator()
|
||||
|
||||
annotated_image = bounding_box_annotator.annotate(
|
||||
annotated_image = box_annotator.annotate(
|
||||
scene=image, detections=detections)
|
||||
annotated_image = label_annotator.annotate(
|
||||
scene=annotated_image, detections=detections)
|
||||
@@ -207,10 +206,10 @@ objects within each, and aggregating the results.
|
||||
slicer = sv.InferenceSlicer(callback = callback)
|
||||
detections = slicer(image)
|
||||
|
||||
bounding_box_annotator = sv.BoundingBoxAnnotator()
|
||||
box_annotator = sv.BoxAnnotator()
|
||||
label_annotator = sv.LabelAnnotator()
|
||||
|
||||
annotated_image = bounding_box_annotator.annotate(
|
||||
annotated_image = box_annotator.annotate(
|
||||
scene=image, detections=detections)
|
||||
annotated_image = label_annotator.annotate(
|
||||
scene=annotated_image, detections=detections)
|
||||
@@ -239,8 +238,8 @@ objects within each, and aggregating the results.
|
||||
with torch.no_grad():
|
||||
outputs = model(**inputs)
|
||||
|
||||
width, height = image.size
|
||||
target_size = torch.tensor([[height, width]])
|
||||
width, height = image_slice.size
|
||||
target_size = torch.tensor([[width, height]])
|
||||
results = processor.post_process_object_detection(
|
||||
outputs=outputs, target_sizes=target_size)[0]
|
||||
return sv.Detections.from_transformers(results)
|
||||
@@ -248,7 +247,7 @@ objects within each, and aggregating the results.
|
||||
slicer = sv.InferenceSlicer(callback = callback)
|
||||
detections = slicer(image)
|
||||
|
||||
bounding_box_annotator = sv.BoundingBoxAnnotator()
|
||||
box_annotator = sv.BoxAnnotator()
|
||||
label_annotator = sv.LabelAnnotator()
|
||||
|
||||
labels = [
|
||||
@@ -257,10 +256,70 @@ objects within each, and aggregating the results.
|
||||
in detections.class_id
|
||||
]
|
||||
|
||||
annotated_image = bounding_box_annotator.annotate(
|
||||
annotated_image = box_annotator.annotate(
|
||||
scene=image, detections=detections)
|
||||
annotated_image = label_annotator.annotate(
|
||||
scene=annotated_image, detections=detections, labels=labels)
|
||||
```
|
||||
|
||||

|
||||
|
||||
## Small Object Segmentation
|
||||
|
||||
[`InferenceSlicer`](/latest/detection/tools/inference_slicer/#supervision.detection.tools.inference_slicer.InferenceSlicer) can perform segmentation tasks too.
|
||||
|
||||
=== "Inference"
|
||||
|
||||
```{ .py hl_lines="6 16 19-20" }
|
||||
import cv2
|
||||
import numpy as np
|
||||
import supervision as sv
|
||||
from inference import get_model
|
||||
|
||||
model = get_model(model_id="yolov8x-seg-640")
|
||||
image = cv2.imread(<SOURCE_IMAGE_PATH>)
|
||||
|
||||
def callback(image_slice: np.ndarray) -> sv.Detections:
|
||||
results = model.infer(image_slice)[0]
|
||||
return sv.Detections.from_inference(results)
|
||||
|
||||
slicer = sv.InferenceSlicer(callback = callback)
|
||||
detections = slicer(image)
|
||||
|
||||
mask_annotator = sv.MaskAnnotator()
|
||||
label_annotator = sv.LabelAnnotator()
|
||||
|
||||
annotated_image = mask_annotator.annotate(
|
||||
scene=image, detections=detections)
|
||||
annotated_image = label_annotator.annotate(
|
||||
scene=annotated_image, detections=detections)
|
||||
```
|
||||
|
||||
=== "Ultralytics"
|
||||
|
||||
```{ .py hl_lines="6 16 19-20" }
|
||||
import cv2
|
||||
import numpy as np
|
||||
import supervision as sv
|
||||
from ultralytics import YOLO
|
||||
|
||||
model = YOLO("yolov8x-seg.pt")
|
||||
image = cv2.imread(<SOURCE_IMAGE_PATH>)
|
||||
|
||||
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)
|
||||
|
||||
mask_annotator = sv.MaskAnnotator()
|
||||
label_annotator = sv.LabelAnnotator()
|
||||
|
||||
annotated_image = mask_annotator.annotate(
|
||||
scene=image, detections=detections)
|
||||
annotated_image = label_annotator.annotate(
|
||||
scene=annotated_image, detections=detections)
|
||||
```
|
||||
|
||||

|
||||
|
||||
@@ -1,6 +1,5 @@
|
||||
---
|
||||
comments: true
|
||||
status: new
|
||||
---
|
||||
|
||||
# Save Detections
|
||||
|
||||
@@ -146,4 +146,13 @@ You can install `supervision` in a
|
||||
|
||||
Master the techniques to selectively filter and focus on objects within a specific zone
|
||||
|
||||
- **Cheatsheet**
|
||||
|
||||
***
|
||||
|
||||
Access a quick reference guide to the most common `supervision` functions
|
||||
|
||||
[:octicons-arrow-right-24: Cheatsheet](https://roboflow.github.io/cheatsheet-supervision/)
|
||||
|
||||
|
||||
</div>
|
||||
|
||||
@@ -1,6 +1,5 @@
|
||||
---
|
||||
comments: true
|
||||
status: new
|
||||
---
|
||||
|
||||
# Annotators
|
||||
@@ -13,7 +12,10 @@ status: new
|
||||
image = ...
|
||||
key_points = sv.KeyPoints(...)
|
||||
|
||||
vertex_annotator = sv.VertexAnnotator(color=sv.Color.GREEN, radius=10)
|
||||
vertex_annotator = sv.VertexAnnotator(
|
||||
color=sv.Color.GREEN,
|
||||
radius=10
|
||||
)
|
||||
annotated_frame = vertex_annotator.annotate(
|
||||
scene=image.copy(),
|
||||
key_points=key_points
|
||||
@@ -34,7 +36,10 @@ status: new
|
||||
image = ...
|
||||
key_points = sv.KeyPoints(...)
|
||||
|
||||
edge_annotator = sv.EdgeAnnotator(color=sv.Color.GREEN, thickness=5)
|
||||
edge_annotator = sv.EdgeAnnotator(
|
||||
color=sv.Color.GREEN,
|
||||
thickness=5
|
||||
)
|
||||
annotated_frame = edge_annotator.annotate(
|
||||
scene=image.copy(),
|
||||
key_points=key_points
|
||||
@@ -47,6 +52,31 @@ status: new
|
||||
|
||||
</div>
|
||||
|
||||
=== "VertexLabelAnnotator"
|
||||
|
||||
```python
|
||||
import supervision as sv
|
||||
|
||||
image = ...
|
||||
key_points = sv.KeyPoints(...)
|
||||
|
||||
vertex_label_annotator = sv.VertexLabelAnnotator(
|
||||
color=sv.Color.GREEN,
|
||||
text_color=sv.Color.BLACK,
|
||||
border_radius=5
|
||||
)
|
||||
annotated_frame = vertex_label_annotator.annotate(
|
||||
scene=image.copy(),
|
||||
key_points=key_points
|
||||
)
|
||||
```
|
||||
|
||||
<div class="result" markdown>
|
||||
|
||||
{ align=center width="800" }
|
||||
|
||||
</div>
|
||||
|
||||
<div class="md-typeset">
|
||||
<h2><a href="#supervision.keypoint.annotators.VertexAnnotator">VertexAnnotator</a></h2>
|
||||
</div>
|
||||
@@ -58,3 +88,9 @@ status: new
|
||||
</div>
|
||||
|
||||
:::supervision.keypoint.annotators.EdgeAnnotator
|
||||
|
||||
<div class="md-typeset">
|
||||
<h2><a href="#supervision.keypoint.annotators.VertexLabelAnnotator">VertexLabelAnnotator</a></h2>
|
||||
</div>
|
||||
|
||||
:::supervision.keypoint.annotators.VertexLabelAnnotator
|
||||
|
||||
Vendored
+333
-193
File diff suppressed because one or more lines are too long
+745
File diff suppressed because one or more lines are too long
+779
File diff suppressed because one or more lines are too long
Vendored
+3
@@ -19,6 +19,9 @@
|
||||
<p class="card repo-card" data-url="/develop/notebooks/object-tracking" data-name="Object Tracking" data-labels="TRACKING, ANNOTATOR" data-version="v0.18.0" data-author="nickherrig"></p>
|
||||
<p class="card repo-card" data-url="/develop/notebooks/occupancy_analytics" data-name="Analyzing Zone Occupancy" data-labels="ANNOTATOR,DETECTION,ZONES" data-version="v0.19.0" data-author="stellasphere"></p>
|
||||
<p class="card repo-card" data-url="/develop/notebooks/evaluating-alignment-of-text-to-image-diffusion-models" data-name="Evaluating Alignment of Text-to-image Diffusion Models" data-labels="ANNOTATORS,YOLO WORLD" data-version="v0.19.0rc5" data-author="iamhatesz"></p>
|
||||
<p class="card repo-card" data-url="/develop/notebooks/serialise-detections-to-csv" data-name="Serialise Detections to a CSV File" data-labels="DETECTIONS,CSV SINK,INFERENCE" data-version="v0.21.0" data-author="onuralpszr"></p>
|
||||
<p class="card repo-card" data-url="/develop/notebooks/serialise-detections-to-json" data-name="Serialise Detections to a JSON File" data-labels="DETECTIONS,JSON SINK,INFERENCE" data-version="v0.21.0" data-author="onuralpszr"></p>
|
||||
|
||||
</div>
|
||||
</div>
|
||||
</section>
|
||||
|
||||
@@ -1,6 +1,5 @@
|
||||
---
|
||||
comments: true
|
||||
status: new
|
||||
---
|
||||
|
||||
# ByteTrack
|
||||
|
||||
+1
-1
@@ -41,7 +41,7 @@ comments: true
|
||||
:::supervision.draw.utils.draw_image
|
||||
|
||||
<div class="md-typeset">
|
||||
<h2><a href="#supervision.draw.utils.calculate_optimal_font_scale">calculate_optimal_font_scale</a></h2>
|
||||
<h2><a href="#supervision.draw.utils.calculate_optimal_text_scale">calculate_optimal_text_scale</a></h2>
|
||||
</div>
|
||||
|
||||
:::supervision.draw.utils.calculate_optimal_text_scale
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
---
|
||||
comments: true
|
||||
status: new
|
||||
---
|
||||
|
||||
# File Utils
|
||||
|
||||
+1
-2
@@ -1,6 +1,5 @@
|
||||
---
|
||||
comments: true
|
||||
status: new
|
||||
---
|
||||
|
||||
# Image Utils
|
||||
@@ -12,7 +11,7 @@ status: new
|
||||
:::supervision.utils.image.crop_image
|
||||
|
||||
<div class="md-typeset">
|
||||
<h2><a href="#supervision.utils.image.scale_image">crop_image</a></h2>
|
||||
<h2><a href="#supervision.utils.image.scale_image">scale_image</a></h2>
|
||||
</div>
|
||||
|
||||
:::supervision.utils.image.scale_image
|
||||
|
||||
@@ -1,6 +1,5 @@
|
||||
---
|
||||
comments: true
|
||||
status: new
|
||||
---
|
||||
|
||||
# Iterables Utils
|
||||
|
||||
@@ -11,7 +11,7 @@ from tqdm import tqdm
|
||||
|
||||
import supervision as sv
|
||||
|
||||
COLORS = sv.ColorPalette.default()
|
||||
COLORS = sv.ColorPalette.DEFAULT
|
||||
|
||||
|
||||
def load_zones_config(file_path: str) -> List[np.ndarray]:
|
||||
@@ -38,15 +38,15 @@ def initiate_annotators(
|
||||
) -> Tuple[
|
||||
List[sv.PolygonZone], List[sv.PolygonZoneAnnotator], List[sv.BoundingBoxAnnotator]
|
||||
]:
|
||||
line_thickness = sv.calculate_dynamic_line_thickness(resolution_wh=resolution_wh)
|
||||
text_scale = sv.calculate_dynamic_text_scale(resolution_wh=resolution_wh)
|
||||
line_thickness = sv.calculate_optimal_line_thickness(resolution_wh=resolution_wh)
|
||||
text_scale = sv.calculate_optimal_text_scale(resolution_wh=resolution_wh)
|
||||
|
||||
zones = []
|
||||
zone_annotators = []
|
||||
box_annotators = []
|
||||
|
||||
for index, polygon in enumerate(polygons):
|
||||
zone = sv.PolygonZone(polygon=polygon, frame_resolution_wh=resolution_wh)
|
||||
zone = sv.PolygonZone(polygon=polygon)
|
||||
zone_annotator = sv.PolygonZoneAnnotator(
|
||||
zone=zone,
|
||||
color=COLORS.by_idx(index),
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
gdown
|
||||
inference
|
||||
supervision==0.19.0
|
||||
inference==0.9.17
|
||||
supervision>=0.20.0
|
||||
tqdm
|
||||
ultralytics
|
||||
|
||||
@@ -9,7 +9,7 @@ from ultralytics import YOLO
|
||||
|
||||
import supervision as sv
|
||||
|
||||
COLORS = sv.ColorPalette.default()
|
||||
COLORS = sv.ColorPalette.DEFAULT
|
||||
|
||||
|
||||
def load_zones_config(file_path: str) -> List[np.ndarray]:
|
||||
@@ -36,15 +36,15 @@ def initiate_annotators(
|
||||
) -> Tuple[
|
||||
List[sv.PolygonZone], List[sv.PolygonZoneAnnotator], List[sv.BoundingBoxAnnotator]
|
||||
]:
|
||||
line_thickness = sv.calculate_dynamic_line_thickness(resolution_wh=resolution_wh)
|
||||
text_scale = sv.calculate_dynamic_text_scale(resolution_wh=resolution_wh)
|
||||
line_thickness = sv.calculate_optimal_line_thickness(resolution_wh=resolution_wh)
|
||||
text_scale = sv.calculate_optimal_text_scale(resolution_wh=resolution_wh)
|
||||
|
||||
zones = []
|
||||
zone_annotators = []
|
||||
box_annotators = []
|
||||
|
||||
for index, polygon in enumerate(polygons):
|
||||
zone = sv.PolygonZone(polygon=polygon, frame_resolution_wh=resolution_wh)
|
||||
zone = sv.PolygonZone(polygon=polygon)
|
||||
zone_annotator = sv.PolygonZoneAnnotator(
|
||||
zone=zone,
|
||||
color=COLORS.by_idx(index),
|
||||
|
||||
@@ -98,11 +98,11 @@ if __name__ == "__main__":
|
||||
frame_rate=video_info.fps, track_thresh=args.confidence_threshold
|
||||
)
|
||||
|
||||
thickness = sv.calculate_dynamic_line_thickness(
|
||||
thickness = sv.calculate_optimal_line_thickness(
|
||||
resolution_wh=video_info.resolution_wh
|
||||
)
|
||||
text_scale = sv.calculate_dynamic_text_scale(resolution_wh=video_info.resolution_wh)
|
||||
bounding_box_annotator = sv.BoundingBoxAnnotator(thickness=thickness)
|
||||
text_scale = sv.calculate_optimal_text_scale(resolution_wh=video_info.resolution_wh)
|
||||
box_annotator = sv.BoxAnnotator(thickness=thickness)
|
||||
label_annotator = sv.LabelAnnotator(
|
||||
text_scale=text_scale,
|
||||
text_thickness=thickness,
|
||||
@@ -116,9 +116,7 @@ if __name__ == "__main__":
|
||||
|
||||
frame_generator = sv.get_video_frames_generator(source_path=args.source_video_path)
|
||||
|
||||
polygon_zone = sv.PolygonZone(
|
||||
polygon=SOURCE, frame_resolution_wh=video_info.resolution_wh
|
||||
)
|
||||
polygon_zone = sv.PolygonZone(polygon=SOURCE)
|
||||
view_transformer = ViewTransformer(source=SOURCE, target=TARGET)
|
||||
|
||||
coordinates = defaultdict(lambda: deque(maxlen=video_info.fps))
|
||||
@@ -156,7 +154,7 @@ if __name__ == "__main__":
|
||||
annotated_frame = trace_annotator.annotate(
|
||||
scene=annotated_frame, detections=detections
|
||||
)
|
||||
annotated_frame = bounding_box_annotator.annotate(
|
||||
annotated_frame = box_annotator.annotate(
|
||||
scene=annotated_frame, detections=detections
|
||||
)
|
||||
annotated_frame = label_annotator.annotate(
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
supervision==0.19.0
|
||||
tqdm==4.66.1
|
||||
supervision>=0.20.0
|
||||
tqdm==4.66.3
|
||||
requests
|
||||
ultralytics==8.0.237
|
||||
super-gradients==3.5.0
|
||||
inference==0.9.8
|
||||
inference==0.9.17
|
||||
|
||||
@@ -76,11 +76,11 @@ if __name__ == "__main__":
|
||||
frame_rate=video_info.fps, track_thresh=args.confidence_threshold
|
||||
)
|
||||
|
||||
thickness = sv.calculate_dynamic_line_thickness(
|
||||
thickness = sv.calculate_optimal_line_thickness(
|
||||
resolution_wh=video_info.resolution_wh
|
||||
)
|
||||
text_scale = sv.calculate_dynamic_text_scale(resolution_wh=video_info.resolution_wh)
|
||||
bounding_box_annotator = sv.BoundingBoxAnnotator(thickness=thickness)
|
||||
text_scale = sv.calculate_optimal_text_scale(resolution_wh=video_info.resolution_wh)
|
||||
box_annotator = sv.BoxAnnotator(thickness=thickness)
|
||||
label_annotator = sv.LabelAnnotator(
|
||||
text_scale=text_scale,
|
||||
text_thickness=thickness,
|
||||
@@ -94,9 +94,7 @@ if __name__ == "__main__":
|
||||
|
||||
frame_generator = sv.get_video_frames_generator(source_path=args.source_video_path)
|
||||
|
||||
polygon_zone = sv.PolygonZone(
|
||||
polygon=SOURCE, frame_resolution_wh=video_info.resolution_wh
|
||||
)
|
||||
polygon_zone = sv.PolygonZone(polygon=SOURCE)
|
||||
view_transformer = ViewTransformer(source=SOURCE, target=TARGET)
|
||||
|
||||
coordinates = defaultdict(lambda: deque(maxlen=video_info.fps))
|
||||
@@ -134,7 +132,7 @@ if __name__ == "__main__":
|
||||
annotated_frame = trace_annotator.annotate(
|
||||
scene=annotated_frame, detections=detections
|
||||
)
|
||||
annotated_frame = bounding_box_annotator.annotate(
|
||||
annotated_frame = box_annotator.annotate(
|
||||
scene=annotated_frame, detections=detections
|
||||
)
|
||||
annotated_frame = label_annotator.annotate(
|
||||
|
||||
@@ -77,11 +77,11 @@ if __name__ == "__main__":
|
||||
frame_rate=video_info.fps, track_thresh=args.confidence_threshold
|
||||
)
|
||||
|
||||
thickness = sv.calculate_dynamic_line_thickness(
|
||||
thickness = sv.calculate_optimal_line_thickness(
|
||||
resolution_wh=video_info.resolution_wh
|
||||
)
|
||||
text_scale = sv.calculate_dynamic_text_scale(resolution_wh=video_info.resolution_wh)
|
||||
bounding_box_annotator = sv.BoundingBoxAnnotator(thickness=thickness)
|
||||
text_scale = sv.calculate_optimal_text_scale(resolution_wh=video_info.resolution_wh)
|
||||
box_annotator = sv.BoxAnnotator(thickness=thickness)
|
||||
label_annotator = sv.LabelAnnotator(
|
||||
text_scale=text_scale,
|
||||
text_thickness=thickness,
|
||||
@@ -95,9 +95,7 @@ if __name__ == "__main__":
|
||||
|
||||
frame_generator = sv.get_video_frames_generator(source_path=args.source_video_path)
|
||||
|
||||
polygon_zone = sv.PolygonZone(
|
||||
polygon=SOURCE, frame_resolution_wh=video_info.resolution_wh
|
||||
)
|
||||
polygon_zone = sv.PolygonZone(polygon=SOURCE)
|
||||
view_transformer = ViewTransformer(source=SOURCE, target=TARGET)
|
||||
|
||||
coordinates = defaultdict(lambda: deque(maxlen=video_info.fps))
|
||||
@@ -134,7 +132,7 @@ if __name__ == "__main__":
|
||||
annotated_frame = trace_annotator.annotate(
|
||||
scene=annotated_frame, detections=detections
|
||||
)
|
||||
annotated_frame = bounding_box_annotator.annotate(
|
||||
annotated_frame = box_annotator.annotate(
|
||||
scene=annotated_frame, detections=detections
|
||||
)
|
||||
annotated_frame = label_annotator.annotate(
|
||||
|
||||
@@ -103,7 +103,7 @@ python scripts/draw_zones.py \
|
||||
```bash
|
||||
python scripts/draw_zones.py \
|
||||
--source_path "data/traffic/video.mp4" \
|
||||
--zone_configuration_path "data/traffic/custom_config.json"
|
||||
--zone_configuration_path "data/traffic/config.json"
|
||||
```
|
||||
|
||||
https://github.com/roboflow/supervision/assets/26109316/9d514c9e-2a61-418b-ae49-6ac1ad6ae5ac
|
||||
@@ -157,7 +157,7 @@ Script to run object detection on a video stream using the Roboflow Inference mo
|
||||
- `--iou_threshold`: IOU threshold for non-max suppression. Default is `0.7`.
|
||||
|
||||
```bash
|
||||
python inference_file_example.py \
|
||||
python inference_stream_example.py \
|
||||
--zone_configuration_path "data/checkout/config.json" \
|
||||
--rtsp_url "rtsp://localhost:8554/live0.stream" \
|
||||
--model_id "yolov8x-640" \
|
||||
@@ -167,7 +167,7 @@ python inference_file_example.py \
|
||||
```
|
||||
|
||||
```bash
|
||||
python inference_file_example.py \
|
||||
python inference_stream_example.py \
|
||||
--zone_configuration_path "data/traffic/config.json" \
|
||||
--rtsp_url "rtsp://localhost:8554/live0.stream" \
|
||||
--model_id "yolov8x-640" \
|
||||
@@ -192,7 +192,7 @@ Script to run object detection on a video file using the Ultralytics YOLOv8 mode
|
||||
- `--iou_threshold`: IOU threshold for non-max suppression. Default is `0.7`.
|
||||
|
||||
```bash
|
||||
python inference_file_example.py \
|
||||
python ultralytics_file_example.py \
|
||||
--zone_configuration_path "data/checkout/config.json" \
|
||||
--source_video_path "data/checkout/video.mp4" \
|
||||
--weights "yolov8x.pt" \
|
||||
@@ -203,7 +203,7 @@ python inference_file_example.py \
|
||||
```
|
||||
|
||||
```bash
|
||||
python inference_file_example.py \
|
||||
python ultralytics_file_example.py \
|
||||
--zone_configuration_path "data/traffic/config.json" \
|
||||
--source_video_path "data/traffic/video.mp4" \
|
||||
--weights "yolov8x.pt" \
|
||||
@@ -226,7 +226,7 @@ Script to run object detection on a video stream using the Ultralytics YOLOv8 mo
|
||||
- `--iou_threshold`: IOU threshold for non-max suppression. Default is `0.7`.
|
||||
|
||||
```bash
|
||||
python inference_file_example.py \
|
||||
python ultralytics_stream_example.py \
|
||||
--zone_configuration_path "data/checkout/config.json" \
|
||||
--rtsp_url "rtsp://localhost:8554/live0.stream" \
|
||||
--weights "yolov8x.pt" \
|
||||
@@ -237,7 +237,7 @@ python inference_file_example.py \
|
||||
```
|
||||
|
||||
```bash
|
||||
python inference_file_example.py \
|
||||
python ultralytics_stream_example.py \
|
||||
--zone_configuration_path "data/traffic/config.json" \
|
||||
--rtsp_url "rtsp://localhost:8554/live0.stream" \
|
||||
--weights "yolov8x.pt" \
|
||||
|
||||
@@ -29,14 +29,10 @@ def main(
|
||||
video_info = sv.VideoInfo.from_video_path(video_path=source_video_path)
|
||||
frames_generator = sv.get_video_frames_generator(source_video_path)
|
||||
|
||||
frame = next(frames_generator)
|
||||
resolution_wh = frame.shape[1], frame.shape[0]
|
||||
|
||||
polygons = load_zones_config(file_path=zone_configuration_path)
|
||||
zones = [
|
||||
sv.PolygonZone(
|
||||
polygon=polygon,
|
||||
frame_resolution_wh=resolution_wh,
|
||||
triggering_anchors=(sv.Position.CENTER,),
|
||||
)
|
||||
for polygon in polygons
|
||||
|
||||
@@ -29,14 +29,10 @@ def main(
|
||||
frames_generator = get_stream_frames_generator(rtsp_url=rtsp_url)
|
||||
fps_monitor = sv.FPSMonitor()
|
||||
|
||||
frame = next(frames_generator)
|
||||
resolution_wh = frame.shape[1], frame.shape[0]
|
||||
|
||||
polygons = load_zones_config(file_path=zone_configuration_path)
|
||||
zones = [
|
||||
sv.PolygonZone(
|
||||
polygon=polygon,
|
||||
frame_resolution_wh=resolution_wh,
|
||||
triggering_anchors=(sv.Position.CENTER,),
|
||||
)
|
||||
for polygon in polygons
|
||||
|
||||
@@ -24,20 +24,15 @@ class CustomSink:
|
||||
self.fps_monitor = sv.FPSMonitor()
|
||||
self.polygons = load_zones_config(file_path=zone_configuration_path)
|
||||
self.timers = [ClockBasedTimer() for _ in self.polygons]
|
||||
self.zones = None
|
||||
self.zones = [
|
||||
sv.PolygonZone(
|
||||
polygon=polygon,
|
||||
triggering_anchors=(sv.Position.CENTER,),
|
||||
)
|
||||
for polygon in self.polygons
|
||||
]
|
||||
|
||||
def on_prediction(self, result: dict, frame: VideoFrame) -> None:
|
||||
if self.zones is None:
|
||||
resolution_wh = frame.image.shape[1], frame.image.shape[0]
|
||||
self.zones = [
|
||||
sv.PolygonZone(
|
||||
polygon=polygon,
|
||||
frame_resolution_wh=resolution_wh,
|
||||
triggering_anchors=(sv.Position.CENTER,),
|
||||
)
|
||||
for polygon in self.polygons
|
||||
]
|
||||
|
||||
self.fps_monitor.tick()
|
||||
fps = self.fps_monitor.fps
|
||||
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
opencv-python
|
||||
supervision
|
||||
supervision>=0.20.0
|
||||
ultralytics
|
||||
inference
|
||||
inference==0.9.17
|
||||
pytube
|
||||
|
||||
@@ -30,14 +30,10 @@ def main(
|
||||
video_info = sv.VideoInfo.from_video_path(video_path=source_video_path)
|
||||
frames_generator = sv.get_video_frames_generator(source_video_path)
|
||||
|
||||
frame = next(frames_generator)
|
||||
resolution_wh = frame.shape[1], frame.shape[0]
|
||||
|
||||
polygons = load_zones_config(file_path=zone_configuration_path)
|
||||
zones = [
|
||||
sv.PolygonZone(
|
||||
polygon=polygon,
|
||||
frame_resolution_wh=resolution_wh,
|
||||
triggering_anchors=(sv.Position.CENTER,),
|
||||
)
|
||||
for polygon in polygons
|
||||
|
||||
@@ -30,14 +30,10 @@ def main(
|
||||
frames_generator = get_stream_frames_generator(rtsp_url=rtsp_url)
|
||||
fps_monitor = sv.FPSMonitor()
|
||||
|
||||
frame = next(frames_generator)
|
||||
resolution_wh = frame.shape[1], frame.shape[0]
|
||||
|
||||
polygons = load_zones_config(file_path=zone_configuration_path)
|
||||
zones = [
|
||||
sv.PolygonZone(
|
||||
polygon=polygon,
|
||||
frame_resolution_wh=resolution_wh,
|
||||
triggering_anchors=(sv.Position.CENTER,),
|
||||
)
|
||||
for polygon in polygons
|
||||
|
||||
@@ -25,20 +25,15 @@ class CustomSink:
|
||||
self.fps_monitor = sv.FPSMonitor()
|
||||
self.polygons = load_zones_config(file_path=zone_configuration_path)
|
||||
self.timers = [ClockBasedTimer() for _ in self.polygons]
|
||||
self.zones = None
|
||||
self.zones = [
|
||||
sv.PolygonZone(
|
||||
polygon=polygon,
|
||||
triggering_anchors=(sv.Position.CENTER,),
|
||||
)
|
||||
for polygon in self.polygons
|
||||
]
|
||||
|
||||
def on_prediction(self, detections: sv.Detections, frame: VideoFrame) -> None:
|
||||
if self.zones is None:
|
||||
resolution_wh = frame.image.shape[1], frame.image.shape[0]
|
||||
self.zones = [
|
||||
sv.PolygonZone(
|
||||
polygon=polygon,
|
||||
frame_resolution_wh=resolution_wh,
|
||||
triggering_anchors=(sv.Position.CENTER,),
|
||||
)
|
||||
for polygon in self.polygons
|
||||
]
|
||||
|
||||
self.fps_monitor.tick()
|
||||
fps = self.fps_monitor.fps
|
||||
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
inference
|
||||
inference==0.9.17
|
||||
supervision==0.19.0
|
||||
tqdm
|
||||
ultralytics
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
import argparse
|
||||
import os
|
||||
from typing import Dict, Iterable, List, Set, Tuple
|
||||
from typing import Dict, Iterable, List, Set
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
@@ -60,13 +60,11 @@ class DetectionsManager:
|
||||
|
||||
def initiate_polygon_zones(
|
||||
polygons: List[np.ndarray],
|
||||
frame_resolution_wh: Tuple[int, int],
|
||||
triggering_anchors: Iterable[sv.Position] = [sv.Position.CENTER],
|
||||
) -> List[sv.PolygonZone]:
|
||||
return [
|
||||
sv.PolygonZone(
|
||||
polygon=polygon,
|
||||
frame_resolution_wh=frame_resolution_wh,
|
||||
triggering_anchors=triggering_anchors,
|
||||
)
|
||||
for polygon in polygons
|
||||
@@ -92,14 +90,10 @@ class VideoProcessor:
|
||||
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.zones_in = initiate_polygon_zones(ZONE_IN_POLYGONS, [sv.Position.CENTER])
|
||||
self.zones_out = initiate_polygon_zones(ZONE_OUT_POLYGONS, [sv.Position.CENTER])
|
||||
|
||||
self.bounding_box_annotator = sv.BoundingBoxAnnotator(color=COLORS)
|
||||
self.box_annotator = sv.BoxAnnotator(color=COLORS)
|
||||
self.label_annotator = sv.LabelAnnotator(
|
||||
color=COLORS, text_color=sv.Color.BLACK
|
||||
)
|
||||
@@ -140,9 +134,7 @@ class VideoProcessor:
|
||||
|
||||
labels = [f"#{tracker_id}" for tracker_id in detections.tracker_id]
|
||||
annotated_frame = self.trace_annotator.annotate(annotated_frame, detections)
|
||||
annotated_frame = self.bounding_box_annotator.annotate(
|
||||
annotated_frame, detections
|
||||
)
|
||||
annotated_frame = self.box_annotator.annotate(annotated_frame, detections)
|
||||
annotated_frame = self.label_annotator.annotate(
|
||||
annotated_frame, detections, labels
|
||||
)
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
gdown
|
||||
inference
|
||||
supervision>=0.19.0
|
||||
inference==0.9.17
|
||||
supervision>=0.20.0
|
||||
tqdm
|
||||
ultralytics
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
import argparse
|
||||
from typing import Dict, Iterable, List, Set, Tuple
|
||||
from typing import Dict, Iterable, List, Set
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
@@ -58,13 +58,11 @@ class DetectionsManager:
|
||||
|
||||
def initiate_polygon_zones(
|
||||
polygons: List[np.ndarray],
|
||||
frame_resolution_wh: Tuple[int, int],
|
||||
triggering_anchors: Iterable[sv.Position] = [sv.Position.CENTER],
|
||||
) -> List[sv.PolygonZone]:
|
||||
return [
|
||||
sv.PolygonZone(
|
||||
polygon=polygon,
|
||||
frame_resolution_wh=frame_resolution_wh,
|
||||
triggering_anchors=triggering_anchors,
|
||||
)
|
||||
for polygon in polygons
|
||||
@@ -89,14 +87,10 @@ class VideoProcessor:
|
||||
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.zones_in = initiate_polygon_zones(ZONE_IN_POLYGONS, [sv.Position.CENTER])
|
||||
self.zones_out = initiate_polygon_zones(ZONE_OUT_POLYGONS, [sv.Position.CENTER])
|
||||
|
||||
self.bounding_box_annotator = sv.BoundingBoxAnnotator(color=COLORS)
|
||||
self.box_annotator = sv.BoxAnnotator(color=COLORS)
|
||||
self.label_annotator = sv.LabelAnnotator(
|
||||
color=COLORS, text_color=sv.Color.BLACK
|
||||
)
|
||||
@@ -137,9 +131,7 @@ class VideoProcessor:
|
||||
|
||||
labels = [f"#{tracker_id}" for tracker_id in detections.tracker_id]
|
||||
annotated_frame = self.trace_annotator.annotate(annotated_frame, detections)
|
||||
annotated_frame = self.bounding_box_annotator.annotate(
|
||||
annotated_frame, detections
|
||||
)
|
||||
annotated_frame = self.box_annotator.annotate(annotated_frame, detections)
|
||||
annotated_frame = self.label_annotator.annotate(
|
||||
annotated_frame, detections, labels
|
||||
)
|
||||
|
||||
+19
-6
@@ -35,19 +35,20 @@ extra_css:
|
||||
|
||||
|
||||
nav:
|
||||
- Home: index.md
|
||||
- How to:
|
||||
- Supervision: index.md
|
||||
- Learn:
|
||||
- Detect and Annotate: how_to/detect_and_annotate.md
|
||||
- Save Detections: how_to/save_detections.md
|
||||
- Filter Detections: how_to/filter_detections.md
|
||||
- Detect Small Objects: how_to/detect_small_objects.md
|
||||
- Track Objects: how_to/track_objects.md
|
||||
- Track Objects on Video: how_to/track_objects.md
|
||||
|
||||
- API:
|
||||
- Reference - Code API:
|
||||
- Detection and Segmentation:
|
||||
- Core: detection/core.md
|
||||
- Annotators: detection/annotators.md
|
||||
- Metrics: detection/metrics.md
|
||||
- Double Detection Filter: detection/double_detection_filter.md
|
||||
- Utils: detection/utils.md
|
||||
- Keypoint Detection:
|
||||
- Core: keypoint/core.md
|
||||
@@ -61,7 +62,9 @@ nav:
|
||||
- Detection Smoother: detection/tools/smoother.md
|
||||
- Save Detections: detection/tools/save_detections.md
|
||||
- Trackers: trackers.md
|
||||
- Datasets: datasets.md
|
||||
- Datasets:
|
||||
- Core: datasets/core.md
|
||||
- Utils: datasets/utils.md
|
||||
- Utils:
|
||||
- Video: utils/video.md
|
||||
- Image: utils/image.md
|
||||
@@ -72,11 +75,12 @@ nav:
|
||||
- Geometry: utils/geometry.md
|
||||
- Assets: assets.md
|
||||
- Cookbooks: cookbooks.md
|
||||
- Cheatsheet: https://roboflow.github.io/cheatsheet-supervision/
|
||||
- Contribute:
|
||||
- Contributing: contributing.md
|
||||
- Code of Conduct: code_of_conduct.md
|
||||
- License: license.md
|
||||
- Changelog:
|
||||
- Release Notes:
|
||||
- Changelog: changelog.md
|
||||
- Deprecated: deprecated.md
|
||||
|
||||
@@ -175,3 +179,12 @@ extra_javascript:
|
||||
- "javascripts/init_kapa_widget.js"
|
||||
- "javascripts/cookbooks-card.js"
|
||||
- "https://cdnjs.cloudflare.com/ajax/libs/dompurify/3.0.8/purify.min.js"
|
||||
|
||||
# Messages shown during document build
|
||||
# Reference: https://www.mkdocs.org/user-guide/configuration/#validation
|
||||
# Values: [warn, info, ignore]
|
||||
validation:
|
||||
nav:
|
||||
absolute_links: ignore
|
||||
links:
|
||||
absolute_links: ignore
|
||||
|
||||
Generated
+1040
-977
File diff suppressed because it is too large
Load Diff
+72
-18
@@ -1,18 +1,26 @@
|
||||
[tool.poetry]
|
||||
name = "supervision"
|
||||
version = "0.20.0"
|
||||
version = "0.22.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"}]
|
||||
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", "Roboflow"]
|
||||
keywords = [
|
||||
"machine-learning",
|
||||
"deep-learning",
|
||||
"vision",
|
||||
"ML",
|
||||
"DL",
|
||||
"AI",
|
||||
"Roboflow",
|
||||
]
|
||||
|
||||
classifiers=[
|
||||
classifiers = [
|
||||
'Development Status :: 4 - Beta',
|
||||
'Intended Audience :: Developers',
|
||||
'Intended Audience :: Education',
|
||||
@@ -27,7 +35,7 @@ classifiers=[
|
||||
'Typing :: Typed',
|
||||
'Operating System :: Microsoft :: Windows',
|
||||
'Operating System :: POSIX :: Linux',
|
||||
'Operating System :: MacOS'
|
||||
'Operating System :: MacOS',
|
||||
]
|
||||
|
||||
[tool.poetry.dependencies]
|
||||
@@ -35,23 +43,23 @@ python = "^3.8"
|
||||
numpy = ">=1.21.2"
|
||||
scipy = [
|
||||
{ version = "1.10.0", python = "<3.9" },
|
||||
{ version = "^1.10.0", python = ">=3.9" }
|
||||
{ version = "^1.10.0", python = ">=3.9" },
|
||||
]
|
||||
matplotlib = ">=3.6.0"
|
||||
pyyaml = ">=5.3"
|
||||
defusedxml = "^0.7.1"
|
||||
opencv-python = { version = ">=4.5.5.64", optional = true }
|
||||
opencv-python-headless = ">=4.5.5.64"
|
||||
requests = { version = ">=2.26.0,<=2.31.0", optional = true }
|
||||
tqdm = { version = ">=4.62.3,<=4.66.2", optional = true }
|
||||
requests = { version = ">=2.26.0,<=2.32.3", optional = true }
|
||||
tqdm = { version = ">=4.62.3,<=4.66.4", optional = true }
|
||||
pillow = ">=9.4"
|
||||
|
||||
[tool.poetry.extras]
|
||||
desktop = ["opencv-python"]
|
||||
assets = ["requests","tqdm"]
|
||||
assets = ["requests", "tqdm"]
|
||||
|
||||
[tool.poetry.group.dev.dependencies]
|
||||
twine = ">=4.0.2,<6.0.0"
|
||||
twine = "^5.1.1"
|
||||
pytest = ">=7.2.2,<9.0.0"
|
||||
wheel = ">=0.40,<0.44"
|
||||
build = ">=0.10,<1.3"
|
||||
@@ -63,11 +71,15 @@ notebook = ">=6.5.3,<8.0.0"
|
||||
ipywidgets = "^8.1.1"
|
||||
jupytext = "^1.16.1"
|
||||
nbconvert = "^7.14.2"
|
||||
docutils = [
|
||||
{ version = "^0.20.1", python = "<3.9" },
|
||||
{ version = "^0.21.1", python = ">=3.9" },
|
||||
]
|
||||
|
||||
|
||||
[tool.poetry.group.docs.dependencies]
|
||||
mkdocs-material = {extras = ["imaging"], version = "^9.5.5"}
|
||||
mkdocstrings = {extras = ["python"], version = ">=0.20,<0.25"}
|
||||
mkdocs-material = { extras = ["imaging"], version = "^9.5.5" }
|
||||
mkdocstrings = { extras = ["python"], version = ">=0.20,<0.26" }
|
||||
mike = "^2.0.0"
|
||||
# For Documentation Development use Python 3.10 or above
|
||||
# Use Latest mkdocs-jupyter min 0.24.6 for Jupyter Notebook Theme support
|
||||
@@ -75,8 +87,6 @@ mkdocs-jupyter = "^0.24.3"
|
||||
mkdocs-git-committers-plugin-2 = "^2.2.3"
|
||||
mkdocs-git-revision-date-localized-plugin = "^1.2.4"
|
||||
|
||||
|
||||
|
||||
[tool.isort]
|
||||
line_length = 88
|
||||
profile = "black"
|
||||
@@ -89,7 +99,6 @@ tests = ["B201", "B301", "B318", "B314", "B303", "B413", "B412", "B410"]
|
||||
check = true
|
||||
imports = ["cv2", "supervision"]
|
||||
|
||||
|
||||
[tool.black]
|
||||
target-version = ["py38"]
|
||||
line-length = 88
|
||||
@@ -148,7 +157,52 @@ indent-width = 4
|
||||
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"]
|
||||
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 = []
|
||||
# Allow unused variables when underscore-prefixed.
|
||||
dummy-variable-rgx = "^(_+|(_+[a-zA-Z0-9_]*[a-zA-Z0-9]+?))$"
|
||||
@@ -163,7 +217,7 @@ docstring-quotes = "double"
|
||||
convention = "google"
|
||||
|
||||
[tool.ruff.lint.per-file-ignores]
|
||||
"__init__.py" = ["E402","F401"]
|
||||
"__init__.py" = ["E402", "F401"]
|
||||
|
||||
[tool.ruff.lint.mccabe]
|
||||
# Flag errors (`C901`) whenever the complexity level exceeds 5.
|
||||
@@ -186,7 +240,7 @@ line-ending = "auto"
|
||||
include-package-data = false
|
||||
|
||||
[tool.setuptools.packages.find]
|
||||
exclude = ["docs*", "test*","examples*"]
|
||||
exclude = ["docs*", "test*", "examples*"]
|
||||
|
||||
[build-system]
|
||||
requires = ["poetry-core"]
|
||||
|
||||
@@ -0,0 +1,19 @@
|
||||
# Release Process
|
||||
|
||||
This doc outlines how supervision is released into production.
|
||||
|
||||
It assumes you already have the code changes, as well as a draft of the release notes.
|
||||
|
||||
1. Make sure you have all required changes were merged into `develop`.
|
||||
2. Create and merge a PR, merging `develop` into `main`, containing:
|
||||
- A commit that updates the project version in `pyproject.toml`.
|
||||
- All changes made during the release.
|
||||
3. Tag the commit with the new supervision version.
|
||||
- Upon pushing the tag, the [PyPi](https://pypi.org/project/supervision/) should update to the new version. Check this!
|
||||
4. Open and merge a PR, merging `main` into `develop`.
|
||||
5. Update the docs by running the [Supervision Release Documentation Workflow 📚](https://github.com/roboflow/supervision/actions/workflows/publish-release-docs.yml) workflow from GitHub.
|
||||
- Select the `main` branch from the dropdown.
|
||||
6. Create a release on GitHub.
|
||||
- Go to releases
|
||||
- Assign the release notes to the tag created in step 3.
|
||||
- Publish the release.
|
||||
+22
-4
@@ -9,6 +9,7 @@ except importlib_metadata.PackageNotFoundError:
|
||||
from supervision.annotators.core import (
|
||||
BlurAnnotator,
|
||||
BoundingBoxAnnotator,
|
||||
BoxAnnotator,
|
||||
BoxCornerAnnotator,
|
||||
CircleAnnotator,
|
||||
ColorAnnotator,
|
||||
@@ -23,6 +24,7 @@ from supervision.annotators.core import (
|
||||
PercentageBarAnnotator,
|
||||
PixelateAnnotator,
|
||||
PolygonAnnotator,
|
||||
RichLabelAnnotator,
|
||||
RoundBoxAnnotator,
|
||||
TraceAnnotator,
|
||||
TriangleAnnotator,
|
||||
@@ -34,9 +36,16 @@ from supervision.dataset.core import (
|
||||
ClassificationDataset,
|
||||
DetectionDataset,
|
||||
)
|
||||
from supervision.detection.annotate import BoxAnnotator
|
||||
from supervision.dataset.utils import mask_to_rle, rle_to_mask
|
||||
from supervision.detection.core import Detections
|
||||
from supervision.detection.line_zone import LineZone, LineZoneAnnotator
|
||||
from supervision.detection.lmm import LMM
|
||||
from supervision.detection.overlap_filter import (
|
||||
OverlapFilter,
|
||||
box_non_max_merge,
|
||||
box_non_max_suppression,
|
||||
mask_non_max_suppression,
|
||||
)
|
||||
from supervision.detection.tools.csv_sink import CSVSink
|
||||
from supervision.detection.tools.inference_slicer import InferenceSlicer
|
||||
from supervision.detection.tools.json_sink import JSONSink
|
||||
@@ -44,17 +53,22 @@ from supervision.detection.tools.polygon_zone import PolygonZone, PolygonZoneAnn
|
||||
from supervision.detection.tools.smoother import DetectionsSmoother
|
||||
from supervision.detection.utils import (
|
||||
box_iou_batch,
|
||||
box_non_max_suppression,
|
||||
calculate_masks_centroids,
|
||||
clip_boxes,
|
||||
contains_holes,
|
||||
contains_multiple_segments,
|
||||
filter_polygons_by_area,
|
||||
mask_iou_batch,
|
||||
mask_non_max_suppression,
|
||||
mask_to_polygons,
|
||||
mask_to_xyxy,
|
||||
move_boxes,
|
||||
move_masks,
|
||||
pad_boxes,
|
||||
polygon_to_mask,
|
||||
polygon_to_xyxy,
|
||||
scale_boxes,
|
||||
xcycwh_to_xyxy,
|
||||
xywh_to_xyxy,
|
||||
)
|
||||
from supervision.draw.color import Color, ColorPalette
|
||||
from supervision.draw.utils import (
|
||||
@@ -69,7 +83,11 @@ from supervision.draw.utils import (
|
||||
)
|
||||
from supervision.geometry.core import Point, Position, Rect
|
||||
from supervision.geometry.utils import get_polygon_center
|
||||
from supervision.keypoint.annotators import EdgeAnnotator, VertexAnnotator
|
||||
from supervision.keypoint.annotators import (
|
||||
EdgeAnnotator,
|
||||
VertexAnnotator,
|
||||
VertexLabelAnnotator,
|
||||
)
|
||||
from supervision.keypoint.core import KeyPoints
|
||||
from supervision.metrics.detection import ConfusionMatrix, MeanAveragePrecision
|
||||
from supervision.tracker.byte_tracker.core import ByteTrack
|
||||
|
||||
+326
-102
@@ -3,19 +3,114 @@ from typing import List, Optional, Tuple, Union
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
from PIL import ImageDraw, ImageFont
|
||||
|
||||
from supervision.annotators.base import BaseAnnotator, ImageType
|
||||
from supervision.annotators.utils import ColorLookup, Trace, resolve_color
|
||||
from supervision.annotators.utils import (
|
||||
ColorLookup,
|
||||
Trace,
|
||||
resolve_color,
|
||||
resolve_text_background_xyxy,
|
||||
)
|
||||
from supervision.config import CLASS_NAME_DATA_FIELD, ORIENTED_BOX_COORDINATES
|
||||
from supervision.detection.core import Detections
|
||||
from supervision.detection.utils import clip_boxes, mask_to_polygons
|
||||
from supervision.draw.color import Color, ColorPalette
|
||||
from supervision.draw.utils import draw_polygon
|
||||
from supervision.geometry.core import Position
|
||||
from supervision.utils.conversion import convert_for_annotation_method
|
||||
from supervision.utils.conversion import (
|
||||
ensure_cv2_image_for_annotation,
|
||||
ensure_pil_image_for_annotation,
|
||||
)
|
||||
from supervision.utils.image import crop_image, overlay_image, scale_image
|
||||
from supervision.utils.internal import deprecated
|
||||
|
||||
|
||||
class BoxAnnotator(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_lookup: ColorLookup = ColorLookup.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_lookup (str): Strategy for mapping colors to annotations.
|
||||
Options are `INDEX`, `CLASS`, `TRACK`.
|
||||
"""
|
||||
self.color: Union[Color, ColorPalette] = color
|
||||
self.thickness: int = thickness
|
||||
self.color_lookup: ColorLookup = color_lookup
|
||||
|
||||
@ensure_cv2_image_for_annotation
|
||||
def annotate(
|
||||
self,
|
||||
scene: ImageType,
|
||||
detections: Detections,
|
||||
custom_color_lookup: Optional[np.ndarray] = None,
|
||||
) -> ImageType:
|
||||
"""
|
||||
Annotates the given scene with bounding boxes based on the provided detections.
|
||||
|
||||
Args:
|
||||
scene (ImageType): The image where bounding boxes will be drawn. `ImageType`
|
||||
is a flexible type, accepting either `numpy.ndarray` or `PIL.Image.Image`.
|
||||
detections (Detections): Object detections to annotate.
|
||||
custom_color_lookup (Optional[np.ndarray]): Custom color lookup array.
|
||||
Allows to override the default color mapping strategy.
|
||||
|
||||
Returns:
|
||||
The annotated image, matching the type of `scene` (`numpy.ndarray`
|
||||
or `PIL.Image.Image`)
|
||||
|
||||
Example:
|
||||
```python
|
||||
import supervision as sv
|
||||
|
||||
image = ...
|
||||
detections = sv.Detections(...)
|
||||
|
||||
box_annotator = sv.BoxAnnotator()
|
||||
annotated_frame = 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)
|
||||
color = resolve_color(
|
||||
color=self.color,
|
||||
detections=detections,
|
||||
detection_idx=detection_idx,
|
||||
color_lookup=self.color_lookup
|
||||
if custom_color_lookup is None
|
||||
else custom_color_lookup,
|
||||
)
|
||||
cv2.rectangle(
|
||||
img=scene,
|
||||
pt1=(x1, y1),
|
||||
pt2=(x2, y2),
|
||||
color=color.as_bgr(),
|
||||
thickness=self.thickness,
|
||||
)
|
||||
return scene
|
||||
|
||||
|
||||
@deprecated(
|
||||
"`BoundingBoxAnnotator` is deprecated and has been renamed to `BoxAnnotator`."
|
||||
" `BoundingBoxAnnotator` will be removed in supervision-0.26.0."
|
||||
)
|
||||
class BoundingBoxAnnotator(BaseAnnotator):
|
||||
"""
|
||||
A class for drawing bounding boxes on an image using provided detections.
|
||||
@@ -39,7 +134,7 @@ class BoundingBoxAnnotator(BaseAnnotator):
|
||||
self.thickness: int = thickness
|
||||
self.color_lookup: ColorLookup = color_lookup
|
||||
|
||||
@convert_for_annotation_method
|
||||
@ensure_cv2_image_for_annotation
|
||||
def annotate(
|
||||
self,
|
||||
scene: ImageType,
|
||||
@@ -120,7 +215,7 @@ class OrientedBoxAnnotator(BaseAnnotator):
|
||||
self.thickness: int = thickness
|
||||
self.color_lookup: ColorLookup = color_lookup
|
||||
|
||||
@convert_for_annotation_method
|
||||
@ensure_cv2_image_for_annotation
|
||||
def annotate(
|
||||
self,
|
||||
scene: ImageType,
|
||||
@@ -166,7 +261,7 @@ class OrientedBoxAnnotator(BaseAnnotator):
|
||||
return scene
|
||||
|
||||
for detection_idx in range(len(detections)):
|
||||
bbox = np.int0(detections.data.get(ORIENTED_BOX_COORDINATES)[detection_idx])
|
||||
bbox = np.intp(detections.data.get(ORIENTED_BOX_COORDINATES)[detection_idx])
|
||||
color = resolve_color(
|
||||
color=self.color,
|
||||
detections=detections,
|
||||
@@ -208,7 +303,7 @@ class MaskAnnotator(BaseAnnotator):
|
||||
self.opacity = opacity
|
||||
self.color_lookup: ColorLookup = color_lookup
|
||||
|
||||
@convert_for_annotation_method
|
||||
@ensure_cv2_image_for_annotation
|
||||
def annotate(
|
||||
self,
|
||||
scene: ImageType,
|
||||
@@ -264,8 +359,10 @@ class MaskAnnotator(BaseAnnotator):
|
||||
mask = detections.mask[detection_idx]
|
||||
colored_mask[mask] = color.as_bgr()
|
||||
|
||||
scene = cv2.addWeighted(colored_mask, self.opacity, scene, 1 - self.opacity, 0)
|
||||
return scene.astype(np.uint8)
|
||||
cv2.addWeighted(
|
||||
colored_mask, self.opacity, scene, 1 - self.opacity, 0, dst=scene
|
||||
)
|
||||
return scene
|
||||
|
||||
|
||||
class PolygonAnnotator(BaseAnnotator):
|
||||
@@ -295,7 +392,7 @@ class PolygonAnnotator(BaseAnnotator):
|
||||
self.thickness: int = thickness
|
||||
self.color_lookup: ColorLookup = color_lookup
|
||||
|
||||
@convert_for_annotation_method
|
||||
@ensure_cv2_image_for_annotation
|
||||
def annotate(
|
||||
self,
|
||||
scene: ImageType,
|
||||
@@ -381,7 +478,7 @@ class ColorAnnotator(BaseAnnotator):
|
||||
self.color_lookup: ColorLookup = color_lookup
|
||||
self.opacity = opacity
|
||||
|
||||
@convert_for_annotation_method
|
||||
@ensure_cv2_image_for_annotation
|
||||
def annotate(
|
||||
self,
|
||||
scene: ImageType,
|
||||
@@ -420,7 +517,7 @@ class ColorAnnotator(BaseAnnotator):
|
||||

|
||||
"""
|
||||
mask_image = scene.copy()
|
||||
scene_with_boxes = scene.copy()
|
||||
for detection_idx in range(len(detections)):
|
||||
x1, y1, x2, y2 = detections.xyxy[detection_idx].astype(int)
|
||||
color = resolve_color(
|
||||
@@ -432,14 +529,15 @@ class ColorAnnotator(BaseAnnotator):
|
||||
else custom_color_lookup,
|
||||
)
|
||||
cv2.rectangle(
|
||||
img=scene,
|
||||
img=scene_with_boxes,
|
||||
pt1=(x1, y1),
|
||||
pt2=(x2, y2),
|
||||
color=color.as_bgr(),
|
||||
thickness=-1,
|
||||
)
|
||||
scene = cv2.addWeighted(
|
||||
scene, self.opacity, mask_image, 1 - self.opacity, gamma=0
|
||||
|
||||
cv2.addWeighted(
|
||||
scene_with_boxes, self.opacity, scene, 1 - self.opacity, gamma=0, dst=scene
|
||||
)
|
||||
return scene
|
||||
|
||||
@@ -475,7 +573,7 @@ class HaloAnnotator(BaseAnnotator):
|
||||
self.color_lookup: ColorLookup = color_lookup
|
||||
self.kernel_size: int = kernel_size
|
||||
|
||||
@convert_for_annotation_method
|
||||
@ensure_cv2_image_for_annotation
|
||||
def annotate(
|
||||
self,
|
||||
scene: ImageType,
|
||||
@@ -540,7 +638,8 @@ class HaloAnnotator(BaseAnnotator):
|
||||
gray = cv2.cvtColor(colored_mask, cv2.COLOR_BGR2GRAY)
|
||||
alpha = self.opacity * gray / gray.max()
|
||||
alpha_mask = alpha[:, :, np.newaxis]
|
||||
scene = np.uint8(scene * (1 - alpha_mask) + colored_mask * self.opacity)
|
||||
blended_scene = np.uint8(scene * (1 - alpha_mask) + colored_mask * self.opacity)
|
||||
np.copyto(scene, blended_scene)
|
||||
return scene
|
||||
|
||||
|
||||
@@ -573,7 +672,7 @@ class EllipseAnnotator(BaseAnnotator):
|
||||
self.end_angle: int = end_angle
|
||||
self.color_lookup: ColorLookup = color_lookup
|
||||
|
||||
@convert_for_annotation_method
|
||||
@ensure_cv2_image_for_annotation
|
||||
def annotate(
|
||||
self,
|
||||
scene: ImageType,
|
||||
@@ -664,7 +763,7 @@ class BoxCornerAnnotator(BaseAnnotator):
|
||||
self.corner_length: int = corner_length
|
||||
self.color_lookup: ColorLookup = color_lookup
|
||||
|
||||
@convert_for_annotation_method
|
||||
@ensure_cv2_image_for_annotation
|
||||
def annotate(
|
||||
self,
|
||||
scene: ImageType,
|
||||
@@ -752,7 +851,7 @@ class CircleAnnotator(BaseAnnotator):
|
||||
self.thickness: int = thickness
|
||||
self.color_lookup: ColorLookup = color_lookup
|
||||
|
||||
@convert_for_annotation_method
|
||||
@ensure_cv2_image_for_annotation
|
||||
def annotate(
|
||||
self,
|
||||
scene: ImageType,
|
||||
@@ -827,6 +926,7 @@ class DotAnnotator(BaseAnnotator):
|
||||
radius: int = 4,
|
||||
position: Position = Position.CENTER,
|
||||
color_lookup: ColorLookup = ColorLookup.CLASS,
|
||||
outline_thickness: int = 0,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
@@ -836,13 +936,15 @@ class DotAnnotator(BaseAnnotator):
|
||||
position (Position): The anchor position for placing the dot.
|
||||
color_lookup (ColorLookup): Strategy for mapping colors to annotations.
|
||||
Options are `INDEX`, `CLASS`, `TRACK`.
|
||||
outline_thickness (int): Thickness of the outline of the dot.
|
||||
"""
|
||||
self.color: Union[Color, ColorPalette] = color
|
||||
self.radius: int = radius
|
||||
self.position: Position = position
|
||||
self.color_lookup: ColorLookup = color_lookup
|
||||
self.outline_thickness = outline_thickness
|
||||
|
||||
@convert_for_annotation_method
|
||||
@ensure_cv2_image_for_annotation
|
||||
def annotate(
|
||||
self,
|
||||
scene: ImageType,
|
||||
@@ -892,11 +994,16 @@ class DotAnnotator(BaseAnnotator):
|
||||
else custom_color_lookup,
|
||||
)
|
||||
center = (int(xy[detection_idx, 0]), int(xy[detection_idx, 1]))
|
||||
|
||||
cv2.circle(scene, center, self.radius, color.as_bgr(), -1)
|
||||
if self.outline_thickness:
|
||||
cv2.circle(
|
||||
scene, center, self.radius, (0, 0, 0), self.outline_thickness
|
||||
)
|
||||
return scene
|
||||
|
||||
|
||||
class LabelAnnotator:
|
||||
class LabelAnnotator(BaseAnnotator):
|
||||
"""
|
||||
A class for annotating labels on an image using provided detections.
|
||||
"""
|
||||
@@ -936,65 +1043,12 @@ class LabelAnnotator:
|
||||
self.text_anchor: Position = text_position
|
||||
self.color_lookup: ColorLookup = color_lookup
|
||||
|
||||
@staticmethod
|
||||
def resolve_text_background_xyxy(
|
||||
center_coordinates: Tuple[int, int],
|
||||
text_wh: Tuple[int, int],
|
||||
position: Position,
|
||||
) -> Tuple[int, int, int, int]:
|
||||
center_x, center_y = center_coordinates
|
||||
text_w, text_h = text_wh
|
||||
|
||||
if position == Position.TOP_LEFT:
|
||||
return center_x, center_y - text_h, center_x + text_w, center_y
|
||||
elif position == Position.TOP_RIGHT:
|
||||
return center_x - text_w, center_y - text_h, center_x, center_y
|
||||
elif position == Position.TOP_CENTER:
|
||||
return (
|
||||
center_x - text_w // 2,
|
||||
center_y - text_h,
|
||||
center_x + text_w // 2,
|
||||
center_y,
|
||||
)
|
||||
elif position == Position.CENTER or position == Position.CENTER_OF_MASS:
|
||||
return (
|
||||
center_x - text_w // 2,
|
||||
center_y - text_h // 2,
|
||||
center_x + text_w // 2,
|
||||
center_y + text_h // 2,
|
||||
)
|
||||
elif position == Position.BOTTOM_LEFT:
|
||||
return center_x, center_y, center_x + text_w, center_y + text_h
|
||||
elif position == Position.BOTTOM_RIGHT:
|
||||
return center_x - text_w, center_y, center_x, center_y + text_h
|
||||
elif position == Position.BOTTOM_CENTER:
|
||||
return (
|
||||
center_x - text_w // 2,
|
||||
center_y,
|
||||
center_x + text_w // 2,
|
||||
center_y + text_h,
|
||||
)
|
||||
elif position == Position.CENTER_LEFT:
|
||||
return (
|
||||
center_x - text_w,
|
||||
center_y - text_h // 2,
|
||||
center_x,
|
||||
center_y + text_h // 2,
|
||||
)
|
||||
elif position == Position.CENTER_RIGHT:
|
||||
return (
|
||||
center_x,
|
||||
center_y - text_h // 2,
|
||||
center_x + text_w,
|
||||
center_y + text_h // 2,
|
||||
)
|
||||
|
||||
@convert_for_annotation_method
|
||||
@ensure_cv2_image_for_annotation
|
||||
def annotate(
|
||||
self,
|
||||
scene: ImageType,
|
||||
detections: Detections,
|
||||
labels: List[str] = None,
|
||||
labels: Optional[List[str]] = None,
|
||||
custom_color_lookup: Optional[np.ndarray] = None,
|
||||
) -> ImageType:
|
||||
"""
|
||||
@@ -1005,7 +1059,7 @@ class LabelAnnotator:
|
||||
`ImageType` is a flexible type, accepting either `numpy.ndarray`
|
||||
or `PIL.Image.Image`.
|
||||
detections (Detections): Object detections to annotate.
|
||||
labels (List[str]): Optional. Custom labels for each detection.
|
||||
labels (Optional[List[str]]): Custom labels for each detection.
|
||||
custom_color_lookup (Optional[np.ndarray]): Custom color lookup array.
|
||||
Allows to override the default color mapping strategy.
|
||||
|
||||
@@ -1015,7 +1069,7 @@ class LabelAnnotator:
|
||||
|
||||
Example:
|
||||
```python
|
||||
import supervision as sv
|
||||
import supervision as sv
|
||||
|
||||
image = ...
|
||||
detections = sv.Detections(...)
|
||||
@@ -1043,12 +1097,9 @@ class LabelAnnotator:
|
||||
).astype(int)
|
||||
if labels is not None and len(labels) != len(detections):
|
||||
raise ValueError(
|
||||
f"The number of labels provided ({len(labels)}) does not match the "
|
||||
f"number of detections ({len(detections)}). Each detection should have "
|
||||
f"a corresponding label. This discrepancy can occur if the labels and "
|
||||
f"detections are not aligned or if an incorrect number of labels has "
|
||||
f"been provided. Please ensure that the labels array has the same "
|
||||
f"length as the Detections object."
|
||||
f"The number of labels ({len(labels)}) does not match the "
|
||||
f"number of detections ({len(detections)}). Each detection "
|
||||
f"should have exactly 1 label."
|
||||
)
|
||||
|
||||
for detection_idx, center_coordinates in enumerate(anchors_coordinates):
|
||||
@@ -1056,9 +1107,11 @@ class LabelAnnotator:
|
||||
color=self.color,
|
||||
detections=detections,
|
||||
detection_idx=detection_idx,
|
||||
color_lookup=self.color_lookup
|
||||
if custom_color_lookup is None
|
||||
else custom_color_lookup,
|
||||
color_lookup=(
|
||||
self.color_lookup
|
||||
if custom_color_lookup is None
|
||||
else custom_color_lookup
|
||||
),
|
||||
)
|
||||
|
||||
if labels is not None:
|
||||
@@ -1078,7 +1131,7 @@ class LabelAnnotator:
|
||||
)[0]
|
||||
text_w_padded = text_w + 2 * self.text_padding
|
||||
text_h_padded = text_h + 2 * self.text_padding
|
||||
text_background_xyxy = self.resolve_text_background_xyxy(
|
||||
text_background_xyxy = resolve_text_background_xyxy(
|
||||
center_coordinates=tuple(center_coordinates),
|
||||
text_wh=(text_w_padded, text_h_padded),
|
||||
position=self.text_anchor,
|
||||
@@ -1148,6 +1201,173 @@ class LabelAnnotator:
|
||||
return scene
|
||||
|
||||
|
||||
class RichLabelAnnotator(BaseAnnotator):
|
||||
"""
|
||||
A class for annotating labels on an image using provided detections,
|
||||
with support for Unicode characters by using a custom font.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
color: Union[Color, ColorPalette] = ColorPalette.DEFAULT,
|
||||
text_color: Color = Color.WHITE,
|
||||
font_path: Optional[str] = None,
|
||||
font_size: int = 10,
|
||||
text_padding: int = 10,
|
||||
text_position: Position = Position.TOP_LEFT,
|
||||
color_lookup: ColorLookup = ColorLookup.CLASS,
|
||||
border_radius: int = 0,
|
||||
):
|
||||
"""
|
||||
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.
|
||||
font_path (Optional[str]): Path to the font file (e.g., ".ttf" or ".otf")
|
||||
to use for rendering text. If `None`, the default PIL font will be used.
|
||||
font_size (int): Font size for the text.
|
||||
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_lookup (ColorLookup): Strategy for mapping colors to annotations.
|
||||
Options are `INDEX`, `CLASS`, `TRACK`.
|
||||
border_radius (int): The radius to apply round edges. If the selected
|
||||
value is higher than the lower dimension, width or height, is clipped.
|
||||
"""
|
||||
self.color = color
|
||||
self.text_color = text_color
|
||||
self.text_padding = text_padding
|
||||
self.text_anchor = text_position
|
||||
self.color_lookup = color_lookup
|
||||
self.border_radius = border_radius
|
||||
if font_path is not None:
|
||||
try:
|
||||
self.font = ImageFont.truetype(font_path, font_size)
|
||||
except OSError:
|
||||
print(f"Font path '{font_path}' not found. Using PIL's default font.")
|
||||
self.font = self._load_default_font(font_size)
|
||||
else:
|
||||
self.font = self._load_default_font(font_size)
|
||||
|
||||
@ensure_pil_image_for_annotation
|
||||
def annotate(
|
||||
self,
|
||||
scene: ImageType,
|
||||
detections: Detections,
|
||||
labels: Optional[List[str]] = None,
|
||||
custom_color_lookup: Optional[np.ndarray] = None,
|
||||
) -> ImageType:
|
||||
"""
|
||||
Annotates the given scene with labels based on the provided
|
||||
detections, with support for Unicode characters.
|
||||
|
||||
Args:
|
||||
scene (ImageType): The image where labels will be drawn.
|
||||
`ImageType` is a flexible type, accepting either `numpy.ndarray`
|
||||
or `PIL.Image.Image`.
|
||||
detections (Detections): Object detections to annotate.
|
||||
labels (Optional[List[str]]): Custom labels for each detection.
|
||||
custom_color_lookup (Optional[np.ndarray]): Custom color lookup array.
|
||||
Allows to override the default color mapping strategy.
|
||||
|
||||
Returns:
|
||||
The annotated image, matching the type of `scene` (`numpy.ndarray`
|
||||
or `PIL.Image.Image`)
|
||||
|
||||
Example:
|
||||
```python
|
||||
import supervision as sv
|
||||
|
||||
image = ...
|
||||
detections = sv.Detections(...)
|
||||
|
||||
labels = [
|
||||
f"{class_name} {confidence:.2f}"
|
||||
for class_name, confidence
|
||||
in zip(detections['class_name'], detections.confidence)
|
||||
]
|
||||
|
||||
rich_label_annotator = sv.RichLabelAnnotator(font_path="path/to/font.ttf")
|
||||
annotated_frame = label_annotator.annotate(
|
||||
scene=image.copy(),
|
||||
detections=detections,
|
||||
labels=labels
|
||||
)
|
||||
```
|
||||
|
||||
"""
|
||||
draw = ImageDraw.Draw(scene)
|
||||
anchors_coordinates = detections.get_anchors_coordinates(
|
||||
anchor=self.text_anchor
|
||||
).astype(int)
|
||||
if labels is not None and len(labels) != len(detections):
|
||||
raise ValueError(
|
||||
f"The number of labels provided ({len(labels)}) does not match the "
|
||||
f"number of detections ({len(detections)}). Each detection should have "
|
||||
f"a corresponding label."
|
||||
)
|
||||
for detection_idx, center_coordinates in enumerate(anchors_coordinates):
|
||||
color = resolve_color(
|
||||
color=self.color,
|
||||
detections=detections,
|
||||
detection_idx=detection_idx,
|
||||
color_lookup=(
|
||||
self.color_lookup
|
||||
if custom_color_lookup is None
|
||||
else custom_color_lookup
|
||||
),
|
||||
)
|
||||
if labels is not None:
|
||||
text = labels[detection_idx]
|
||||
elif detections[CLASS_NAME_DATA_FIELD] is not None:
|
||||
text = detections[CLASS_NAME_DATA_FIELD][detection_idx]
|
||||
elif detections.class_id is not None:
|
||||
text = str(detections.class_id[detection_idx])
|
||||
else:
|
||||
text = str(detection_idx)
|
||||
|
||||
left, top, right, bottom = draw.textbbox((0, 0), text, font=self.font)
|
||||
text_width = right - left
|
||||
text_height = bottom - top
|
||||
text_w_padded = text_width + 2 * self.text_padding
|
||||
text_h_padded = text_height + 2 * self.text_padding
|
||||
text_background_xyxy = resolve_text_background_xyxy(
|
||||
center_coordinates=tuple(center_coordinates),
|
||||
text_wh=(text_w_padded, text_h_padded),
|
||||
position=self.text_anchor,
|
||||
)
|
||||
|
||||
text_x = text_background_xyxy[0] + self.text_padding - left
|
||||
text_y = text_background_xyxy[1] + self.text_padding - top
|
||||
|
||||
draw.rounded_rectangle(
|
||||
text_background_xyxy,
|
||||
radius=self.border_radius,
|
||||
fill=color.as_rgb(),
|
||||
outline=None,
|
||||
)
|
||||
draw.text(
|
||||
xy=(text_x, text_y),
|
||||
text=text,
|
||||
font=self.font,
|
||||
fill=self.text_color.as_rgb(),
|
||||
)
|
||||
return scene
|
||||
|
||||
@staticmethod
|
||||
def _load_default_font(size):
|
||||
"""
|
||||
PIL either loads a font that accepts a size (e.g. on my machine)
|
||||
or raises an error saying `load_default` does not accept arguments
|
||||
(e.g. in Colab).
|
||||
"""
|
||||
try:
|
||||
font = ImageFont.load_default(size)
|
||||
except TypeError:
|
||||
font = ImageFont.load_default()
|
||||
return font
|
||||
|
||||
|
||||
class BlurAnnotator(BaseAnnotator):
|
||||
"""
|
||||
A class for blurring regions in an image using provided detections.
|
||||
@@ -1160,7 +1380,7 @@ class BlurAnnotator(BaseAnnotator):
|
||||
"""
|
||||
self.kernel_size: int = kernel_size
|
||||
|
||||
@convert_for_annotation_method
|
||||
@ensure_cv2_image_for_annotation
|
||||
def annotate(
|
||||
self,
|
||||
scene: ImageType,
|
||||
@@ -1209,7 +1429,7 @@ class BlurAnnotator(BaseAnnotator):
|
||||
return scene
|
||||
|
||||
|
||||
class TraceAnnotator:
|
||||
class TraceAnnotator(BaseAnnotator):
|
||||
"""
|
||||
A class for drawing trace paths on an image based on detection coordinates.
|
||||
|
||||
@@ -1245,7 +1465,7 @@ class TraceAnnotator:
|
||||
self.thickness = thickness
|
||||
self.color_lookup: ColorLookup = color_lookup
|
||||
|
||||
@convert_for_annotation_method
|
||||
@ensure_cv2_image_for_annotation
|
||||
def annotate(
|
||||
self,
|
||||
scene: ImageType,
|
||||
@@ -1318,7 +1538,7 @@ class TraceAnnotator:
|
||||
return scene
|
||||
|
||||
|
||||
class HeatMapAnnotator:
|
||||
class HeatMapAnnotator(BaseAnnotator):
|
||||
"""
|
||||
A class for drawing heatmaps on an image based on provided detections.
|
||||
Heat accumulates over time and is drawn as a semi-transparent overlay
|
||||
@@ -1352,7 +1572,7 @@ class HeatMapAnnotator:
|
||||
self.top_hue = top_hue
|
||||
self.low_hue = low_hue
|
||||
|
||||
@convert_for_annotation_method
|
||||
@ensure_cv2_image_for_annotation
|
||||
def annotate(self, scene: ImageType, detections: Detections) -> ImageType:
|
||||
"""
|
||||
Annotates the scene with a heatmap based on the provided detections.
|
||||
@@ -1428,7 +1648,7 @@ class PixelateAnnotator(BaseAnnotator):
|
||||
"""
|
||||
self.pixel_size: int = pixel_size
|
||||
|
||||
@convert_for_annotation_method
|
||||
@ensure_cv2_image_for_annotation
|
||||
def annotate(
|
||||
self,
|
||||
scene: ImageType,
|
||||
@@ -1499,6 +1719,7 @@ class TriangleAnnotator(BaseAnnotator):
|
||||
height: int = 10,
|
||||
position: Position = Position.TOP_CENTER,
|
||||
color_lookup: ColorLookup = ColorLookup.CLASS,
|
||||
outline_thickness: int = 0,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
@@ -1509,14 +1730,16 @@ class TriangleAnnotator(BaseAnnotator):
|
||||
position (Position): The anchor position for placing the triangle.
|
||||
color_lookup (ColorLookup): Strategy for mapping colors to annotations.
|
||||
Options are `INDEX`, `CLASS`, `TRACK`.
|
||||
outline_thickness (int): Thickness of the outline of the triangle.
|
||||
"""
|
||||
self.color: Union[Color, ColorPalette] = color
|
||||
self.base: int = base
|
||||
self.height: int = height
|
||||
self.position: Position = position
|
||||
self.color_lookup: ColorLookup = color_lookup
|
||||
self.outline_thickness: int = outline_thickness
|
||||
|
||||
@convert_for_annotation_method
|
||||
@ensure_cv2_image_for_annotation
|
||||
def annotate(
|
||||
self,
|
||||
scene: ImageType,
|
||||
@@ -1576,7 +1799,10 @@ class TriangleAnnotator(BaseAnnotator):
|
||||
)
|
||||
|
||||
cv2.fillPoly(scene, [vertices], color.as_bgr())
|
||||
|
||||
if self.outline_thickness:
|
||||
cv2.polylines(
|
||||
scene, [vertices], True, (0, 0, 0), thickness=self.outline_thickness
|
||||
)
|
||||
return scene
|
||||
|
||||
|
||||
@@ -1612,7 +1838,7 @@ class RoundBoxAnnotator(BaseAnnotator):
|
||||
raise ValueError("roundness attribute must be float between (0, 1.0]")
|
||||
self.roundness: float = roundness
|
||||
|
||||
@convert_for_annotation_method
|
||||
@ensure_cv2_image_for_annotation
|
||||
def annotate(
|
||||
self,
|
||||
scene: ImageType,
|
||||
@@ -1749,7 +1975,7 @@ class PercentageBarAnnotator(BaseAnnotator):
|
||||
if border_thickness is None:
|
||||
self.border_thickness = int(0.15 * self.height)
|
||||
|
||||
@convert_for_annotation_method
|
||||
@ensure_cv2_image_for_annotation
|
||||
def annotate(
|
||||
self,
|
||||
scene: ImageType,
|
||||
@@ -1922,7 +2148,7 @@ class CropAnnotator(BaseAnnotator):
|
||||
self.border_thickness: int = border_thickness
|
||||
self.border_color_lookup: ColorLookup = border_color_lookup
|
||||
|
||||
@convert_for_annotation_method
|
||||
@ensure_cv2_image_for_annotation
|
||||
def annotate(
|
||||
self,
|
||||
scene: ImageType,
|
||||
@@ -1974,9 +2200,7 @@ class CropAnnotator(BaseAnnotator):
|
||||
(x1, y1), (x2, y2) = self.calculate_crop_coordinates(
|
||||
anchor=anchor, crop_wh=crop_wh, position=self.position
|
||||
)
|
||||
scene = overlay_image(
|
||||
scene=scene, inserted_image=resized_crop, anchor=(x1, y1)
|
||||
)
|
||||
scene = overlay_image(image=scene, overlay=resized_crop, anchor=(x1, y1))
|
||||
color = resolve_color(
|
||||
color=self.border_color,
|
||||
detections=detections,
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
from enum import Enum
|
||||
from typing import Optional, Union
|
||||
from typing import Optional, Tuple, Union
|
||||
|
||||
import numpy as np
|
||||
|
||||
@@ -34,14 +34,14 @@ def resolve_color_idx(
|
||||
) -> int:
|
||||
if detection_idx >= len(detections):
|
||||
raise ValueError(
|
||||
f"Detection index {detection_idx}"
|
||||
f"Detection index {detection_idx} "
|
||||
f"is out of bounds for detections of length {len(detections)}"
|
||||
)
|
||||
|
||||
if isinstance(color_lookup, np.ndarray):
|
||||
if len(color_lookup) != len(detections):
|
||||
raise ValueError(
|
||||
f"Length of color lookup {len(color_lookup)}"
|
||||
f"Length of color lookup {len(color_lookup)} "
|
||||
f"does not match length of detections {len(detections)}"
|
||||
)
|
||||
return color_lookup[detection_idx]
|
||||
@@ -50,19 +50,72 @@ def resolve_color_idx(
|
||||
elif color_lookup == ColorLookup.CLASS:
|
||||
if detections.class_id is None:
|
||||
raise ValueError(
|
||||
"Could not resolve color by class because"
|
||||
"Could not resolve color by class because "
|
||||
"Detections do not have class_id"
|
||||
)
|
||||
return detections.class_id[detection_idx]
|
||||
elif color_lookup == ColorLookup.TRACK:
|
||||
if detections.tracker_id is None:
|
||||
raise ValueError(
|
||||
"Could not resolve color by track because"
|
||||
"Could not resolve color by track because "
|
||||
"Detections do not have tracker_id"
|
||||
)
|
||||
return detections.tracker_id[detection_idx]
|
||||
|
||||
|
||||
def resolve_text_background_xyxy(
|
||||
center_coordinates: Tuple[int, int],
|
||||
text_wh: Tuple[int, int],
|
||||
position: Position,
|
||||
) -> Tuple[int, int, int, int]:
|
||||
center_x, center_y = center_coordinates
|
||||
text_w, text_h = text_wh
|
||||
|
||||
if position == Position.TOP_LEFT:
|
||||
return center_x, center_y - text_h, center_x + text_w, center_y
|
||||
elif position == Position.TOP_RIGHT:
|
||||
return center_x - text_w, center_y - text_h, center_x, center_y
|
||||
elif position == Position.TOP_CENTER:
|
||||
return (
|
||||
center_x - text_w // 2,
|
||||
center_y - text_h,
|
||||
center_x + text_w // 2,
|
||||
center_y,
|
||||
)
|
||||
elif position == Position.CENTER or position == Position.CENTER_OF_MASS:
|
||||
return (
|
||||
center_x - text_w // 2,
|
||||
center_y - text_h // 2,
|
||||
center_x + text_w // 2,
|
||||
center_y + text_h // 2,
|
||||
)
|
||||
elif position == Position.BOTTOM_LEFT:
|
||||
return center_x, center_y, center_x + text_w, center_y + text_h
|
||||
elif position == Position.BOTTOM_RIGHT:
|
||||
return center_x - text_w, center_y, center_x, center_y + text_h
|
||||
elif position == Position.BOTTOM_CENTER:
|
||||
return (
|
||||
center_x - text_w // 2,
|
||||
center_y,
|
||||
center_x + text_w // 2,
|
||||
center_y + text_h,
|
||||
)
|
||||
elif position == Position.CENTER_LEFT:
|
||||
return (
|
||||
center_x - text_w,
|
||||
center_y - text_h // 2,
|
||||
center_x,
|
||||
center_y + text_h // 2,
|
||||
)
|
||||
elif position == Position.CENTER_RIGHT:
|
||||
return (
|
||||
center_x,
|
||||
center_y - text_h // 2,
|
||||
center_x + text_w,
|
||||
center_y + text_h // 2,
|
||||
)
|
||||
|
||||
|
||||
def get_color_by_index(color: Union[Color, ColorPalette], idx: int) -> Color:
|
||||
if isinstance(color, ColorPalette):
|
||||
return color.by_idx(idx)
|
||||
|
||||
+384
-149
@@ -3,8 +3,9 @@ from __future__ import annotations
|
||||
import os
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass
|
||||
from itertools import chain
|
||||
from pathlib import Path
|
||||
from typing import Dict, Iterator, List, Optional, Tuple
|
||||
from typing import Dict, Iterator, List, Optional, Tuple, Union
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
@@ -31,9 +32,10 @@ from supervision.dataset.utils import (
|
||||
train_test_split,
|
||||
)
|
||||
from supervision.detection.core import Detections
|
||||
from supervision.utils.internal import deprecated, warn_deprecated
|
||||
from supervision.utils.iterables import find_duplicates
|
||||
|
||||
|
||||
@dataclass
|
||||
class BaseDataset(ABC):
|
||||
@abstractmethod
|
||||
def __len__(self) -> int:
|
||||
@@ -46,30 +48,95 @@ class BaseDataset(ABC):
|
||||
pass
|
||||
|
||||
|
||||
@dataclass
|
||||
class DetectionDataset(BaseDataset):
|
||||
"""
|
||||
Dataclass containing information about object detection dataset.
|
||||
Contains information about a detection dataset. Handles lazy image loading
|
||||
and annotation retrieval, dataset splitting, conversions into multiple
|
||||
formats.
|
||||
|
||||
Attributes:
|
||||
classes (List[str]): List containing dataset class names.
|
||||
images (Dict[str, np.ndarray]): Dictionary mapping image name to image.
|
||||
images (Union[List[str], Dict[str, np.ndarray]]):
|
||||
Accepts a list of image paths, or dictionaries of loaded cv2 images
|
||||
with paths as keys. If you pass a list of paths, the dataset will
|
||||
lazily load images on demand, which is much more memory-efficient.
|
||||
annotations (Dict[str, Detections]): Dictionary mapping
|
||||
image name to annotations.
|
||||
image path to annotations. The dictionary keys match
|
||||
match the keys in `images` or entries in the list of
|
||||
image paths.
|
||||
"""
|
||||
|
||||
classes: List[str]
|
||||
images: Dict[str, np.ndarray]
|
||||
annotations: Dict[str, Detections]
|
||||
def __init__(
|
||||
self,
|
||||
classes: List[str],
|
||||
images: Union[List[str], Dict[str, np.ndarray]],
|
||||
annotations: Dict[str, Detections],
|
||||
) -> None:
|
||||
self.classes = classes
|
||||
|
||||
if set(images) != set(annotations):
|
||||
raise ValueError(
|
||||
"The keys of the images and annotations dictionaries must match."
|
||||
)
|
||||
self.annotations = annotations
|
||||
|
||||
# Eliminate duplicates while preserving order
|
||||
self.image_paths = list(dict.fromkeys(images))
|
||||
|
||||
self._images_in_memory: Dict[str, np.ndarray] = {}
|
||||
if isinstance(images, dict):
|
||||
self._images_in_memory = images
|
||||
warn_deprecated(
|
||||
"Passing a `Dict[str, np.ndarray]` into `DetectionDataset` is "
|
||||
"deprecated and will be removed in `supervision-0.26.0`. Use "
|
||||
"a list of paths `List[str]` instead."
|
||||
)
|
||||
# TODO: when supervision-0.26.0 is released, and Dict[str, np.ndarray]
|
||||
# for images is no longer supported, also simplify the rest of
|
||||
# the code. E.g. list(images) is no longer needed, and merge can
|
||||
# be simplified.
|
||||
|
||||
@property
|
||||
@deprecated(
|
||||
"`DetectionDataset.images` property is deprecated and will be removed in "
|
||||
"`supervision-0.26.0`. Iterate with `for path, image, annotation in dataset:` "
|
||||
"instead."
|
||||
)
|
||||
def images(self) -> Dict[str, np.ndarray]:
|
||||
"""
|
||||
Load all images to memory and return them as a dictionary.
|
||||
|
||||
!!! warning
|
||||
|
||||
Only use this when you need all images at once.
|
||||
It is much more memory-efficient to initialize dataset with
|
||||
image paths and use `for path, image, annotation in dataset:`.
|
||||
"""
|
||||
if self._images_in_memory:
|
||||
return self._images_in_memory
|
||||
|
||||
images = {image_path: cv2.imread(image_path) for image_path in self.image_paths}
|
||||
return images
|
||||
|
||||
def _get_image(self, image_path: str) -> np.ndarray:
|
||||
"""Assumes that image is in dataset"""
|
||||
if self._images_in_memory:
|
||||
return self._images_in_memory[image_path]
|
||||
return cv2.imread(image_path)
|
||||
|
||||
def __len__(self) -> int:
|
||||
"""
|
||||
Return the number of images in the dataset.
|
||||
return len(self._images_in_memory) or len(self.image_paths)
|
||||
|
||||
Returns:
|
||||
int: The number of images.
|
||||
def __getitem__(self, i: int) -> Tuple[str, np.ndarray, Detections]:
|
||||
"""
|
||||
return len(self.images)
|
||||
Returns:
|
||||
Tuple[str, np.ndarray, Detections]: The image path, image data,
|
||||
and its corresponding annotation at index i.
|
||||
"""
|
||||
image_path = self.image_paths[i]
|
||||
image = self._get_image(image_path)
|
||||
annotation = self.annotations[image_path]
|
||||
return image_path, image, annotation
|
||||
|
||||
def __iter__(self) -> Iterator[Tuple[str, np.ndarray, Detections]]:
|
||||
"""
|
||||
@@ -77,25 +144,33 @@ class DetectionDataset(BaseDataset):
|
||||
|
||||
Yields:
|
||||
Iterator[Tuple[str, np.ndarray, Detections]]:
|
||||
An iterator that yields tuples containing the image name,
|
||||
An iterator that yields tuples containing the image path,
|
||||
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)
|
||||
for i in range(len(self)):
|
||||
image_path, image, annotation = self[i]
|
||||
yield image_path, image, annotation
|
||||
|
||||
def __eq__(self, other):
|
||||
def __eq__(self, other) -> bool:
|
||||
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]:
|
||||
if self.image_paths != other.image_paths:
|
||||
return False
|
||||
|
||||
if self._images_in_memory or other._images_in_memory:
|
||||
if not np.array_equal(
|
||||
list(self._images_in_memory.values()),
|
||||
list(other._images_in_memory.values()),
|
||||
):
|
||||
return False
|
||||
|
||||
if self.annotations != other.annotations:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def split(
|
||||
@@ -116,38 +191,140 @@ class DetectionDataset(BaseDataset):
|
||||
Tuple[DetectionDataset, DetectionDataset]: A tuple containing
|
||||
the training and testing datasets.
|
||||
|
||||
Example:
|
||||
Examples:
|
||||
```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)
|
||||
```
|
||||
"""
|
||||
|
||||
image_names = list(self.images.keys())
|
||||
train_names, test_names = train_test_split(
|
||||
data=image_names,
|
||||
train_paths, test_paths = train_test_split(
|
||||
data=self.image_paths,
|
||||
train_ratio=split_ratio,
|
||||
random_state=random_state,
|
||||
shuffle=shuffle,
|
||||
)
|
||||
|
||||
train_input: Union[List[str], Dict[str, np.ndarray]]
|
||||
test_input: Union[List[str], Dict[str, np.ndarray]]
|
||||
if self._images_in_memory:
|
||||
train_input = {path: self._images_in_memory[path] for path in train_paths}
|
||||
test_input = {path: self._images_in_memory[path] for path in test_paths}
|
||||
else:
|
||||
train_input = train_paths
|
||||
test_input = test_paths
|
||||
train_annotations = {path: self.annotations[path] for path in train_paths}
|
||||
test_annotations = {path: self.annotations[path] for path in test_paths}
|
||||
|
||||
train_dataset = DetectionDataset(
|
||||
classes=self.classes,
|
||||
images={name: self.images[name] for name in train_names},
|
||||
annotations={name: self.annotations[name] for name in train_names},
|
||||
images=train_input,
|
||||
annotations=train_annotations,
|
||||
)
|
||||
test_dataset = DetectionDataset(
|
||||
classes=self.classes,
|
||||
images={name: self.images[name] for name in test_names},
|
||||
annotations={name: self.annotations[name] for name in test_names},
|
||||
images=test_input,
|
||||
annotations=test_annotations,
|
||||
)
|
||||
return train_dataset, test_dataset
|
||||
|
||||
@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.
|
||||
|
||||
Examples:
|
||||
```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']
|
||||
```
|
||||
"""
|
||||
|
||||
def is_in_memory(dataset: DetectionDataset) -> bool:
|
||||
return len(dataset._images_in_memory) > 0 or len(dataset.image_paths) == 0
|
||||
|
||||
def is_lazy(dataset: DetectionDataset) -> bool:
|
||||
return len(dataset._images_in_memory) == 0
|
||||
|
||||
all_in_memory = all([is_in_memory(dataset) for dataset in dataset_list])
|
||||
all_lazy = all([is_lazy(dataset) for dataset in dataset_list])
|
||||
if not all_in_memory and not all_lazy:
|
||||
raise ValueError(
|
||||
"Merging lazy and in-memory DetectionDatasets is not supported."
|
||||
)
|
||||
|
||||
images_in_memory = {}
|
||||
for dataset in dataset_list:
|
||||
images_in_memory.update(dataset._images_in_memory)
|
||||
|
||||
image_paths = list(
|
||||
chain.from_iterable(dataset.image_paths for dataset in dataset_list)
|
||||
)
|
||||
image_paths_unique = list(dict.fromkeys(image_paths))
|
||||
if len(image_paths) != len(image_paths_unique):
|
||||
duplicates = find_duplicates(image_paths)
|
||||
raise ValueError(
|
||||
f"Image paths {duplicates} are not unique across datasets."
|
||||
)
|
||||
image_paths = image_paths_unique
|
||||
|
||||
classes = merge_class_lists(
|
||||
class_lists=[dataset.classes for dataset in dataset_list]
|
||||
)
|
||||
|
||||
annotations = {}
|
||||
for dataset in dataset_list:
|
||||
annotations.update(dataset.annotations)
|
||||
for dataset in dataset_list:
|
||||
class_index_mapping = build_class_index_mapping(
|
||||
source_classes=dataset.classes, target_classes=classes
|
||||
)
|
||||
for image_path in dataset.image_paths:
|
||||
annotations[image_path] = map_detections_class_id(
|
||||
source_to_target_mapping=class_index_mapping,
|
||||
detections=annotations[image_path],
|
||||
)
|
||||
|
||||
return cls(
|
||||
classes=classes,
|
||||
images=images_in_memory or image_paths,
|
||||
annotations=annotations,
|
||||
)
|
||||
|
||||
def as_pascal_voc(
|
||||
self,
|
||||
images_directory_path: Optional[str] = None,
|
||||
@@ -181,25 +358,22 @@ class DetectionDataset(BaseDataset):
|
||||
"""
|
||||
if images_directory_path:
|
||||
save_dataset_images(
|
||||
images_directory_path=images_directory_path, images=self.images
|
||||
dataset=self,
|
||||
images_directory_path=images_directory_path,
|
||||
)
|
||||
if annotations_directory_path:
|
||||
Path(annotations_directory_path).mkdir(parents=True, exist_ok=True)
|
||||
|
||||
for image_path, image in self.images.items():
|
||||
detections = self.annotations[image_path]
|
||||
|
||||
if annotations_directory_path:
|
||||
for image_path, image, annotations in self:
|
||||
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,
|
||||
detections=annotations,
|
||||
classes=self.classes,
|
||||
filename=image_name,
|
||||
image_shape=image.shape,
|
||||
image_shape=image.shape, # type: ignore
|
||||
min_image_area_percentage=min_image_area_percentage,
|
||||
max_image_area_percentage=max_image_area_percentage,
|
||||
approximation_percentage=approximation_percentage,
|
||||
@@ -229,7 +403,7 @@ class DetectionDataset(BaseDataset):
|
||||
DetectionDataset: A DetectionDataset instance containing
|
||||
the loaded images and annotations.
|
||||
|
||||
Example:
|
||||
Examples:
|
||||
```python
|
||||
import roboflow
|
||||
from roboflow import Roboflow
|
||||
@@ -252,13 +426,15 @@ class DetectionDataset(BaseDataset):
|
||||
```
|
||||
"""
|
||||
|
||||
classes, images, annotations = load_pascal_voc_annotations(
|
||||
classes, image_paths, annotations = load_pascal_voc_annotations(
|
||||
images_directory_path=images_directory_path,
|
||||
annotations_directory_path=annotations_directory_path,
|
||||
force_masks=force_masks,
|
||||
)
|
||||
|
||||
return DetectionDataset(classes=classes, images=images, annotations=annotations)
|
||||
return DetectionDataset(
|
||||
classes=classes, images=image_paths, annotations=annotations
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_yolo(
|
||||
@@ -267,6 +443,7 @@ class DetectionDataset(BaseDataset):
|
||||
annotations_directory_path: str,
|
||||
data_yaml_path: str,
|
||||
force_masks: bool = False,
|
||||
is_obb: bool = False,
|
||||
) -> DetectionDataset:
|
||||
"""
|
||||
Creates a Dataset instance from YOLO formatted data.
|
||||
@@ -281,12 +458,15 @@ class DetectionDataset(BaseDataset):
|
||||
force_masks (bool, optional): If True, forces
|
||||
masks to be loaded for all annotations,
|
||||
regardless of whether they are present.
|
||||
is_obb (bool, optional): If True, loads the annotations in OBB format.
|
||||
OBB annotations are defined as `[class_id, x, y, x, y, x, y, x, y]`,
|
||||
where pairs of [x, y] are box corners.
|
||||
|
||||
Returns:
|
||||
DetectionDataset: A DetectionDataset instance
|
||||
containing the loaded images and annotations.
|
||||
|
||||
Example:
|
||||
Examples:
|
||||
```python
|
||||
import roboflow
|
||||
from roboflow import Roboflow
|
||||
@@ -308,13 +488,16 @@ class DetectionDataset(BaseDataset):
|
||||
# ['dog', 'person']
|
||||
```
|
||||
"""
|
||||
classes, images, annotations = load_yolo_annotations(
|
||||
classes, image_paths, annotations = load_yolo_annotations(
|
||||
images_directory_path=images_directory_path,
|
||||
annotations_directory_path=annotations_directory_path,
|
||||
data_yaml_path=data_yaml_path,
|
||||
force_masks=force_masks,
|
||||
is_obb=is_obb,
|
||||
)
|
||||
return DetectionDataset(
|
||||
classes=classes, images=image_paths, annotations=annotations
|
||||
)
|
||||
return DetectionDataset(classes=classes, images=images, annotations=annotations)
|
||||
|
||||
def as_yolo(
|
||||
self,
|
||||
@@ -355,13 +538,12 @@ class DetectionDataset(BaseDataset):
|
||||
"""
|
||||
if images_directory_path is not None:
|
||||
save_dataset_images(
|
||||
images_directory_path=images_directory_path, images=self.images
|
||||
dataset=self, images_directory_path=images_directory_path
|
||||
)
|
||||
if annotations_directory_path is not None:
|
||||
save_yolo_annotations(
|
||||
dataset=self,
|
||||
annotations_directory_path=annotations_directory_path,
|
||||
images=self.images,
|
||||
annotations=self.annotations,
|
||||
min_image_area_percentage=min_image_area_percentage,
|
||||
max_image_area_percentage=max_image_area_percentage,
|
||||
approximation_percentage=approximation_percentage,
|
||||
@@ -391,7 +573,7 @@ class DetectionDataset(BaseDataset):
|
||||
DetectionDataset: A DetectionDataset instance containing
|
||||
the loaded images and annotations.
|
||||
|
||||
Example:
|
||||
Examples:
|
||||
```python
|
||||
import roboflow
|
||||
from roboflow import Roboflow
|
||||
@@ -431,6 +613,20 @@ class DetectionDataset(BaseDataset):
|
||||
Exports the dataset to COCO format. This method saves the
|
||||
images and their corresponding annotations in COCO format.
|
||||
|
||||
!!! tip
|
||||
|
||||
The format of the mask is determined automatically based on its structure:
|
||||
|
||||
- If a mask contains multiple disconnected components or holes, it will be
|
||||
saved using the Run-Length Encoding (RLE) format for efficient storage and
|
||||
processing.
|
||||
- If a mask consists of a single, contiguous region without any holes, it
|
||||
will be encoded as a polygon, preserving the outline of the object.
|
||||
|
||||
This automatic selection ensures that the masks are stored in the most
|
||||
appropriate and space-efficient format, complying with COCO dataset
|
||||
standards.
|
||||
|
||||
Args:
|
||||
images_directory_path (Optional[str]): The path to the directory
|
||||
where the images should be saved.
|
||||
@@ -451,103 +647,134 @@ class DetectionDataset(BaseDataset):
|
||||
"""
|
||||
if images_directory_path is not None:
|
||||
save_dataset_images(
|
||||
images_directory_path=images_directory_path, images=self.images
|
||||
dataset=self, images_directory_path=images_directory_path
|
||||
)
|
||||
if annotations_path is not None:
|
||||
save_coco_annotations(
|
||||
dataset=self,
|
||||
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.
|
||||
Contains information about a classification dataset, handles lazy image
|
||||
loading, dataset splitting.
|
||||
|
||||
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
|
||||
images (Union[List[str], Dict[str, np.ndarray]]):
|
||||
List of image paths or dictionary mapping image name to image data.
|
||||
annotations (Dict[str, Classifications]): Dictionary mapping
|
||||
image name to annotations.
|
||||
"""
|
||||
|
||||
classes: List[str]
|
||||
images: Dict[str, np.ndarray]
|
||||
annotations: Dict[str, Classifications]
|
||||
def __init__(
|
||||
self,
|
||||
classes: List[str],
|
||||
images: Union[List[str], Dict[str, np.ndarray]],
|
||||
annotations: Dict[str, Classifications],
|
||||
) -> None:
|
||||
self.classes = classes
|
||||
|
||||
if set(images) != set(annotations):
|
||||
raise ValueError(
|
||||
"The keys of the images and annotations dictionaries must match."
|
||||
)
|
||||
self.annotations = annotations
|
||||
|
||||
# Eliminate duplicates while preserving order
|
||||
self.image_paths = list(dict.fromkeys(images))
|
||||
|
||||
self._images_in_memory: Dict[str, np.ndarray] = {}
|
||||
if isinstance(images, dict):
|
||||
self._images_in_memory = images
|
||||
warn_deprecated(
|
||||
"Passing a `Dict[str, np.ndarray]` into `ClassificationDataset` is "
|
||||
"deprecated and will be removed in a future release. Use "
|
||||
"a list of paths `List[str]` instead."
|
||||
)
|
||||
|
||||
@property
|
||||
@deprecated(
|
||||
"`DetectionDataset.images` property is deprecated and will be removed in "
|
||||
"`supervision-0.26.0`. Iterate with `for path, image, annotation in dataset:` "
|
||||
"instead."
|
||||
)
|
||||
def images(self) -> Dict[str, np.ndarray]:
|
||||
"""
|
||||
Load all images to memory and return them as a dictionary.
|
||||
|
||||
!!! warning
|
||||
|
||||
Only use this when you need all images at once.
|
||||
It is much more memory-efficient to initialize dataset with
|
||||
image paths and use `for path, image, annotation in dataset:`.
|
||||
"""
|
||||
if self._images_in_memory:
|
||||
return self._images_in_memory
|
||||
|
||||
images = {image_path: cv2.imread(image_path) for image_path in self.image_paths}
|
||||
return images
|
||||
|
||||
def _get_image(self, image_path: str) -> np.ndarray:
|
||||
"""Assumes that image is in dataset"""
|
||||
if self._images_in_memory:
|
||||
return self._images_in_memory[image_path]
|
||||
return cv2.imread(image_path)
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.images)
|
||||
return len(self._images_in_memory) or len(self.image_paths)
|
||||
|
||||
def __getitem__(self, i: int) -> Tuple[str, np.ndarray, Classifications]:
|
||||
"""
|
||||
Returns:
|
||||
Tuple[str, np.ndarray, Classifications]: The image path, image data,
|
||||
and its corresponding annotation at index i.
|
||||
"""
|
||||
image_path = self.image_paths[i]
|
||||
image = self._get_image(image_path)
|
||||
annotation = self.annotations[image_path]
|
||||
return image_path, image, annotation
|
||||
|
||||
def __iter__(self) -> Iterator[Tuple[str, np.ndarray, Classifications]]:
|
||||
"""
|
||||
Iterate over the images and annotations in the dataset.
|
||||
|
||||
Yields:
|
||||
Iterator[Tuple[str, np.ndarray, Detections]]:
|
||||
An iterator that yields tuples containing the image path,
|
||||
the image data, and its corresponding annotation.
|
||||
"""
|
||||
for i in range(len(self)):
|
||||
image_path, image, annotation = self[i]
|
||||
yield image_path, image, annotation
|
||||
|
||||
def __eq__(self, other) -> bool:
|
||||
if not isinstance(other, ClassificationDataset):
|
||||
return False
|
||||
|
||||
if set(self.classes) != set(other.classes):
|
||||
return False
|
||||
|
||||
if self.image_paths != other.image_paths:
|
||||
return False
|
||||
|
||||
if self._images_in_memory or other._images_in_memory:
|
||||
if not np.array_equal(
|
||||
list(self._images_in_memory.values()),
|
||||
list(other._images_in_memory.values()),
|
||||
):
|
||||
return False
|
||||
|
||||
if self.annotations != other.annotations:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def split(
|
||||
self, split_ratio=0.8, random_state=None, shuffle: bool = True
|
||||
@@ -567,35 +794,45 @@ class ClassificationDataset(BaseDataset):
|
||||
Tuple[ClassificationDataset, ClassificationDataset]: A tuple containing
|
||||
the training and testing datasets.
|
||||
|
||||
Example:
|
||||
Examples:
|
||||
```python
|
||||
import supervision as sv
|
||||
|
||||
cd = sv.ClassificationDataset(...)
|
||||
train_cd,test_cd = cd.split(split_ratio=0.7,
|
||||
random_state=42,shuffle=True)
|
||||
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_paths, test_paths = train_test_split(
|
||||
data=self.image_paths,
|
||||
train_ratio=split_ratio,
|
||||
random_state=random_state,
|
||||
shuffle=shuffle,
|
||||
)
|
||||
|
||||
train_input: Union[List[str], Dict[str, np.ndarray]]
|
||||
test_input: Union[List[str], Dict[str, np.ndarray]]
|
||||
if self._images_in_memory:
|
||||
train_input = {path: self._images_in_memory[path] for path in train_paths}
|
||||
test_input = {path: self._images_in_memory[path] for path in test_paths}
|
||||
else:
|
||||
train_input = train_paths
|
||||
test_input = test_paths
|
||||
train_annotations = {path: self.annotations[path] for path in train_paths}
|
||||
test_annotations = {path: self.annotations[path] for path in test_paths}
|
||||
|
||||
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},
|
||||
images=train_input,
|
||||
annotations=train_annotations,
|
||||
)
|
||||
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},
|
||||
images=test_input,
|
||||
annotations=test_annotations,
|
||||
)
|
||||
|
||||
return train_dataset, test_dataset
|
||||
|
||||
def as_folder_structure(self, root_directory_path: str) -> None:
|
||||
@@ -611,18 +848,16 @@ class ClassificationDataset(BaseDataset):
|
||||
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
|
||||
for image_save_path, image, annotation in self:
|
||||
image_name = Path(image_save_path).name
|
||||
class_id = (
|
||||
classification.class_id[0]
|
||||
if classification.confidence is None
|
||||
else classification.get_top_k(1)[0][0]
|
||||
annotation.class_id[0]
|
||||
if annotation.confidence is None
|
||||
else annotation.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)
|
||||
image_save_path = os.path.join(root_directory_path, class_name, image_name)
|
||||
cv2.imwrite(image_save_path, image)
|
||||
|
||||
@classmethod
|
||||
def from_folder_structure(cls, root_directory_path: str) -> ClassificationDataset:
|
||||
@@ -635,7 +870,7 @@ class ClassificationDataset(BaseDataset):
|
||||
Returns:
|
||||
ClassificationDataset: The dataset.
|
||||
|
||||
Example:
|
||||
Examples:
|
||||
```python
|
||||
import roboflow
|
||||
from roboflow import Roboflow
|
||||
@@ -655,7 +890,7 @@ class ClassificationDataset(BaseDataset):
|
||||
classes = os.listdir(root_directory_path)
|
||||
classes = sorted(set(classes))
|
||||
|
||||
images = {}
|
||||
image_paths = []
|
||||
annotations = {}
|
||||
|
||||
for class_name in classes:
|
||||
@@ -663,13 +898,13 @@ class ClassificationDataset(BaseDataset):
|
||||
|
||||
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)
|
||||
image_paths.append(image_path)
|
||||
annotations[image_path] = Classifications(
|
||||
class_id=np.array([class_id]),
|
||||
)
|
||||
|
||||
return cls(
|
||||
classes=classes,
|
||||
images=images,
|
||||
images=image_paths,
|
||||
annotations=annotations,
|
||||
)
|
||||
|
||||
@@ -1,19 +1,28 @@
|
||||
import os
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Tuple
|
||||
from typing import TYPE_CHECKING, Dict, List, Tuple
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import numpy.typing as npt
|
||||
|
||||
from supervision.dataset.utils import (
|
||||
approximate_mask_with_polygons,
|
||||
map_detections_class_id,
|
||||
mask_to_rle,
|
||||
rle_to_mask,
|
||||
)
|
||||
from supervision.detection.core import Detections
|
||||
from supervision.detection.utils import polygon_to_mask
|
||||
from supervision.detection.utils import (
|
||||
contains_holes,
|
||||
contains_multiple_segments,
|
||||
polygon_to_mask,
|
||||
)
|
||||
from supervision.utils.file import read_json_file, save_json_file
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from supervision.dataset.core import DetectionDataset
|
||||
|
||||
|
||||
def coco_categories_to_classes(coco_categories: List[dict]) -> List[str]:
|
||||
return [
|
||||
@@ -57,13 +66,24 @@ def group_coco_annotations_by_image_id(
|
||||
return annotations
|
||||
|
||||
|
||||
def _polygons_to_masks(
|
||||
polygons: List[np.ndarray], resolution_wh: Tuple[int, int]
|
||||
) -> np.ndarray:
|
||||
def coco_annotations_to_masks(
|
||||
image_annotations: List[dict], resolution_wh: Tuple[int, int]
|
||||
) -> npt.NDArray[np.bool_]:
|
||||
return np.array(
|
||||
[
|
||||
polygon_to_mask(polygon=polygon, resolution_wh=resolution_wh)
|
||||
for polygon in polygons
|
||||
rle_to_mask(
|
||||
rle=np.array(image_annotation["segmentation"]["counts"]),
|
||||
resolution_wh=resolution_wh,
|
||||
)
|
||||
if image_annotation["iscrowd"]
|
||||
else polygon_to_mask(
|
||||
polygon=np.reshape(
|
||||
np.asarray(image_annotation["segmentation"], dtype=np.int32),
|
||||
(-1, 2),
|
||||
),
|
||||
resolution_wh=resolution_wh,
|
||||
)
|
||||
for image_annotation in image_annotations
|
||||
],
|
||||
dtype=bool,
|
||||
)
|
||||
@@ -83,13 +103,9 @@ def coco_annotations_to_detections(
|
||||
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)
|
||||
mask = coco_annotations_to_masks(
|
||||
image_annotations=image_annotations, resolution_wh=resolution_wh
|
||||
)
|
||||
return Detections(
|
||||
class_id=np.asarray(class_ids, dtype=int), xyxy=xyxy, mask=mask
|
||||
)
|
||||
@@ -108,24 +124,35 @@ def detections_to_coco_annotations(
|
||||
coco_annotations = []
|
||||
for xyxy, mask, _, class_id, _, _ in detections:
|
||||
box_width, box_height = xyxy[2] - xyxy[0], xyxy[3] - xyxy[1]
|
||||
polygon = []
|
||||
segmentation = []
|
||||
iscrowd = 0
|
||||
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()
|
||||
)
|
||||
iscrowd = contains_holes(mask=mask) or contains_multiple_segments(mask=mask)
|
||||
|
||||
if iscrowd:
|
||||
segmentation = {
|
||||
"counts": mask_to_rle(mask=mask),
|
||||
"size": list(mask.shape[:2]),
|
||||
}
|
||||
else:
|
||||
segmentation = [
|
||||
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,
|
||||
"segmentation": segmentation,
|
||||
"iscrowd": iscrowd,
|
||||
}
|
||||
coco_annotations.append(coco_annotation)
|
||||
annotation_id += 1
|
||||
@@ -136,7 +163,7 @@ 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]]:
|
||||
) -> Tuple[List[str], List[str], 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(
|
||||
@@ -147,7 +174,7 @@ def load_coco_annotations(
|
||||
coco_annotations=coco_data["annotations"]
|
||||
)
|
||||
|
||||
images = {}
|
||||
images = []
|
||||
annotations = {}
|
||||
|
||||
for coco_image in coco_images:
|
||||
@@ -159,7 +186,6 @@ def load_coco_annotations(
|
||||
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),
|
||||
@@ -170,23 +196,20 @@ def load_coco_annotations(
|
||||
detections=annotation,
|
||||
)
|
||||
|
||||
images[image_path] = image
|
||||
images.append(image_path)
|
||||
annotations[image_path] = annotation
|
||||
|
||||
return classes, images, annotations
|
||||
|
||||
|
||||
def save_coco_annotations(
|
||||
dataset: "DetectionDataset",
|
||||
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,
|
||||
@@ -197,10 +220,10 @@ def save_coco_annotations(
|
||||
|
||||
coco_annotations = []
|
||||
coco_images = []
|
||||
coco_categories = classes_to_coco_categories(classes=classes)
|
||||
coco_categories = classes_to_coco_categories(classes=dataset.classes)
|
||||
|
||||
image_id, annotation_id = 1, 1
|
||||
for image_path, image in images.items():
|
||||
for image_path, image, annotation in dataset:
|
||||
image_height, image_width, _ = image.shape
|
||||
image_name = f"{Path(image_path).stem}{Path(image_path).suffix}"
|
||||
coco_image = {
|
||||
@@ -213,10 +236,8 @@ def save_coco_annotations(
|
||||
}
|
||||
|
||||
coco_images.append(coco_image)
|
||||
detections = annotations[image_path]
|
||||
|
||||
coco_annotation, annotation_id = detections_to_coco_annotations(
|
||||
detections=detections,
|
||||
detections=annotation,
|
||||
image_id=image_id,
|
||||
annotation_id=annotation_id,
|
||||
min_image_area_percentage=min_image_area_percentage,
|
||||
@@ -228,7 +249,7 @@ def save_coco_annotations(
|
||||
image_id += 1
|
||||
|
||||
annotation_dict = {
|
||||
"info": info,
|
||||
"info": {},
|
||||
"licenses": licenses,
|
||||
"categories": coco_categories,
|
||||
"images": coco_images,
|
||||
|
||||
@@ -5,7 +5,8 @@ from xml.etree.ElementTree import Element, SubElement
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
from defusedxml.ElementTree import fromstring, parse, tostring
|
||||
from defusedxml.ElementTree import parse, tostring
|
||||
from defusedxml.minidom import parseString
|
||||
|
||||
from supervision.dataset.utils import approximate_mask_with_polygons
|
||||
from supervision.detection.core import Detections
|
||||
@@ -129,8 +130,7 @@ def detections_to_pascal_voc(
|
||||
annotation.append(next_object)
|
||||
|
||||
# Generate XML string
|
||||
xml_string = fromstring(tostring(annotation)).toprettyxml(indent=" ")
|
||||
|
||||
xml_string = parseString(tostring(annotation)).toprettyxml(indent=" ")
|
||||
return xml_string
|
||||
|
||||
|
||||
@@ -138,7 +138,7 @@ def load_pascal_voc_annotations(
|
||||
images_directory_path: str,
|
||||
annotations_directory_path: str,
|
||||
force_masks: bool = False,
|
||||
) -> Tuple[List[str], Dict[str, np.ndarray], Dict[str, Detections]]:
|
||||
) -> Tuple[List[str], List[str], Dict[str, Detections]]:
|
||||
"""
|
||||
Loads PASCAL VOC XML annotations and returns the image name,
|
||||
a Detections instance, and a list of class names.
|
||||
@@ -151,44 +151,39 @@ def load_pascal_voc_annotations(
|
||||
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], List[str], Dict[str, Detections]]: A tuple with a list
|
||||
of class names, a list of paths to images, and a dictionary with image
|
||||
paths as keys and corresponding Detections instances as values.
|
||||
"""
|
||||
|
||||
image_paths = list_files_with_extensions(
|
||||
directory=images_directory_path, extensions=["jpg", "jpeg", "png"]
|
||||
)
|
||||
image_paths = [
|
||||
str(path)
|
||||
for path in list_files_with_extensions(
|
||||
directory=images_directory_path, extensions=["jpg", "jpeg", "png"]
|
||||
)
|
||||
]
|
||||
|
||||
classes = []
|
||||
images = {}
|
||||
classes: List[str] = []
|
||||
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")
|
||||
image_stem = Path(image_path).stem
|
||||
annotation_path = os.path.join(annotations_directory_path, f"{image_stem}.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()
|
||||
|
||||
image = cv2.imread(image_path)
|
||||
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
|
||||
return classes, image_paths, annotations
|
||||
|
||||
|
||||
def detections_from_xml_obj(
|
||||
|
||||
@@ -1,10 +1,11 @@
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Optional, Tuple
|
||||
from typing import TYPE_CHECKING, Dict, List, Optional, Tuple
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from supervision.config import ORIENTED_BOX_COORDINATES
|
||||
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
|
||||
@@ -16,6 +17,9 @@ from supervision.utils.file import (
|
||||
save_yaml_file,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from supervision.dataset.core import DetectionDataset
|
||||
|
||||
|
||||
def _parse_box(values: List[str]) -> np.ndarray:
|
||||
x_center, y_center, width, height = values
|
||||
@@ -70,12 +74,15 @@ def _image_name_to_annotation_name(image_name: str) -> str:
|
||||
|
||||
|
||||
def yolo_annotations_to_detections(
|
||||
lines: List[str], resolution_wh: Tuple[int, int], with_masks: bool
|
||||
lines: List[str],
|
||||
resolution_wh: Tuple[int, int],
|
||||
with_masks: bool,
|
||||
is_obb: bool = False,
|
||||
) -> Detections:
|
||||
if len(lines) == 0:
|
||||
return Detections.empty()
|
||||
|
||||
class_id, relative_xyxy, relative_polygon = [], [], []
|
||||
class_id, relative_xyxy, relative_polygon, relative_xyxyxyxy = [], [], [], []
|
||||
w, h = resolution_wh
|
||||
for line in lines:
|
||||
values = line.split()
|
||||
@@ -88,21 +95,30 @@ def yolo_annotations_to_detections(
|
||||
elif len(values) > 5:
|
||||
polygon = _parse_polygon(values=values[1:])
|
||||
relative_xyxy.append(polygon_to_xyxy(polygon=polygon))
|
||||
if is_obb:
|
||||
relative_xyxyxyxy.append(np.array(values[1:]))
|
||||
if with_masks:
|
||||
relative_polygon.append(polygon)
|
||||
|
||||
class_id = np.array(class_id, dtype=int)
|
||||
relative_xyxy = np.array(relative_xyxy, dtype=np.float32)
|
||||
xyxy = relative_xyxy * np.array([w, h, w, h], dtype=np.float32)
|
||||
data = {}
|
||||
|
||||
if is_obb:
|
||||
relative_xyxyxyxy = np.array(relative_xyxyxyxy, dtype=np.float32)
|
||||
xyxyxyxy = relative_xyxyxyxy.reshape(-1, 4, 2)
|
||||
xyxyxyxy *= np.array([w, h], dtype=np.float32)
|
||||
data[ORIENTED_BOX_COORDINATES] = xyxyxyxy
|
||||
|
||||
if not with_masks:
|
||||
return Detections(class_id=class_id, xyxy=xyxy)
|
||||
return Detections(class_id=class_id, xyxy=xyxy, data=data)
|
||||
|
||||
polygons = [
|
||||
(polygon * np.array(resolution_wh)).astype(int) for polygon in relative_polygon
|
||||
]
|
||||
mask = _polygons_to_masks(polygons=polygons, resolution_wh=resolution_wh)
|
||||
return Detections(class_id=class_id, xyxy=xyxy, mask=mask)
|
||||
return Detections(class_id=class_id, xyxy=xyxy, data=data, mask=mask)
|
||||
|
||||
|
||||
def load_yolo_annotations(
|
||||
@@ -110,7 +126,8 @@ def load_yolo_annotations(
|
||||
annotations_directory_path: str,
|
||||
data_yaml_path: str,
|
||||
force_masks: bool = False,
|
||||
) -> Tuple[List[str], Dict[str, np.ndarray], Dict[str, Detections]]:
|
||||
is_obb: bool = False,
|
||||
) -> Tuple[List[str], List[str], Dict[str, Detections]]:
|
||||
"""
|
||||
Loads YOLO annotations and returns class names, images,
|
||||
and their corresponding detections.
|
||||
@@ -123,32 +140,34 @@ def load_yolo_annotations(
|
||||
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.
|
||||
is_obb (bool, optional): If True, loads the annotations in OBB format.
|
||||
OBB annotations are defined as `[class_id, x, y, x, y, x, y, x, y]`,
|
||||
where pairs of [x, y] are box corners.
|
||||
|
||||
Returns:
|
||||
Tuple[List[str], Dict[str, np.ndarray], Dict[str, Detections]]:
|
||||
Tuple[List[str], List[str], 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"]
|
||||
)
|
||||
image_paths = [
|
||||
str(path)
|
||||
for path in list_files_with_extensions(
|
||||
directory=images_directory_path, extensions=["jpg", "jpeg", "png"]
|
||||
)
|
||||
]
|
||||
|
||||
classes = _extract_class_names(file_path=data_yaml_path)
|
||||
images = {}
|
||||
annotations = {}
|
||||
|
||||
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
|
||||
|
||||
image = cv2.imread(image_path)
|
||||
lines = read_txt_file(file_path=annotation_path, skip_empty=True)
|
||||
h, w, _ = image.shape
|
||||
resolution_wh = (w, h)
|
||||
@@ -156,12 +175,13 @@ def load_yolo_annotations(
|
||||
with_masks = _with_mask(lines=lines)
|
||||
with_masks = force_masks if force_masks else with_masks
|
||||
annotation = yolo_annotations_to_detections(
|
||||
lines=lines, resolution_wh=resolution_wh, with_masks=with_masks
|
||||
lines=lines,
|
||||
resolution_wh=resolution_wh,
|
||||
with_masks=with_masks,
|
||||
is_obb=is_obb,
|
||||
)
|
||||
|
||||
images[image_path] = image
|
||||
annotations[image_path] = annotation
|
||||
return classes, images, annotations
|
||||
return classes, image_paths, annotations
|
||||
|
||||
|
||||
def object_to_yolo(
|
||||
@@ -195,6 +215,9 @@ def detections_to_yolo_annotations(
|
||||
) -> List[str]:
|
||||
annotation = []
|
||||
for xyxy, mask, _, class_id, _, _ in detections:
|
||||
if class_id is None:
|
||||
raise ValueError("Class ID is required for YOLO annotations.")
|
||||
|
||||
if mask is not None:
|
||||
polygons = approximate_mask_with_polygons(
|
||||
mask=mask,
|
||||
@@ -220,24 +243,22 @@ def detections_to_yolo_annotations(
|
||||
|
||||
|
||||
def save_yolo_annotations(
|
||||
dataset: "DetectionDataset",
|
||||
annotations_directory_path: str,
|
||||
images: Dict[str, np.ndarray],
|
||||
annotations: Dict[str, Detections],
|
||||
min_image_area_percentage: float = 0.0,
|
||||
max_image_area_percentage: float = 1.0,
|
||||
approximation_percentage: float = 0.75,
|
||||
) -> None:
|
||||
Path(annotations_directory_path).mkdir(parents=True, exist_ok=True)
|
||||
for image_path, image in images.items():
|
||||
detections = annotations[image_path]
|
||||
for image_path, image, annotation in dataset:
|
||||
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
|
||||
)
|
||||
lines = detections_to_yolo_annotations(
|
||||
detections=detections,
|
||||
image_shape=image.shape,
|
||||
detections=annotation,
|
||||
image_shape=image.shape, # type: ignore
|
||||
min_image_area_percentage=min_image_area_percentage,
|
||||
max_image_area_percentage=max_image_area_percentage,
|
||||
approximation_percentage=approximation_percentage,
|
||||
|
||||
@@ -1,11 +1,13 @@
|
||||
import copy
|
||||
import os
|
||||
import random
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Optional, Tuple, TypeVar
|
||||
from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, TypeVar, Union
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import numpy.typing as npt
|
||||
|
||||
from supervision.detection.core import Detections
|
||||
from supervision.detection.utils import (
|
||||
@@ -14,6 +16,9 @@ from supervision.detection.utils import (
|
||||
mask_to_polygons,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from supervision.dataset.core import DetectionDataset
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
@@ -58,6 +63,7 @@ def merge_class_lists(class_lists: List[List[str]]) -> List[str]:
|
||||
def build_class_index_mapping(
|
||||
source_classes: List[str], target_classes: List[str]
|
||||
) -> Dict[int, int]:
|
||||
"""Returns the index map of source classes -> target classes."""
|
||||
index_mapping = {}
|
||||
|
||||
for i, class_name in enumerate(source_classes):
|
||||
@@ -93,14 +99,16 @@ def map_detections_class_id(
|
||||
|
||||
|
||||
def save_dataset_images(
|
||||
images_directory_path: str, images: Dict[str, np.ndarray]
|
||||
dataset: "DetectionDataset", images_directory_path: str
|
||||
) -> None:
|
||||
Path(images_directory_path).mkdir(parents=True, exist_ok=True)
|
||||
|
||||
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)
|
||||
for image_path in dataset.image_paths:
|
||||
final_path = os.path.join(images_directory_path, Path(image_path).name)
|
||||
if image_path in dataset._images_in_memory:
|
||||
image = dataset._images_in_memory[image_path]
|
||||
cv2.imwrite(final_path, image)
|
||||
else:
|
||||
shutil.copyfile(image_path, final_path)
|
||||
|
||||
|
||||
def train_test_split(
|
||||
@@ -129,3 +137,123 @@ def train_test_split(
|
||||
|
||||
split_index = int(len(data) * train_ratio)
|
||||
return data[:split_index], data[split_index:]
|
||||
|
||||
|
||||
def rle_to_mask(
|
||||
rle: Union[npt.NDArray[np.int_], List[int]], resolution_wh: Tuple[int, int]
|
||||
) -> npt.NDArray[np.bool_]:
|
||||
"""
|
||||
Converts run-length encoding (RLE) to a binary mask.
|
||||
|
||||
Args:
|
||||
rle (Union[npt.NDArray[np.int_], List[int]]): The 1D RLE array, the format
|
||||
used in the COCO dataset (column-wise encoding, values of an array with
|
||||
even indices represent the number of pixels assigned as background,
|
||||
values of an array with odd indices represent the number of pixels
|
||||
assigned as foreground object).
|
||||
resolution_wh (Tuple[int, int]): The width (w) and height (h)
|
||||
of the desired binary mask.
|
||||
|
||||
Returns:
|
||||
The generated 2D Boolean mask of shape `(h, w)`, where the foreground object is
|
||||
marked with `True`'s and the rest is filled with `False`'s.
|
||||
|
||||
Raises:
|
||||
AssertionError: If the sum of pixels encoded in RLE differs from the
|
||||
number of pixels in the expected mask (computed based on resolution_wh).
|
||||
|
||||
Examples:
|
||||
```python
|
||||
import supervision as sv
|
||||
|
||||
sv.rle_to_mask([5, 2, 2, 2, 5], (4, 4))
|
||||
# array([
|
||||
# [False, False, False, False],
|
||||
# [False, True, True, False],
|
||||
# [False, True, True, False],
|
||||
# [False, False, False, False],
|
||||
# ])
|
||||
```
|
||||
"""
|
||||
if isinstance(rle, list):
|
||||
rle = np.array(rle, dtype=int)
|
||||
|
||||
width, height = resolution_wh
|
||||
|
||||
assert width * height == np.sum(rle), (
|
||||
"the sum of the number of pixels in the RLE must be the same "
|
||||
"as the number of pixels in the expected mask"
|
||||
)
|
||||
|
||||
zero_one_values = np.zeros(shape=(rle.size, 1), dtype=np.uint8)
|
||||
zero_one_values[1::2] = 1
|
||||
|
||||
decoded_rle = np.repeat(zero_one_values, rle, axis=0)
|
||||
decoded_rle = np.append(
|
||||
decoded_rle, np.zeros(width * height - len(decoded_rle), dtype=np.uint8)
|
||||
)
|
||||
return decoded_rle.reshape((height, width), order="F")
|
||||
|
||||
|
||||
def mask_to_rle(mask: npt.NDArray[np.bool_]) -> List[int]:
|
||||
"""
|
||||
Converts a binary mask into a run-length encoding (RLE).
|
||||
|
||||
Args:
|
||||
mask (npt.NDArray[np.bool_]): 2D binary mask where `True` indicates foreground
|
||||
object and `False` indicates background.
|
||||
|
||||
Returns:
|
||||
The run-length encoded mask. Values of a list with even indices
|
||||
represent the number of pixels assigned as background (`False`), values
|
||||
of a list with odd indices represent the number of pixels assigned
|
||||
as foreground object (`True`).
|
||||
|
||||
Raises:
|
||||
AssertionError: If input mask is not 2D or is empty.
|
||||
|
||||
Examples:
|
||||
```python
|
||||
import numpy as np
|
||||
import supervision as sv
|
||||
|
||||
mask = np.array([
|
||||
[True, True, True, True],
|
||||
[True, True, True, True],
|
||||
[True, True, True, True],
|
||||
[True, True, True, True],
|
||||
])
|
||||
sv.mask_to_rle(mask)
|
||||
# [0, 16]
|
||||
|
||||
mask = np.array([
|
||||
[False, False, False, False],
|
||||
[False, True, True, False],
|
||||
[False, True, True, False],
|
||||
[False, False, False, False],
|
||||
])
|
||||
sv.mask_to_rle(mask)
|
||||
# [5, 2, 2, 2, 5]
|
||||
```
|
||||
|
||||
{ align=center width="800" }
|
||||
""" # noqa E501 // docs
|
||||
assert mask.ndim == 2, "Input mask must be 2D"
|
||||
assert mask.size != 0, "Input mask cannot be empty"
|
||||
|
||||
on_value_change_indices = np.where(
|
||||
mask.ravel(order="F") != np.roll(mask.ravel(order="F"), 1)
|
||||
)[0]
|
||||
|
||||
on_value_change_indices = np.append(on_value_change_indices, mask.size)
|
||||
# need to add 0 at the beginning when the same value is in the first and
|
||||
# last element of the flattened mask
|
||||
if on_value_change_indices[0] != 0:
|
||||
on_value_change_indices = np.insert(on_value_change_indices, 0, 0)
|
||||
|
||||
rle = np.diff(on_value_change_indices)
|
||||
|
||||
if mask[0][0] == 1:
|
||||
rle = np.insert(rle, 0, 0)
|
||||
|
||||
return list(rle)
|
||||
|
||||
@@ -1,155 +0,0 @@
|
||||
from typing import List, Optional, Union
|
||||
|
||||
import cv2
|
||||
|
||||
from supervision.annotators.base import ImageType
|
||||
from supervision.detection.core import Detections
|
||||
from supervision.draw.color import Color, ColorPalette
|
||||
from supervision.utils.conversion import convert_for_annotation_method
|
||||
from supervision.utils.internal import deprecated
|
||||
|
||||
|
||||
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
|
||||
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
|
||||
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
color: Union[Color, ColorPalette] = ColorPalette.DEFAULT,
|
||||
thickness: int = 2,
|
||||
text_color: Color = Color.BLACK,
|
||||
text_scale: float = 0.5,
|
||||
text_thickness: int = 1,
|
||||
text_padding: int = 10,
|
||||
):
|
||||
self.color: Union[Color, ColorPalette] = color
|
||||
self.thickness: int = thickness
|
||||
self.text_color: Color = text_color
|
||||
self.text_scale: float = text_scale
|
||||
self.text_thickness: int = text_thickness
|
||||
self.text_padding: int = text_padding
|
||||
|
||||
@deprecated(
|
||||
"`BoxAnnotator` is deprecated and will be removed in "
|
||||
"`supervision-0.22.0`. Use `BoundingBoxAnnotator` and `LabelAnnotator` instead"
|
||||
)
|
||||
@convert_for_annotation_method
|
||||
def annotate(
|
||||
self,
|
||||
scene: ImageType,
|
||||
detections: Detections,
|
||||
labels: Optional[List[str]] = None,
|
||||
skip_label: bool = False,
|
||||
) -> ImageType:
|
||||
"""
|
||||
Draws bounding boxes on the frame using the detections provided.
|
||||
|
||||
Args:
|
||||
scene (ImageType): The image on which the bounding boxes will be drawn.
|
||||
`ImageType` is a flexible type, accepting either `numpy.ndarray`
|
||||
or `PIL.Image.Image`.
|
||||
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:
|
||||
ImageType: The image with the bounding boxes drawn on it, matching the
|
||||
type of `scene` (`numpy.ndarray` or `PIL.Image.Image`)
|
||||
|
||||
Example:
|
||||
```python
|
||||
import supervision as sv
|
||||
|
||||
classes = ['person', ...]
|
||||
image = ...
|
||||
detections = sv.Detections(...)
|
||||
|
||||
box_annotator = sv.BoxAnnotator()
|
||||
labels = [
|
||||
f"{classes[class_id]} {confidence:0.2f}"
|
||||
for _, _, confidence, class_id, _ in detections
|
||||
]
|
||||
annotated_frame = box_annotator.annotate(
|
||||
scene=image.copy(),
|
||||
detections=detections,
|
||||
labels=labels
|
||||
)
|
||||
```
|
||||
"""
|
||||
font = cv2.FONT_HERSHEY_SIMPLEX
|
||||
for i in range(len(detections)):
|
||||
x1, y1, x2, y2 = detections.xyxy[i].astype(int)
|
||||
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
|
||||
)
|
||||
cv2.rectangle(
|
||||
img=scene,
|
||||
pt1=(x1, y1),
|
||||
pt2=(x2, y2),
|
||||
color=color.as_bgr(),
|
||||
thickness=self.thickness,
|
||||
)
|
||||
if skip_label:
|
||||
continue
|
||||
|
||||
text = (
|
||||
f"{class_id}"
|
||||
if (labels is None or len(detections) != len(labels))
|
||||
else labels[i]
|
||||
)
|
||||
|
||||
text_width, text_height = cv2.getTextSize(
|
||||
text=text,
|
||||
fontFace=font,
|
||||
fontScale=self.text_scale,
|
||||
thickness=self.text_thickness,
|
||||
)[0]
|
||||
|
||||
text_x = x1 + self.text_padding
|
||||
text_y = y1 - self.text_padding
|
||||
|
||||
text_background_x1 = x1
|
||||
text_background_y1 = y1 - 2 * self.text_padding - text_height
|
||||
|
||||
text_background_x2 = x1 + 2 * self.text_padding + text_width
|
||||
text_background_y2 = y1
|
||||
|
||||
cv2.rectangle(
|
||||
img=scene,
|
||||
pt1=(text_background_x1, text_background_y1),
|
||||
pt2=(text_background_x2, text_background_y2),
|
||||
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
|
||||
+296
-47
@@ -7,20 +7,30 @@ from typing import Any, Dict, Iterator, List, Optional, Tuple, Union
|
||||
import numpy as np
|
||||
|
||||
from supervision.config import CLASS_NAME_DATA_FIELD, ORIENTED_BOX_COORDINATES
|
||||
from supervision.detection.utils import (
|
||||
from supervision.detection.lmm import (
|
||||
LMM,
|
||||
from_florence_2,
|
||||
from_paligemma,
|
||||
validate_lmm_parameters,
|
||||
)
|
||||
from supervision.detection.overlap_filter import (
|
||||
box_non_max_merge,
|
||||
box_non_max_suppression,
|
||||
mask_non_max_suppression,
|
||||
)
|
||||
from supervision.detection.utils import (
|
||||
box_iou_batch,
|
||||
calculate_masks_centroids,
|
||||
extract_ultralytics_masks,
|
||||
get_data_item,
|
||||
is_data_equal,
|
||||
mask_non_max_suppression,
|
||||
mask_to_xyxy,
|
||||
merge_data,
|
||||
process_roboflow_result,
|
||||
xywh_to_xyxy,
|
||||
)
|
||||
from supervision.geometry.core import Position
|
||||
from supervision.utils.internal import deprecated
|
||||
from supervision.utils.internal import get_instance_variables
|
||||
from supervision.validators import validate_detections_fields
|
||||
|
||||
|
||||
@@ -240,7 +250,7 @@ class Detections:
|
||||
Class names values can be accessed using `detections["class_name"]`.
|
||||
""" # noqa: E501 // docs
|
||||
|
||||
if "obb" in ultralytics_results and ultralytics_results.obb is not None:
|
||||
if hasattr(ultralytics_results, "obb") and ultralytics_results.obb is not None:
|
||||
class_id = ultralytics_results.obb.cls.cpu().numpy().astype(int)
|
||||
class_names = np.array([ultralytics_results.names[i] for i in class_id])
|
||||
oriented_box_coordinates = ultralytics_results.obb.xyxyxyxy.cpu().numpy()
|
||||
@@ -418,6 +428,9 @@ class Detections:
|
||||
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),
|
||||
mask=mmdet_results.pred_instances.masks.cpu().numpy()
|
||||
if "masks" in mmdet_results.pred_instances
|
||||
else None,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
@@ -532,6 +545,9 @@ class Detections:
|
||||
return cls(
|
||||
xyxy=detectron2_results["instances"].pred_boxes.tensor.cpu().numpy(),
|
||||
confidence=detectron2_results["instances"].scores.cpu().numpy(),
|
||||
mask=detectron2_results["instances"].pred_masks.cpu().numpy()
|
||||
if hasattr(detectron2_results["instances"], "pred_masks")
|
||||
else None,
|
||||
class_id=detectron2_results["instances"]
|
||||
.pred_classes.cpu()
|
||||
.numpy()
|
||||
@@ -592,45 +608,6 @@ class Detections:
|
||||
data=data,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
@deprecated(
|
||||
"`Detections.from_roboflow` is deprecated and will be removed in "
|
||||
"`supervision-0.22.0`. Use `Detections.from_inference` instead."
|
||||
)
|
||||
def from_roboflow(cls, roboflow_result: Union[dict, Any]) -> Detections:
|
||||
"""
|
||||
!!! failure "Deprecated"
|
||||
|
||||
`Detections.from_roboflow` is deprecated and will be removed in
|
||||
`supervision-0.22.0`. Use `Detections.from_inference` instead.
|
||||
|
||||
Create a Detections object from the [Roboflow](https://roboflow.com/)
|
||||
API inference result or the [Inference](https://inference.roboflow.com/)
|
||||
package results.
|
||||
|
||||
Args:
|
||||
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.
|
||||
|
||||
Example:
|
||||
```python
|
||||
import cv2
|
||||
import supervision as sv
|
||||
from inference import get_model
|
||||
|
||||
image = cv2.imread(<SOURCE_IMAGE_PATH>)
|
||||
model = get_model(model_id="yolov8s-640")
|
||||
|
||||
result = model.infer(image)[0]
|
||||
detections = sv.Detections.from_roboflow(result)
|
||||
```
|
||||
"""
|
||||
return cls.from_inference(roboflow_result)
|
||||
|
||||
@classmethod
|
||||
def from_sam(cls, sam_result: List[dict]) -> Detections:
|
||||
"""
|
||||
@@ -670,7 +647,7 @@ class Detections:
|
||||
if np.asarray(xywh).shape[0] == 0:
|
||||
return cls.empty()
|
||||
|
||||
xyxy = xywh_to_xyxy(boxes_xywh=xywh)
|
||||
xyxy = xywh_to_xyxy(xywh=xywh)
|
||||
return cls(xyxy=xyxy, mask=mask)
|
||||
|
||||
@classmethod
|
||||
@@ -802,6 +779,69 @@ class Detections:
|
||||
class_id=paddledet_result["bbox"][:, 0].astype(int),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_lmm(
|
||||
cls, lmm: Union[LMM, str], result: Union[str, dict], **kwargs
|
||||
) -> Detections:
|
||||
"""
|
||||
Creates a Detections object from the given result string based on the specified
|
||||
Large Multimodal Model (LMM).
|
||||
|
||||
Args:
|
||||
lmm (Union[LMM, str]): The type of LMM (Large Multimodal Model) to use.
|
||||
result (str): The result string containing the detection data.
|
||||
**kwargs: Additional keyword arguments required by the specified LMM.
|
||||
|
||||
Returns:
|
||||
Detections: A new Detections object.
|
||||
|
||||
Raises:
|
||||
ValueError: If the LMM is invalid, required arguments are missing, or
|
||||
disallowed arguments are provided.
|
||||
ValueError: If the specified LMM is not supported.
|
||||
|
||||
Examples:
|
||||
```python
|
||||
import supervision as sv
|
||||
|
||||
paligemma_result = "<loc0256><loc0256><loc0768><loc0768> cat"
|
||||
detections = sv.Detections.from_lmm(
|
||||
sv.LMM.PALIGEMMA,
|
||||
paligemma_result,
|
||||
resolution_wh=(1000, 1000),
|
||||
classes=['cat', 'dog']
|
||||
)
|
||||
detections.xyxy
|
||||
# array([[250., 250., 750., 750.]])
|
||||
|
||||
detections.class_id
|
||||
# array([0])
|
||||
```
|
||||
"""
|
||||
lmm = validate_lmm_parameters(lmm, result, kwargs)
|
||||
|
||||
if lmm == LMM.PALIGEMMA:
|
||||
assert isinstance(result, str)
|
||||
xyxy, class_id, class_name = from_paligemma(result, **kwargs)
|
||||
data = {CLASS_NAME_DATA_FIELD: class_name}
|
||||
return cls(xyxy=xyxy, class_id=class_id, data=data)
|
||||
|
||||
if lmm == LMM.FLORENCE_2:
|
||||
assert isinstance(result, dict)
|
||||
xyxy, labels, mask, xyxyxyxy = from_florence_2(result, **kwargs)
|
||||
if len(xyxy) == 0:
|
||||
return cls.empty()
|
||||
|
||||
data = {}
|
||||
if labels is not None:
|
||||
data[CLASS_NAME_DATA_FIELD] = labels
|
||||
if xyxyxyxy is not None:
|
||||
data[ORIENTED_BOX_COORDINATES] = xyxyxyxy
|
||||
|
||||
return cls(xyxy=xyxy, mask=mask, data=data)
|
||||
|
||||
raise ValueError(f"Unsupported LMM: {lmm}")
|
||||
|
||||
@classmethod
|
||||
def empty(cls) -> Detections:
|
||||
"""
|
||||
@@ -824,6 +864,14 @@ class Detections:
|
||||
class_id=np.array([], dtype=int),
|
||||
)
|
||||
|
||||
def is_empty(self) -> bool:
|
||||
"""
|
||||
Returns `True` if the `Detections` object is considered empty.
|
||||
"""
|
||||
empty_detections = Detections.empty()
|
||||
empty_detections.data = self.data
|
||||
return self == empty_detections
|
||||
|
||||
@classmethod
|
||||
def merge(cls, detections_list: List[Detections]) -> Detections:
|
||||
"""
|
||||
@@ -831,9 +879,14 @@ class Detections:
|
||||
|
||||
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`.
|
||||
into a single Detections object.
|
||||
|
||||
For example, if merging Detections with 3 and 4 detected objects, this method
|
||||
will return a Detections with 7 objects (7 entries in `xyxy`, `mask`, etc).
|
||||
|
||||
!!! Note
|
||||
|
||||
When merging, empty `Detections` objects are ignored.
|
||||
|
||||
Args:
|
||||
detections_list (List[Detections]): A list of Detections objects to merge.
|
||||
@@ -873,6 +926,10 @@ class Detections:
|
||||
array([0.1, 0.2, 0.3])
|
||||
```
|
||||
"""
|
||||
detections_list = [
|
||||
detections for detections in detections_list if not detections.is_empty()
|
||||
]
|
||||
|
||||
if len(detections_list) == 0:
|
||||
return Detections.empty()
|
||||
|
||||
@@ -1147,3 +1204,195 @@ class Detections:
|
||||
)
|
||||
|
||||
return self[indices]
|
||||
|
||||
def with_nmm(
|
||||
self, threshold: float = 0.5, class_agnostic: bool = False
|
||||
) -> Detections:
|
||||
"""
|
||||
Perform non-maximum merging on the current set of object detections.
|
||||
|
||||
Args:
|
||||
threshold (float, optional): The intersection-over-union threshold
|
||||
to use for non-maximum merging. Defaults to 0.5.
|
||||
class_agnostic (bool, optional): Whether to perform class-agnostic
|
||||
non-maximum merging. 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 merging.
|
||||
|
||||
Raises:
|
||||
AssertionError: If `confidence` is None or `class_id` is None and
|
||||
class_agnostic is False.
|
||||
|
||||
{ align=center width="800" }
|
||||
""" # noqa: E501 // docs
|
||||
if len(self) == 0:
|
||||
return self
|
||||
|
||||
assert (
|
||||
self.confidence is not None
|
||||
), "Detections confidence must be given for NMM to be executed."
|
||||
|
||||
if class_agnostic:
|
||||
predictions = np.hstack((self.xyxy, self.confidence.reshape(-1, 1)))
|
||||
else:
|
||||
assert self.class_id is not None, (
|
||||
"Detections class_id must be given for NMM to be executed. If you"
|
||||
" intended to perform class agnostic NMM set class_agnostic=True."
|
||||
)
|
||||
predictions = np.hstack(
|
||||
(
|
||||
self.xyxy,
|
||||
self.confidence.reshape(-1, 1),
|
||||
self.class_id.reshape(-1, 1),
|
||||
)
|
||||
)
|
||||
|
||||
merge_groups = box_non_max_merge(
|
||||
predictions=predictions, iou_threshold=threshold
|
||||
)
|
||||
|
||||
result = []
|
||||
for merge_group in merge_groups:
|
||||
unmerged_detections = [self[i] for i in merge_group]
|
||||
merged_detections = merge_inner_detections_objects(
|
||||
unmerged_detections, threshold
|
||||
)
|
||||
result.append(merged_detections)
|
||||
|
||||
return Detections.merge(result)
|
||||
|
||||
|
||||
def merge_inner_detection_object_pair(
|
||||
detections_1: Detections, detections_2: Detections
|
||||
) -> Detections:
|
||||
"""
|
||||
Merges two Detections object into a single Detections object.
|
||||
Assumes each Detections contains exactly one object.
|
||||
|
||||
A `winning` detection is determined based on the confidence score of the two
|
||||
input detections. This winning detection is then used to specify which
|
||||
`class_id`, `tracker_id`, and `data` to include in the merged Detections object.
|
||||
|
||||
The resulting `confidence` of the merged object is calculated by the weighted
|
||||
contribution of ea detection to the merged object.
|
||||
The bounding boxes and masks of the two input detections are merged into a
|
||||
single bounding box and mask, respectively.
|
||||
|
||||
Args:
|
||||
detections_1 (Detections):
|
||||
The first Detections object
|
||||
detections_2 (Detections):
|
||||
The second Detections object
|
||||
|
||||
Returns:
|
||||
Detections: A new Detections object, with merged attributes.
|
||||
|
||||
Raises:
|
||||
ValueError: If the input Detections objects do not have exactly 1 detected
|
||||
object.
|
||||
|
||||
Example:
|
||||
```python
|
||||
import cv2
|
||||
import supervision as sv
|
||||
from inference import get_model
|
||||
|
||||
image = cv2.imread(<SOURCE_IMAGE_PATH>)
|
||||
model = get_model(model_id="yolov8s-640")
|
||||
|
||||
result = model.infer(image)[0]
|
||||
detections = sv.Detections.from_inference(result)
|
||||
|
||||
merged_detections = merge_object_detection_pair(
|
||||
detections[0], detections[1])
|
||||
```
|
||||
"""
|
||||
if len(detections_1) != 1 or len(detections_2) != 1:
|
||||
raise ValueError("Both Detections should have exactly 1 detected object.")
|
||||
|
||||
validate_fields_both_defined_or_none(detections_1, detections_2)
|
||||
|
||||
xyxy_1 = detections_1.xyxy[0]
|
||||
xyxy_2 = detections_2.xyxy[0]
|
||||
if detections_1.confidence is None and detections_2.confidence is None:
|
||||
merged_confidence = None
|
||||
else:
|
||||
detection_1_area = (xyxy_1[2] - xyxy_1[0]) * (xyxy_1[3] - xyxy_1[1])
|
||||
detections_2_area = (xyxy_2[2] - xyxy_2[0]) * (xyxy_2[3] - xyxy_2[1])
|
||||
merged_confidence = (
|
||||
detection_1_area * detections_1.confidence[0]
|
||||
+ detections_2_area * detections_2.confidence[0]
|
||||
) / (detection_1_area + detections_2_area)
|
||||
merged_confidence = np.array([merged_confidence])
|
||||
|
||||
merged_x1, merged_y1 = np.minimum(xyxy_1[:2], xyxy_2[:2])
|
||||
merged_x2, merged_y2 = np.maximum(xyxy_1[2:], xyxy_2[2:])
|
||||
merged_xyxy = np.array([[merged_x1, merged_y1, merged_x2, merged_y2]])
|
||||
|
||||
if detections_1.mask is None and detections_2.mask is None:
|
||||
merged_mask = None
|
||||
else:
|
||||
merged_mask = np.logical_or(detections_1.mask, detections_2.mask)
|
||||
|
||||
if detections_1.confidence is None and detections_2.confidence is None:
|
||||
winning_detection = detections_1
|
||||
elif detections_1.confidence[0] >= detections_2.confidence[0]:
|
||||
winning_detection = detections_1
|
||||
else:
|
||||
winning_detection = detections_2
|
||||
|
||||
return Detections(
|
||||
xyxy=merged_xyxy,
|
||||
mask=merged_mask,
|
||||
confidence=merged_confidence,
|
||||
class_id=winning_detection.class_id,
|
||||
tracker_id=winning_detection.tracker_id,
|
||||
data=winning_detection.data,
|
||||
)
|
||||
|
||||
|
||||
def merge_inner_detections_objects(
|
||||
detections: List[Detections], threshold=0.5
|
||||
) -> Detections:
|
||||
"""
|
||||
Given N detections each of length 1 (exactly one object inside), combine them into a
|
||||
single detection object of length 1. The contained inner object will be the merged
|
||||
result of all the input detections.
|
||||
|
||||
For example, this lets you merge N boxes into one big box, N masks into one mask,
|
||||
etc.
|
||||
"""
|
||||
detections_1 = detections[0]
|
||||
for detections_2 in detections[1:]:
|
||||
box_iou = box_iou_batch(detections_1.xyxy, detections_2.xyxy)[0]
|
||||
if box_iou < threshold:
|
||||
break
|
||||
detections_1 = merge_inner_detection_object_pair(detections_1, detections_2)
|
||||
return detections_1
|
||||
|
||||
|
||||
def validate_fields_both_defined_or_none(
|
||||
detections_1: Detections, detections_2: Detections
|
||||
) -> None:
|
||||
"""
|
||||
Verify that for each optional field in the Detections, both instances either have
|
||||
the field set to None or both have it set to non-None values.
|
||||
|
||||
`data` field is ignored.
|
||||
|
||||
Raises:
|
||||
ValueError: If one field is None and the other is not, for any of the fields.
|
||||
"""
|
||||
attributes = get_instance_variables(detections_1)
|
||||
for attribute in attributes:
|
||||
value_1 = getattr(detections_1, attribute)
|
||||
value_2 = getattr(detections_2, attribute)
|
||||
|
||||
if (value_1 is None) != (value_2 is None):
|
||||
raise ValueError(
|
||||
f"Field '{attribute}' should be consistently None or not None in both "
|
||||
"Detections."
|
||||
)
|
||||
|
||||
@@ -1,12 +1,15 @@
|
||||
import warnings
|
||||
from typing import Dict, Iterable, Optional, Tuple
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from supervision.detection.core import Detections
|
||||
from supervision.detection.utils import cross_product
|
||||
from supervision.draw.color import Color
|
||||
from supervision.draw.utils import draw_text
|
||||
from supervision.geometry.core import Point, Position, Vector
|
||||
from supervision.utils.internal import SupervisionWarnings
|
||||
|
||||
|
||||
class LineZone:
|
||||
@@ -81,6 +84,8 @@ class LineZone:
|
||||
self.in_count: int = 0
|
||||
self.out_count: int = 0
|
||||
self.triggering_anchors = triggering_anchors
|
||||
if not list(self.triggering_anchors):
|
||||
raise ValueError("Triggering anchors cannot be empty.")
|
||||
|
||||
@staticmethod
|
||||
def calculate_region_of_interest_limits(vector: Vector) -> Tuple[Vector, Vector]:
|
||||
@@ -140,6 +145,15 @@ class LineZone:
|
||||
if len(detections) == 0:
|
||||
return crossed_in, crossed_out
|
||||
|
||||
if detections.tracker_id is None:
|
||||
warnings.warn(
|
||||
"Line zone counting skipped. LineZone requires tracker_id. Refer to "
|
||||
"https://supervision.roboflow.com/latest/trackers for more "
|
||||
"information.",
|
||||
category=SupervisionWarnings,
|
||||
)
|
||||
return crossed_in, crossed_out
|
||||
|
||||
all_anchors = np.array(
|
||||
[
|
||||
detections.get_anchors_coordinates(anchor)
|
||||
@@ -147,31 +161,23 @@ class LineZone:
|
||||
]
|
||||
)
|
||||
|
||||
cross_products_1 = cross_product(all_anchors, self.limits[0])
|
||||
cross_products_2 = cross_product(all_anchors, self.limits[1])
|
||||
in_limits = (cross_products_1 > 0) == (cross_products_2 > 0)
|
||||
in_limits = np.all(in_limits, axis=0)
|
||||
|
||||
triggers = cross_product(all_anchors, self.vector) < 0
|
||||
has_any_left_trigger = np.any(triggers, axis=0)
|
||||
has_any_right_trigger = np.any(~triggers, axis=0)
|
||||
is_uniformly_triggered = ~(has_any_left_trigger & has_any_right_trigger)
|
||||
for i, tracker_id in enumerate(detections.tracker_id):
|
||||
if tracker_id is None:
|
||||
if not in_limits[i]:
|
||||
continue
|
||||
|
||||
box_anchors = [Point(x=x, y=y) for x, y in all_anchors[:, i, :]]
|
||||
|
||||
in_limits = all(
|
||||
[
|
||||
self.is_point_in_limits(point=anchor, limits=self.limits)
|
||||
for anchor in box_anchors
|
||||
]
|
||||
)
|
||||
|
||||
if not in_limits:
|
||||
if not is_uniformly_triggered[i]:
|
||||
continue
|
||||
|
||||
triggers = [
|
||||
self.vector.cross_product(point=anchor) < 0 for anchor in box_anchors
|
||||
]
|
||||
|
||||
if len(set(triggers)) == 2:
|
||||
continue
|
||||
|
||||
tracker_state = triggers[0]
|
||||
|
||||
tracker_state = has_any_left_trigger[i]
|
||||
if tracker_id not in self.tracker_state:
|
||||
self.tracker_state[tracker_id] = tracker_state
|
||||
continue
|
||||
|
||||
@@ -0,0 +1,184 @@
|
||||
import re
|
||||
from enum import Enum
|
||||
from typing import Any, Dict, List, Optional, Tuple, Union
|
||||
|
||||
import numpy as np
|
||||
|
||||
from supervision.detection.utils import polygon_to_mask, polygon_to_xyxy
|
||||
|
||||
|
||||
class LMM(Enum):
|
||||
PALIGEMMA = "paligemma"
|
||||
FLORENCE_2 = "florence_2"
|
||||
|
||||
|
||||
RESULT_TYPES: Dict[LMM, type] = {LMM.PALIGEMMA: str, LMM.FLORENCE_2: dict}
|
||||
|
||||
REQUIRED_ARGUMENTS: Dict[LMM, List[str]] = {
|
||||
LMM.PALIGEMMA: ["resolution_wh"],
|
||||
LMM.FLORENCE_2: ["resolution_wh"],
|
||||
}
|
||||
|
||||
ALLOWED_ARGUMENTS: Dict[LMM, List[str]] = {
|
||||
LMM.PALIGEMMA: ["resolution_wh", "classes"],
|
||||
LMM.FLORENCE_2: ["resolution_wh"],
|
||||
}
|
||||
|
||||
SUPPORTED_TASKS_FLORENCE_2 = [
|
||||
"<OD>",
|
||||
"<CAPTION_TO_PHRASE_GROUNDING>",
|
||||
"<DENSE_REGION_CAPTION>",
|
||||
"<REGION_PROPOSAL>",
|
||||
"<OCR_WITH_REGION>",
|
||||
"<REFERRING_EXPRESSION_SEGMENTATION>",
|
||||
"<REGION_TO_SEGMENTATION>",
|
||||
"<OPEN_VOCABULARY_DETECTION>",
|
||||
"<REGION_TO_CATEGORY>",
|
||||
"<REGION_TO_DESCRIPTION>",
|
||||
]
|
||||
|
||||
|
||||
def validate_lmm_parameters(
|
||||
lmm: Union[LMM, str], result: Any, kwargs: Dict[str, Any]
|
||||
) -> LMM:
|
||||
if isinstance(lmm, str):
|
||||
try:
|
||||
lmm = LMM(lmm.lower())
|
||||
except ValueError:
|
||||
raise ValueError(
|
||||
f"Invalid lmm value: {lmm}. Must be one of {[e.value for e in LMM]}"
|
||||
)
|
||||
|
||||
if not isinstance(result, RESULT_TYPES[lmm]):
|
||||
raise ValueError(
|
||||
f"Invalid LMM result type: {type(result)}. Must be {RESULT_TYPES[lmm]}"
|
||||
)
|
||||
|
||||
required_args = REQUIRED_ARGUMENTS.get(lmm, [])
|
||||
for arg in required_args:
|
||||
if arg not in kwargs:
|
||||
raise ValueError(f"Missing required argument: {arg}")
|
||||
|
||||
allowed_args = ALLOWED_ARGUMENTS.get(lmm, [])
|
||||
for arg in kwargs:
|
||||
if arg not in allowed_args:
|
||||
raise ValueError(f"Argument {arg} is not allowed for {lmm.name}")
|
||||
|
||||
return lmm
|
||||
|
||||
|
||||
def from_paligemma(
|
||||
result: str, resolution_wh: Tuple[int, int], classes: Optional[List[str]] = None
|
||||
) -> Tuple[np.ndarray, Optional[np.ndarray], np.ndarray]:
|
||||
w, h = resolution_wh
|
||||
pattern = re.compile(
|
||||
r"(?<!<loc\d{4}>)<loc(\d{4})><loc(\d{4})><loc(\d{4})><loc(\d{4})> ([\w\s\-]+)"
|
||||
)
|
||||
matches = pattern.findall(result)
|
||||
matches = np.array(matches) if matches else np.empty((0, 5))
|
||||
|
||||
xyxy, class_name = matches[:, [1, 0, 3, 2]], matches[:, 4]
|
||||
xyxy = xyxy.astype(int) / 1024 * np.array([w, h, w, h])
|
||||
class_name = np.char.strip(class_name.astype(str))
|
||||
class_id = None
|
||||
|
||||
if classes is not None:
|
||||
mask = np.array([name in classes for name in class_name]).astype(bool)
|
||||
xyxy, class_name = xyxy[mask], class_name[mask]
|
||||
class_id = np.array([classes.index(name) for name in class_name])
|
||||
|
||||
return xyxy, class_id, class_name
|
||||
|
||||
|
||||
def from_florence_2(
|
||||
result: dict, resolution_wh: Tuple[int, int]
|
||||
) -> Tuple[
|
||||
np.ndarray, Optional[np.ndarray], Optional[np.ndarray], Optional[np.ndarray]
|
||||
]:
|
||||
"""
|
||||
Parse results from the Florence 2 multi-model model.
|
||||
https://huggingface.co/microsoft/Florence-2-large
|
||||
|
||||
Parameters:
|
||||
result: dict containing the model output
|
||||
|
||||
Returns:
|
||||
xyxy (np.ndarray): An array of shape `(n, 4)` containing
|
||||
the bounding boxes coordinates in format `[x1, y1, x2, y2]`
|
||||
labels: (Optional[np.ndarray]): An array of shape `(n,)` containing
|
||||
the class labels for each bounding box
|
||||
masks: (Optional[np.ndarray]): An array of shape `(n, h, w)` containing
|
||||
the segmentation masks for each bounding box
|
||||
obb_boxes: (Optional[np.ndarray]): An array of shape `(n, 4, 2)` containing
|
||||
oriented bounding boxes.
|
||||
"""
|
||||
assert len(result) == 1, f"Expected result with a single element. Got: {result}"
|
||||
task = list(result.keys())[0]
|
||||
if task not in SUPPORTED_TASKS_FLORENCE_2:
|
||||
raise ValueError(
|
||||
f"{task} not supported. Supported tasks are: {SUPPORTED_TASKS_FLORENCE_2}"
|
||||
)
|
||||
result = result[task]
|
||||
|
||||
if task in ["<OD>", "<CAPTION_TO_PHRASE_GROUNDING>", "<DENSE_REGION_CAPTION>"]:
|
||||
xyxy = np.array(result["bboxes"], dtype=np.float32)
|
||||
labels = np.array(result["labels"])
|
||||
return xyxy, labels, None, None
|
||||
|
||||
if task == "<REGION_PROPOSAL>":
|
||||
xyxy = np.array(result["bboxes"], dtype=np.float32)
|
||||
# provides labels, but they are ["", "", "", ...]
|
||||
return xyxy, None, None, None
|
||||
|
||||
if task == "<OCR_WITH_REGION>":
|
||||
xyxyxyxy = np.array(result["quad_boxes"], dtype=np.float32)
|
||||
xyxyxyxy = xyxyxyxy.reshape(-1, 4, 2)
|
||||
xyxy = np.array([polygon_to_xyxy(polygon) for polygon in xyxyxyxy])
|
||||
labels = np.array(result["labels"])
|
||||
return xyxy, labels, None, xyxyxyxy
|
||||
|
||||
if task in ["<REFERRING_EXPRESSION_SEGMENTATION>", "<REGION_TO_SEGMENTATION>"]:
|
||||
xyxy_list = []
|
||||
masks_list = []
|
||||
for polygons_of_same_class in result["polygons"]:
|
||||
for polygon in polygons_of_same_class:
|
||||
polygon = np.reshape(polygon, (-1, 2)).astype(np.int32)
|
||||
mask = polygon_to_mask(polygon, resolution_wh).astype(bool)
|
||||
masks_list.append(mask)
|
||||
xyxy = polygon_to_xyxy(polygon)
|
||||
xyxy_list.append(xyxy)
|
||||
# per-class labels also provided, but they are ["", "", "", ...]
|
||||
# when we figure out how to set class names, we can do
|
||||
# zip(result["labels"], result["polygons"])
|
||||
xyxy = np.array(xyxy_list, dtype=np.float32)
|
||||
masks = np.array(masks_list)
|
||||
return xyxy, None, masks, None
|
||||
|
||||
if task == "<OPEN_VOCABULARY_DETECTION>":
|
||||
xyxy = np.array(result["bboxes"], dtype=np.float32)
|
||||
labels = np.array(result["bboxes_labels"])
|
||||
# Also has "polygons" and "polygons_labels", but they don't seem to be used
|
||||
return xyxy, labels, None, None
|
||||
|
||||
if task in ["<REGION_TO_CATEGORY>", "<REGION_TO_DESCRIPTION>"]:
|
||||
assert isinstance(
|
||||
result, str
|
||||
), f"Expected string as <REGION_TO_CATEGORY> result, got {type(result)}"
|
||||
|
||||
if result == "No object detected.":
|
||||
return np.empty((0, 4), dtype=np.float32), np.array([]), None, None
|
||||
|
||||
pattern = re.compile(r"<loc_(\d+)><loc_(\d+)><loc_(\d+)><loc_(\d+)>")
|
||||
match = pattern.search(result)
|
||||
assert (
|
||||
match is not None
|
||||
), f"Expected string to end in location tags, but got {result}"
|
||||
|
||||
w, h = resolution_wh
|
||||
xyxy = np.array([match.groups()], dtype=np.float32)
|
||||
xyxy *= np.array([w, h, w, h]) / 1000
|
||||
result_string = result[: match.start()]
|
||||
labels = np.array([result_string])
|
||||
return xyxy, labels, None, None
|
||||
|
||||
assert False, f"Unimplemented task: {task}"
|
||||
@@ -0,0 +1,263 @@
|
||||
from enum import Enum
|
||||
from typing import List, Union
|
||||
|
||||
import numpy as np
|
||||
import numpy.typing as npt
|
||||
|
||||
from supervision.detection.utils import box_iou_batch, mask_iou_batch
|
||||
|
||||
|
||||
def resize_masks(masks: np.ndarray, max_dimension: int = 640) -> np.ndarray:
|
||||
"""
|
||||
Resize all masks in the array to have a maximum dimension of max_dimension,
|
||||
maintaining aspect ratio.
|
||||
|
||||
Args:
|
||||
masks (np.ndarray): 3D array of binary masks with shape (N, H, W).
|
||||
max_dimension (int): The maximum dimension for the resized masks.
|
||||
|
||||
Returns:
|
||||
np.ndarray: Array of resized masks.
|
||||
"""
|
||||
max_height = np.max(masks.shape[1])
|
||||
max_width = np.max(masks.shape[2])
|
||||
scale = min(max_dimension / max_height, max_dimension / max_width)
|
||||
|
||||
new_height = int(scale * max_height)
|
||||
new_width = int(scale * max_width)
|
||||
|
||||
x = np.linspace(0, max_width - 1, new_width).astype(int)
|
||||
y = np.linspace(0, max_height - 1, new_height).astype(int)
|
||||
xv, yv = np.meshgrid(x, y)
|
||||
|
||||
resized_masks = masks[:, yv, xv]
|
||||
|
||||
resized_masks = resized_masks.reshape(masks.shape[0], new_height, new_width)
|
||||
return resized_masks
|
||||
|
||||
|
||||
def mask_non_max_suppression(
|
||||
predictions: np.ndarray,
|
||||
masks: np.ndarray,
|
||||
iou_threshold: float = 0.5,
|
||||
mask_dimension: int = 640,
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
Perform Non-Maximum Suppression (NMS) on segmentation predictions.
|
||||
|
||||
Args:
|
||||
predictions (np.ndarray): A 2D 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)`. Shape: `(N, 5)` or
|
||||
`(N, 6)`, where N is the number of predictions.
|
||||
masks (np.ndarray): A 3D array of binary masks corresponding to the predictions.
|
||||
Shape: `(N, H, W)`, where N is the number of predictions, and H, W are the
|
||||
dimensions of each mask.
|
||||
iou_threshold (float, optional): The intersection-over-union threshold
|
||||
to use for non-maximum suppression.
|
||||
mask_dimension (int, optional): The dimension to which the masks should be
|
||||
resized before computing IOU values. Defaults to 640.
|
||||
|
||||
Returns:
|
||||
np.ndarray: A boolean array indicating which predictions to keep after
|
||||
non-maximum suppression.
|
||||
|
||||
Raises:
|
||||
AssertionError: If `iou_threshold` is not within the closed
|
||||
range from `0` to `1`.
|
||||
"""
|
||||
assert 0 <= iou_threshold <= 1, (
|
||||
"Value of `iou_threshold` must be in the closed range from 0 to 1, "
|
||||
f"{iou_threshold} given."
|
||||
)
|
||||
rows, columns = predictions.shape
|
||||
|
||||
if columns == 5:
|
||||
predictions = np.c_[predictions, np.zeros(rows)]
|
||||
|
||||
sort_index = predictions[:, 4].argsort()[::-1]
|
||||
predictions = predictions[sort_index]
|
||||
masks = masks[sort_index]
|
||||
masks_resized = resize_masks(masks, mask_dimension)
|
||||
ious = mask_iou_batch(masks_resized, masks_resized)
|
||||
categories = predictions[:, 5]
|
||||
|
||||
keep = np.ones(rows, dtype=bool)
|
||||
for i in range(rows):
|
||||
if keep[i]:
|
||||
condition = (ious[i] > iou_threshold) & (categories[i] == categories)
|
||||
keep[i + 1 :] = np.where(condition[i + 1 :], False, keep[i + 1 :])
|
||||
|
||||
return keep[sort_index.argsort()]
|
||||
|
||||
|
||||
def box_non_max_suppression(
|
||||
predictions: np.ndarray, iou_threshold: float = 0.5
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
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.
|
||||
|
||||
Returns:
|
||||
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`.
|
||||
"""
|
||||
assert 0 <= iou_threshold <= 1, (
|
||||
"Value of `iou_threshold` must be in the closed range from 0 to 1, "
|
||||
f"{iou_threshold} given."
|
||||
)
|
||||
rows, columns = predictions.shape
|
||||
|
||||
# add column #5 - category filled with zeros for agnostic nms
|
||||
if columns == 5:
|
||||
predictions = np.c_[predictions, np.zeros(rows)]
|
||||
|
||||
# sort predictions column #4 - score
|
||||
sort_index = np.flip(predictions[:, 4].argsort())
|
||||
predictions = predictions[sort_index]
|
||||
|
||||
boxes = predictions[:, :4]
|
||||
categories = predictions[:, 5]
|
||||
ious = box_iou_batch(boxes, boxes)
|
||||
ious = ious - np.eye(rows)
|
||||
|
||||
keep = np.ones(rows, dtype=bool)
|
||||
|
||||
for index, (iou, category) in enumerate(zip(ious, categories)):
|
||||
if not keep[index]:
|
||||
continue
|
||||
|
||||
# drop detections with iou > iou_threshold and
|
||||
# same category as current detections
|
||||
condition = (iou > iou_threshold) & (categories == category)
|
||||
keep = keep & ~condition
|
||||
|
||||
return keep[sort_index.argsort()]
|
||||
|
||||
|
||||
def group_overlapping_boxes(
|
||||
predictions: npt.NDArray[np.float64], iou_threshold: float = 0.5
|
||||
) -> List[List[int]]:
|
||||
"""
|
||||
Apply greedy version of non-maximum merging to avoid detecting too many
|
||||
overlapping bounding boxes for a given object.
|
||||
|
||||
Args:
|
||||
predictions (npt.NDArray[np.float64]): An array of shape `(n, 5)` containing
|
||||
the bounding boxes coordinates in format `[x1, y1, x2, y2]`
|
||||
and the confidence scores.
|
||||
iou_threshold (float, optional): The intersection-over-union threshold
|
||||
to use for non-maximum suppression. Defaults to 0.5.
|
||||
|
||||
Returns:
|
||||
List[List[int]]: Groups of prediction indices be merged.
|
||||
Each group may have 1 or more elements.
|
||||
"""
|
||||
merge_groups: List[List[int]] = []
|
||||
|
||||
scores = predictions[:, 4]
|
||||
order = scores.argsort()
|
||||
|
||||
while len(order) > 0:
|
||||
idx = int(order[-1])
|
||||
|
||||
order = order[:-1]
|
||||
if len(order) == 0:
|
||||
merge_groups.append([idx])
|
||||
break
|
||||
|
||||
merge_candidate = np.expand_dims(predictions[idx], axis=0)
|
||||
ious = box_iou_batch(predictions[order][:, :4], merge_candidate[:, :4])
|
||||
ious = ious.flatten()
|
||||
|
||||
above_threshold = ious >= iou_threshold
|
||||
merge_group = [idx] + np.flip(order[above_threshold]).tolist()
|
||||
merge_groups.append(merge_group)
|
||||
order = order[~above_threshold]
|
||||
return merge_groups
|
||||
|
||||
|
||||
def box_non_max_merge(
|
||||
predictions: npt.NDArray[np.float64],
|
||||
iou_threshold: float = 0.5,
|
||||
) -> List[List[int]]:
|
||||
"""
|
||||
Apply greedy version of non-maximum merging per category to avoid detecting
|
||||
too many overlapping bounding boxes for a given object.
|
||||
|
||||
Args:
|
||||
predictions (npt.NDArray[np.float64]): An array of shape `(n, 5)` or `(n, 6)`
|
||||
containing the bounding boxes coordinates in format `[x1, y1, x2, y2]`,
|
||||
the confidence scores and class_ids. Omit class_id column to allow
|
||||
detections of different classes to be merged.
|
||||
iou_threshold (float, optional): The intersection-over-union threshold
|
||||
to use for non-maximum suppression. Defaults to 0.5.
|
||||
|
||||
Returns:
|
||||
List[List[int]]: Groups of prediction indices be merged.
|
||||
Each group may have 1 or more elements.
|
||||
"""
|
||||
if predictions.shape[1] == 5:
|
||||
return group_overlapping_boxes(predictions, iou_threshold)
|
||||
|
||||
category_ids = predictions[:, 5]
|
||||
merge_groups = []
|
||||
for category_id in np.unique(category_ids):
|
||||
curr_indices = np.where(category_ids == category_id)[0]
|
||||
merge_class_groups = group_overlapping_boxes(
|
||||
predictions[curr_indices], iou_threshold
|
||||
)
|
||||
|
||||
for merge_class_group in merge_class_groups:
|
||||
merge_groups.append(curr_indices[merge_class_group].tolist())
|
||||
|
||||
for merge_group in merge_groups:
|
||||
if len(merge_group) == 0:
|
||||
raise ValueError(
|
||||
f"Empty group detected when non-max-merging "
|
||||
f"detections: {merge_groups}"
|
||||
)
|
||||
return merge_groups
|
||||
|
||||
|
||||
class OverlapFilter(Enum):
|
||||
"""
|
||||
Enum specifying the strategy for filtering overlapping detections.
|
||||
|
||||
Attributes:
|
||||
NONE: Do not filter detections based on overlap.
|
||||
NON_MAX_SUPPRESSION: Filter detections using non-max suppression. This means,
|
||||
detections that overlap by more than a set threshold will be discarded,
|
||||
except for the one with the highest confidence.
|
||||
NON_MAX_MERGE: Merge detections with non-max merging. This means,
|
||||
detections that overlap by more than a set threshold will be merged
|
||||
into a single detection.
|
||||
"""
|
||||
|
||||
NONE = "none"
|
||||
NON_MAX_SUPPRESSION = "non_max_suppression"
|
||||
NON_MAX_MERGE = "non_max_merge"
|
||||
|
||||
|
||||
def validate_overlap_filter(
|
||||
strategy: Union[OverlapFilter, str],
|
||||
) -> OverlapFilter:
|
||||
if isinstance(strategy, str):
|
||||
try:
|
||||
strategy = OverlapFilter(strategy.lower())
|
||||
except ValueError:
|
||||
raise ValueError(
|
||||
f"Invalid strategy value: {strategy}. Must be one of "
|
||||
f"{[e.value for e in OverlapFilter]}"
|
||||
)
|
||||
return strategy
|
||||
@@ -1,23 +1,42 @@
|
||||
import warnings
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
from typing import Callable, Optional, Tuple
|
||||
from typing import Callable, Optional, Tuple, Union
|
||||
|
||||
import numpy as np
|
||||
|
||||
from supervision.detection.core import Detections
|
||||
from supervision.detection.utils import move_boxes
|
||||
from supervision.detection.overlap_filter import OverlapFilter, validate_overlap_filter
|
||||
from supervision.detection.utils import move_boxes, move_masks
|
||||
from supervision.utils.image import crop_image
|
||||
from supervision.utils.internal import SupervisionWarnings
|
||||
|
||||
|
||||
def move_detections(detections: Detections, offset: np.array) -> Detections:
|
||||
def move_detections(
|
||||
detections: Detections,
|
||||
offset: np.ndarray,
|
||||
resolution_wh: Optional[Tuple[int, int]] = None,
|
||||
) -> Detections:
|
||||
"""
|
||||
Args:
|
||||
detections (sv.Detections): Detections object to be moved.
|
||||
offset (np.array): An array of shape `(2,)` containing offset values in format
|
||||
offset (np.ndarray): An array of shape `(2,)` containing offset values in format
|
||||
is `[dx, dy]`.
|
||||
resolution_wh (Tuple[int, int]): The width and height of the desired mask
|
||||
resolution. Required for segmentation detections.
|
||||
|
||||
Returns:
|
||||
(sv.Detections) repositioned Detections object.
|
||||
"""
|
||||
detections.xyxy = move_boxes(xyxy=detections.xyxy, offset=offset)
|
||||
if detections.mask is not None:
|
||||
if resolution_wh is None:
|
||||
raise ValueError(
|
||||
"Resolution width and height are required for moving segmentation "
|
||||
"detections. This should be the same as (width, height) of image shape."
|
||||
)
|
||||
detections.mask = move_masks(
|
||||
masks=detections.mask, offset=offset, resolution_wh=resolution_wh
|
||||
)
|
||||
return detections
|
||||
|
||||
|
||||
@@ -34,8 +53,10 @@ class InferenceSlicer:
|
||||
`(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.
|
||||
overlap_filter_strategy (Union[OverlapFilter, str]): Strategy for
|
||||
filtering or merging overlapping detections in slices.
|
||||
iou_threshold (float): Intersection over Union (IoU) threshold
|
||||
used when filtering by overlap.
|
||||
callback (Callable): A function that performs inference on a given image
|
||||
slice and returns detections.
|
||||
thread_workers (int): Number of threads for parallel execution.
|
||||
@@ -52,12 +73,18 @@ class InferenceSlicer:
|
||||
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,
|
||||
overlap_filter_strategy: Union[
|
||||
OverlapFilter, str
|
||||
] = OverlapFilter.NON_MAX_SUPPRESSION,
|
||||
iou_threshold: float = 0.5,
|
||||
thread_workers: int = 1,
|
||||
):
|
||||
overlap_filter_strategy = validate_overlap_filter(overlap_filter_strategy)
|
||||
|
||||
self.slice_wh = slice_wh
|
||||
self.overlap_ratio_wh = overlap_ratio_wh
|
||||
self.iou_threshold = iou_threshold
|
||||
self.overlap_filter_strategy = overlap_filter_strategy
|
||||
self.callback = callback
|
||||
self.thread_workers = thread_workers
|
||||
|
||||
@@ -88,7 +115,10 @@ class InferenceSlicer:
|
||||
result = model(image_slice)[0]
|
||||
return sv.Detections.from_ultralytics(result)
|
||||
|
||||
slicer = sv.InferenceSlicer(callback = callback)
|
||||
slicer = sv.InferenceSlicer(
|
||||
callback=callback,
|
||||
overlap_filter_strategy=sv.OverlapFilter.NON_MAX_SUPPRESSION,
|
||||
)
|
||||
|
||||
detections = slicer(image)
|
||||
```
|
||||
@@ -108,9 +138,19 @@ class InferenceSlicer:
|
||||
for future in as_completed(futures):
|
||||
detections_list.append(future.result())
|
||||
|
||||
return Detections.merge(detections_list=detections_list).with_nms(
|
||||
threshold=self.iou_threshold
|
||||
)
|
||||
merged = Detections.merge(detections_list=detections_list)
|
||||
if self.overlap_filter_strategy == OverlapFilter.NONE:
|
||||
return merged
|
||||
elif self.overlap_filter_strategy == OverlapFilter.NON_MAX_SUPPRESSION:
|
||||
return merged.with_nms(threshold=self.iou_threshold)
|
||||
elif self.overlap_filter_strategy == OverlapFilter.NON_MAX_MERGE:
|
||||
return merged.with_nmm(threshold=self.iou_threshold)
|
||||
else:
|
||||
warnings.warn(
|
||||
f"Invalid overlap filter strategy: {self.overlap_filter_strategy}",
|
||||
category=SupervisionWarnings,
|
||||
)
|
||||
return merged
|
||||
|
||||
def _run_callback(self, image, offset) -> Detections:
|
||||
"""
|
||||
@@ -126,7 +166,10 @@ class InferenceSlicer:
|
||||
"""
|
||||
image_slice = crop_image(image=image, xyxy=offset)
|
||||
detections = self.callback(image_slice)
|
||||
detections = move_detections(detections=detections, offset=offset[:2])
|
||||
resolution_wh = (image.shape[1], image.shape[0])
|
||||
detections = move_detections(
|
||||
detections=detections, offset=offset[:2], resolution_wh=resolution_wh
|
||||
)
|
||||
|
||||
return detections
|
||||
|
||||
|
||||
@@ -54,6 +54,8 @@ class PolygonZone:
|
||||
|
||||
self.polygon = polygon.astype(int)
|
||||
self.triggering_anchors = triggering_anchors
|
||||
if not list(self.triggering_anchors):
|
||||
raise ValueError("Triggering anchors cannot be empty.")
|
||||
|
||||
self.current_count = 0
|
||||
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
import warnings
|
||||
from collections import defaultdict, deque
|
||||
from copy import deepcopy
|
||||
from typing import Optional
|
||||
@@ -5,6 +6,7 @@ from typing import Optional
|
||||
import numpy as np
|
||||
|
||||
from supervision.detection.core import Detections
|
||||
from supervision.utils.internal import SupervisionWarnings
|
||||
|
||||
|
||||
class DetectionsSmoother:
|
||||
@@ -39,7 +41,7 @@ class DetectionsSmoother:
|
||||
tracker = sv.ByteTrack(frame_rate=video_info.fps)
|
||||
smoother = sv.DetectionsSmoother()
|
||||
|
||||
annotator = sv.BoundingBoxAnnotator()
|
||||
box_annotator = sv.BoxAnnotator()
|
||||
|
||||
with sv.VideoSink(<TARGET_FILE_PATH>, video_info=video_info) as sink:
|
||||
for frame in frame_generator:
|
||||
@@ -48,7 +50,7 @@ class DetectionsSmoother:
|
||||
detections = tracker.update_with_detections(detections)
|
||||
detections = smoother.update_with_detections(detections)
|
||||
|
||||
annotated_frame = bounding_box_annotator.annotate(frame.copy(), detections)
|
||||
annotated_frame = box_annotator.annotate(frame.copy(), detections)
|
||||
sink.write_frame(annotated_frame)
|
||||
```
|
||||
""" # noqa: E501 // docs
|
||||
@@ -70,16 +72,16 @@ class DetectionsSmoother:
|
||||
"""
|
||||
|
||||
if detections.tracker_id is None:
|
||||
print(
|
||||
warnings.warn(
|
||||
"Smoothing skipped. DetectionsSmoother requires tracker_id. Refer to "
|
||||
"https://supervision.roboflow.com/latest/trackers for more information."
|
||||
"https://supervision.roboflow.com/latest/trackers for more "
|
||||
"information.",
|
||||
category=SupervisionWarnings,
|
||||
)
|
||||
return detections
|
||||
|
||||
for detection_idx in range(len(detections)):
|
||||
tracker_id = detections.tracker_id[detection_idx]
|
||||
if tracker_id is None:
|
||||
continue
|
||||
|
||||
self.tracks[tracker_id].append(detections[detection_idx])
|
||||
|
||||
|
||||
+352
-173
@@ -3,8 +3,10 @@ from typing import Dict, List, Optional, Tuple, Union
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import numpy.typing as npt
|
||||
|
||||
from supervision.config import CLASS_NAME_DATA_FIELD
|
||||
from supervision.geometry.core import Vector
|
||||
|
||||
MIN_POLYGON_POINT_COUNT = 3
|
||||
|
||||
@@ -56,7 +58,9 @@ def box_iou_batch(boxes_true: np.ndarray, boxes_detection: np.ndarray) -> np.nda
|
||||
bottom_right = np.minimum(boxes_true[:, None, 2:], boxes_detection[:, 2:])
|
||||
|
||||
area_inter = np.prod(np.clip(bottom_right - top_left, a_min=0, a_max=None), 2)
|
||||
return area_inter / (area_true[:, None] + area_detection - area_inter)
|
||||
ious = area_inter / (area_true[:, None] + area_detection - area_inter)
|
||||
ious = np.nan_to_num(ious)
|
||||
return ious
|
||||
|
||||
|
||||
def _mask_iou_batch_split(
|
||||
@@ -136,144 +140,6 @@ def mask_iou_batch(
|
||||
return np.vstack(ious)
|
||||
|
||||
|
||||
def resize_masks(masks: np.ndarray, max_dimension: int = 640) -> np.ndarray:
|
||||
"""
|
||||
Resize all masks in the array to have a maximum dimension of max_dimension,
|
||||
maintaining aspect ratio.
|
||||
|
||||
Args:
|
||||
masks (np.ndarray): 3D array of binary masks with shape (N, H, W).
|
||||
max_dimension (int): The maximum dimension for the resized masks.
|
||||
|
||||
Returns:
|
||||
np.ndarray: Array of resized masks.
|
||||
"""
|
||||
max_height = np.max(masks.shape[1])
|
||||
max_width = np.max(masks.shape[2])
|
||||
scale = min(max_dimension / max_height, max_dimension / max_width)
|
||||
|
||||
new_height = int(scale * max_height)
|
||||
new_width = int(scale * max_width)
|
||||
|
||||
x = np.linspace(0, max_width - 1, new_width).astype(int)
|
||||
y = np.linspace(0, max_height - 1, new_height).astype(int)
|
||||
xv, yv = np.meshgrid(x, y)
|
||||
|
||||
resized_masks = masks[:, yv, xv]
|
||||
|
||||
resized_masks = resized_masks.reshape(masks.shape[0], new_height, new_width)
|
||||
return resized_masks
|
||||
|
||||
|
||||
def mask_non_max_suppression(
|
||||
predictions: np.ndarray,
|
||||
masks: np.ndarray,
|
||||
iou_threshold: float = 0.5,
|
||||
mask_dimension: int = 640,
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
Perform Non-Maximum Suppression (NMS) on segmentation predictions.
|
||||
|
||||
Args:
|
||||
predictions (np.ndarray): A 2D 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)`. Shape: `(N, 5)` or
|
||||
`(N, 6)`, where N is the number of predictions.
|
||||
masks (np.ndarray): A 3D array of binary masks corresponding to the predictions.
|
||||
Shape: `(N, H, W)`, where N is the number of predictions, and H, W are the
|
||||
dimensions of each mask.
|
||||
iou_threshold (float, optional): The intersection-over-union threshold
|
||||
to use for non-maximum suppression.
|
||||
mask_dimension (int, optional): The dimension to which the masks should be
|
||||
resized before computing IOU values. Defaults to 640.
|
||||
|
||||
Returns:
|
||||
np.ndarray: A boolean array indicating which predictions to keep after
|
||||
non-maximum suppression.
|
||||
|
||||
Raises:
|
||||
AssertionError: If `iou_threshold` is not within the closed
|
||||
range from `0` to `1`.
|
||||
"""
|
||||
assert 0 <= iou_threshold <= 1, (
|
||||
"Value of `iou_threshold` must be in the closed range from 0 to 1, "
|
||||
f"{iou_threshold} given."
|
||||
)
|
||||
rows, columns = predictions.shape
|
||||
|
||||
if columns == 5:
|
||||
predictions = np.c_[predictions, np.zeros(rows)]
|
||||
|
||||
sort_index = predictions[:, 4].argsort()[::-1]
|
||||
predictions = predictions[sort_index]
|
||||
masks = masks[sort_index]
|
||||
masks_resized = resize_masks(masks, mask_dimension)
|
||||
ious = mask_iou_batch(masks_resized, masks_resized)
|
||||
categories = predictions[:, 5]
|
||||
|
||||
keep = np.ones(rows, dtype=bool)
|
||||
for i in range(rows):
|
||||
if keep[i]:
|
||||
condition = (ious[i] > iou_threshold) & (categories[i] == categories)
|
||||
keep[i + 1 :] = np.where(condition[i + 1 :], False, keep[i + 1 :])
|
||||
|
||||
return keep[sort_index.argsort()]
|
||||
|
||||
|
||||
def box_non_max_suppression(
|
||||
predictions: np.ndarray, iou_threshold: float = 0.5
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
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.
|
||||
|
||||
Returns:
|
||||
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`.
|
||||
"""
|
||||
assert 0 <= iou_threshold <= 1, (
|
||||
"Value of `iou_threshold` must be in the closed range from 0 to 1, "
|
||||
f"{iou_threshold} given."
|
||||
)
|
||||
rows, columns = predictions.shape
|
||||
|
||||
# add column #5 - category filled with zeros for agnostic nms
|
||||
if columns == 5:
|
||||
predictions = np.c_[predictions, np.zeros(rows)]
|
||||
|
||||
# sort predictions column #4 - score
|
||||
sort_index = np.flip(predictions[:, 4].argsort())
|
||||
predictions = predictions[sort_index]
|
||||
|
||||
boxes = predictions[:, :4]
|
||||
categories = predictions[:, 5]
|
||||
ious = box_iou_batch(boxes, boxes)
|
||||
ious = ious - np.eye(rows)
|
||||
|
||||
keep = np.ones(rows, dtype=bool)
|
||||
|
||||
for index, (iou, category) in enumerate(zip(ious, categories)):
|
||||
if not keep[index]:
|
||||
continue
|
||||
|
||||
# drop detections with iou > iou_threshold and
|
||||
# same category as current detections
|
||||
condition = (iou > iou_threshold) & (categories == category)
|
||||
keep = keep & ~condition
|
||||
|
||||
return keep[sort_index.argsort()]
|
||||
|
||||
|
||||
def clip_boxes(xyxy: np.ndarray, resolution_wh: Tuple[int, int]) -> np.ndarray:
|
||||
"""
|
||||
Clips bounding boxes coordinates to fit within the frame resolution.
|
||||
@@ -289,6 +155,25 @@ def clip_boxes(xyxy: np.ndarray, resolution_wh: Tuple[int, int]) -> np.ndarray:
|
||||
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.
|
||||
|
||||
Examples:
|
||||
```python
|
||||
import numpy as np
|
||||
import supervision as sv
|
||||
|
||||
xyxy = np.array([
|
||||
[10, 20, 300, 200],
|
||||
[15, 25, 350, 450],
|
||||
[-10, -20, 30, 40]
|
||||
])
|
||||
|
||||
sv.clip_boxes(xyxy=xyxy, resolution_wh=(320, 240))
|
||||
# array([
|
||||
# [ 10, 20, 300, 200],
|
||||
# [ 15, 25, 320, 240],
|
||||
# [ 0, 0, 30, 40]
|
||||
# ])
|
||||
```
|
||||
"""
|
||||
result = np.copy(xyxy)
|
||||
width, height = resolution_wh
|
||||
@@ -297,10 +182,124 @@ def clip_boxes(xyxy: np.ndarray, resolution_wh: Tuple[int, int]) -> np.ndarray:
|
||||
return result
|
||||
|
||||
|
||||
def xywh_to_xyxy(boxes_xywh: np.ndarray) -> np.ndarray:
|
||||
xyxy = boxes_xywh.copy()
|
||||
xyxy[:, 2] = boxes_xywh[:, 0] + boxes_xywh[:, 2]
|
||||
xyxy[:, 3] = boxes_xywh[:, 1] + boxes_xywh[:, 3]
|
||||
def pad_boxes(xyxy: np.ndarray, px: int, py: Optional[int] = None) -> np.ndarray:
|
||||
"""
|
||||
Pads bounding boxes coordinates with a constant padding.
|
||||
|
||||
Args:
|
||||
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)`.
|
||||
px (int): The padding value to be added to both the left and right sides of
|
||||
each bounding box.
|
||||
py (Optional[int]): The padding value to be added to both the top and bottom
|
||||
sides of each bounding box. If not provided, `px` will be used for both
|
||||
dimensions.
|
||||
|
||||
Returns:
|
||||
np.ndarray: A numpy array of shape `(N, 4)` where each row corresponds to a
|
||||
bounding box with coordinates padded according to the provided padding
|
||||
values.
|
||||
|
||||
Examples:
|
||||
```python
|
||||
import numpy as np
|
||||
import supervision as sv
|
||||
|
||||
xyxy = np.array([
|
||||
[10, 20, 30, 40],
|
||||
[15, 25, 35, 45]
|
||||
])
|
||||
|
||||
sv.pad_boxes(xyxy=xyxy, px=5, py=10)
|
||||
# array([
|
||||
# [ 5, 10, 35, 50],
|
||||
# [10, 15, 40, 55]
|
||||
# ])
|
||||
```
|
||||
"""
|
||||
if py is None:
|
||||
py = px
|
||||
|
||||
result = xyxy.copy()
|
||||
result[:, [0, 1]] -= [px, py]
|
||||
result[:, [2, 3]] += [px, py]
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def xywh_to_xyxy(xywh: np.ndarray) -> np.ndarray:
|
||||
"""
|
||||
Converts bounding box coordinates from `(x, y, width, height)`
|
||||
format to `(x_min, y_min, x_max, y_max)` format.
|
||||
|
||||
Args:
|
||||
xywh (np.ndarray): A numpy array of shape `(N, 4)` where each row
|
||||
corresponds to a bounding box in the format `(x, y, width, height)`.
|
||||
|
||||
Returns:
|
||||
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)`.
|
||||
|
||||
Examples:
|
||||
```python
|
||||
import numpy as np
|
||||
import supervision as sv
|
||||
|
||||
xywh = np.array([
|
||||
[10, 20, 30, 40],
|
||||
[15, 25, 35, 45]
|
||||
])
|
||||
|
||||
sv.xywh_to_xyxy(xywh=xywh)
|
||||
# array([
|
||||
# [10, 20, 40, 60],
|
||||
# [15, 25, 50, 70]
|
||||
# ])
|
||||
```
|
||||
"""
|
||||
xyxy = xywh.copy()
|
||||
xyxy[:, 2] = xywh[:, 0] + xywh[:, 2]
|
||||
xyxy[:, 3] = xywh[:, 1] + xywh[:, 3]
|
||||
return xyxy
|
||||
|
||||
|
||||
def xcycwh_to_xyxy(xcycwh: np.ndarray) -> np.ndarray:
|
||||
"""
|
||||
Converts bounding box coordinates from `(center_x, center_y, width, height)`
|
||||
format to `(x_min, y_min, x_max, y_max)` format.
|
||||
|
||||
Args:
|
||||
xcycwh (np.ndarray): A numpy array of shape `(N, 4)` where each row
|
||||
corresponds to a bounding box in the format `(center_x, center_y, width,
|
||||
height)`.
|
||||
|
||||
Returns:
|
||||
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)`.
|
||||
|
||||
Examples:
|
||||
```python
|
||||
import numpy as np
|
||||
import supervision as sv
|
||||
|
||||
xcycwh = np.array([
|
||||
[50, 50, 20, 30],
|
||||
[30, 40, 10, 15]
|
||||
])
|
||||
|
||||
sv.xcycwh_to_xyxy(xcycwh=xcycwh)
|
||||
# array([
|
||||
# [40, 35, 60, 65],
|
||||
# [25, 32.5, 35, 47.5]
|
||||
# ])
|
||||
```
|
||||
"""
|
||||
xyxy = xcycwh.copy()
|
||||
xyxy[:, 0] = xcycwh[:, 0] - xcycwh[:, 2] / 2
|
||||
xyxy[:, 1] = xcycwh[:, 1] - xcycwh[:, 3] / 2
|
||||
xyxy[:, 2] = xcycwh[:, 0] + xcycwh[:, 2] / 2
|
||||
xyxy[:, 3] = xcycwh[:, 1] + xcycwh[:, 3] / 2
|
||||
return xyxy
|
||||
|
||||
|
||||
@@ -317,7 +316,7 @@ def mask_to_xyxy(masks: np.ndarray) -> np.ndarray:
|
||||
`(x_min, y_min, x_max, y_max)` for each mask
|
||||
"""
|
||||
n = masks.shape[0]
|
||||
bboxes = np.zeros((n, 4), dtype=int)
|
||||
xyxy = np.zeros((n, 4), dtype=int)
|
||||
|
||||
for i, mask in enumerate(masks):
|
||||
rows, cols = np.where(mask)
|
||||
@@ -325,9 +324,9 @@ def mask_to_xyxy(masks: np.ndarray) -> np.ndarray:
|
||||
if len(rows) > 0 and len(cols) > 0:
|
||||
x_min, x_max = np.min(cols), np.max(cols)
|
||||
y_min, y_max = np.min(rows), np.max(rows)
|
||||
bboxes[i, :] = [x_min, y_min, x_max, y_max]
|
||||
xyxy[i, :] = [x_min, y_min, x_max, y_max]
|
||||
|
||||
return bboxes
|
||||
return xyxy
|
||||
|
||||
|
||||
def mask_to_polygons(mask: np.ndarray) -> List[np.ndarray]:
|
||||
@@ -500,7 +499,7 @@ def process_roboflow_result(
|
||||
np.ndarray,
|
||||
Optional[np.ndarray],
|
||||
Optional[np.ndarray],
|
||||
Dict[str, List[np.ndarray]],
|
||||
Dict[str, Union[List[np.ndarray], np.ndarray]],
|
||||
]:
|
||||
if not roboflow_result["predictions"]:
|
||||
return (
|
||||
@@ -563,59 +562,103 @@ def process_roboflow_result(
|
||||
return xyxy, confidence, class_id, masks, tracker_id, data
|
||||
|
||||
|
||||
def move_boxes(xyxy: np.ndarray, offset: np.ndarray) -> np.ndarray:
|
||||
def move_boxes(
|
||||
xyxy: npt.NDArray[np.float64], offset: npt.NDArray[np.int32]
|
||||
) -> npt.NDArray[np.float64]:
|
||||
"""
|
||||
Parameters:
|
||||
xyxy (np.ndarray): An array of shape `(n, 4)` containing the bounding boxes
|
||||
coordinates in format `[x1, y1, x2, y2]`
|
||||
xyxy (npt.NDArray[np.float64]): 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.
|
||||
npt.NDArray[np.float64]: Repositioned bounding boxes.
|
||||
|
||||
Example:
|
||||
Examples:
|
||||
```python
|
||||
import numpy as np
|
||||
import supervision as sv
|
||||
|
||||
boxes = np.array([[10, 10, 20, 20], [30, 30, 40, 40]])
|
||||
xyxy = np.array([
|
||||
[10, 10, 20, 20],
|
||||
[30, 30, 40, 40]
|
||||
])
|
||||
offset = np.array([5, 5])
|
||||
moved_box = sv.move_boxes(boxes, offset)
|
||||
print(moved_box)
|
||||
# np.array([
|
||||
|
||||
sv.move_boxes(xyxy=xyxy, offset=offset)
|
||||
# array([
|
||||
# [15, 15, 25, 25],
|
||||
# [35, 35, 45, 45]
|
||||
# [35, 35, 45, 45]
|
||||
# ])
|
||||
```
|
||||
"""
|
||||
return xyxy + np.hstack([offset, offset])
|
||||
|
||||
|
||||
def scale_boxes(xyxy: np.ndarray, factor: float) -> np.ndarray:
|
||||
def move_masks(
|
||||
masks: npt.NDArray[np.bool_],
|
||||
offset: npt.NDArray[np.int32],
|
||||
resolution_wh: Tuple[int, int],
|
||||
) -> npt.NDArray[np.bool_]:
|
||||
"""
|
||||
Offset the masks in an array by the specified (x, y) amount.
|
||||
|
||||
Args:
|
||||
masks (npt.NDArray[np.bool_]): A 3D array of binary masks corresponding to the
|
||||
predictions. Shape: `(N, H, W)`, where N is the number of predictions, and
|
||||
H, W are the dimensions of each mask.
|
||||
offset (npt.NDArray[np.int32]): An array of shape `(2,)` containing non-negative
|
||||
int values `[dx, dy]`.
|
||||
resolution_wh (Tuple[int, int]): The width and height of the desired mask
|
||||
resolution.
|
||||
|
||||
Returns:
|
||||
(npt.NDArray[np.bool_]) repositioned masks, optionally padded to the specified
|
||||
shape.
|
||||
"""
|
||||
|
||||
if offset[0] < 0 or offset[1] < 0:
|
||||
raise ValueError(f"Offset values must be non-negative integers. Got: {offset}")
|
||||
|
||||
mask_array = np.full((masks.shape[0], resolution_wh[1], resolution_wh[0]), False)
|
||||
mask_array[
|
||||
:,
|
||||
offset[1] : masks.shape[1] + offset[1],
|
||||
offset[0] : masks.shape[2] + offset[0],
|
||||
] = masks
|
||||
|
||||
return mask_array
|
||||
|
||||
|
||||
def scale_boxes(
|
||||
xyxy: npt.NDArray[np.float64], factor: float
|
||||
) -> npt.NDArray[np.float64]:
|
||||
"""
|
||||
Scale the dimensions of bounding boxes.
|
||||
|
||||
Parameters:
|
||||
xyxy (np.ndarray): An array of shape `(n, 4)` containing the bounding boxes
|
||||
coordinates in format `[x1, y1, x2, y2]`
|
||||
xyxy (npt.NDArray[np.float64]): An array of shape `(n, 4)` containing the
|
||||
bounding boxes coordinates in format `[x1, y1, x2, y2]`
|
||||
factor (float): A float value representing the factor by which the box
|
||||
dimensions are scaled. A factor greater than 1 enlarges the boxes, while a
|
||||
factor less than 1 shrinks them.
|
||||
|
||||
Returns:
|
||||
np.ndarray: Scaled bounding boxes.
|
||||
npt.NDArray[np.float64]: Scaled bounding boxes.
|
||||
|
||||
Example:
|
||||
Examples:
|
||||
```python
|
||||
import numpy as np
|
||||
import supervision as sv
|
||||
|
||||
boxes = np.array([[10, 10, 20, 20], [30, 30, 40, 40]])
|
||||
factor = 1.5
|
||||
scaled_bb = sv.scale_boxes(boxes, factor)
|
||||
print(scaled_bb)
|
||||
# np.array([
|
||||
xyxy = np.array([
|
||||
[10, 10, 20, 20],
|
||||
[30, 30, 40, 40]
|
||||
])
|
||||
|
||||
sv.scale_boxes(xyxy=xyxy, factor=1.5)
|
||||
# array([
|
||||
# [ 7.5, 7.5, 22.5, 22.5],
|
||||
# [27.5, 27.5, 42.5, 42.5]
|
||||
# ])
|
||||
@@ -672,17 +715,19 @@ def is_data_equal(data_a: Dict[str, np.ndarray], data_b: Dict[str, np.ndarray])
|
||||
|
||||
|
||||
def merge_data(
|
||||
data_list: List[Dict[str, Union[np.ndarray, List]]],
|
||||
) -> Dict[str, Union[np.ndarray, List]]:
|
||||
data_list: List[Dict[str, Union[npt.NDArray[np.generic], List]]],
|
||||
) -> Dict[str, Union[npt.NDArray[np.generic], List]]:
|
||||
"""
|
||||
Merges the data payloads of a list of Detections instances.
|
||||
|
||||
Args:
|
||||
data_list: The data payloads of the instances.
|
||||
data_list: The data payloads of the Detections instances. Each data payload
|
||||
is a dictionary with the same keys, and the values are either lists or
|
||||
npt.NDArray[np.generic].
|
||||
|
||||
Returns:
|
||||
A single data payload containing the merged data, preserving the original data
|
||||
types (list or np.ndarray).
|
||||
types (list or npt.NDArray[np.generic]).
|
||||
|
||||
Raises:
|
||||
ValueError: If data values within a single object have different lengths or if
|
||||
@@ -703,9 +748,8 @@ def merge_data(
|
||||
)
|
||||
|
||||
merged_data = {key: [] for key in all_keys_sets[0]}
|
||||
|
||||
for data in data_list:
|
||||
for key in merged_data:
|
||||
for key in data:
|
||||
merged_data[key].append(data[key])
|
||||
|
||||
for key in merged_data:
|
||||
@@ -766,3 +810,138 @@ def get_data_item(
|
||||
raise TypeError(f"Unsupported data type for key '{key}': {type(value)}")
|
||||
|
||||
return subset_data
|
||||
|
||||
|
||||
def contains_holes(mask: npt.NDArray[np.bool_]) -> bool:
|
||||
"""
|
||||
Checks if the binary mask contains holes (background pixels fully enclosed by
|
||||
foreground pixels).
|
||||
|
||||
Args:
|
||||
mask (npt.NDArray[np.bool_]): 2D binary mask where `True` indicates foreground
|
||||
object and `False` indicates background.
|
||||
|
||||
Returns:
|
||||
True if holes are detected, False otherwise.
|
||||
|
||||
Examples:
|
||||
```python
|
||||
import numpy as np
|
||||
import supervision as sv
|
||||
|
||||
mask = np.array([
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 1, 1, 1, 0],
|
||||
[0, 1, 0, 1, 0],
|
||||
[0, 1, 1, 1, 0],
|
||||
[0, 0, 0, 0, 0]
|
||||
]).astype(bool)
|
||||
|
||||
sv.contains_holes(mask=mask)
|
||||
# True
|
||||
|
||||
mask = np.array([
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 1, 1, 1, 0],
|
||||
[0, 1, 1, 1, 0],
|
||||
[0, 1, 1, 1, 0],
|
||||
[0, 0, 0, 0, 0]
|
||||
]).astype(bool)
|
||||
|
||||
sv.contains_holes(mask=mask)
|
||||
# False
|
||||
```
|
||||
|
||||
{ align=center width="800" }
|
||||
""" # noqa E501 // docs
|
||||
mask_uint8 = mask.astype(np.uint8)
|
||||
_, hierarchy = cv2.findContours(mask_uint8, cv2.RETR_CCOMP, cv2.CHAIN_APPROX_SIMPLE)
|
||||
|
||||
if hierarchy is not None:
|
||||
parent_contour_index = 3
|
||||
for h in hierarchy[0]:
|
||||
if h[parent_contour_index] != -1:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def contains_multiple_segments(
|
||||
mask: npt.NDArray[np.bool_], connectivity: int = 4
|
||||
) -> bool:
|
||||
"""
|
||||
Checks if the binary mask contains multiple unconnected foreground segments.
|
||||
|
||||
Args:
|
||||
mask (npt.NDArray[np.bool_]): 2D binary mask where `True` indicates foreground
|
||||
object and `False` indicates background.
|
||||
connectivity (int) : Default: 4 is 4-way connectivity, which means that
|
||||
foreground pixels are the part of the same segment/component
|
||||
if their edges touch.
|
||||
Alternatively: 8 for 8-way connectivity, when foreground pixels are
|
||||
connected by their edges or corners touch.
|
||||
|
||||
Returns:
|
||||
True when the mask contains multiple not connected components, False otherwise.
|
||||
|
||||
Raises:
|
||||
ValueError: If connectivity(int) parameter value is not 4 or 8.
|
||||
|
||||
Examples:
|
||||
```python
|
||||
import numpy as np
|
||||
import supervision as sv
|
||||
|
||||
mask = np.array([
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[0, 1, 1, 0, 1, 1],
|
||||
[0, 1, 1, 0, 1, 1],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[0, 1, 1, 1, 0, 0],
|
||||
[0, 1, 1, 1, 0, 0]
|
||||
]).astype(bool)
|
||||
|
||||
sv.contains_multiple_segments(mask=mask, connectivity=4)
|
||||
# True
|
||||
|
||||
mask = np.array([
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[0, 1, 1, 1, 1, 1],
|
||||
[0, 1, 1, 1, 1, 1],
|
||||
[0, 1, 1, 1, 1, 1],
|
||||
[0, 1, 1, 1, 1, 1],
|
||||
[0, 0, 0, 0, 0, 0]
|
||||
]).astype(bool)
|
||||
|
||||
sv.contains_multiple_segments(mask=mask, connectivity=4)
|
||||
# False
|
||||
```
|
||||
|
||||
{ align=center width="800" }
|
||||
""" # noqa E501 // docs
|
||||
if connectivity != 4 and connectivity != 8:
|
||||
raise ValueError(
|
||||
"Incorrect connectivity value. Possible connectivity values: 4 or 8."
|
||||
)
|
||||
mask_uint8 = mask.astype(np.uint8)
|
||||
labels = np.zeros_like(mask_uint8, dtype=np.int32)
|
||||
number_of_labels, _ = cv2.connectedComponents(
|
||||
mask_uint8, labels, connectivity=connectivity
|
||||
)
|
||||
return number_of_labels > 2
|
||||
|
||||
|
||||
def cross_product(anchors: np.ndarray, vector: Vector) -> np.ndarray:
|
||||
"""
|
||||
Get array of cross products of each anchor with a vector.
|
||||
Args:
|
||||
anchors: Array of anchors of shape (number of anchors, detections, 2)
|
||||
vector: Vector to calculate cross product with
|
||||
|
||||
Returns:
|
||||
Array of cross products of shape (number of anchors, detections)
|
||||
"""
|
||||
vector_at_zero = np.array(
|
||||
[vector.end.x - vector.start.x, vector.end.y - vector.start.y]
|
||||
)
|
||||
vector_start = np.array([vector.start.x, vector.start.y])
|
||||
return np.cross(vector_at_zero, anchors - vector_start)
|
||||
|
||||
@@ -5,7 +5,7 @@ from typing import List, Tuple, Union
|
||||
|
||||
import matplotlib.pyplot as plt
|
||||
|
||||
from supervision.utils.internal import classproperty, deprecated
|
||||
from supervision.utils.internal import classproperty
|
||||
|
||||
DEFAULT_COLOR_PALETTE = [
|
||||
"A351FB",
|
||||
@@ -255,46 +255,6 @@ class Color:
|
||||
def ROBOFLOW(cls) -> Color:
|
||||
return Color.from_hex("#A351FB")
|
||||
|
||||
@classmethod
|
||||
@deprecated(
|
||||
"`Color.white()` is deprecated and will be removed in "
|
||||
"`supervision-0.22.0`. Use `Color.WHITE` instead."
|
||||
)
|
||||
def white(cls) -> Color:
|
||||
return Color.from_hex(color_hex="#ffffff")
|
||||
|
||||
@classmethod
|
||||
@deprecated(
|
||||
"`Color.black()` is deprecated and will be removed in "
|
||||
"`supervision-0.22.0`. Use `Color.BLACK` instead."
|
||||
)
|
||||
def black(cls) -> Color:
|
||||
return Color.from_hex(color_hex="#000000")
|
||||
|
||||
@classmethod
|
||||
@deprecated(
|
||||
"`Color.red()` is deprecated and will be removed in "
|
||||
"`supervision-0.22.0`. Use `Color.RED` instead."
|
||||
)
|
||||
def red(cls) -> Color:
|
||||
return Color.from_hex(color_hex="#ff0000")
|
||||
|
||||
@classmethod
|
||||
@deprecated(
|
||||
"`Color.green()` is deprecated and will be removed in "
|
||||
"`supervision-0.22.0`. Use `Color.GREEN` instead."
|
||||
)
|
||||
def green(cls) -> Color:
|
||||
return Color.from_hex(color_hex="#00ff00")
|
||||
|
||||
@classmethod
|
||||
@deprecated(
|
||||
"`Color.blue()` is deprecated and will be removed in "
|
||||
"`supervision-0.22.0`. Use `Color.BLUE` instead."
|
||||
)
|
||||
def blue(cls) -> Color:
|
||||
return Color.from_hex(color_hex="#0000ff")
|
||||
|
||||
|
||||
@dataclass
|
||||
class ColorPalette:
|
||||
@@ -346,33 +306,6 @@ class ColorPalette:
|
||||
def LEGACY(cls) -> ColorPalette:
|
||||
return ColorPalette.from_hex(color_hex_list=LEGACY_COLOR_PALETTE)
|
||||
|
||||
@classmethod
|
||||
@deprecated(
|
||||
"`ColorPalette.default()` is deprecated and will be removed in "
|
||||
"`supervision-0.22.0`. Use `Color.DEFAULT` instead."
|
||||
)
|
||||
def default(cls) -> ColorPalette:
|
||||
"""
|
||||
!!! failure "Deprecated"
|
||||
|
||||
`ColorPalette.default()` is deprecated and will be removed in
|
||||
`supervision-0.22.0`. Use `Color.DEFAULT` instead.
|
||||
|
||||
Returns a default color palette.
|
||||
|
||||
Returns:
|
||||
ColorPalette: A ColorPalette instance with default colors.
|
||||
|
||||
Example:
|
||||
```python
|
||||
import supervision as sv
|
||||
|
||||
sv.ColorPalette.default()
|
||||
# ColorPalette(colors=[Color(r=255, g=64, b=64), Color(r=255, g=161, b=160), ...])
|
||||
```
|
||||
""" # noqa: E501 // docs
|
||||
return ColorPalette.from_hex(color_hex_list=DEFAULT_COLOR_PALETTE)
|
||||
|
||||
@classmethod
|
||||
def from_hex(cls, color_hex_list: List[str]) -> ColorPalette:
|
||||
"""
|
||||
|
||||
@@ -81,6 +81,58 @@ def draw_filled_rectangle(scene: np.ndarray, rect: Rect, color: Color) -> np.nda
|
||||
return scene
|
||||
|
||||
|
||||
def draw_rounded_rectangle(
|
||||
scene: np.ndarray,
|
||||
rect: Rect,
|
||||
color: Color,
|
||||
border_radius: int,
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
Draws a rounded rectangle on an image.
|
||||
|
||||
Parameters:
|
||||
scene (np.ndarray): The image on which the rounded rectangle will be drawn.
|
||||
rect (Rect): The rectangle to be drawn.
|
||||
color (Color): The color of the rounded rectangle.
|
||||
border_radius (int): The radius of the corner rounding.
|
||||
|
||||
Returns:
|
||||
np.ndarray: The image with the rounded rectangle drawn on it.
|
||||
"""
|
||||
x1, y1, x2, y2 = rect.as_xyxy_int_tuple()
|
||||
width, height = x2 - x1, y2 - y1
|
||||
border_radius = min(border_radius, min(width, height) // 2)
|
||||
|
||||
rectangle_coordinates = [
|
||||
((x1 + border_radius, y1), (x2 - border_radius, y2)),
|
||||
((x1, y1 + border_radius), (x2, y2 - border_radius)),
|
||||
]
|
||||
circle_centers = [
|
||||
(x1 + border_radius, y1 + border_radius),
|
||||
(x2 - border_radius, y1 + border_radius),
|
||||
(x1 + border_radius, y2 - border_radius),
|
||||
(x2 - border_radius, y2 - border_radius),
|
||||
]
|
||||
|
||||
for coordinates in rectangle_coordinates:
|
||||
cv2.rectangle(
|
||||
img=scene,
|
||||
pt1=coordinates[0],
|
||||
pt2=coordinates[1],
|
||||
color=color.as_bgr(),
|
||||
thickness=-1,
|
||||
)
|
||||
for center in circle_centers:
|
||||
cv2.circle(
|
||||
img=scene,
|
||||
center=center,
|
||||
radius=border_radius,
|
||||
color=color.as_bgr(),
|
||||
thickness=-1,
|
||||
)
|
||||
return scene
|
||||
|
||||
|
||||
def draw_polygon(
|
||||
scene: np.ndarray, polygon: np.ndarray, color: Color, thickness: int = 2
|
||||
) -> np.ndarray:
|
||||
@@ -243,8 +295,8 @@ def calculate_optimal_text_scale(resolution_wh: Tuple[int, int]) -> float:
|
||||
Calculate font scale based on the resolution of an image.
|
||||
|
||||
Parameters:
|
||||
resolution_wh (Tuple[int, int]): A tuple representing the width and height
|
||||
of the image.
|
||||
resolution_wh (Tuple[int, int]): A tuple representing the width and height
|
||||
of the image.
|
||||
|
||||
Returns:
|
||||
float: The calculated font scale factor.
|
||||
|
||||
@@ -98,6 +98,11 @@ class Rect:
|
||||
width: float
|
||||
height: float
|
||||
|
||||
@classmethod
|
||||
def from_xyxy(cls, xyxy: Tuple[float, float, float, float]) -> Rect:
|
||||
x1, y1, x2, y2 = xyxy
|
||||
return cls(x=x1, y=y1, width=x2 - x1, height=y2 - y1)
|
||||
|
||||
@property
|
||||
def top_left(self) -> Point:
|
||||
return Point(x=self.x, y=self.y)
|
||||
@@ -113,3 +118,11 @@ class Rect:
|
||||
width=self.width + 2 * padding,
|
||||
height=self.height + 2 * padding,
|
||||
)
|
||||
|
||||
def as_xyxy_int_tuple(self) -> Tuple[int, int, int, int]:
|
||||
return (
|
||||
int(self.x),
|
||||
int(self.y),
|
||||
int(self.x + self.width),
|
||||
int(self.y + self.height),
|
||||
)
|
||||
|
||||
@@ -1,15 +1,17 @@
|
||||
from abc import ABC, abstractmethod
|
||||
from logging import warn
|
||||
from typing import List, Optional, Tuple
|
||||
from typing import List, Optional, Tuple, Union
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from supervision import Rect, pad_boxes
|
||||
from supervision.annotators.base import ImageType
|
||||
from supervision.draw.color import Color
|
||||
from supervision.draw.utils import draw_rounded_rectangle
|
||||
from supervision.keypoint.core import KeyPoints
|
||||
from supervision.keypoint.skeletons import SKELETONS_BY_VERTEX_COUNT
|
||||
from supervision.utils.conversion import convert_for_annotation_method
|
||||
from supervision.utils.conversion import ensure_cv2_image_for_annotation
|
||||
|
||||
|
||||
class BaseKeyPointAnnotator(ABC):
|
||||
@@ -39,15 +41,15 @@ class VertexAnnotator(BaseKeyPointAnnotator):
|
||||
self.color = color
|
||||
self.radius = radius
|
||||
|
||||
@convert_for_annotation_method
|
||||
@ensure_cv2_image_for_annotation
|
||||
def annotate(self, scene: ImageType, key_points: KeyPoints) -> ImageType:
|
||||
"""
|
||||
Annotates the given scene with skeleton vertices based on the provided key
|
||||
points. It draws circles at each key point location.
|
||||
|
||||
Args:
|
||||
scene (ImageType): The image where bounding boxes will be drawn. `ImageType`
|
||||
is a flexible type, accepting either `numpy.ndarray` or
|
||||
scene (ImageType): The image where skeleton vertices will be drawn.
|
||||
`ImageType` is a flexible type, accepting either `numpy.ndarray` or
|
||||
`PIL.Image.Image`.
|
||||
key_points (KeyPoints): A collection of key points where each key point
|
||||
consists of x and y coordinates.
|
||||
@@ -63,7 +65,10 @@ class VertexAnnotator(BaseKeyPointAnnotator):
|
||||
image = ...
|
||||
key_points = sv.KeyPoints(...)
|
||||
|
||||
vertex_annotator = sv.VertexAnnotator(color=sv.Color.GREEN, radius=10)
|
||||
vertex_annotator = sv.VertexAnnotator(
|
||||
color=sv.Color.GREEN,
|
||||
radius=10
|
||||
)
|
||||
annotated_frame = vertex_annotator.annotate(
|
||||
scene=image.copy(),
|
||||
key_points=key_points
|
||||
@@ -112,14 +117,14 @@ class EdgeAnnotator(BaseKeyPointAnnotator):
|
||||
self.thickness = thickness
|
||||
self.edges = edges
|
||||
|
||||
@convert_for_annotation_method
|
||||
@ensure_cv2_image_for_annotation
|
||||
def annotate(self, scene: ImageType, key_points: KeyPoints) -> ImageType:
|
||||
"""
|
||||
Annotates the given scene by drawing lines between specified key points to form
|
||||
edges.
|
||||
|
||||
Args:
|
||||
scene (ImageType): The image where bounding boxes will be drawn. `ImageType`
|
||||
scene (ImageType): The image where skeleton edges will be drawn. `ImageType`
|
||||
is a flexible type, accepting either `numpy.ndarray` or
|
||||
`PIL.Image.Image`.
|
||||
key_points (KeyPoints): A collection of key points where each key point
|
||||
@@ -137,7 +142,10 @@ class EdgeAnnotator(BaseKeyPointAnnotator):
|
||||
image = ...
|
||||
key_points = sv.KeyPoints(...)
|
||||
|
||||
edge_annotator = sv.EdgeAnnotator(color=sv.Color.GREEN, thickness=5)
|
||||
edge_annotator = sv.EdgeAnnotator(
|
||||
color=sv.Color.GREEN,
|
||||
thickness=5
|
||||
)
|
||||
annotated_frame = edge_annotator.annotate(
|
||||
scene=image.copy(),
|
||||
key_points=key_points
|
||||
@@ -175,3 +183,236 @@ class EdgeAnnotator(BaseKeyPointAnnotator):
|
||||
)
|
||||
|
||||
return scene
|
||||
|
||||
|
||||
class VertexLabelAnnotator:
|
||||
"""
|
||||
A class that draws labels of skeleton vertices on images. It uses specified key
|
||||
points to determine the locations where the vertices should be drawn.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
color: Union[Color, List[Color]] = Color.ROBOFLOW,
|
||||
text_color: Color = Color.WHITE,
|
||||
text_scale: float = 0.5,
|
||||
text_thickness: int = 1,
|
||||
text_padding: int = 10,
|
||||
border_radius: int = 0,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
color (Union[Color, List[Color]], optional): The color to use for each
|
||||
keypoint label. If a list is provided, the colors will be used in order
|
||||
for each keypoint.
|
||||
text_color (Color, optional): The color to use for the labels.
|
||||
text_scale (float, optional): The scale of the text.
|
||||
text_thickness (int, optional): The thickness of the text.
|
||||
text_padding (int, optional): The padding around the text.
|
||||
border_radius (int, optional): The radius of the rounded corners of the
|
||||
boxes. Set to a high value to produce circles.
|
||||
"""
|
||||
self.border_radius: int = border_radius
|
||||
self.color: Union[Color, List[Color]] = 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
|
||||
|
||||
def annotate(
|
||||
self, scene: ImageType, key_points: KeyPoints, labels: List[str] = None
|
||||
) -> ImageType:
|
||||
"""
|
||||
A class that draws labels of skeleton vertices on images. It uses specified key
|
||||
points to determine the locations where the vertices should be drawn.
|
||||
|
||||
Args:
|
||||
scene (ImageType): The image where vertex labels will be drawn. `ImageType`
|
||||
is a flexible type, accepting either `numpy.ndarray` or
|
||||
`PIL.Image.Image`.
|
||||
key_points (KeyPoints): A collection of key points where each key point
|
||||
consists of x and y coordinates.
|
||||
labels (List[str], optional): A list of labels to be displayed on the
|
||||
annotated image. If not provided, keypoint indices will be used.
|
||||
|
||||
Returns:
|
||||
The annotated image, matching the type of `scene` (`numpy.ndarray`
|
||||
or `PIL.Image.Image`)
|
||||
|
||||
Example:
|
||||
```python
|
||||
import supervision as sv
|
||||
|
||||
image = ...
|
||||
key_points = sv.KeyPoints(...)
|
||||
|
||||
vertex_label_annotator = sv.VertexLabelAnnotator(
|
||||
color=sv.Color.GREEN,
|
||||
text_color=sv.Color.BLACK,
|
||||
border_radius=5
|
||||
)
|
||||
annotated_frame = vertex_label_annotator.annotate(
|
||||
scene=image.copy(),
|
||||
key_points=key_points
|
||||
)
|
||||
```
|
||||
|
||||

|
||||
|
||||
!!! tip
|
||||
|
||||
`VertexLabelAnnotator` allows to customize the color of each keypoint label
|
||||
values.
|
||||
|
||||
Example:
|
||||
```python
|
||||
import supervision as sv
|
||||
|
||||
image = ...
|
||||
key_points = sv.KeyPoints(...)
|
||||
|
||||
LABELS = [
|
||||
"nose", "left eye", "right eye", "left ear",
|
||||
"right ear", "left shoulder", "right shoulder", "left elbow",
|
||||
"right elbow", "left wrist", "right wrist", "left hip",
|
||||
"right hip", "left knee", "right knee", "left ankle",
|
||||
"right ankle"
|
||||
]
|
||||
|
||||
COLORS = [
|
||||
"#FF6347", "#FF6347", "#FF6347", "#FF6347",
|
||||
"#FF6347", "#FF1493", "#00FF00", "#FF1493",
|
||||
"#00FF00", "#FF1493", "#00FF00", "#FFD700",
|
||||
"#00BFFF", "#FFD700", "#00BFFF", "#FFD700",
|
||||
"#00BFFF"
|
||||
]
|
||||
COLORS = [sv.Color.from_hex(color_hex=c) for c in COLORS]
|
||||
|
||||
vertex_label_annotator = sv.VertexLabelAnnotator(
|
||||
color=COLORS,
|
||||
text_color=sv.Color.BLACK,
|
||||
border_radius=5
|
||||
)
|
||||
annotated_frame = vertex_label_annotator.annotate(
|
||||
scene=image.copy(),
|
||||
key_points=key_points,
|
||||
labels=labels
|
||||
)
|
||||
```
|
||||

|
||||
"""
|
||||
font = cv2.FONT_HERSHEY_SIMPLEX
|
||||
|
||||
skeletons_count, points_count, _ = key_points.xy.shape
|
||||
if skeletons_count == 0:
|
||||
return scene
|
||||
|
||||
anchors = key_points.xy.reshape(points_count * skeletons_count, 2).astype(int)
|
||||
mask = np.all(anchors != 0, axis=1)
|
||||
|
||||
if not np.any(mask):
|
||||
return scene
|
||||
|
||||
colors = self.preprocess_and_validate_colors(
|
||||
colors=self.color,
|
||||
points_count=points_count,
|
||||
skeletons_count=skeletons_count,
|
||||
)
|
||||
|
||||
labels = self.preprocess_and_validate_labels(
|
||||
labels=labels, points_count=points_count, skeletons_count=skeletons_count
|
||||
)
|
||||
|
||||
anchors = anchors[mask]
|
||||
colors = colors[mask]
|
||||
labels = labels[mask]
|
||||
|
||||
xyxy = np.array(
|
||||
[
|
||||
self.get_text_bounding_box(
|
||||
text=label,
|
||||
font=font,
|
||||
text_scale=self.text_scale,
|
||||
text_thickness=self.text_thickness,
|
||||
center_coordinates=tuple(anchor),
|
||||
)
|
||||
for anchor, label in zip(anchors, labels)
|
||||
]
|
||||
)
|
||||
|
||||
xyxy_padded = pad_boxes(xyxy=xyxy, px=self.text_padding)
|
||||
|
||||
for text, color, box, box_padded in zip(labels, colors, xyxy, xyxy_padded):
|
||||
draw_rounded_rectangle(
|
||||
scene=scene,
|
||||
rect=Rect.from_xyxy(box_padded),
|
||||
color=color,
|
||||
border_radius=self.border_radius,
|
||||
)
|
||||
cv2.putText(
|
||||
img=scene,
|
||||
text=text,
|
||||
org=(box[0], box[1] + self.text_padding),
|
||||
fontFace=font,
|
||||
fontScale=self.text_scale,
|
||||
color=self.text_color.as_rgb(),
|
||||
thickness=self.text_thickness,
|
||||
lineType=cv2.LINE_AA,
|
||||
)
|
||||
|
||||
return scene
|
||||
|
||||
@staticmethod
|
||||
def get_text_bounding_box(
|
||||
text: str,
|
||||
font: int,
|
||||
text_scale: float,
|
||||
text_thickness: int,
|
||||
center_coordinates: Tuple[int, int],
|
||||
) -> Tuple[int, int, int, int]:
|
||||
text_w, text_h = cv2.getTextSize(
|
||||
text=text,
|
||||
fontFace=font,
|
||||
fontScale=text_scale,
|
||||
thickness=text_thickness,
|
||||
)[0]
|
||||
center_x, center_y = center_coordinates
|
||||
return (
|
||||
center_x - text_w // 2,
|
||||
center_y - text_h // 2,
|
||||
center_x + text_w // 2,
|
||||
center_y + text_h // 2,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def preprocess_and_validate_labels(
|
||||
labels: Optional[List[str]], points_count: int, skeletons_count: int
|
||||
) -> np.array:
|
||||
if labels and len(labels) != points_count:
|
||||
raise ValueError(
|
||||
f"Number of labels ({len(labels)}) must match number of key points "
|
||||
f"({points_count})."
|
||||
)
|
||||
if labels is None:
|
||||
labels = [str(i) for i in range(points_count)]
|
||||
|
||||
return np.array(labels * skeletons_count)
|
||||
|
||||
@staticmethod
|
||||
def preprocess_and_validate_colors(
|
||||
colors: Optional[Union[Color, List[Color]]],
|
||||
points_count: int,
|
||||
skeletons_count: int,
|
||||
) -> np.array:
|
||||
if isinstance(colors, list) and len(colors) != points_count:
|
||||
raise ValueError(
|
||||
f"Number of colors ({len(colors)}) must match number of key points "
|
||||
f"({points_count})."
|
||||
)
|
||||
return (
|
||||
np.array(colors * skeletons_count)
|
||||
if isinstance(colors, list)
|
||||
else np.array([colors] * points_count * skeletons_count)
|
||||
)
|
||||
|
||||
+414
-35
@@ -1,5 +1,6 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from contextlib import suppress
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Dict, Iterator, List, Optional, Tuple, Union
|
||||
|
||||
@@ -17,12 +18,13 @@ class KeyPoints:
|
||||
The `sv.KeyPoints` class in the Supervision library standardizes results from
|
||||
various keypoint detection and pose estimation models into a consistent format. This
|
||||
class simplifies data manipulation and filtering, providing a uniform API for
|
||||
integration with Supervision annotators.
|
||||
integration with Supervision [keypoints annotators](/keypoint/annotators).
|
||||
|
||||
=== "Ultralytics"
|
||||
|
||||
Use [`sv.KeyPoints.from_ultralytics`](/keypoint/core/#supervision.keypoint.core.KeyPoints.from_ultralytics)
|
||||
method, which accepts model results.
|
||||
method, which accepts [YOLOv8](https://github.com/ultralytics/ultralytics)
|
||||
pose result.
|
||||
|
||||
```python
|
||||
import cv2
|
||||
@@ -31,10 +33,60 @@ class KeyPoints:
|
||||
|
||||
image = cv2.imread(<SOURCE_IMAGE_PATH>)
|
||||
model = YOLO('yolov8s-pose.pt')
|
||||
|
||||
result = model(image)[0]
|
||||
key_points = sv.KeyPoints.from_ultralytics(result)
|
||||
```
|
||||
|
||||
=== "Inference"
|
||||
|
||||
Use [`sv.KeyPoints.from_inference`](/keypoint/core/#supervision.keypoint.core.KeyPoints.from_inference)
|
||||
method, which accepts [Inference](https://inference.roboflow.com/) pose result.
|
||||
|
||||
```python
|
||||
import cv2
|
||||
import supervision as sv
|
||||
from inference import get_model
|
||||
|
||||
image = cv2.imread(<SOURCE_IMAGE_PATH>)
|
||||
model = get_model(model_id=<POSE_MODEL_ID>, api_key=<ROBOFLOW_API_KEY>)
|
||||
|
||||
result = model.infer(image)[0]
|
||||
key_points = sv.KeyPoints.from_inference(result)
|
||||
```
|
||||
|
||||
=== "MediaPipe"
|
||||
|
||||
Use [`sv.KeyPoints.from_mediapipe`](/keypoint/core/#supervision.keypoint.core.KeyPoints.from_mediapipe)
|
||||
method, which accepts [MediaPipe](https://github.com/google-ai-edge/mediapipe)
|
||||
pose result.
|
||||
|
||||
```python
|
||||
import cv2
|
||||
import mediapipe as mp
|
||||
import supervision as sv
|
||||
|
||||
image = cv2.imread(<SOURCE_IMAGE_PATH>)
|
||||
image_height, image_width, _ = image.shape
|
||||
mediapipe_image = mp.Image(
|
||||
image_format=mp.ImageFormat.SRGB,
|
||||
data=cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
|
||||
|
||||
options = mp.tasks.vision.PoseLandmarkerOptions(
|
||||
base_options=mp.tasks.BaseOptions(
|
||||
model_asset_path="pose_landmarker_heavy.task"
|
||||
),
|
||||
running_mode=mp.tasks.vision.RunningMode.IMAGE,
|
||||
num_poses=2)
|
||||
|
||||
PoseLandmarker = mp.tasks.vision.PoseLandmarker
|
||||
with PoseLandmarker.create_from_options(options) as landmarker:
|
||||
pose_landmarker_result = landmarker.detect(mediapipe_image)
|
||||
|
||||
key_points = sv.KeyPoints.from_mediapipe(
|
||||
pose_landmarker_result, (image_width, image_height))
|
||||
```
|
||||
|
||||
Attributes:
|
||||
xy (np.ndarray): An array of shape `(n, 2)` containing
|
||||
the bounding boxes coordinates in format `[x1, y1]`
|
||||
@@ -62,7 +114,7 @@ class KeyPoints:
|
||||
|
||||
def __len__(self) -> int:
|
||||
"""
|
||||
Returns the number of keypoints in the keypoints object.
|
||||
Returns the number of keypoints in the `sv.KeyPoints` object.
|
||||
"""
|
||||
return len(self.xy)
|
||||
|
||||
@@ -100,20 +152,228 @@ class KeyPoints:
|
||||
]
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_inference(cls, inference_result: Union[dict, Any]) -> KeyPoints:
|
||||
"""
|
||||
Create a `sv.KeyPoints` object from the [Roboflow](https://roboflow.com/)
|
||||
API inference result or the [Inference](https://inference.roboflow.com/)
|
||||
package results.
|
||||
|
||||
Args:
|
||||
inference_result (dict, any): The result from the
|
||||
Roboflow API or Inference package containing predictions with keypoints.
|
||||
|
||||
Returns:
|
||||
A `sv.KeyPoints` object containing the keypoint coordinates, class IDs,
|
||||
and class names, and confidences of each keypoint.
|
||||
|
||||
Examples:
|
||||
```python
|
||||
import cv2
|
||||
import supervision as sv
|
||||
from inference import get_model
|
||||
|
||||
image = cv2.imread(<SOURCE_IMAGE_PATH>)
|
||||
model = get_model(model_id=<POSE_MODEL_ID>, api_key=<ROBOFLOW_API_KEY>)
|
||||
|
||||
result = model.infer(image)[0]
|
||||
key_points = sv.KeyPoints.from_inference(result)
|
||||
```
|
||||
|
||||
```python
|
||||
import cv2
|
||||
import supervision as sv
|
||||
from inference_sdk import InferenceHTTPClient
|
||||
|
||||
image = cv2.imread(<SOURCE_IMAGE_PATH>)
|
||||
client = InferenceHTTPClient(
|
||||
api_url="https://detect.roboflow.com",
|
||||
api_key=<ROBOFLOW_API_KEY>
|
||||
)
|
||||
|
||||
result = client.infer(image, model_id=<POSE_MODEL_ID>)
|
||||
key_points = sv.KeyPoints.from_inference(result)
|
||||
```
|
||||
"""
|
||||
if isinstance(inference_result, list):
|
||||
raise ValueError(
|
||||
"from_inference() operates on a single result at a time."
|
||||
"You can retrieve it like so: inference_result = model.infer(image)[0]"
|
||||
)
|
||||
|
||||
with suppress(AttributeError):
|
||||
inference_result = inference_result.dict(exclude_none=True, by_alias=True)
|
||||
|
||||
if not inference_result.get("predictions"):
|
||||
return cls.empty()
|
||||
|
||||
xy = []
|
||||
confidence = []
|
||||
class_id = []
|
||||
class_names = []
|
||||
|
||||
for prediction in inference_result["predictions"]:
|
||||
prediction_xy = []
|
||||
prediction_confidence = []
|
||||
for keypoint in prediction["keypoints"]:
|
||||
prediction_xy.append([keypoint["x"], keypoint["y"]])
|
||||
prediction_confidence.append(keypoint["confidence"])
|
||||
xy.append(prediction_xy)
|
||||
confidence.append(prediction_confidence)
|
||||
|
||||
class_id.append(prediction["class_id"])
|
||||
class_names.append(prediction["class"])
|
||||
|
||||
data = {CLASS_NAME_DATA_FIELD: np.array(class_names)}
|
||||
|
||||
return cls(
|
||||
xy=np.array(xy, dtype=np.float32),
|
||||
confidence=np.array(confidence, dtype=np.float32),
|
||||
class_id=np.array(class_id, dtype=int),
|
||||
data=data,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_mediapipe(
|
||||
cls, mediapipe_results, resolution_wh: Tuple[int, int]
|
||||
) -> KeyPoints:
|
||||
"""
|
||||
Creates a `sv.KeyPoints` instance from a
|
||||
[MediaPipe](https://github.com/google-ai-edge/mediapipe)
|
||||
pose landmark detection inference result.
|
||||
|
||||
Args:
|
||||
mediapipe_results (Union[PoseLandmarkerResult, FaceLandmarkerResult, SolutionOutputs]):
|
||||
The output results from Mediapipe. It support pose and face landmarks
|
||||
from `PoseLandmaker`, `FaceLandmarker` and the legacy ones
|
||||
from `Pose` and `FaceMesh`.
|
||||
resolution_wh (Tuple[int, int]): A tuple of the form `(width, height)`
|
||||
representing the resolution of the frame.
|
||||
|
||||
Returns:
|
||||
A `sv.KeyPoints` object containing the keypoint coordinates and
|
||||
confidences of each keypoint.
|
||||
|
||||
!!! tip
|
||||
Before you start, download model bundles from the
|
||||
[MediaPipe website](https://ai.google.dev/edge/mediapipe/solutions/vision/pose_landmarker/index#models).
|
||||
|
||||
Examples:
|
||||
```python
|
||||
import cv2
|
||||
import mediapipe as mp
|
||||
import supervision as sv
|
||||
|
||||
image = cv2.imread(<SOURCE_IMAGE_PATH>)
|
||||
image_height, image_width, _ = image.shape
|
||||
mediapipe_image = mp.Image(
|
||||
image_format=mp.ImageFormat.SRGB,
|
||||
data=cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
|
||||
|
||||
options = mp.tasks.vision.PoseLandmarkerOptions(
|
||||
base_options=mp.tasks.BaseOptions(
|
||||
model_asset_path="pose_landmarker_heavy.task"
|
||||
),
|
||||
running_mode=mp.tasks.vision.RunningMode.IMAGE,
|
||||
num_poses=2)
|
||||
|
||||
PoseLandmarker = mp.tasks.vision.PoseLandmarker
|
||||
with PoseLandmarker.create_from_options(options) as landmarker:
|
||||
pose_landmarker_result = landmarker.detect(mediapipe_image)
|
||||
|
||||
key_points = sv.KeyPoints.from_mediapipe(
|
||||
pose_landmarker_result, (image_width, image_height))
|
||||
```
|
||||
|
||||
```python
|
||||
import cv2
|
||||
import mediapipe as mp
|
||||
import supervision as sv
|
||||
|
||||
image = cv2.imread(<SOURCE_IMAGE_PATH>)
|
||||
image_height, image_width, _ = image.shape
|
||||
mediapipe_image = mp.Image(
|
||||
image_format=mp.ImageFormat.SRGB,
|
||||
data=cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
|
||||
|
||||
options = mp.tasks.vision.FaceLandmarkerOptions(
|
||||
base_options=mp.tasks.BaseOptions(
|
||||
model_asset_path="face_landmarker.task"
|
||||
),
|
||||
output_face_blendshapes=True,
|
||||
output_facial_transformation_matrixes=True,
|
||||
num_faces=2)
|
||||
|
||||
FaceLandmarker = mp.tasks.vision.FaceLandmarker
|
||||
with FaceLandmarker.create_from_options(options) as landmarker:
|
||||
face_landmarker_result = landmarker.detect(mediapipe_image)
|
||||
|
||||
key_points = sv.KeyPoints.from_mediapipe(
|
||||
face_landmarker_result, (image_width, image_height))
|
||||
```
|
||||
""" # noqa: E501 // docs
|
||||
if hasattr(mediapipe_results, "pose_landmarks"):
|
||||
results = mediapipe_results.pose_landmarks
|
||||
if not isinstance(mediapipe_results.pose_landmarks, list):
|
||||
if mediapipe_results.pose_landmarks is None:
|
||||
results = []
|
||||
else:
|
||||
results = [
|
||||
[
|
||||
landmark
|
||||
for landmark in mediapipe_results.pose_landmarks.landmark
|
||||
]
|
||||
]
|
||||
elif hasattr(mediapipe_results, "face_landmarks"):
|
||||
results = mediapipe_results.face_landmarks
|
||||
elif hasattr(mediapipe_results, "multi_face_landmarks"):
|
||||
if mediapipe_results.multi_face_landmarks is None:
|
||||
results = []
|
||||
else:
|
||||
results = [
|
||||
face_landmark.landmark
|
||||
for face_landmark in mediapipe_results.multi_face_landmarks
|
||||
]
|
||||
|
||||
if len(results) == 0:
|
||||
return cls.empty()
|
||||
|
||||
xy = []
|
||||
confidence = []
|
||||
for pose in results:
|
||||
prediction_xy = []
|
||||
prediction_confidence = []
|
||||
for landmark in pose:
|
||||
keypoint_xy = [
|
||||
landmark.x * resolution_wh[0],
|
||||
landmark.y * resolution_wh[1],
|
||||
]
|
||||
prediction_xy.append(keypoint_xy)
|
||||
prediction_confidence.append(landmark.visibility)
|
||||
|
||||
xy.append(prediction_xy)
|
||||
confidence.append(prediction_confidence)
|
||||
|
||||
return cls(
|
||||
xy=np.array(xy, dtype=np.float32),
|
||||
confidence=np.array(confidence, dtype=np.float32),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_ultralytics(cls, ultralytics_results) -> KeyPoints:
|
||||
"""
|
||||
Creates a Keypoints instance from a
|
||||
[YOLOv8](https://github.com/ultralytics/ultralytics) inference result.
|
||||
Creates a `sv.KeyPoints` instance from a
|
||||
[YOLOv8](https://github.com/ultralytics/ultralytics) pose inference result.
|
||||
|
||||
Args:
|
||||
ultralytics_results (ultralytics.engine.results.Keypoints):
|
||||
The output Results instance from YOLOv8
|
||||
|
||||
Returns:
|
||||
KeyPoints: A new Keypoints object.
|
||||
A `sv.KeyPoints` object containing the keypoint coordinates, class IDs,
|
||||
and class names, and confidences of each keypoint.
|
||||
|
||||
Example:
|
||||
Examples:
|
||||
```python
|
||||
import cv2
|
||||
import supervision as sv
|
||||
@@ -121,8 +381,9 @@ class KeyPoints:
|
||||
|
||||
image = cv2.imread(<SOURCE_IMAGE_PATH>)
|
||||
model = YOLO('yolov8s-pose.pt')
|
||||
|
||||
result = model(image)[0]
|
||||
keypoints = sv.KeyPoints.from_ultralytics(result)
|
||||
key_points = sv.KeyPoints.from_ultralytics(result)
|
||||
```
|
||||
"""
|
||||
if ultralytics_results.keypoints.xy.numel() == 0:
|
||||
@@ -136,38 +397,155 @@ class KeyPoints:
|
||||
data = {CLASS_NAME_DATA_FIELD: class_names}
|
||||
return cls(xy, class_id, confidence, data)
|
||||
|
||||
def __getitem__(
|
||||
self, index: Union[int, slice, List[int], np.ndarray, str]
|
||||
) -> Union["KeyPoints", List, np.ndarray, None]:
|
||||
@classmethod
|
||||
def from_yolo_nas(cls, yolo_nas_results) -> KeyPoints:
|
||||
"""
|
||||
Get a subset of the KeyPoints object or access an item from its data field.
|
||||
|
||||
When provided with an integer, slice, list of integers, or a numpy array, this
|
||||
method returns a new KeyPoints object that represents a subset of the original
|
||||
keypoints. When provided with a string, it accesses the corresponding item in
|
||||
the data dictionary.
|
||||
Create a `sv.KeyPoints` instance from a [YOLO-NAS](https://github.com/Deci-AI/super-gradients/blob/master/YOLONAS-POSE.md)
|
||||
pose inference results.
|
||||
|
||||
Args:
|
||||
index (Union[int, slice, List[int], np.ndarray, str]): The index, indices,
|
||||
or key to access a subset of the KeyPoints or an item from the data.
|
||||
yolo_nas_results (ImagePoseEstimationPrediction): The output object from
|
||||
YOLO NAS.
|
||||
|
||||
Returns:
|
||||
Union[KeyPoints, Any]: A subset of the KeyPoints object or an item from
|
||||
the data field.
|
||||
A `sv.KeyPoints` object containing the keypoint coordinates, class IDs,
|
||||
and class names, and confidences of each keypoint.
|
||||
|
||||
Examples:
|
||||
```python
|
||||
import cv2
|
||||
import torch
|
||||
import supervision as sv
|
||||
import super_gradients
|
||||
|
||||
image = cv2.imread(<SOURCE_IMAGE_PATH>)
|
||||
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
model = super_gradients.training.models.get(
|
||||
"yolo_nas_pose_s", pretrained_weights="coco_pose").to(device)
|
||||
|
||||
results = model.predict(image, conf=0.1)
|
||||
key_points = sv.KeyPoints.from_yolo_nas(results)
|
||||
```
|
||||
""" # noqa: E501 // docs
|
||||
if len(yolo_nas_results.prediction.poses) == 0:
|
||||
return cls.empty()
|
||||
|
||||
xy = yolo_nas_results.prediction.poses[:, :, :2]
|
||||
confidence = yolo_nas_results.prediction.poses[:, :, 2]
|
||||
|
||||
# yolo_nas_results treats params differently.
|
||||
# prediction.labels may not exist, whereas class_names might be None
|
||||
if hasattr(yolo_nas_results.prediction, "labels"):
|
||||
class_id = yolo_nas_results.prediction.labels # np.array[int]
|
||||
else:
|
||||
class_id = None
|
||||
|
||||
data = {}
|
||||
if class_id is not None and yolo_nas_results.class_names is not None:
|
||||
class_names = []
|
||||
for c_id in class_id:
|
||||
name = yolo_nas_results.class_names[c_id] # tuple[str]
|
||||
class_names.append(name)
|
||||
data[CLASS_NAME_DATA_FIELD] = class_names
|
||||
|
||||
return cls(
|
||||
xy=xy,
|
||||
confidence=confidence,
|
||||
class_id=class_id,
|
||||
data=data,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_detectron2(cls, detectron2_results) -> KeyPoints:
|
||||
"""
|
||||
Create a `sv.KeyPoints` 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.
|
||||
|
||||
Returns:
|
||||
A `sv.KeyPoints` object containing the keypoint coordinates, class IDs,
|
||||
and class names, and confidences of each keypoint.
|
||||
|
||||
Example:
|
||||
```python
|
||||
import cv2
|
||||
import supervision as sv
|
||||
from detectron2.engine import DefaultPredictor
|
||||
from detectron2.config import get_cfg
|
||||
|
||||
|
||||
image = cv2.imread(<SOURCE_IMAGE_PATH>)
|
||||
cfg = get_cfg()
|
||||
cfg.merge_from_file(<CONFIG_PATH>)
|
||||
cfg.MODEL.WEIGHTS = <WEIGHTS_PATH>
|
||||
predictor = DefaultPredictor(cfg)
|
||||
|
||||
result = predictor(image)
|
||||
keypoints = sv.KeyPoints.from_detectron2(result)
|
||||
```
|
||||
"""
|
||||
|
||||
if hasattr(detectron2_results["instances"], "pred_keypoints"):
|
||||
if detectron2_results["instances"].pred_keypoints.cpu().numpy().size == 0:
|
||||
return cls.empty()
|
||||
return cls(
|
||||
xy=detectron2_results["instances"]
|
||||
.pred_keypoints.cpu()
|
||||
.numpy()[:, :, :2],
|
||||
confidence=detectron2_results["instances"]
|
||||
.pred_keypoints.cpu()
|
||||
.numpy()[:, :, 2],
|
||||
class_id=detectron2_results["instances"]
|
||||
.pred_classes.cpu()
|
||||
.numpy()
|
||||
.astype(int),
|
||||
)
|
||||
else:
|
||||
return cls.empty()
|
||||
|
||||
def __getitem__(
|
||||
self, index: Union[int, slice, List[int], np.ndarray, str]
|
||||
) -> Union[KeyPoints, List, np.ndarray, None]:
|
||||
"""
|
||||
Get a subset of the `sv.KeyPoints` object or access an item from its data field.
|
||||
|
||||
When provided with an integer, slice, list of integers, or a numpy array, this
|
||||
method returns a new `sv.KeyPoints` object that represents a subset of the
|
||||
original `sv.KeyPoints`. When provided with a string, it accesses the
|
||||
corresponding item in the data dictionary.
|
||||
|
||||
Args:
|
||||
index (Union[int, slice, List[int], np.ndarray, str]): The index, indices,
|
||||
or key to access a subset of the `sv.KeyPoints` or an item from the
|
||||
data.
|
||||
|
||||
Returns:
|
||||
A subset of the `sv.KeyPoints` object or an item from the data field.
|
||||
|
||||
Examples:
|
||||
```python
|
||||
import supervision as sv
|
||||
|
||||
keypoints = sv.KeyPoints()
|
||||
key_points = sv.KeyPoints()
|
||||
|
||||
first_detection = keypoints[0]
|
||||
first_10_keypoints = keypoints[0:10]
|
||||
some_keypoints = keypoints[[0, 2, 4]]
|
||||
class_0_keypoints = keypoints[keypoints.class_id == 0]
|
||||
high_confidence_keypoints = keypoints[keypoints.confidence > 0.5]
|
||||
# access the first keypoint using an integer index
|
||||
key_points[0]
|
||||
|
||||
feature_vector = keypoints['feature_vector']
|
||||
# access the first 10 keypoints using index slice
|
||||
key_points[0:10]
|
||||
|
||||
# access selected keypoints using a list of indices
|
||||
key_points[[0, 2, 4]]
|
||||
|
||||
# access keypoints with selected class_id
|
||||
key_points[key_points.class_id == 0]
|
||||
|
||||
# access keypoints with confidence greater than 0.5
|
||||
key_points[key_points.confidence > 0.5]
|
||||
```
|
||||
"""
|
||||
if isinstance(index, str):
|
||||
@@ -183,13 +561,13 @@ class KeyPoints:
|
||||
|
||||
def __setitem__(self, key: str, value: Union[np.ndarray, List]):
|
||||
"""
|
||||
Set a value in the data dictionary of the KeyPoints object.
|
||||
Set a value in the data dictionary of the `sv.KeyPoints` object.
|
||||
|
||||
Args:
|
||||
key (str): The key in the data dictionary to set.
|
||||
value (Union[np.ndarray, List]): The value to set for the key.
|
||||
|
||||
Example:
|
||||
Examples:
|
||||
```python
|
||||
import cv2
|
||||
import supervision as sv
|
||||
@@ -201,7 +579,7 @@ class KeyPoints:
|
||||
result = model(image)[0]
|
||||
keypoints = sv.KeyPoints.from_ultralytics(result)
|
||||
|
||||
keypoints['names'] = [
|
||||
keypoints['class_name'] = [
|
||||
model.model.names[class_id]
|
||||
for class_id
|
||||
in keypoints.class_id
|
||||
@@ -222,12 +600,13 @@ class KeyPoints:
|
||||
Create an empty Keypoints object with no keypoints.
|
||||
|
||||
Returns:
|
||||
(KeyPoints): An empty Keypoints object.
|
||||
An empty `sv.KeyPoints` object.
|
||||
|
||||
Example:
|
||||
Examples:
|
||||
```python
|
||||
from supervision import Keypoints
|
||||
empty_keypoints = Keypoints.empty()
|
||||
import supervision as sv
|
||||
|
||||
key_points = sv.KeyPoints.empty()
|
||||
```
|
||||
"""
|
||||
return cls(xy=np.empty((0, 0, 2), dtype=np.float32))
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -12,8 +12,9 @@ from supervision.utils.internal import deprecated_parameter
|
||||
|
||||
class STrack(BaseTrack):
|
||||
shared_kalman = KalmanFilter()
|
||||
_external_count = 0
|
||||
|
||||
def __init__(self, tlwh, score, class_ids):
|
||||
def __init__(self, tlwh, score, class_ids, minimum_consecutive_frames):
|
||||
# wait activate
|
||||
self._tlwh = np.asarray(tlwh, dtype=np.float32)
|
||||
self.kalman_filter = None
|
||||
@@ -24,6 +25,10 @@ class STrack(BaseTrack):
|
||||
self.class_ids = class_ids
|
||||
self.tracklet_len = 0
|
||||
|
||||
self.external_track_id = -1
|
||||
|
||||
self.minimum_consecutive_frames = minimum_consecutive_frames
|
||||
|
||||
def predict(self):
|
||||
mean_state = self.mean.copy()
|
||||
if self.state != TrackState.Tracked:
|
||||
@@ -53,7 +58,7 @@ class STrack(BaseTrack):
|
||||
def activate(self, kalman_filter, frame_id):
|
||||
"""Start a new tracklet"""
|
||||
self.kalman_filter = kalman_filter
|
||||
self.track_id = self.next_id()
|
||||
self.internal_track_id = self.next_id()
|
||||
self.mean, self.covariance = self.kalman_filter.initiate(
|
||||
self.tlwh_to_xyah(self._tlwh)
|
||||
)
|
||||
@@ -62,6 +67,10 @@ class STrack(BaseTrack):
|
||||
self.state = TrackState.Tracked
|
||||
if frame_id == 1:
|
||||
self.is_activated = True
|
||||
|
||||
if self.minimum_consecutive_frames == 1:
|
||||
self.external_track_id = self.next_external_id()
|
||||
|
||||
self.frame_id = frame_id
|
||||
self.start_frame = frame_id
|
||||
|
||||
@@ -71,10 +80,10 @@ class STrack(BaseTrack):
|
||||
)
|
||||
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.internal_track_id = self.next_id()
|
||||
self.score = new_track.score
|
||||
|
||||
def update(self, new_track, frame_id):
|
||||
@@ -93,7 +102,10 @@ class STrack(BaseTrack):
|
||||
self.mean, self.covariance, self.tlwh_to_xyah(new_tlwh)
|
||||
)
|
||||
self.state = TrackState.Tracked
|
||||
self.is_activated = True
|
||||
if self.tracklet_len == self.minimum_consecutive_frames:
|
||||
self.is_activated = True
|
||||
if self.external_track_id == -1:
|
||||
self.external_track_id = self.next_external_id()
|
||||
|
||||
self.score = new_track.score
|
||||
|
||||
@@ -131,6 +143,15 @@ class STrack(BaseTrack):
|
||||
def to_xyah(self):
|
||||
return self.tlwh_to_xyah(self.tlwh)
|
||||
|
||||
@staticmethod
|
||||
def next_external_id():
|
||||
STrack._external_count += 1
|
||||
return STrack._external_count
|
||||
|
||||
@staticmethod
|
||||
def reset_external_counter():
|
||||
STrack._external_count = 0
|
||||
|
||||
@staticmethod
|
||||
def tlbr_to_tlwh(tlbr):
|
||||
ret = np.asarray(tlbr).copy()
|
||||
@@ -144,7 +165,9 @@ class STrack(BaseTrack):
|
||||
return ret
|
||||
|
||||
def __repr__(self):
|
||||
return "OT_{}_({}-{})".format(self.track_id, self.start_frame, self.end_frame)
|
||||
return "OT_{}_({}-{})".format(
|
||||
self.internal_track_id, self.start_frame, self.end_frame
|
||||
)
|
||||
|
||||
|
||||
def detections2boxes(detections: Detections) -> np.ndarray:
|
||||
@@ -186,6 +209,10 @@ class ByteTrack:
|
||||
Increasing minimum_matching_threshold improves accuracy but risks fragmentation.
|
||||
Decreasing it improves completeness but risks false positives and drift.
|
||||
frame_rate (int, optional): The frame rate of the video.
|
||||
minimum_consecutive_frames (int, optional): Number of consecutive frames that an object must
|
||||
be tracked before it is considered a 'valid' track.
|
||||
Increasing minimum_consecutive_frames prevents the creation of accidental tracks from
|
||||
false detection or double detection, but risks missing shorter tracks.
|
||||
""" # noqa: E501 // docs
|
||||
|
||||
@deprecated_parameter(
|
||||
@@ -218,6 +245,7 @@ class ByteTrack:
|
||||
lost_track_buffer: int = 30,
|
||||
minimum_matching_threshold: float = 0.8,
|
||||
frame_rate: int = 30,
|
||||
minimum_consecutive_frames: int = 1,
|
||||
):
|
||||
self.track_activation_threshold = track_activation_threshold
|
||||
self.minimum_matching_threshold = minimum_matching_threshold
|
||||
@@ -225,6 +253,7 @@ class ByteTrack:
|
||||
self.frame_id = 0
|
||||
self.det_thresh = self.track_activation_threshold + 0.1
|
||||
self.max_time_lost = int(frame_rate / 30.0 * lost_track_buffer)
|
||||
self.minimum_consecutive_frames = minimum_consecutive_frames
|
||||
self.kalman_filter = KalmanFilter()
|
||||
|
||||
self.tracked_tracks: List[STrack] = []
|
||||
@@ -247,7 +276,7 @@ class ByteTrack:
|
||||
model = YOLO(<MODEL_PATH>)
|
||||
tracker = sv.ByteTrack()
|
||||
|
||||
bounding_box_annotator = sv.BoundingBoxAnnotator()
|
||||
box_annotator = sv.BoxAnnotator()
|
||||
label_annotator = sv.LabelAnnotator()
|
||||
|
||||
def callback(frame: np.ndarray, index: int) -> np.ndarray:
|
||||
@@ -257,7 +286,7 @@ class ByteTrack:
|
||||
|
||||
labels = [f"#{tracker_id}" for tracker_id in detections.tracker_id]
|
||||
|
||||
annotated_frame = bounding_box_annotator.annotate(
|
||||
annotated_frame = box_annotator.annotate(
|
||||
scene=frame.copy(), detections=detections)
|
||||
annotated_frame = label_annotator.annotate(
|
||||
scene=annotated_frame, detections=detections, labels=labels)
|
||||
@@ -285,11 +314,14 @@ class ByteTrack:
|
||||
matches, _, _ = matching.linear_assignment(iou_costs, 0.5)
|
||||
detections.tracker_id = np.full(len(detections), -1, dtype=int)
|
||||
for i_detection, i_track in matches:
|
||||
detections.tracker_id[i_detection] = int(tracks[i_track].track_id)
|
||||
detections.tracker_id[i_detection] = int(
|
||||
tracks[i_track].external_track_id
|
||||
)
|
||||
|
||||
return detections[detections.tracker_id != -1]
|
||||
|
||||
else:
|
||||
detections = Detections.empty()
|
||||
detections.tracker_id = np.array([], dtype=int)
|
||||
|
||||
return detections
|
||||
@@ -308,6 +340,7 @@ class ByteTrack:
|
||||
self.lost_tracks: List[STrack] = []
|
||||
self.removed_tracks: List[STrack] = []
|
||||
BaseTrack.reset_counter()
|
||||
STrack.reset_external_counter()
|
||||
|
||||
def update_with_tensors(self, tensors: np.ndarray) -> List[STrack]:
|
||||
"""
|
||||
@@ -345,7 +378,7 @@ class ByteTrack:
|
||||
if len(dets) > 0:
|
||||
"""Detections"""
|
||||
detections = [
|
||||
STrack(STrack.tlbr_to_tlwh(tlbr), s, c)
|
||||
STrack(STrack.tlbr_to_tlwh(tlbr), s, c, self.minimum_consecutive_frames)
|
||||
for (tlbr, s, c) in zip(dets, scores_keep, class_ids_keep)
|
||||
]
|
||||
else:
|
||||
@@ -387,7 +420,7 @@ class ByteTrack:
|
||||
if len(dets_second) > 0:
|
||||
"""Detections"""
|
||||
detections_second = [
|
||||
STrack(STrack.tlbr_to_tlwh(tlbr), s, c)
|
||||
STrack(STrack.tlbr_to_tlwh(tlbr), s, c, self.minimum_consecutive_frames)
|
||||
for (tlbr, s, c) in zip(dets_second, scores_second, class_ids_second)
|
||||
]
|
||||
else:
|
||||
@@ -454,7 +487,7 @@ class ByteTrack:
|
||||
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.removed_tracks = removed_stracks
|
||||
self.tracked_tracks, self.lost_tracks = remove_duplicate_tracks(
|
||||
self.tracked_tracks, self.lost_tracks
|
||||
)
|
||||
@@ -468,22 +501,22 @@ def joint_tracks(
|
||||
) -> List[STrack]:
|
||||
"""
|
||||
Joins two lists of tracks, ensuring that the resulting list does not
|
||||
contain tracks with duplicate track_id values.
|
||||
contain tracks with duplicate internal_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).
|
||||
track_list_a: First list of tracks (with internal_track_id attribute).
|
||||
track_list_b: Second list of tracks (with internal_track_id attribute).
|
||||
|
||||
Returns:
|
||||
Combined list of tracks from track_list_a and track_list_b
|
||||
without duplicate track_id values.
|
||||
without duplicate internal_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)
|
||||
if track.internal_track_id not in seen_track_ids:
|
||||
seen_track_ids.add(track.internal_track_id)
|
||||
result.append(track)
|
||||
|
||||
return result
|
||||
@@ -492,17 +525,17 @@ def joint_tracks(
|
||||
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.
|
||||
that share the same internal_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
|
||||
track_list_a: List of tracks (with internal_track_id attribute).
|
||||
track_list_b: List of tracks (with internal_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}
|
||||
tracks = {track.internal_track_id: track for track in track_list_a}
|
||||
track_ids_b = {track.internal_track_id for track in track_list_b}
|
||||
|
||||
for track_id in track_ids_b:
|
||||
tracks.pop(track_id, None)
|
||||
|
||||
@@ -8,11 +8,13 @@ from PIL import Image
|
||||
from supervision.annotators.base import ImageType
|
||||
|
||||
|
||||
def convert_for_annotation_method(annotate_func):
|
||||
def ensure_cv2_image_for_annotation(annotate_func):
|
||||
"""
|
||||
Decorates `BaseAnnotator.annotate` implementations, converts scene to
|
||||
an image type used internally by the annotators, converts back when annotation
|
||||
is complete.
|
||||
|
||||
Assumes the annotators modify the scene in-place.
|
||||
"""
|
||||
|
||||
@wraps(annotate_func)
|
||||
@@ -21,19 +23,22 @@ def convert_for_annotation_method(annotate_func):
|
||||
return annotate_func(self, scene, *args, **kwargs)
|
||||
|
||||
if isinstance(scene, Image.Image):
|
||||
scene = pillow_to_cv2(scene)
|
||||
annotated = annotate_func(self, scene, *args, **kwargs)
|
||||
return cv2_to_pillow(image=annotated)
|
||||
scene_np = pillow_to_cv2(scene)
|
||||
annotated_np = annotate_func(self, scene_np, *args, **kwargs)
|
||||
scene.paste(cv2_to_pillow(annotated_np))
|
||||
return scene
|
||||
|
||||
raise ValueError(f"Unsupported image type: {type(scene)}")
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
def convert_for_image_processing(image_processing_fun):
|
||||
def ensure_cv2_image_for_processing(image_processing_fun):
|
||||
"""
|
||||
Decorates image processing functions that accept np.ndarray, converting `image` to
|
||||
np.ndarray, converts back when processing is complete.
|
||||
|
||||
Assumes the annotators do NOT modify the scene in-place.
|
||||
"""
|
||||
|
||||
@wraps(image_processing_fun)
|
||||
@@ -44,13 +49,37 @@ def convert_for_image_processing(image_processing_fun):
|
||||
if isinstance(image, Image.Image):
|
||||
scene = pillow_to_cv2(image)
|
||||
annotated = image_processing_fun(scene, *args, **kwargs)
|
||||
return cv2_to_pillow(image=annotated)
|
||||
return cv2_to_pillow(annotated)
|
||||
|
||||
raise ValueError(f"Unsupported image type: {type(image)}")
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
def ensure_pil_image_for_annotation(annotate_func):
|
||||
"""
|
||||
Decorates image processing functions that accept np.ndarray, converting `image` to
|
||||
PIL image, converts back when processing is complete.
|
||||
|
||||
Assumes the annotators modify the scene in-place.
|
||||
"""
|
||||
|
||||
@wraps(annotate_func)
|
||||
def wrapper(self, scene: ImageType, *args, **kwargs):
|
||||
if isinstance(scene, np.ndarray):
|
||||
scene_pil = cv2_to_pillow(scene)
|
||||
annotated_pil = annotate_func(self, scene_pil, *args, **kwargs)
|
||||
np.copyto(scene, pillow_to_cv2(annotated_pil))
|
||||
return scene
|
||||
|
||||
if isinstance(scene, Image.Image):
|
||||
return annotate_func(self, scene, *args, **kwargs)
|
||||
|
||||
raise ValueError(f"Unsupported image type: {type(scene)}")
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
def images_to_cv2(images: List[ImageType]) -> List[np.ndarray]:
|
||||
"""
|
||||
Converts images provided either as Pillow images or OpenCV
|
||||
@@ -67,7 +96,7 @@ def images_to_cv2(images: List[ImageType]) -> List[np.ndarray]:
|
||||
result = []
|
||||
for image in images:
|
||||
if issubclass(type(image), Image.Image):
|
||||
image = pillow_to_cv2(image=image)
|
||||
image = pillow_to_cv2(image)
|
||||
result.append(image)
|
||||
return result
|
||||
|
||||
|
||||
+19
-19
@@ -1,6 +1,6 @@
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import List, Optional, Union
|
||||
from typing import Dict, List, Optional, Union
|
||||
|
||||
import numpy as np
|
||||
import yaml
|
||||
@@ -57,20 +57,20 @@ def list_files_with_extensions(
|
||||
return files_with_extensions
|
||||
|
||||
|
||||
def read_txt_file(file_path: str, skip_empty: bool = False) -> List[str]:
|
||||
def read_txt_file(file_path: Union[str, Path], skip_empty: bool = False) -> List[str]:
|
||||
"""
|
||||
Read a text file and return a list of strings without newline characters.
|
||||
Optionally skip empty lines.
|
||||
|
||||
Args:
|
||||
file_path (str): The path to the text file.
|
||||
file_path (Union[str, Path]): The file path as a string or Path object.
|
||||
skip_empty (bool): If True, skip lines that are empty or contain only
|
||||
whitespace. Default is False.
|
||||
|
||||
Returns:
|
||||
List[str]: A list of strings representing the lines in the text file.
|
||||
"""
|
||||
with open(file_path, "r") as file:
|
||||
with open(str(file_path), "r") as file:
|
||||
if skip_empty:
|
||||
lines = [line.rstrip("\n") for line in file if line.strip()]
|
||||
else:
|
||||
@@ -79,71 +79,71 @@ def read_txt_file(file_path: str, skip_empty: bool = False) -> List[str]:
|
||||
return lines
|
||||
|
||||
|
||||
def save_text_file(lines: List[str], file_path: str):
|
||||
def save_text_file(lines: List[str], file_path: Union[str, Path]) -> None:
|
||||
"""
|
||||
Write a list of strings to a text file, each string on a new line.
|
||||
|
||||
Args:
|
||||
lines (List[str]): The list of strings to be written to the file.
|
||||
file_path (str): The path to the text file.
|
||||
file_path (Union[str, Path]): The file path as a string or Path object.
|
||||
"""
|
||||
with open(file_path, "w") as file:
|
||||
with open(str(file_path), "w") as file:
|
||||
for line in lines:
|
||||
file.write(line + "\n")
|
||||
|
||||
|
||||
def read_json_file(file_path: str) -> dict:
|
||||
def read_json_file(file_path: Union[str, Path]) -> Dict:
|
||||
"""
|
||||
Read a json file and return a dict.
|
||||
|
||||
Args:
|
||||
file_path (str): The path to the json file.
|
||||
file_path (Union[str, Path]): The file path as a string or Path object.
|
||||
|
||||
Returns:
|
||||
dict: A dict of annotations information
|
||||
"""
|
||||
with open(file_path, "r") as file:
|
||||
with open(str(file_path), "r") as file:
|
||||
data = json.load(file)
|
||||
return data
|
||||
|
||||
|
||||
def save_json_file(data: dict, file_path: str, indent: int = 3) -> None:
|
||||
def save_json_file(data: dict, file_path: Union[str, Path], 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.
|
||||
file_path (Union[str, Path]): The file path as a string or Path object.
|
||||
"""
|
||||
with open(file_path, "w") as fp:
|
||||
with open(str(file_path), "w") as fp:
|
||||
json.dump(data, fp, cls=NumpyJsonEncoder, indent=indent)
|
||||
|
||||
|
||||
def read_yaml_file(file_path: str) -> dict:
|
||||
def read_yaml_file(file_path: Union[str, Path]) -> Dict:
|
||||
"""
|
||||
Read a yaml file and return a dict.
|
||||
|
||||
Args:
|
||||
file_path (str): The path to the yaml file.
|
||||
file_path (Union[str, Path]): The file path as a string or Path object.
|
||||
|
||||
Returns:
|
||||
dict: A dict of content information
|
||||
"""
|
||||
with open(file_path, "r") as file:
|
||||
with open(str(file_path), "r") as file:
|
||||
data = yaml.safe_load(file)
|
||||
return data
|
||||
|
||||
|
||||
def save_yaml_file(data: dict, file_path: str) -> None:
|
||||
def save_yaml_file(data: dict, file_path: Union[str, Path]) -> 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.
|
||||
file_path (Union[str, Path]): The file path as a string or Path object.
|
||||
"""
|
||||
|
||||
with open(file_path, "w") as outfile:
|
||||
with open(str(file_path), "w") as outfile:
|
||||
yaml.dump(data, outfile, sort_keys=False, default_flow_style=None)
|
||||
|
||||
@@ -14,8 +14,8 @@ from supervision.draw.color import Color, unify_to_bgr
|
||||
from supervision.draw.utils import calculate_optimal_text_scale, draw_text
|
||||
from supervision.geometry.core import Point
|
||||
from supervision.utils.conversion import (
|
||||
convert_for_image_processing,
|
||||
cv2_to_pillow,
|
||||
ensure_cv2_image_for_processing,
|
||||
images_to_cv2,
|
||||
)
|
||||
from supervision.utils.iterables import create_batches, fill
|
||||
@@ -25,7 +25,7 @@ RelativePosition = Literal["top", "bottom"]
|
||||
MAX_COLUMNS_FOR_SINGLE_ROW_GRID = 3
|
||||
|
||||
|
||||
@convert_for_image_processing
|
||||
@ensure_cv2_image_for_processing
|
||||
def crop_image(
|
||||
image: ImageType,
|
||||
xyxy: Union[npt.NDArray[int], List[int], Tuple[int, int, int, int]],
|
||||
@@ -86,7 +86,7 @@ def crop_image(
|
||||
return image[y_min:y_max, x_min:x_max]
|
||||
|
||||
|
||||
@convert_for_image_processing
|
||||
@ensure_cv2_image_for_processing
|
||||
def scale_image(image: ImageType, scale_factor: float) -> ImageType:
|
||||
"""
|
||||
Scales the given image based on the given scale factor.
|
||||
@@ -143,7 +143,7 @@ def scale_image(image: ImageType, scale_factor: float) -> ImageType:
|
||||
return cv2.resize(image, (width_new, height_new), interpolation=cv2.INTER_LINEAR)
|
||||
|
||||
|
||||
@convert_for_image_processing
|
||||
@ensure_cv2_image_for_processing
|
||||
def resize_image(
|
||||
image: ImageType,
|
||||
resolution_wh: Tuple[int, int],
|
||||
@@ -216,7 +216,7 @@ def resize_image(
|
||||
return cv2.resize(image, (width_new, height_new), interpolation=cv2.INTER_LINEAR)
|
||||
|
||||
|
||||
@convert_for_image_processing
|
||||
@ensure_cv2_image_for_processing
|
||||
def letterbox_image(
|
||||
image: ImageType,
|
||||
resolution_wh: Tuple[int, int],
|
||||
|
||||
@@ -1,7 +1,8 @@
|
||||
import functools
|
||||
import inspect
|
||||
import os
|
||||
import warnings
|
||||
from typing import Callable
|
||||
from typing import Any, Callable, Set
|
||||
|
||||
|
||||
class SupervisionWarnings(Warning):
|
||||
@@ -30,6 +31,16 @@ else:
|
||||
warnings.simplefilter("always", SupervisionWarnings)
|
||||
|
||||
|
||||
def warn_deprecated(message: str):
|
||||
"""
|
||||
Issue a warning that a function is deprecated.
|
||||
|
||||
Args:
|
||||
message (str): The message to display when the function is called.
|
||||
"""
|
||||
warnings.warn(message, category=SupervisionWarnings, stacklevel=2)
|
||||
|
||||
|
||||
def deprecated_parameter(
|
||||
old_parameter: str,
|
||||
new_parameter: str,
|
||||
@@ -81,15 +92,13 @@ def deprecated_parameter(
|
||||
else:
|
||||
function_name = func.__name__
|
||||
|
||||
warnings.warn(
|
||||
warn_deprecated(
|
||||
message=warning_message.format(
|
||||
function_name=function_name,
|
||||
old_parameter=old_parameter,
|
||||
new_parameter=new_parameter,
|
||||
**message_kwargs,
|
||||
),
|
||||
category=SupervisionWarnings,
|
||||
stacklevel=2,
|
||||
)
|
||||
)
|
||||
|
||||
kwargs[new_parameter] = map_function(kwargs.pop(old_parameter))
|
||||
@@ -105,11 +114,7 @@ def deprecated(reason: str):
|
||||
def decorator(func):
|
||||
@functools.wraps(func)
|
||||
def wrapper(*args, **kwargs):
|
||||
warnings.warn(
|
||||
f"{func.__name__} is deprecated: {reason}",
|
||||
category=SupervisionWarnings,
|
||||
stacklevel=2,
|
||||
)
|
||||
warn_deprecated(f"{func.__name__} is deprecated: {reason}")
|
||||
return func(*args, **kwargs)
|
||||
|
||||
return wrapper
|
||||
@@ -141,3 +146,42 @@ class classproperty(property):
|
||||
The result of calling the function stored in 'fget' with 'owner_cls'.
|
||||
"""
|
||||
return self.fget(owner_cls)
|
||||
|
||||
|
||||
def get_instance_variables(instance: Any, include_properties=False) -> Set[str]:
|
||||
"""
|
||||
Get the public variables of a class instance.
|
||||
|
||||
Args:
|
||||
instance (Any): The instance of a class
|
||||
include_properties (bool): Whether to include properties in the result
|
||||
|
||||
Usage:
|
||||
```python
|
||||
detections = Detections(xyxy=np.array([1,2,3,4]))
|
||||
variables = get_class_variables(detections)
|
||||
# ["xyxy", "mask", "confidence", ..., "data"]
|
||||
```
|
||||
"""
|
||||
if isinstance(instance, type):
|
||||
raise ValueError("Only class instances are supported, not classes.")
|
||||
|
||||
fields = set(
|
||||
(
|
||||
name
|
||||
for name, val in inspect.getmembers(instance)
|
||||
if not callable(val) and not name.startswith("_")
|
||||
)
|
||||
)
|
||||
|
||||
if not include_properties:
|
||||
properties = set(
|
||||
(
|
||||
name
|
||||
for name, val in inspect.getmembers(instance.__class__)
|
||||
if isinstance(val, property)
|
||||
)
|
||||
)
|
||||
fields -= properties
|
||||
|
||||
return fields
|
||||
|
||||
@@ -68,3 +68,17 @@ def fill(sequence: List[V], desired_size: int, content: V) -> List[V]:
|
||||
missing_size = max(0, desired_size - len(sequence))
|
||||
sequence.extend([content] * missing_size)
|
||||
return sequence
|
||||
|
||||
|
||||
def find_duplicates(sequence: List) -> List:
|
||||
"""
|
||||
Find all duplicate elements in the input sequence.
|
||||
"""
|
||||
seen = set()
|
||||
duplicates = set()
|
||||
for element in sequence:
|
||||
if element in seen:
|
||||
duplicates.add(element)
|
||||
else:
|
||||
seen.add(element)
|
||||
return list(duplicates)
|
||||
|
||||
@@ -8,8 +8,6 @@ from typing import Callable, Generator, Optional, Tuple
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from supervision.utils.internal import deprecated
|
||||
|
||||
|
||||
@dataclass
|
||||
class VideoInfo:
|
||||
@@ -239,24 +237,6 @@ class FPSMonitor:
|
||||
""" # noqa: E501 // docs
|
||||
self.all_timestamps = deque(maxlen=sample_size)
|
||||
|
||||
@deprecated(
|
||||
"`FPSMonitor.__call__` is deprecated and will be removed in "
|
||||
"`supervision-0.22.0`. Use `FPSMonitor.fps` instead."
|
||||
)
|
||||
def __call__(self) -> float:
|
||||
"""
|
||||
!!! failure "Deprecated"
|
||||
|
||||
`FPSMonitor.__call__` is deprecated and will be removed in
|
||||
`supervision-0.22.0`. Use `FPSMonitor.fps` instead.
|
||||
|
||||
Computes and returns the average FPS based on the stored time stamps.
|
||||
|
||||
Returns:
|
||||
float: The average FPS. Returns 0.0 if no time stamps are stored.
|
||||
"""
|
||||
return self.fps
|
||||
|
||||
@property
|
||||
def fps(self) -> float:
|
||||
"""
|
||||
|
||||
@@ -62,8 +62,7 @@ def validate_keypoint_confidence(confidence: Any, n: int, m: int) -> None:
|
||||
if not is_valid:
|
||||
raise ValueError(
|
||||
f"confidence must be a 1D np.ndarray with shape {expected_shape}, but "
|
||||
"got"
|
||||
f"shape {actual_shape}"
|
||||
f"got shape {actual_shape}"
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
from contextlib import ExitStack as DoesNotRaise
|
||||
from typing import Dict, List, Tuple
|
||||
from typing import Dict, List, Tuple, Union
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
@@ -10,24 +10,30 @@ from supervision.dataset.formats.coco import (
|
||||
classes_to_coco_categories,
|
||||
coco_annotations_to_detections,
|
||||
coco_categories_to_classes,
|
||||
detections_to_coco_annotations,
|
||||
group_coco_annotations_by_image_id,
|
||||
)
|
||||
|
||||
|
||||
def mock_cock_coco_annotation(
|
||||
def mock_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,
|
||||
segmentation: Union[List[list], Dict] = None,
|
||||
iscrowd: bool = False,
|
||||
) -> dict:
|
||||
if not segmentation:
|
||||
segmentation = []
|
||||
return {
|
||||
"id": annotation_id,
|
||||
"image_id": image_id,
|
||||
"category_id": category_id,
|
||||
"bbox": list(bbox),
|
||||
"area": area,
|
||||
"iscrowd": 0,
|
||||
"segmentation": segmentation,
|
||||
"iscrowd": int(iscrowd),
|
||||
}
|
||||
|
||||
|
||||
@@ -101,74 +107,46 @@ def test_classes_to_coco_categories_and_back_to_classes(
|
||||
[
|
||||
([], {}, 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
|
||||
)
|
||||
]
|
||||
},
|
||||
[mock_coco_annotation(annotation_id=0, image_id=0, category_id=0)],
|
||||
{0: [mock_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),
|
||||
mock_coco_annotation(annotation_id=0, image_id=0, category_id=0),
|
||||
mock_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
|
||||
)
|
||||
],
|
||||
0: [mock_coco_annotation(annotation_id=0, image_id=0, category_id=0)],
|
||||
1: [mock_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),
|
||||
mock_coco_annotation(annotation_id=0, image_id=0, category_id=0),
|
||||
mock_coco_annotation(annotation_id=1, image_id=1, category_id=1),
|
||||
mock_coco_annotation(annotation_id=2, image_id=1, category_id=2),
|
||||
mock_coco_annotation(annotation_id=3, image_id=2, category_id=3),
|
||||
mock_coco_annotation(annotation_id=4, image_id=3, category_id=1),
|
||||
mock_coco_annotation(annotation_id=5, image_id=3, category_id=2),
|
||||
mock_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
|
||||
),
|
||||
mock_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
|
||||
),
|
||||
mock_coco_annotation(annotation_id=1, image_id=1, category_id=1),
|
||||
mock_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
|
||||
),
|
||||
mock_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
|
||||
),
|
||||
mock_coco_annotation(annotation_id=4, image_id=3, category_id=1),
|
||||
mock_coco_annotation(annotation_id=5, image_id=3, category_id=2),
|
||||
mock_coco_annotation(annotation_id=5, image_id=3, category_id=3),
|
||||
],
|
||||
},
|
||||
DoesNotRaise(),
|
||||
@@ -195,7 +173,7 @@ def test_group_coco_annotations_by_image_id(
|
||||
), # empty image annotations
|
||||
(
|
||||
[
|
||||
mock_cock_coco_annotation(
|
||||
mock_coco_annotation(
|
||||
category_id=0, bbox=(0, 0, 100, 100), area=100 * 100
|
||||
)
|
||||
],
|
||||
@@ -209,10 +187,10 @@ def test_group_coco_annotations_by_image_id(
|
||||
), # single image annotations
|
||||
(
|
||||
[
|
||||
mock_cock_coco_annotation(
|
||||
mock_coco_annotation(
|
||||
category_id=0, bbox=(0, 0, 100, 100), area=100 * 100
|
||||
),
|
||||
mock_cock_coco_annotation(
|
||||
mock_coco_annotation(
|
||||
category_id=0, bbox=(100, 100, 100, 100), area=100 * 100
|
||||
),
|
||||
],
|
||||
@@ -226,6 +204,156 @@ def test_group_coco_annotations_by_image_id(
|
||||
),
|
||||
DoesNotRaise(),
|
||||
), # two image annotations
|
||||
(
|
||||
[
|
||||
mock_coco_annotation(
|
||||
category_id=0,
|
||||
bbox=(0, 0, 5, 5),
|
||||
area=5 * 5,
|
||||
segmentation=[[0, 0, 2, 0, 2, 2, 4, 2, 4, 4, 0, 4]],
|
||||
)
|
||||
],
|
||||
(5, 5),
|
||||
True,
|
||||
Detections(
|
||||
xyxy=np.array([[0, 0, 5, 5]], dtype=np.float32),
|
||||
class_id=np.array([0], dtype=int),
|
||||
mask=np.array(
|
||||
[
|
||||
[
|
||||
[1, 1, 1, 0, 0],
|
||||
[1, 1, 1, 0, 0],
|
||||
[1, 1, 1, 1, 1],
|
||||
[1, 1, 1, 1, 1],
|
||||
[1, 1, 1, 1, 1],
|
||||
]
|
||||
]
|
||||
),
|
||||
),
|
||||
DoesNotRaise(),
|
||||
), # single image annotations with mask as polygon
|
||||
(
|
||||
[
|
||||
mock_coco_annotation(
|
||||
category_id=0,
|
||||
bbox=(0, 0, 5, 5),
|
||||
area=5 * 5,
|
||||
segmentation={
|
||||
"size": [5, 5],
|
||||
"counts": [0, 15, 2, 3, 2, 3],
|
||||
},
|
||||
iscrowd=True,
|
||||
)
|
||||
],
|
||||
(5, 5),
|
||||
True,
|
||||
Detections(
|
||||
xyxy=np.array([[0, 0, 5, 5]], dtype=np.float32),
|
||||
class_id=np.array([0], dtype=int),
|
||||
mask=np.array(
|
||||
[
|
||||
[
|
||||
[1, 1, 1, 0, 0],
|
||||
[1, 1, 1, 0, 0],
|
||||
[1, 1, 1, 1, 1],
|
||||
[1, 1, 1, 1, 1],
|
||||
[1, 1, 1, 1, 1],
|
||||
]
|
||||
]
|
||||
),
|
||||
),
|
||||
DoesNotRaise(),
|
||||
), # single image annotations with mask, RLE segmentation mask
|
||||
(
|
||||
[
|
||||
mock_coco_annotation(
|
||||
category_id=0,
|
||||
bbox=(0, 0, 5, 5),
|
||||
area=5 * 5,
|
||||
segmentation=[[0, 0, 2, 0, 2, 2, 4, 2, 4, 4, 0, 4]],
|
||||
),
|
||||
mock_coco_annotation(
|
||||
category_id=0,
|
||||
bbox=(3, 0, 2, 2),
|
||||
area=2 * 2,
|
||||
segmentation={
|
||||
"size": [5, 5],
|
||||
"counts": [15, 2, 3, 2, 3],
|
||||
},
|
||||
iscrowd=True,
|
||||
),
|
||||
],
|
||||
(5, 5),
|
||||
True,
|
||||
Detections(
|
||||
xyxy=np.array([[0, 0, 5, 5], [3, 0, 5, 2]], dtype=np.float32),
|
||||
class_id=np.array([0, 0], dtype=int),
|
||||
mask=np.array(
|
||||
[
|
||||
[
|
||||
[1, 1, 1, 0, 0],
|
||||
[1, 1, 1, 0, 0],
|
||||
[1, 1, 1, 1, 1],
|
||||
[1, 1, 1, 1, 1],
|
||||
[1, 1, 1, 1, 1],
|
||||
],
|
||||
[
|
||||
[0, 0, 0, 1, 1],
|
||||
[0, 0, 0, 1, 1],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
],
|
||||
]
|
||||
),
|
||||
),
|
||||
DoesNotRaise(),
|
||||
), # two image annotations with mask, one mask as polygon ans second as RLE
|
||||
(
|
||||
[
|
||||
mock_coco_annotation(
|
||||
category_id=0,
|
||||
bbox=(3, 0, 2, 2),
|
||||
area=2 * 2,
|
||||
segmentation={
|
||||
"size": [5, 5],
|
||||
"counts": [15, 2, 3, 2, 3],
|
||||
},
|
||||
iscrowd=True,
|
||||
),
|
||||
mock_coco_annotation(
|
||||
category_id=1,
|
||||
bbox=(0, 0, 5, 5),
|
||||
area=5 * 5,
|
||||
segmentation=[[0, 0, 2, 0, 2, 2, 4, 2, 4, 4, 0, 4]],
|
||||
),
|
||||
],
|
||||
(5, 5),
|
||||
True,
|
||||
Detections(
|
||||
xyxy=np.array([[3, 0, 5, 2], [0, 0, 5, 5]], dtype=np.float32),
|
||||
class_id=np.array([0, 1], dtype=int),
|
||||
mask=np.array(
|
||||
[
|
||||
[
|
||||
[0, 0, 0, 1, 1],
|
||||
[0, 0, 0, 1, 1],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
],
|
||||
[
|
||||
[1, 1, 1, 0, 0],
|
||||
[1, 1, 1, 0, 0],
|
||||
[1, 1, 1, 1, 1],
|
||||
[1, 1, 1, 1, 1],
|
||||
[1, 1, 1, 1, 1],
|
||||
],
|
||||
]
|
||||
),
|
||||
),
|
||||
DoesNotRaise(),
|
||||
), # two image annotations with mask, first mask as RLE and second as polygon
|
||||
],
|
||||
)
|
||||
def test_coco_annotations_to_detections(
|
||||
@@ -301,3 +429,131 @@ def test_build_coco_class_index_mapping(
|
||||
coco_categories=coco_categories, target_classes=target_classes
|
||||
)
|
||||
assert result == expected_result
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"detections, image_id, annotation_id, expected_result, exception",
|
||||
[
|
||||
(
|
||||
Detections(
|
||||
xyxy=np.array([[0, 0, 100, 100]], dtype=np.float32),
|
||||
class_id=np.array([0], dtype=int),
|
||||
),
|
||||
0,
|
||||
0,
|
||||
[
|
||||
mock_coco_annotation(
|
||||
category_id=0, bbox=(0, 0, 100, 100), area=100 * 100
|
||||
)
|
||||
],
|
||||
DoesNotRaise(),
|
||||
), # no segmentation mask
|
||||
(
|
||||
Detections(
|
||||
xyxy=np.array([[0, 0, 4, 5]], dtype=np.float32),
|
||||
class_id=np.array([0], dtype=int),
|
||||
mask=np.array(
|
||||
[
|
||||
[
|
||||
[1, 1, 1, 1, 0],
|
||||
[1, 1, 1, 1, 0],
|
||||
[1, 1, 1, 1, 0],
|
||||
[1, 1, 1, 1, 0],
|
||||
[1, 1, 1, 1, 0],
|
||||
]
|
||||
]
|
||||
),
|
||||
),
|
||||
0,
|
||||
0,
|
||||
[
|
||||
mock_coco_annotation(
|
||||
category_id=0,
|
||||
bbox=(0, 0, 4, 5),
|
||||
area=4 * 5,
|
||||
segmentation=[[0, 0, 0, 4, 3, 4, 3, 0]],
|
||||
)
|
||||
],
|
||||
DoesNotRaise(),
|
||||
), # segmentation mask in single component,no holes in mask,
|
||||
# expects polygon mask
|
||||
(
|
||||
Detections(
|
||||
xyxy=np.array([[0, 0, 5, 5]], dtype=np.float32),
|
||||
class_id=np.array([0], dtype=int),
|
||||
mask=np.array(
|
||||
[
|
||||
[
|
||||
[1, 1, 1, 0, 0],
|
||||
[1, 1, 1, 0, 0],
|
||||
[1, 1, 1, 0, 0],
|
||||
[0, 0, 0, 1, 1],
|
||||
[0, 0, 0, 1, 1],
|
||||
]
|
||||
]
|
||||
),
|
||||
),
|
||||
0,
|
||||
0,
|
||||
[
|
||||
mock_coco_annotation(
|
||||
category_id=0,
|
||||
bbox=(0, 0, 5, 5),
|
||||
area=5 * 5,
|
||||
segmentation={
|
||||
"size": [5, 5],
|
||||
"counts": [0, 3, 2, 3, 2, 3, 5, 2, 3, 2],
|
||||
},
|
||||
iscrowd=True,
|
||||
)
|
||||
],
|
||||
DoesNotRaise(),
|
||||
), # segmentation mask with 2 components, no holes in mask, expects RLE mask
|
||||
(
|
||||
Detections(
|
||||
xyxy=np.array([[0, 0, 5, 5]], dtype=np.float32),
|
||||
class_id=np.array([0], dtype=int),
|
||||
mask=np.array(
|
||||
[
|
||||
[
|
||||
[0, 1, 1, 1, 1],
|
||||
[0, 1, 1, 1, 1],
|
||||
[1, 1, 0, 0, 1],
|
||||
[1, 1, 0, 0, 1],
|
||||
[1, 1, 1, 1, 1],
|
||||
]
|
||||
]
|
||||
),
|
||||
),
|
||||
0,
|
||||
0,
|
||||
[
|
||||
mock_coco_annotation(
|
||||
category_id=0,
|
||||
bbox=(0, 0, 5, 5),
|
||||
area=5 * 5,
|
||||
segmentation={
|
||||
"size": [5, 5],
|
||||
"counts": [2, 10, 2, 3, 2, 6],
|
||||
},
|
||||
iscrowd=True,
|
||||
)
|
||||
],
|
||||
DoesNotRaise(),
|
||||
), # seg mask in single component, with holes in mask, expects RLE mask
|
||||
],
|
||||
)
|
||||
def test_detections_to_coco_annotations(
|
||||
detections: Detections,
|
||||
image_id: int,
|
||||
annotation_id: int,
|
||||
expected_result: List[Dict],
|
||||
exception: Exception,
|
||||
) -> None:
|
||||
with exception:
|
||||
result, _ = detections_to_coco_annotations(
|
||||
detections=detections,
|
||||
image_id=image_id,
|
||||
annotation_id=annotation_id,
|
||||
)
|
||||
assert result == expected_result
|
||||
|
||||
@@ -2,13 +2,17 @@ from contextlib import ExitStack as DoesNotRaise
|
||||
from test.test_utils import mock_detections
|
||||
from typing import Dict, List, Optional, Tuple, TypeVar
|
||||
|
||||
import numpy as np
|
||||
import numpy.typing as npt
|
||||
import pytest
|
||||
|
||||
from supervision import Detections
|
||||
from supervision.dataset.utils import (
|
||||
build_class_index_mapping,
|
||||
map_detections_class_id,
|
||||
mask_to_rle,
|
||||
merge_class_lists,
|
||||
rle_to_mask,
|
||||
train_test_split,
|
||||
)
|
||||
|
||||
@@ -229,3 +233,131 @@ def test_map_detections_class_id(
|
||||
source_to_target_mapping=source_to_target_mapping, detections=detections
|
||||
)
|
||||
assert result == expected_result
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"mask, expected_rle, exception",
|
||||
[
|
||||
(
|
||||
np.zeros((3, 3)).astype(bool),
|
||||
[9],
|
||||
DoesNotRaise(),
|
||||
), # mask with background only (mask with only False values)
|
||||
(
|
||||
np.ones((3, 3)).astype(bool),
|
||||
[0, 9],
|
||||
DoesNotRaise(),
|
||||
), # mask with foreground only (mask with only True values)
|
||||
(
|
||||
np.array(
|
||||
[
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 1, 1, 1, 0],
|
||||
[0, 1, 0, 1, 0],
|
||||
[0, 1, 1, 1, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
]
|
||||
).astype(bool),
|
||||
[6, 3, 2, 1, 1, 1, 2, 3, 6],
|
||||
DoesNotRaise(),
|
||||
), # mask where foreground object has hole
|
||||
(
|
||||
np.array(
|
||||
[
|
||||
[1, 0, 1, 0, 1],
|
||||
[1, 0, 1, 0, 1],
|
||||
[1, 0, 1, 0, 1],
|
||||
[1, 0, 1, 0, 1],
|
||||
[1, 0, 1, 0, 1],
|
||||
]
|
||||
).astype(bool),
|
||||
[0, 5, 5, 5, 5, 5],
|
||||
DoesNotRaise(),
|
||||
), # mask where foreground consists of 3 separate components
|
||||
(
|
||||
np.array([[[]]]).astype(bool),
|
||||
None,
|
||||
pytest.raises(AssertionError),
|
||||
), # raises AssertionError because mask dimentionality is not 2D
|
||||
(
|
||||
np.array([[]]).astype(bool),
|
||||
None,
|
||||
pytest.raises(AssertionError),
|
||||
), # raises AssertionError because mask is empty
|
||||
],
|
||||
)
|
||||
def test_mask_to_rle(
|
||||
mask: npt.NDArray[np.bool_], expected_rle: List[int], exception: Exception
|
||||
) -> None:
|
||||
with exception:
|
||||
result = mask_to_rle(mask=mask)
|
||||
assert result == expected_rle
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"rle, resolution_wh, expected_mask, exception",
|
||||
[
|
||||
(
|
||||
np.array([9]),
|
||||
[3, 3],
|
||||
np.zeros((3, 3)).astype(bool),
|
||||
DoesNotRaise(),
|
||||
), # mask with background only (mask with only False values); rle as array
|
||||
(
|
||||
[9],
|
||||
[3, 3],
|
||||
np.zeros((3, 3)).astype(bool),
|
||||
DoesNotRaise(),
|
||||
), # mask with background only (mask with only False values); rle as list
|
||||
(
|
||||
np.array([0, 9]),
|
||||
[3, 3],
|
||||
np.ones((3, 3)).astype(bool),
|
||||
DoesNotRaise(),
|
||||
), # mask with foreground only (mask with only True values)
|
||||
(
|
||||
np.array([6, 3, 2, 1, 1, 1, 2, 3, 6]),
|
||||
[5, 5],
|
||||
np.array(
|
||||
[
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 1, 1, 1, 0],
|
||||
[0, 1, 0, 1, 0],
|
||||
[0, 1, 1, 1, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
]
|
||||
).astype(bool),
|
||||
DoesNotRaise(),
|
||||
), # mask where foreground object has hole
|
||||
(
|
||||
np.array([0, 5, 5, 5, 5, 5]),
|
||||
[5, 5],
|
||||
np.array(
|
||||
[
|
||||
[1, 0, 1, 0, 1],
|
||||
[1, 0, 1, 0, 1],
|
||||
[1, 0, 1, 0, 1],
|
||||
[1, 0, 1, 0, 1],
|
||||
[1, 0, 1, 0, 1],
|
||||
]
|
||||
).astype(bool),
|
||||
DoesNotRaise(),
|
||||
), # mask where foreground consists of 3 separate components
|
||||
(
|
||||
np.array([0, 5, 5, 5, 5, 5]),
|
||||
[2, 2],
|
||||
None,
|
||||
pytest.raises(AssertionError),
|
||||
), # raises AssertionError because number of pixels in RLE does not match
|
||||
# number of pixels in expected mask (width x height).
|
||||
],
|
||||
)
|
||||
def test_rle_to_mask(
|
||||
rle: npt.NDArray[np.int_],
|
||||
resolution_wh: Tuple[int, int],
|
||||
expected_mask: npt.NDArray[np.bool_],
|
||||
exception: Exception,
|
||||
) -> None:
|
||||
with exception:
|
||||
result = rle_to_mask(rle=rle, resolution_wh=resolution_wh)
|
||||
assert np.all(result == expected_mask)
|
||||
|
||||
+300
-32
@@ -5,7 +5,7 @@ from typing import List, Optional, Union
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from supervision.detection.core import Detections
|
||||
from supervision.detection.core import Detections, merge_inner_detection_object_pair
|
||||
from supervision.geometry.core import Position
|
||||
|
||||
PREDICTIONS = np.array(
|
||||
@@ -30,6 +30,84 @@ DETECTIONS = Detections(
|
||||
)
|
||||
|
||||
|
||||
# Merge test
|
||||
TEST_MASK = np.zeros((1000, 1000), dtype=bool)
|
||||
TEST_MASK[300:351, 200:251] = True
|
||||
TEST_DET_1 = Detections(
|
||||
xyxy=np.array([[10, 10, 20, 20], [30, 30, 40, 40], [50, 50, 60, 60]]),
|
||||
mask=np.array([TEST_MASK, TEST_MASK, TEST_MASK]),
|
||||
confidence=np.array([0.1, 0.2, 0.3]),
|
||||
class_id=np.array([1, 2, 3]),
|
||||
tracker_id=np.array([1, 2, 3]),
|
||||
data={
|
||||
"some_key": [1, 2, 3],
|
||||
"other_key": [["1", "2"], ["3", "4"], ["5", "6"]],
|
||||
},
|
||||
)
|
||||
TEST_DET_2 = Detections(
|
||||
xyxy=np.array([[70, 70, 80, 80], [90, 90, 100, 100]]),
|
||||
mask=np.array([TEST_MASK, TEST_MASK]),
|
||||
confidence=np.array([0.4, 0.5]),
|
||||
class_id=np.array([4, 5]),
|
||||
tracker_id=np.array([4, 5]),
|
||||
data={
|
||||
"some_key": [4, 5],
|
||||
"other_key": [["7", "8"], ["9", "10"]],
|
||||
},
|
||||
)
|
||||
TEST_DET_1_2 = Detections(
|
||||
xyxy=np.array(
|
||||
[
|
||||
[10, 10, 20, 20],
|
||||
[30, 30, 40, 40],
|
||||
[50, 50, 60, 60],
|
||||
[70, 70, 80, 80],
|
||||
[90, 90, 100, 100],
|
||||
]
|
||||
),
|
||||
mask=np.array([TEST_MASK, TEST_MASK, TEST_MASK, TEST_MASK, TEST_MASK]),
|
||||
confidence=np.array([0.1, 0.2, 0.3, 0.4, 0.5]),
|
||||
class_id=np.array([1, 2, 3, 4, 5]),
|
||||
tracker_id=np.array([1, 2, 3, 4, 5]),
|
||||
data={
|
||||
"some_key": [1, 2, 3, 4, 5],
|
||||
"other_key": [["1", "2"], ["3", "4"], ["5", "6"], ["7", "8"], ["9", "10"]],
|
||||
},
|
||||
)
|
||||
TEST_DET_ZERO_LENGTH = Detections(
|
||||
xyxy=np.empty((0, 4), dtype=np.float32),
|
||||
mask=np.empty((0, *TEST_MASK.shape), dtype=bool),
|
||||
confidence=np.empty((0,)),
|
||||
class_id=np.empty((0,)),
|
||||
tracker_id=np.empty((0,)),
|
||||
data={
|
||||
"some_key": [],
|
||||
"other_key": [],
|
||||
},
|
||||
)
|
||||
TEST_DET_NONE = Detections(
|
||||
xyxy=np.empty((0, 4), dtype=np.float32),
|
||||
)
|
||||
TEST_DET_DIFFERENT_FIELDS = Detections(
|
||||
xyxy=np.array([[88, 88, 99, 99]]),
|
||||
mask=np.array([np.logical_not(TEST_MASK)]),
|
||||
confidence=None,
|
||||
class_id=None,
|
||||
tracker_id=np.array([9]),
|
||||
data={"some_key": [9], "other_key": [["11", "12"]]},
|
||||
)
|
||||
TEST_DET_DIFFERENT_DATA = Detections(
|
||||
xyxy=np.array([[88, 88, 99, 99]]),
|
||||
mask=np.array([np.logical_not(TEST_MASK)]),
|
||||
confidence=np.array([0.9]),
|
||||
class_id=np.array([9]),
|
||||
tracker_id=np.array([9]),
|
||||
data={
|
||||
"never_seen_key": [9],
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"detections, index, expected_result, exception",
|
||||
[
|
||||
@@ -148,52 +226,73 @@ def test_getitem(
|
||||
DoesNotRaise(),
|
||||
), # single empty detections
|
||||
(
|
||||
[mock_detections(xyxy=[[10, 10, 20, 20]])],
|
||||
mock_detections(xyxy=[[10, 10, 20, 20]]),
|
||||
[Detections.empty(), Detections.empty()],
|
||||
Detections.empty(),
|
||||
DoesNotRaise(),
|
||||
), # single detection with xyxy field
|
||||
), # two empty detections
|
||||
(
|
||||
[TEST_DET_1],
|
||||
TEST_DET_1,
|
||||
DoesNotRaise(),
|
||||
), # single detection with fields
|
||||
(
|
||||
[TEST_DET_NONE],
|
||||
TEST_DET_NONE,
|
||||
DoesNotRaise(),
|
||||
), # Single weakly-defined detection
|
||||
(
|
||||
[TEST_DET_1, TEST_DET_2],
|
||||
TEST_DET_1_2,
|
||||
DoesNotRaise(),
|
||||
), # Fields with same keys
|
||||
(
|
||||
[TEST_DET_1, Detections.empty()],
|
||||
TEST_DET_1,
|
||||
DoesNotRaise(),
|
||||
), # single detection with fields
|
||||
(
|
||||
[
|
||||
mock_detections(xyxy=[[10, 10, 20, 20]]),
|
||||
mock_detections(xyxy=np.empty((0, 4), dtype=np.float32)),
|
||||
TEST_DET_1,
|
||||
TEST_DET_ZERO_LENGTH,
|
||||
],
|
||||
mock_detections(xyxy=[[10, 10, 20, 20]]),
|
||||
TEST_DET_1,
|
||||
DoesNotRaise(),
|
||||
), # single detection with xyxy field + empty detection
|
||||
), # Single detection and empty-array fields
|
||||
(
|
||||
[
|
||||
mock_detections(xyxy=[[10, 10, 20, 20]]),
|
||||
mock_detections(xyxy=[[20, 20, 30, 30]]),
|
||||
TEST_DET_1,
|
||||
TEST_DET_NONE,
|
||||
],
|
||||
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], [20, 20, 30, 30]]),
|
||||
None,
|
||||
pytest.raises(ValueError),
|
||||
), # detection with xyxy, class_id fields + detection with xyxy field
|
||||
), # Empty detection, but not Detections.empty()
|
||||
# Errors: Non-zero-length differently defined keys & data
|
||||
(
|
||||
[TEST_DET_1, TEST_DET_DIFFERENT_FIELDS],
|
||||
None,
|
||||
pytest.raises(ValueError),
|
||||
), # Non-empty detections with different fields
|
||||
(
|
||||
[TEST_DET_1, TEST_DET_DIFFERENT_DATA],
|
||||
None,
|
||||
pytest.raises(ValueError),
|
||||
), # Non-empty detections with different data keys
|
||||
(
|
||||
[
|
||||
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(),
|
||||
), # two detections with xyxy, class_id fields
|
||||
(
|
||||
[
|
||||
mock_detections(xyxy=[[10, 10, 20, 20]], data={"test": [1]}),
|
||||
mock_detections(xyxy=[[20, 20, 30, 30]], data={"test": [2]}),
|
||||
mock_detections(
|
||||
xyxy=[[10, 10, 20, 20]],
|
||||
class_id=[1],
|
||||
mask=[np.zeros((4, 4), dtype=bool)],
|
||||
),
|
||||
Detections.empty(),
|
||||
],
|
||||
mock_detections(
|
||||
xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]], data={"test": [1, 2]}
|
||||
xyxy=[[10, 10, 20, 20]],
|
||||
class_id=[1],
|
||||
mask=[np.zeros((4, 4), dtype=bool)],
|
||||
),
|
||||
DoesNotRaise(),
|
||||
), # two detections with xyxy, data fields
|
||||
), # Segmentation + Empty
|
||||
],
|
||||
)
|
||||
def test_merge(
|
||||
@@ -337,3 +436,172 @@ def test_equal(
|
||||
detections_a: Detections, detections_b: Detections, expected_result: bool
|
||||
) -> None:
|
||||
assert (detections_a == detections_b) == expected_result
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"detection_1, detection_2, expected_result, exception",
|
||||
[
|
||||
(
|
||||
mock_detections(
|
||||
xyxy=[[10, 10, 30, 30]],
|
||||
),
|
||||
mock_detections(
|
||||
xyxy=[[10, 10, 30, 30]],
|
||||
),
|
||||
mock_detections(
|
||||
xyxy=[[10, 10, 30, 30]],
|
||||
),
|
||||
DoesNotRaise(),
|
||||
), # Merge with self
|
||||
(
|
||||
mock_detections(
|
||||
xyxy=[[10, 10, 30, 30]],
|
||||
),
|
||||
Detections.empty(),
|
||||
None,
|
||||
pytest.raises(ValueError),
|
||||
), # merge with empty: error
|
||||
(
|
||||
mock_detections(
|
||||
xyxy=[[10, 10, 30, 30]],
|
||||
),
|
||||
mock_detections(
|
||||
xyxy=[[10, 10, 30, 30], [40, 40, 60, 60]],
|
||||
),
|
||||
None,
|
||||
pytest.raises(ValueError),
|
||||
), # merge with 2+ objects: error
|
||||
(
|
||||
mock_detections(
|
||||
xyxy=[[10, 10, 30, 30]],
|
||||
confidence=[0.1],
|
||||
class_id=[1],
|
||||
mask=[np.array([[1, 1, 0], [1, 1, 0], [0, 0, 0]], dtype=bool)],
|
||||
tracker_id=[1],
|
||||
data={"key_1": [1]},
|
||||
),
|
||||
mock_detections(
|
||||
xyxy=[[20, 20, 40, 40]],
|
||||
confidence=[0.1],
|
||||
class_id=[2],
|
||||
mask=[np.array([[0, 0, 0], [0, 1, 1], [0, 1, 1]], dtype=bool)],
|
||||
tracker_id=[2],
|
||||
data={"key_2": [2]},
|
||||
),
|
||||
mock_detections(
|
||||
xyxy=[[10, 10, 40, 40]],
|
||||
confidence=[0.1],
|
||||
class_id=[1],
|
||||
mask=[np.array([[1, 1, 0], [1, 1, 1], [0, 1, 1]], dtype=bool)],
|
||||
tracker_id=[1],
|
||||
data={"key_1": [1]},
|
||||
),
|
||||
DoesNotRaise(),
|
||||
), # Same confidence - merge box & mask, tie-break to detection_1
|
||||
(
|
||||
mock_detections(
|
||||
xyxy=[[0, 0, 20, 20]],
|
||||
confidence=[0.1],
|
||||
class_id=[1],
|
||||
mask=[np.array([[1, 1, 0], [1, 1, 0], [0, 0, 0]], dtype=bool)],
|
||||
tracker_id=[1],
|
||||
data={"key_1": [1]},
|
||||
),
|
||||
mock_detections(
|
||||
xyxy=[[10, 10, 50, 50]],
|
||||
confidence=[0.2],
|
||||
class_id=[2],
|
||||
mask=[np.array([[0, 0, 0], [0, 1, 1], [0, 1, 1]], dtype=bool)],
|
||||
tracker_id=[2],
|
||||
data={"key_2": [2]},
|
||||
),
|
||||
mock_detections(
|
||||
xyxy=[[0, 0, 50, 50]],
|
||||
confidence=[(1 * 0.1 + 4 * 0.2) / 5],
|
||||
class_id=[2],
|
||||
mask=[np.array([[1, 1, 0], [1, 1, 1], [0, 1, 1]], dtype=bool)],
|
||||
tracker_id=[2],
|
||||
data={"key_2": [2]},
|
||||
),
|
||||
DoesNotRaise(),
|
||||
), # Different confidence, different area
|
||||
(
|
||||
mock_detections(
|
||||
xyxy=[[10, 10, 30, 30]],
|
||||
confidence=None,
|
||||
class_id=[1],
|
||||
mask=[np.array([[1, 1, 0], [1, 1, 0], [0, 0, 0]], dtype=bool)],
|
||||
tracker_id=[1],
|
||||
data={"key_1": [1]},
|
||||
),
|
||||
mock_detections(
|
||||
xyxy=[[20, 20, 40, 40]],
|
||||
confidence=None,
|
||||
class_id=[2],
|
||||
mask=[np.array([[0, 0, 0], [0, 1, 1], [0, 1, 1]], dtype=bool)],
|
||||
tracker_id=[2],
|
||||
data={"key_2": [2]},
|
||||
),
|
||||
mock_detections(
|
||||
xyxy=[[10, 10, 40, 40]],
|
||||
confidence=None,
|
||||
class_id=[1],
|
||||
mask=[np.array([[1, 1, 0], [1, 1, 1], [0, 1, 1]], dtype=bool)],
|
||||
tracker_id=[1],
|
||||
data={"key_1": [1]},
|
||||
),
|
||||
DoesNotRaise(),
|
||||
), # No confidence at all
|
||||
(
|
||||
mock_detections(
|
||||
xyxy=[[0, 0, 20, 20]],
|
||||
confidence=None,
|
||||
),
|
||||
mock_detections(
|
||||
xyxy=[[10, 10, 30, 30]],
|
||||
confidence=[0.2],
|
||||
),
|
||||
None,
|
||||
pytest.raises(ValueError),
|
||||
), # confidence: None + [x]
|
||||
(
|
||||
mock_detections(
|
||||
xyxy=[[0, 0, 20, 20]],
|
||||
mask=[np.array([[1, 1, 0], [1, 1, 0], [0, 0, 0]], dtype=bool)],
|
||||
),
|
||||
mock_detections(
|
||||
xyxy=[[10, 10, 30, 30]],
|
||||
mask=None,
|
||||
),
|
||||
None,
|
||||
pytest.raises(ValueError),
|
||||
), # mask: None + [x]
|
||||
(
|
||||
mock_detections(xyxy=[[0, 0, 20, 20]], tracker_id=[1]),
|
||||
mock_detections(
|
||||
xyxy=[[10, 10, 30, 30]],
|
||||
tracker_id=None,
|
||||
),
|
||||
None,
|
||||
pytest.raises(ValueError),
|
||||
), # tracker_id: None + []
|
||||
(
|
||||
mock_detections(xyxy=[[0, 0, 20, 20]], class_id=[1]),
|
||||
mock_detections(
|
||||
xyxy=[[10, 10, 30, 30]],
|
||||
class_id=None,
|
||||
),
|
||||
None,
|
||||
pytest.raises(ValueError),
|
||||
), # class_id: None + []
|
||||
],
|
||||
)
|
||||
def test_merge_inner_detection_object_pair(
|
||||
detection_1: Detections,
|
||||
detection_2: Detections,
|
||||
expected_result: Optional[Detections],
|
||||
exception: Exception,
|
||||
):
|
||||
with exception:
|
||||
result = merge_inner_detection_object_pair(detection_1, detection_2)
|
||||
assert result == expected_result
|
||||
|
||||
@@ -1,10 +1,11 @@
|
||||
from contextlib import ExitStack as DoesNotRaise
|
||||
from typing import Optional, Tuple
|
||||
from test.test_utils import mock_detections
|
||||
from typing import List, Optional, Tuple
|
||||
|
||||
import pytest
|
||||
|
||||
from supervision import LineZone
|
||||
from supervision.geometry.core import Point, Vector
|
||||
from supervision.geometry.core import Point, Position, Vector
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
@@ -70,3 +71,409 @@ def test_calculate_region_of_interest_limits(
|
||||
with exception:
|
||||
result = LineZone.calculate_region_of_interest_limits(vector=vector)
|
||||
assert result == expected_result
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"vector, xyxy_sequence, expected_crossed_in, expected_crossed_out",
|
||||
[
|
||||
( # Vertical line, simple crossing
|
||||
Vector(Point(0, 0), Point(0, 10)),
|
||||
[
|
||||
[4, 4, 6, 6],
|
||||
[4 - 10, 4, 6 - 10, 6],
|
||||
[4, 4, 6, 6],
|
||||
[4 - 10, 4, 6 - 10, 6],
|
||||
],
|
||||
[False, False, True, False],
|
||||
[False, True, False, True],
|
||||
),
|
||||
( # Vertical line reversed, simple crossing
|
||||
Vector(Point(0, 10), Point(0, 0)),
|
||||
[
|
||||
[4, 4, 6, 6],
|
||||
[4 - 10, 4, 6 - 10, 6],
|
||||
[4, 4, 6, 6],
|
||||
[4 - 10, 4, 6 - 10, 6],
|
||||
],
|
||||
[False, True, False, True],
|
||||
[False, False, True, False],
|
||||
),
|
||||
( # Horizontal line, simple crossing
|
||||
Vector(Point(0, 0), Point(10, 0)),
|
||||
[
|
||||
[4, 4, 6, 6],
|
||||
[4, 4 - 10, 6, 6 - 10],
|
||||
[4, 4, 6, 6],
|
||||
[4, 4 - 10, 6, 6 - 10],
|
||||
],
|
||||
[False, True, False, True],
|
||||
[False, False, True, False],
|
||||
),
|
||||
( # Horizontal line reversed, simple crossing
|
||||
Vector(Point(10, 0), Point(0, 0)),
|
||||
[
|
||||
[4, 4, 6, 6],
|
||||
[4, 4 - 10, 6, 6 - 10],
|
||||
[4, 4, 6, 6],
|
||||
[4, 4 - 10, 6, 6 - 10],
|
||||
],
|
||||
[False, False, True, False],
|
||||
[False, True, False, True],
|
||||
),
|
||||
( # Diagonal line, simple crossing
|
||||
Vector(Point(5, 0), Point(0, 5)),
|
||||
[
|
||||
[0, 0, 2, 2],
|
||||
[0 + 10, 0 + 10, 2 + 10, 2 + 10],
|
||||
[0, 0, 2, 2],
|
||||
[0 + 10, 0 + 10, 2 + 10, 2 + 10],
|
||||
],
|
||||
[False, True, False, True],
|
||||
[False, False, True, False],
|
||||
),
|
||||
( # Crossing beside - right side
|
||||
Vector(Point(0, 0), Point(10, 0)),
|
||||
[
|
||||
[20, 4, 24, 6],
|
||||
[20, 4 - 10, 24, 6 - 10],
|
||||
[20, 4, 24, 6],
|
||||
[20, 4 - 10, 24, 6 - 10],
|
||||
],
|
||||
[False, False, False, False],
|
||||
[False, False, False, False],
|
||||
),
|
||||
( # Horizontal line, simple crossing, far away
|
||||
Vector(Point(0, 0), Point(10, 0)),
|
||||
[
|
||||
[4, 1e32, 6, 1e32 + 2],
|
||||
[4, -1e32, 6, -1e32 + 2],
|
||||
[4, 1e32, 6, 1e32 + 2],
|
||||
[4, -1e32, 6, -1e32 + 2],
|
||||
],
|
||||
[False, True, False, True],
|
||||
[False, False, True, False],
|
||||
),
|
||||
( # Crossing beside - left side
|
||||
Vector(Point(0, 0), Point(10, 0)),
|
||||
[
|
||||
[-20, 4, -24, 6],
|
||||
[-20, 4 - 10, -24, 6 - 10],
|
||||
[-20, 4, -24, 6],
|
||||
[-20, 4 - 10, -24, 6 - 10],
|
||||
],
|
||||
[False, False, False, False],
|
||||
[False, False, False, False],
|
||||
),
|
||||
( # Move above
|
||||
Vector(Point(0, 0), Point(10, 0)),
|
||||
[
|
||||
[-4, 4, -2, 6],
|
||||
[-4 + 20, 4, -2 + 20, 6],
|
||||
[-4, 4, -2, 6],
|
||||
[-4 + 20, 4, -2 + 20, 6],
|
||||
],
|
||||
[False, False, False, False],
|
||||
[False, False, False, False],
|
||||
),
|
||||
( # Move below
|
||||
Vector(Point(0, 0), Point(10, 0)),
|
||||
[
|
||||
[-4, -6, -2, -4],
|
||||
[-4 + 20, -6, -2 + 20, -4],
|
||||
[-4, -6, -2, -4],
|
||||
[-4 + 20, -6, -2 + 20, -4],
|
||||
],
|
||||
[False, False, False, False],
|
||||
[False, False, False, False],
|
||||
),
|
||||
( # Move into line partway
|
||||
Vector(Point(0, 0), Point(10, 0)),
|
||||
[
|
||||
[4, 4, 6, 6],
|
||||
[4 + 5, 4, 6 + 5, 6],
|
||||
[4, 4, 6, 6],
|
||||
[4 + 5, 4, 6 + 5, 6],
|
||||
],
|
||||
[False, False, False, False],
|
||||
[False, False, False, False],
|
||||
),
|
||||
( # V-shaped crossing from outside limits - not supported.
|
||||
Vector(Point(0, 0), Point(10, 0)),
|
||||
[[-3, 6, -1, 8], [4, -6, 6, -4], [11, 6, 13, 8]],
|
||||
[False, False, False],
|
||||
[False, False, False],
|
||||
),
|
||||
( # Diagonal movement, from within limits to outside - not supported
|
||||
Vector(Point(0, 0), Point(10, 0)),
|
||||
[[4, 1, 6, 3], [11, 1 - 20, 13, 3 - 20]],
|
||||
[False, False],
|
||||
[False, False],
|
||||
),
|
||||
( # Diagonal movement, from outside limits to within - not supported
|
||||
Vector(Point(0, 0), Point(10, 0)),
|
||||
[
|
||||
[11, 21, 13, 23],
|
||||
[4, -3, 6, -1],
|
||||
],
|
||||
[False, False],
|
||||
[False, False],
|
||||
),
|
||||
( # Diagonal crossing, from outside to outside limits - not supported.
|
||||
Vector(Point(0, 0), Point(10, 0)),
|
||||
[
|
||||
[-4, 4, -2, 8],
|
||||
[-4 + 16, -4, -2 + 16, -6],
|
||||
[-4, 4, -2, 8],
|
||||
[-4 + 16, -4, -2 + 16, -6],
|
||||
],
|
||||
[False, False, False, False],
|
||||
[False, False, False, False],
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_line_zone_one_detection_default_anchors(
|
||||
vector: Vector,
|
||||
xyxy_sequence: List[List[float]],
|
||||
expected_crossed_in: List[bool],
|
||||
expected_crossed_out: List[bool],
|
||||
) -> None:
|
||||
line_zone = LineZone(start=vector.start, end=vector.end)
|
||||
|
||||
crossed_in_list = []
|
||||
crossed_out_list = []
|
||||
for i, bbox in enumerate(xyxy_sequence):
|
||||
detections = mock_detections(
|
||||
xyxy=[bbox],
|
||||
tracker_id=[0],
|
||||
)
|
||||
crossed_in, crossed_out = line_zone.trigger(detections)
|
||||
crossed_in_list.append(crossed_in[0])
|
||||
crossed_out_list.append(crossed_out[0])
|
||||
|
||||
assert (
|
||||
crossed_in_list == expected_crossed_in
|
||||
), f"expected {expected_crossed_in}, got {crossed_in_list}"
|
||||
assert (
|
||||
crossed_out_list == expected_crossed_out
|
||||
), f"expected {expected_crossed_out}, got {crossed_out_list}"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"vector, xyxy_sequence, triggering_anchors, expected_crossed_in, "
|
||||
"expected_crossed_out",
|
||||
[
|
||||
( # Scrape line, left side, corner anchors
|
||||
Vector(Point(0, 0), Point(10, 0)),
|
||||
[
|
||||
[-2, 4, 2, 6],
|
||||
[-2, 4 - 10, 2, 6 - 10],
|
||||
[-2, 4, 2, 6],
|
||||
[-2, 4 - 10, 2, 6 - 10],
|
||||
],
|
||||
[
|
||||
Position.TOP_LEFT,
|
||||
Position.BOTTOM_LEFT,
|
||||
Position.TOP_RIGHT,
|
||||
Position.BOTTOM_RIGHT,
|
||||
],
|
||||
[False, False, False, False],
|
||||
[False, False, False, False],
|
||||
),
|
||||
( # Scrape line, left side, right anchors
|
||||
Vector(Point(0, 0), Point(10, 0)),
|
||||
[
|
||||
[-2, 4, 2, 6],
|
||||
[-2, 4 - 10, 2, 6 - 10],
|
||||
[-2, 4, 2, 6],
|
||||
[-2, 4 - 10, 2, 6 - 10],
|
||||
],
|
||||
[Position.TOP_RIGHT, Position.BOTTOM_RIGHT],
|
||||
[False, True, False, True],
|
||||
[False, False, True, False],
|
||||
),
|
||||
( # Scrape line, left side, center anchor (along line point)
|
||||
Vector(Point(0, 0), Point(10, 0)),
|
||||
[
|
||||
[-2, 4, 2, 6],
|
||||
[-2, 4 - 10, 2, 6 - 10],
|
||||
[-2, 4, 2, 6],
|
||||
[-2, 4 - 10, 2, 6 - 10],
|
||||
],
|
||||
[Position.CENTER],
|
||||
[False, True, False, True],
|
||||
[False, False, True, False],
|
||||
),
|
||||
( # Scrape line, right side, corner anchors
|
||||
Vector(Point(0, 0), Point(10, 0)),
|
||||
[
|
||||
[8, 4, 12, 6],
|
||||
[8, 4 - 10, 12, 6 - 10],
|
||||
[8, 4, 12, 6],
|
||||
[8, 4 - 10, 12, 6 - 10],
|
||||
],
|
||||
[
|
||||
Position.TOP_LEFT,
|
||||
Position.BOTTOM_LEFT,
|
||||
Position.TOP_RIGHT,
|
||||
Position.BOTTOM_RIGHT,
|
||||
],
|
||||
[False, False, False, False],
|
||||
[False, False, False, False],
|
||||
),
|
||||
( # Scrape line, right side, left anchors
|
||||
Vector(Point(0, 0), Point(10, 0)),
|
||||
[
|
||||
[8, 4, 12, 6],
|
||||
[8, 4 - 10, 12, 6 - 10],
|
||||
[8, 4, 12, 6],
|
||||
[8, 4 - 10, 12, 6 - 10],
|
||||
],
|
||||
[Position.TOP_LEFT, Position.BOTTOM_LEFT],
|
||||
[False, True, False, True],
|
||||
[False, False, True, False],
|
||||
),
|
||||
( # Scrape line, right side, center anchor (along line point)
|
||||
Vector(Point(0, 0), Point(10, 0)),
|
||||
[
|
||||
[8, 4, 12, 6],
|
||||
[8, 4 - 10, 12, 6 - 10],
|
||||
[8, 4, 12, 6],
|
||||
[8, 4 - 10, 12, 6 - 10],
|
||||
],
|
||||
[Position.CENTER],
|
||||
[False, True, False, True],
|
||||
[False, False, True, False],
|
||||
),
|
||||
( # Simple crossing, one anchor
|
||||
Vector(Point(0, 0), Point(10, 0)),
|
||||
[
|
||||
[4, 4, 6, 6],
|
||||
[4, 4 - 10, 6, 6 - 10],
|
||||
[4, 4, 6, 6],
|
||||
[4, 4 - 10, 6, 6 - 10],
|
||||
],
|
||||
[Position.CENTER],
|
||||
[False, True, False, True],
|
||||
[False, False, True, False],
|
||||
),
|
||||
( # Simple crossing, all box anchors
|
||||
Vector(Point(0, 0), Point(10, 0)),
|
||||
[
|
||||
[4, 4, 6, 6],
|
||||
[4, 4 - 10, 6, 6 - 10],
|
||||
[4, 4, 6, 6],
|
||||
[4, 4 - 10, 6, 6 - 10],
|
||||
],
|
||||
[
|
||||
Position.CENTER,
|
||||
Position.CENTER_LEFT,
|
||||
Position.CENTER_RIGHT,
|
||||
Position.TOP_CENTER,
|
||||
Position.TOP_LEFT,
|
||||
Position.TOP_RIGHT,
|
||||
Position.BOTTOM_LEFT,
|
||||
Position.BOTTOM_CENTER,
|
||||
Position.BOTTOM_RIGHT,
|
||||
],
|
||||
[False, True, False, True],
|
||||
[False, False, True, False],
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_line_zone_one_detection(
|
||||
vector: Vector,
|
||||
xyxy_sequence: List[List[float]],
|
||||
triggering_anchors: List[Position],
|
||||
expected_crossed_in: List[bool],
|
||||
expected_crossed_out: List[bool],
|
||||
) -> None:
|
||||
line_zone = LineZone(
|
||||
start=vector.start, end=vector.end, triggering_anchors=triggering_anchors
|
||||
)
|
||||
|
||||
crossed_in_list = []
|
||||
crossed_out_list = []
|
||||
for i, bbox in enumerate(xyxy_sequence):
|
||||
detections = mock_detections(
|
||||
xyxy=[bbox],
|
||||
tracker_id=[0],
|
||||
)
|
||||
crossed_in, crossed_out = line_zone.trigger(detections)
|
||||
crossed_in_list.append(crossed_in[0])
|
||||
crossed_out_list.append(crossed_out[0])
|
||||
|
||||
assert (
|
||||
crossed_in_list == expected_crossed_in
|
||||
), f"expected {expected_crossed_in}, got {crossed_in_list}"
|
||||
assert (
|
||||
crossed_out_list == expected_crossed_out
|
||||
), f"expected {expected_crossed_out}, got {crossed_out_list}"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"vector, xyxy_sequence, anchors, expected_crossed_in, "
|
||||
"expected_crossed_out, exception",
|
||||
[
|
||||
( # One stays, one crosses
|
||||
Vector(Point(0, 0), Point(10, 0)),
|
||||
[
|
||||
[[4, 4, 6, 6], [4, 4, 6, 6]],
|
||||
[[4, 4, 6, 6], [4, 4 - 10, 6, 6 - 10]],
|
||||
[[4, 4, 6, 6], [4, 4, 6, 6]],
|
||||
[[4, 4, 6, 6], [4, 4 - 10, 6, 6 - 10]],
|
||||
],
|
||||
[
|
||||
Position.TOP_LEFT,
|
||||
Position.TOP_RIGHT,
|
||||
Position.BOTTOM_LEFT,
|
||||
Position.BOTTOM_RIGHT,
|
||||
],
|
||||
[[False, False], [False, True], [False, False], [False, True]],
|
||||
[[False, False], [False, False], [False, True], [False, False]],
|
||||
DoesNotRaise(),
|
||||
),
|
||||
( # Both cross at the same time
|
||||
Vector(Point(0, 0), Point(10, 0)),
|
||||
[
|
||||
[[4, 4, 6, 6], [4, 4, 6, 6]],
|
||||
[[4, 4 - 10, 6, 6 - 10], [4, 4 - 10, 6, 6 - 10]],
|
||||
[[4, 4, 6, 6], [4, 4, 6, 6]],
|
||||
[[4, 4 - 10, 6, 6 - 10], [4, 4 - 10, 6, 6 - 10]],
|
||||
],
|
||||
[
|
||||
Position.TOP_LEFT,
|
||||
Position.TOP_RIGHT,
|
||||
Position.BOTTOM_LEFT,
|
||||
Position.BOTTOM_RIGHT,
|
||||
],
|
||||
[[False, False], [True, True], [False, False], [True, True]],
|
||||
[[False, False], [False, False], [True, True], [False, False]],
|
||||
DoesNotRaise(),
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_line_zone_multiple_detections(
|
||||
vector: Vector,
|
||||
xyxy_sequence: List[List[List[float]]],
|
||||
anchors: List[Position],
|
||||
expected_crossed_in: List[List[bool]],
|
||||
expected_crossed_out: List[List[bool]],
|
||||
exception: Exception,
|
||||
) -> None:
|
||||
with exception:
|
||||
line_zone = LineZone(
|
||||
start=vector.start, end=vector.end, triggering_anchors=anchors
|
||||
)
|
||||
crossed_in_list = []
|
||||
crossed_out_list = []
|
||||
for bboxes in xyxy_sequence:
|
||||
detections = mock_detections(
|
||||
xyxy=bboxes,
|
||||
tracker_id=[i for i in range(0, len(bboxes))],
|
||||
)
|
||||
crossed_in, crossed_out = line_zone.trigger(detections)
|
||||
crossed_in_list.append(list(crossed_in))
|
||||
crossed_out_list.append(list(crossed_out))
|
||||
|
||||
assert crossed_in_list == expected_crossed_in
|
||||
assert crossed_out_list == expected_crossed_out
|
||||
|
||||
@@ -0,0 +1,151 @@
|
||||
from typing import List, Optional, Tuple
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from supervision.detection.lmm import from_paligemma
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"result, resolution_wh, classes, expected_results",
|
||||
[
|
||||
(
|
||||
"",
|
||||
(1000, 1000),
|
||||
None,
|
||||
(np.empty((0, 4)), None, np.empty(0).astype(str)),
|
||||
), # empty response
|
||||
(
|
||||
"",
|
||||
(1000, 1000),
|
||||
["cat", "dog"],
|
||||
(np.empty((0, 4)), None, np.empty(0).astype(str)),
|
||||
), # empty response with classes
|
||||
(
|
||||
"\n",
|
||||
(1000, 1000),
|
||||
None,
|
||||
(np.empty((0, 4)), None, np.empty(0).astype(str)),
|
||||
), # new line response
|
||||
(
|
||||
"the quick brown fox jumps over the lazy dog.",
|
||||
(1000, 1000),
|
||||
None,
|
||||
(np.empty((0, 4)), None, np.empty(0).astype(str)),
|
||||
), # response with no location
|
||||
(
|
||||
"<loc0256><loc0768><loc0768> cat",
|
||||
(1000, 1000),
|
||||
None,
|
||||
(np.empty((0, 4)), None, np.empty(0).astype(str)),
|
||||
), # response with missing location
|
||||
(
|
||||
"<loc0256><loc0256><loc0768><loc0768><loc0768> cat",
|
||||
(1000, 1000),
|
||||
None,
|
||||
(np.empty((0, 4)), None, np.empty(0).astype(str)),
|
||||
), # response with extra location
|
||||
(
|
||||
"<loc0256><loc0256><loc0768><loc0768>",
|
||||
(1000, 1000),
|
||||
None,
|
||||
(np.empty((0, 4)), None, np.empty(0).astype(str)),
|
||||
), # response with no class
|
||||
(
|
||||
"<loc0256><loc0256><loc0768><loc0768> catt",
|
||||
(1000, 1000),
|
||||
["cat", "dog"],
|
||||
(np.empty((0, 4)), np.empty(0), np.empty(0).astype(str)),
|
||||
), # response with invalid class
|
||||
(
|
||||
"<loc0256><loc0256><loc0768><loc0768> cat",
|
||||
(1000, 1000),
|
||||
None,
|
||||
(
|
||||
np.array([[250.0, 250.0, 750.0, 750.0]]),
|
||||
None,
|
||||
np.array(["cat"]).astype(str),
|
||||
),
|
||||
), # correct response; no classes
|
||||
(
|
||||
"<loc0256><loc0256><loc0768><loc0768> black cat",
|
||||
(1000, 1000),
|
||||
None,
|
||||
(
|
||||
np.array([[250.0, 250.0, 750.0, 750.0]]),
|
||||
None,
|
||||
np.array(["black cat"]).astype(np.dtype("U")),
|
||||
),
|
||||
), # correct response; class name with space; no classes
|
||||
(
|
||||
"<loc0256><loc0256><loc0768><loc0768> black-cat",
|
||||
(1000, 1000),
|
||||
None,
|
||||
(
|
||||
np.array([[250.0, 250.0, 750.0, 750.0]]),
|
||||
None,
|
||||
np.array(["black-cat"]).astype(np.dtype("U")),
|
||||
),
|
||||
), # correct response; class name with hyphen; no classes
|
||||
(
|
||||
"<loc0256><loc0256><loc0768><loc0768> black_cat",
|
||||
(1000, 1000),
|
||||
None,
|
||||
(
|
||||
np.array([[250.0, 250.0, 750.0, 750.0]]),
|
||||
None,
|
||||
np.array(["black_cat"]).astype(np.dtype("U")),
|
||||
),
|
||||
), # correct response; class name with underscore; no classes
|
||||
(
|
||||
"<loc0256><loc0256><loc0768><loc0768> cat ;",
|
||||
(1000, 1000),
|
||||
["cat", "dog"],
|
||||
(
|
||||
np.array([[250.0, 250.0, 750.0, 750.0]]),
|
||||
np.array([0]),
|
||||
np.array(["cat"]).astype(str),
|
||||
),
|
||||
), # correct response; with classes
|
||||
(
|
||||
"<loc0256><loc0256><loc0768><loc0768> cat ; <loc0256><loc0256><loc0768><loc0768> dog", # noqa: E501
|
||||
(1000, 1000),
|
||||
["cat", "dog"],
|
||||
(
|
||||
np.array([[250.0, 250.0, 750.0, 750.0], [250.0, 250.0, 750.0, 750.0]]),
|
||||
np.array([0, 1]),
|
||||
np.array(["cat", "dog"]).astype(np.dtype("U")),
|
||||
),
|
||||
), # correct response; with classes
|
||||
(
|
||||
"<loc0256><loc0256><loc0768><loc0768> cat ; <loc0256><loc0256><loc0768> cat", # noqa: E501
|
||||
(1000, 1000),
|
||||
["cat", "dog"],
|
||||
(
|
||||
np.array([[250.0, 250.0, 750.0, 750.0]]),
|
||||
np.array([0]),
|
||||
np.array(["cat"]).astype(str),
|
||||
),
|
||||
), # partially correct response; with classes
|
||||
(
|
||||
"<loc0256><loc0256><loc0768><loc0768> cat ; <loc0256><loc0256><loc0768><loc0768><loc0768> cat", # noqa: E501
|
||||
(1000, 1000),
|
||||
["cat", "dog"],
|
||||
(
|
||||
np.array([[250.0, 250.0, 750.0, 750.0]]),
|
||||
np.array([0]),
|
||||
np.array(["cat"]).astype(str),
|
||||
),
|
||||
), # partially correct response; with classes
|
||||
],
|
||||
)
|
||||
def test_from_paligemma(
|
||||
result: str,
|
||||
resolution_wh: Tuple[int, int],
|
||||
classes: Optional[List[str]],
|
||||
expected_results: Tuple[np.ndarray, Optional[np.ndarray], np.ndarray],
|
||||
) -> None:
|
||||
result = from_paligemma(result=result, resolution_wh=resolution_wh, classes=classes)
|
||||
np.testing.assert_array_equal(result[0], expected_results[0])
|
||||
np.testing.assert_array_equal(result[1], expected_results[1])
|
||||
np.testing.assert_array_equal(result[2], expected_results[2])
|
||||
@@ -0,0 +1,291 @@
|
||||
from contextlib import ExitStack as DoesNotRaise
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from supervision.detection.lmm import from_florence_2
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"florence_result, resolution_wh, expected_results, exception",
|
||||
[
|
||||
( # Object detection: empty
|
||||
{"<OD>": {"bboxes": [], "labels": []}},
|
||||
(10, 10),
|
||||
(np.array([], dtype=np.float32), np.array([]), None, None),
|
||||
DoesNotRaise(),
|
||||
),
|
||||
( # Object detection: two detections
|
||||
{
|
||||
"<OD>": {
|
||||
"bboxes": [[4, 4, 6, 6], [5, 5, 7, 7]],
|
||||
"labels": ["car", "door"],
|
||||
}
|
||||
},
|
||||
(10, 10),
|
||||
(
|
||||
np.array([[4, 4, 6, 6], [5, 5, 7, 7]], dtype=np.float32),
|
||||
np.array(["car", "door"]),
|
||||
None,
|
||||
None,
|
||||
),
|
||||
DoesNotRaise(),
|
||||
),
|
||||
( # Caption: unsupported
|
||||
{"<CAPTION>": "A green car parked in front of a yellow building."},
|
||||
(10, 10),
|
||||
None,
|
||||
pytest.raises(ValueError),
|
||||
),
|
||||
( # Detailed Caption: unsupported
|
||||
{
|
||||
"<DETAILED_CAPTION>": "The image shows a blue Volkswagen Beetle parked "
|
||||
"in front of a yellow building with two brown doors, surrounded by "
|
||||
"trees and a clear blue sky."
|
||||
},
|
||||
(10, 10),
|
||||
None,
|
||||
pytest.raises(ValueError),
|
||||
),
|
||||
( # More Detailed Caption: unsupported
|
||||
{
|
||||
"<MORE_DETAILED_CAPTION>": "The image shows a vintage Volkswagen "
|
||||
"Beetle car parked on a "
|
||||
"cobblestone street in front of a yellow building with two wooden "
|
||||
"doors. The car is painted in a bright turquoise color and has a "
|
||||
"white stripe running along the side. It has two doors on either side "
|
||||
"of the car, one on top of the other, and a small window on the "
|
||||
"front. The building appears to be old and dilapidated, with peeling "
|
||||
"paint and crumbling walls. The sky is blue and there are trees in "
|
||||
"the background."
|
||||
},
|
||||
(10, 10),
|
||||
None,
|
||||
pytest.raises(ValueError),
|
||||
),
|
||||
( # Caption to Phrase Grounding: empty
|
||||
{"<CAPTION_TO_PHRASE_GROUNDING>": {"bboxes": [], "labels": []}},
|
||||
(10, 10),
|
||||
(np.array([], dtype=np.float32), np.array([]), None, None),
|
||||
DoesNotRaise(),
|
||||
),
|
||||
( # Caption to Phrase Grounding: two detections
|
||||
{
|
||||
"<CAPTION_TO_PHRASE_GROUNDING>": {
|
||||
"bboxes": [[4, 4, 6, 6], [5, 5, 7, 7]],
|
||||
"labels": ["a green car", "a yellow building"],
|
||||
}
|
||||
},
|
||||
(10, 10),
|
||||
(
|
||||
np.array([[4, 4, 6, 6], [5, 5, 7, 7]], dtype=np.float32),
|
||||
np.array(["a green car", "a yellow building"]),
|
||||
None,
|
||||
None,
|
||||
),
|
||||
DoesNotRaise(),
|
||||
),
|
||||
( # Dense Region caption: empty
|
||||
{"<DENSE_REGION_CAPTION>": {"bboxes": [], "labels": []}},
|
||||
(10, 10),
|
||||
(np.array([], dtype=np.float32), np.array([]), None, None),
|
||||
DoesNotRaise(),
|
||||
),
|
||||
( # Caption to Phrase Grounding: two detections
|
||||
{
|
||||
"<DENSE_REGION_CAPTION>": {
|
||||
"bboxes": [[4, 4, 6, 6], [5, 5, 7, 7]],
|
||||
"labels": ["a green car", "a yellow building"],
|
||||
}
|
||||
},
|
||||
(10, 10),
|
||||
(
|
||||
np.array([[4, 4, 6, 6], [5, 5, 7, 7]], dtype=np.float32),
|
||||
np.array(["a green car", "a yellow building"]),
|
||||
None,
|
||||
None,
|
||||
),
|
||||
DoesNotRaise(),
|
||||
),
|
||||
( # Region proposal
|
||||
{
|
||||
"<REGION_PROPOSAL>": {
|
||||
"bboxes": [[4, 4, 6, 6], [5, 5, 7, 7]],
|
||||
"labels": ["", ""],
|
||||
}
|
||||
},
|
||||
(10, 10),
|
||||
(
|
||||
np.array([[4, 4, 6, 6], [5, 5, 7, 7]], dtype=np.float32),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
),
|
||||
DoesNotRaise(),
|
||||
),
|
||||
( # Referring Expression Segmentation
|
||||
{
|
||||
"<REFERRING_EXPRESSION_SEGMENTATION>": {
|
||||
"polygons": [[[1, 1, 2, 1, 2, 2, 1, 2]]],
|
||||
"labels": [""],
|
||||
}
|
||||
},
|
||||
(10, 10),
|
||||
(
|
||||
np.array([[1.0, 1.0, 2.0, 2.0]], dtype=np.float32),
|
||||
None,
|
||||
np.array(
|
||||
[
|
||||
[
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 1, 1, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 1, 1, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
]
|
||||
],
|
||||
dtype=bool,
|
||||
),
|
||||
None,
|
||||
),
|
||||
DoesNotRaise(),
|
||||
),
|
||||
( # Referring Expression Segmentation
|
||||
{
|
||||
"<REFERRING_EXPRESSION_SEGMENTATION>": {
|
||||
"polygons": [[[1, 1, 2, 1, 2, 2, 1, 2]]],
|
||||
"labels": [""],
|
||||
}
|
||||
},
|
||||
(10, 10),
|
||||
(
|
||||
np.array([[1.0, 1.0, 2.0, 2.0]], dtype=np.float32),
|
||||
None,
|
||||
np.array(
|
||||
[
|
||||
[
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 1, 1, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 1, 1, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
]
|
||||
],
|
||||
dtype=bool,
|
||||
),
|
||||
None,
|
||||
),
|
||||
DoesNotRaise(),
|
||||
),
|
||||
( # OCR: unsupported
|
||||
{"<OCR>": "A"},
|
||||
(10, 10),
|
||||
None,
|
||||
pytest.raises(ValueError),
|
||||
),
|
||||
( # OCR with Region: obb boxes
|
||||
{
|
||||
"<OCR_WITH_REGION>": {
|
||||
"quad_boxes": [[2, 2, 6, 4, 5, 6, 1, 5], [4, 4, 5, 5, 4, 6, 3, 5]],
|
||||
"labels": ["some text", "other text"],
|
||||
}
|
||||
},
|
||||
(10, 10),
|
||||
(
|
||||
np.array([[1, 2, 6, 6], [3, 4, 5, 6]], dtype=np.float32),
|
||||
np.array(["some text", "other text"]),
|
||||
None,
|
||||
np.array(
|
||||
[[[2, 2], [6, 4], [5, 6], [1, 5]], [[4, 4], [5, 5], [4, 6], [3, 5]]]
|
||||
),
|
||||
),
|
||||
DoesNotRaise(),
|
||||
),
|
||||
( # Open Vocabulary Detection
|
||||
{
|
||||
"<OPEN_VOCABULARY_DETECTION>": {
|
||||
"bboxes": [[4, 4, 6, 6], [5, 5, 7, 7]],
|
||||
"bboxes_labels": ["cat", "cat"],
|
||||
"polygon": [],
|
||||
"polygons_labels": [],
|
||||
}
|
||||
},
|
||||
(10, 10),
|
||||
(
|
||||
np.array([[4, 4, 6, 6], [5, 5, 7, 7]], dtype=np.float32),
|
||||
np.array(["cat", "cat"]),
|
||||
None,
|
||||
None,
|
||||
),
|
||||
DoesNotRaise(),
|
||||
),
|
||||
( # Region to Category: empty
|
||||
{"<REGION_TO_CATEGORY>": "No object detected."},
|
||||
(10, 10),
|
||||
(np.empty((0, 4), dtype=np.float32), np.array([]), None, None),
|
||||
DoesNotRaise(),
|
||||
),
|
||||
( # Region to Category: detected
|
||||
{"<REGION_TO_CATEGORY>": "some object<loc_300><loc_400><loc_500><loc_600>"},
|
||||
(10, 10),
|
||||
(
|
||||
np.array([[3, 4, 5, 6]], dtype=np.float32),
|
||||
np.array(["some object"]),
|
||||
None,
|
||||
None,
|
||||
),
|
||||
DoesNotRaise(),
|
||||
),
|
||||
( # Region to Description: empty
|
||||
{"<REGION_TO_DESCRIPTION>": "No object detected."},
|
||||
(10, 10),
|
||||
(np.empty((0, 4), dtype=np.float32), np.array([]), None, None),
|
||||
DoesNotRaise(),
|
||||
),
|
||||
( # Region to Description: detected
|
||||
{"<REGION_TO_DESCRIPTION>": "descr<loc_300><loc_400><loc_500><loc_600>"},
|
||||
(10, 10),
|
||||
(
|
||||
np.array([[3, 4, 5, 6]], dtype=np.float32),
|
||||
np.array(["descr"]),
|
||||
None,
|
||||
None,
|
||||
),
|
||||
DoesNotRaise(),
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_florence_2(
|
||||
florence_result: dict,
|
||||
resolution_wh: Tuple[int, int],
|
||||
expected_results: Tuple[
|
||||
np.ndarray, Optional[np.ndarray], Optional[np.ndarray], Optional[np.ndarray]
|
||||
],
|
||||
exception: Exception,
|
||||
) -> None:
|
||||
with exception:
|
||||
result = from_florence_2(florence_result, resolution_wh)
|
||||
np.testing.assert_array_equal(result[0], expected_results[0])
|
||||
if expected_results[1] is None:
|
||||
assert result[1] is None
|
||||
else:
|
||||
np.testing.assert_array_equal(result[1], expected_results[1])
|
||||
if expected_results[2] is None:
|
||||
assert result[2] is None
|
||||
else:
|
||||
np.testing.assert_array_equal(result[2], expected_results[2])
|
||||
if expected_results[3] is None:
|
||||
assert result[3] is None
|
||||
else:
|
||||
np.testing.assert_array_equal(result[3], expected_results[3])
|
||||
@@ -0,0 +1,449 @@
|
||||
from contextlib import ExitStack as DoesNotRaise
|
||||
from typing import List, Optional
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from supervision.detection.overlap_filter import (
|
||||
box_non_max_suppression,
|
||||
group_overlapping_boxes,
|
||||
mask_non_max_suppression,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"predictions, iou_threshold, expected_result, exception",
|
||||
[
|
||||
(
|
||||
np.empty(shape=(0, 5), dtype=float),
|
||||
0.5,
|
||||
[],
|
||||
DoesNotRaise(),
|
||||
),
|
||||
(
|
||||
np.array([[0, 0, 10, 10, 1.0]]),
|
||||
0.5,
|
||||
[[0]],
|
||||
DoesNotRaise(),
|
||||
),
|
||||
(
|
||||
np.array([[0, 0, 10, 10, 1.0], [0, 0, 9, 9, 1.0]]),
|
||||
0.5,
|
||||
[[1, 0]],
|
||||
DoesNotRaise(),
|
||||
), # High overlap, tie-break to second det
|
||||
(
|
||||
np.array([[0, 0, 10, 10, 1.0], [0, 0, 9, 9, 0.99]]),
|
||||
0.5,
|
||||
[[0, 1]],
|
||||
DoesNotRaise(),
|
||||
), # High overlap, merge to high confidence
|
||||
(
|
||||
np.array([[0, 0, 10, 10, 0.99], [0, 0, 9, 9, 1.0]]),
|
||||
0.5,
|
||||
[[1, 0]],
|
||||
DoesNotRaise(),
|
||||
), # (test symmetry) High overlap, merge to high confidence
|
||||
(
|
||||
np.array([[0, 0, 10, 10, 0.90], [0, 0, 9, 9, 1.0]]),
|
||||
0.5,
|
||||
[[1, 0]],
|
||||
DoesNotRaise(),
|
||||
), # (test symmetry) High overlap, merge to high confidence
|
||||
(
|
||||
np.array([[0, 0, 10, 10, 1.0], [0, 0, 9, 9, 1.0]]),
|
||||
1.0,
|
||||
[[1], [0]],
|
||||
DoesNotRaise(),
|
||||
), # High IOU required
|
||||
(
|
||||
np.array([[0, 0, 10, 10, 1.0], [0, 0, 9, 9, 1.0]]),
|
||||
0.0,
|
||||
[[1, 0]],
|
||||
DoesNotRaise(),
|
||||
), # No IOU required
|
||||
(
|
||||
np.array([[0, 0, 10, 10, 1.0], [0, 0, 5, 5, 0.9]]),
|
||||
0.25,
|
||||
[[0, 1]],
|
||||
DoesNotRaise(),
|
||||
), # Below IOU requirement
|
||||
(
|
||||
np.array([[0, 0, 10, 10, 1.0], [0, 0, 5, 5, 0.9]]),
|
||||
0.26,
|
||||
[[0], [1]],
|
||||
DoesNotRaise(),
|
||||
), # Above IOU requirement
|
||||
(
|
||||
np.array([[0, 0, 10, 10, 1.0], [0, 0, 9, 9, 1.0], [0, 0, 8, 8, 1.0]]),
|
||||
0.5,
|
||||
[[2, 1, 0]],
|
||||
DoesNotRaise(),
|
||||
), # 3 boxes
|
||||
(
|
||||
np.array(
|
||||
[
|
||||
[0, 0, 10, 10, 1.0],
|
||||
[0, 0, 9, 9, 1.0],
|
||||
[5, 5, 10, 10, 1.0],
|
||||
[6, 6, 10, 10, 1.0],
|
||||
[9, 9, 10, 10, 1.0],
|
||||
]
|
||||
),
|
||||
0.5,
|
||||
[[4], [3, 2], [1, 0]],
|
||||
DoesNotRaise(),
|
||||
), # 5 boxes, 2 merges, 1 separate
|
||||
(
|
||||
np.array(
|
||||
[
|
||||
[0, 0, 2, 1, 1.0],
|
||||
[1, 0, 3, 1, 1.0],
|
||||
[2, 0, 4, 1, 1.0],
|
||||
[3, 0, 5, 1, 1.0],
|
||||
[4, 0, 6, 1, 1.0],
|
||||
]
|
||||
),
|
||||
0.33,
|
||||
[[4, 3], [2, 1], [0]],
|
||||
DoesNotRaise(),
|
||||
), # sequential merge, half overlap
|
||||
(
|
||||
np.array(
|
||||
[
|
||||
[0, 0, 2, 1, 0.9],
|
||||
[1, 0, 3, 1, 0.9],
|
||||
[2, 0, 4, 1, 1.0],
|
||||
[3, 0, 5, 1, 0.9],
|
||||
[4, 0, 6, 1, 0.9],
|
||||
]
|
||||
),
|
||||
0.33,
|
||||
[[2, 3, 1], [4], [0]],
|
||||
DoesNotRaise(),
|
||||
), # confidence
|
||||
],
|
||||
)
|
||||
def test_group_overlapping_boxes(
|
||||
predictions: np.ndarray,
|
||||
iou_threshold: float,
|
||||
expected_result: List[List[int]],
|
||||
exception: Exception,
|
||||
) -> None:
|
||||
with exception:
|
||||
result = group_overlapping_boxes(
|
||||
predictions=predictions, iou_threshold=iou_threshold
|
||||
)
|
||||
|
||||
assert result == expected_result
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"predictions, iou_threshold, expected_result, exception",
|
||||
[
|
||||
(
|
||||
np.empty(shape=(0, 5)),
|
||||
0.5,
|
||||
np.array([]),
|
||||
DoesNotRaise(),
|
||||
), # single box with no category
|
||||
(
|
||||
np.array([[10.0, 10.0, 40.0, 40.0, 0.8]]),
|
||||
0.5,
|
||||
np.array([True]),
|
||||
DoesNotRaise(),
|
||||
), # single box with no category
|
||||
(
|
||||
np.array([[10.0, 10.0, 40.0, 40.0, 0.8, 0]]),
|
||||
0.5,
|
||||
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],
|
||||
]
|
||||
),
|
||||
0.5,
|
||||
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],
|
||||
]
|
||||
),
|
||||
0.5,
|
||||
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],
|
||||
]
|
||||
),
|
||||
0.5,
|
||||
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],
|
||||
]
|
||||
),
|
||||
0.5,
|
||||
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],
|
||||
]
|
||||
),
|
||||
0.5,
|
||||
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],
|
||||
]
|
||||
),
|
||||
0.5,
|
||||
np.array([False, True, True]),
|
||||
DoesNotRaise(),
|
||||
), # three boxes with different category
|
||||
],
|
||||
)
|
||||
def test_box_non_max_suppression(
|
||||
predictions: np.ndarray,
|
||||
iou_threshold: float,
|
||||
expected_result: Optional[np.ndarray],
|
||||
exception: Exception,
|
||||
) -> None:
|
||||
with exception:
|
||||
result = box_non_max_suppression(
|
||||
predictions=predictions, iou_threshold=iou_threshold
|
||||
)
|
||||
assert np.array_equal(result, expected_result)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"predictions, masks, iou_threshold, expected_result, exception",
|
||||
[
|
||||
(
|
||||
np.empty((0, 6)),
|
||||
np.empty((0, 5, 5)),
|
||||
0.5,
|
||||
np.array([]),
|
||||
DoesNotRaise(),
|
||||
), # empty predictions and masks
|
||||
(
|
||||
np.array([[0, 0, 0, 0, 0.8]]),
|
||||
np.array(
|
||||
[
|
||||
[
|
||||
[False, False, False, False, False],
|
||||
[False, True, True, True, False],
|
||||
[False, True, True, True, False],
|
||||
[False, True, True, True, False],
|
||||
[False, False, False, False, False],
|
||||
]
|
||||
]
|
||||
),
|
||||
0.5,
|
||||
np.array([True]),
|
||||
DoesNotRaise(),
|
||||
), # single mask with no category
|
||||
(
|
||||
np.array([[0, 0, 0, 0, 0.8, 0]]),
|
||||
np.array(
|
||||
[
|
||||
[
|
||||
[False, False, False, False, False],
|
||||
[False, True, True, True, False],
|
||||
[False, True, True, True, False],
|
||||
[False, True, True, True, False],
|
||||
[False, False, False, False, False],
|
||||
]
|
||||
]
|
||||
),
|
||||
0.5,
|
||||
np.array([True]),
|
||||
DoesNotRaise(),
|
||||
), # single mask with category
|
||||
(
|
||||
np.array([[0, 0, 0, 0, 0.8], [0, 0, 0, 0, 0.9]]),
|
||||
np.array(
|
||||
[
|
||||
[
|
||||
[False, False, False, False, False],
|
||||
[False, True, True, False, False],
|
||||
[False, True, True, False, False],
|
||||
[False, False, False, False, False],
|
||||
[False, False, False, False, False],
|
||||
],
|
||||
[
|
||||
[False, False, False, False, False],
|
||||
[False, False, False, False, False],
|
||||
[False, False, False, True, True],
|
||||
[False, False, False, True, True],
|
||||
[False, False, False, False, False],
|
||||
],
|
||||
]
|
||||
),
|
||||
0.5,
|
||||
np.array([True, True]),
|
||||
DoesNotRaise(),
|
||||
), # two masks non-overlapping with no category
|
||||
(
|
||||
np.array([[0, 0, 0, 0, 0.8], [0, 0, 0, 0, 0.9]]),
|
||||
np.array(
|
||||
[
|
||||
[
|
||||
[False, False, False, False, False],
|
||||
[False, True, True, True, False],
|
||||
[False, True, True, True, False],
|
||||
[False, True, True, True, False],
|
||||
[False, False, False, False, False],
|
||||
],
|
||||
[
|
||||
[False, False, False, False, False],
|
||||
[False, False, True, True, True],
|
||||
[False, False, True, True, True],
|
||||
[False, False, True, True, True],
|
||||
[False, False, False, False, False],
|
||||
],
|
||||
]
|
||||
),
|
||||
0.4,
|
||||
np.array([False, True]),
|
||||
DoesNotRaise(),
|
||||
), # two masks partially overlapping with no category
|
||||
(
|
||||
np.array([[0, 0, 0, 0, 0.8, 0], [0, 0, 0, 0, 0.9, 1]]),
|
||||
np.array(
|
||||
[
|
||||
[
|
||||
[False, False, False, False, False],
|
||||
[False, True, True, True, False],
|
||||
[False, True, True, True, False],
|
||||
[False, True, True, True, False],
|
||||
[False, False, False, False, False],
|
||||
],
|
||||
[
|
||||
[False, False, False, False, False],
|
||||
[False, False, True, True, True],
|
||||
[False, False, True, True, True],
|
||||
[False, False, True, True, True],
|
||||
[False, False, False, False, False],
|
||||
],
|
||||
]
|
||||
),
|
||||
0.5,
|
||||
np.array([True, True]),
|
||||
DoesNotRaise(),
|
||||
), # two masks partially overlapping with different category
|
||||
(
|
||||
np.array(
|
||||
[
|
||||
[0, 0, 0, 0, 0.8],
|
||||
[0, 0, 0, 0, 0.85],
|
||||
[0, 0, 0, 0, 0.9],
|
||||
]
|
||||
),
|
||||
np.array(
|
||||
[
|
||||
[
|
||||
[False, False, False, False, False],
|
||||
[False, True, True, False, False],
|
||||
[False, True, True, False, False],
|
||||
[False, False, False, False, False],
|
||||
[False, False, False, False, False],
|
||||
],
|
||||
[
|
||||
[False, False, False, False, False],
|
||||
[False, True, True, False, False],
|
||||
[False, True, True, False, False],
|
||||
[False, False, False, False, False],
|
||||
[False, False, False, False, False],
|
||||
],
|
||||
[
|
||||
[False, False, False, False, False],
|
||||
[False, False, False, True, True],
|
||||
[False, False, False, True, True],
|
||||
[False, False, False, False, False],
|
||||
[False, False, False, False, False],
|
||||
],
|
||||
]
|
||||
),
|
||||
0.5,
|
||||
np.array([False, True, True]),
|
||||
DoesNotRaise(),
|
||||
), # three masks with no category
|
||||
(
|
||||
np.array(
|
||||
[
|
||||
[0, 0, 0, 0, 0.8, 0],
|
||||
[0, 0, 0, 0, 0.85, 1],
|
||||
[0, 0, 0, 0, 0.9, 2],
|
||||
]
|
||||
),
|
||||
np.array(
|
||||
[
|
||||
[
|
||||
[False, False, False, False, False],
|
||||
[False, True, True, False, False],
|
||||
[False, True, True, False, False],
|
||||
[False, False, False, False, False],
|
||||
[False, False, False, False, False],
|
||||
],
|
||||
[
|
||||
[False, False, False, False, False],
|
||||
[False, True, True, False, False],
|
||||
[False, True, True, False, False],
|
||||
[False, True, True, False, False],
|
||||
[False, False, False, False, False],
|
||||
],
|
||||
[
|
||||
[False, False, False, False, False],
|
||||
[False, True, True, False, False],
|
||||
[False, True, True, False, False],
|
||||
[False, False, False, False, False],
|
||||
[False, False, False, False, False],
|
||||
],
|
||||
]
|
||||
),
|
||||
0.5,
|
||||
np.array([True, True, True]),
|
||||
DoesNotRaise(),
|
||||
), # three masks with different category
|
||||
],
|
||||
)
|
||||
def test_mask_non_max_suppression(
|
||||
predictions: np.ndarray,
|
||||
masks: np.ndarray,
|
||||
iou_threshold: float,
|
||||
expected_result: Optional[np.ndarray],
|
||||
exception: Exception,
|
||||
) -> None:
|
||||
with exception:
|
||||
result = mask_non_max_suppression(
|
||||
predictions=predictions, masks=masks, iou_threshold=iou_threshold
|
||||
)
|
||||
assert np.array_equal(result, expected_result)
|
||||
@@ -92,3 +92,19 @@ def test_polygon_zone_trigger(
|
||||
with exception:
|
||||
in_zone = polygon_zone.trigger(detections)
|
||||
assert np.all(in_zone == expected_results)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"polygon, triggering_anchors, exception",
|
||||
[
|
||||
(POLYGON, [sv.Position.CENTER], DoesNotRaise()),
|
||||
(
|
||||
POLYGON,
|
||||
[],
|
||||
pytest.raises(ValueError),
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_polygon_zone_initialization(polygon, triggering_anchors, exception):
|
||||
with exception:
|
||||
sv.PolygonZone(polygon, FRAME_RESOLUTION, triggering_anchors=triggering_anchors)
|
||||
|
||||
+249
-314
@@ -2,337 +2,29 @@ from contextlib import ExitStack as DoesNotRaise
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
import numpy as np
|
||||
import numpy.typing as npt
|
||||
import pytest
|
||||
|
||||
from supervision.config import CLASS_NAME_DATA_FIELD
|
||||
from supervision.detection.utils import (
|
||||
box_non_max_suppression,
|
||||
calculate_masks_centroids,
|
||||
clip_boxes,
|
||||
contains_holes,
|
||||
contains_multiple_segments,
|
||||
filter_polygons_by_area,
|
||||
get_data_item,
|
||||
mask_non_max_suppression,
|
||||
merge_data,
|
||||
move_boxes,
|
||||
process_roboflow_result,
|
||||
scale_boxes,
|
||||
xcycwh_to_xyxy,
|
||||
xywh_to_xyxy,
|
||||
)
|
||||
|
||||
TEST_MASK = np.zeros((1, 1000, 1000), dtype=bool)
|
||||
TEST_MASK[:, 300:351, 200:251] = True
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"predictions, iou_threshold, expected_result, exception",
|
||||
[
|
||||
(
|
||||
np.empty(shape=(0, 5)),
|
||||
0.5,
|
||||
np.array([]),
|
||||
DoesNotRaise(),
|
||||
), # single box with no category
|
||||
(
|
||||
np.array([[10.0, 10.0, 40.0, 40.0, 0.8]]),
|
||||
0.5,
|
||||
np.array([True]),
|
||||
DoesNotRaise(),
|
||||
), # single box with no category
|
||||
(
|
||||
np.array([[10.0, 10.0, 40.0, 40.0, 0.8, 0]]),
|
||||
0.5,
|
||||
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],
|
||||
]
|
||||
),
|
||||
0.5,
|
||||
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],
|
||||
]
|
||||
),
|
||||
0.5,
|
||||
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],
|
||||
]
|
||||
),
|
||||
0.5,
|
||||
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],
|
||||
]
|
||||
),
|
||||
0.5,
|
||||
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],
|
||||
]
|
||||
),
|
||||
0.5,
|
||||
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],
|
||||
]
|
||||
),
|
||||
0.5,
|
||||
np.array([False, True, True]),
|
||||
DoesNotRaise(),
|
||||
), # three boxes with different category
|
||||
],
|
||||
)
|
||||
def test_box_non_max_suppression(
|
||||
predictions: np.ndarray,
|
||||
iou_threshold: float,
|
||||
expected_result: Optional[np.ndarray],
|
||||
exception: Exception,
|
||||
) -> None:
|
||||
with exception:
|
||||
result = box_non_max_suppression(
|
||||
predictions=predictions, iou_threshold=iou_threshold
|
||||
)
|
||||
assert np.array_equal(result, expected_result)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"predictions, masks, iou_threshold, expected_result, exception",
|
||||
[
|
||||
(
|
||||
np.empty((0, 6)),
|
||||
np.empty((0, 5, 5)),
|
||||
0.5,
|
||||
np.array([]),
|
||||
DoesNotRaise(),
|
||||
), # empty predictions and masks
|
||||
(
|
||||
np.array([[0, 0, 0, 0, 0.8]]),
|
||||
np.array(
|
||||
[
|
||||
[
|
||||
[False, False, False, False, False],
|
||||
[False, True, True, True, False],
|
||||
[False, True, True, True, False],
|
||||
[False, True, True, True, False],
|
||||
[False, False, False, False, False],
|
||||
]
|
||||
]
|
||||
),
|
||||
0.5,
|
||||
np.array([True]),
|
||||
DoesNotRaise(),
|
||||
), # single mask with no category
|
||||
(
|
||||
np.array([[0, 0, 0, 0, 0.8, 0]]),
|
||||
np.array(
|
||||
[
|
||||
[
|
||||
[False, False, False, False, False],
|
||||
[False, True, True, True, False],
|
||||
[False, True, True, True, False],
|
||||
[False, True, True, True, False],
|
||||
[False, False, False, False, False],
|
||||
]
|
||||
]
|
||||
),
|
||||
0.5,
|
||||
np.array([True]),
|
||||
DoesNotRaise(),
|
||||
), # single mask with category
|
||||
(
|
||||
np.array([[0, 0, 0, 0, 0.8], [0, 0, 0, 0, 0.9]]),
|
||||
np.array(
|
||||
[
|
||||
[
|
||||
[False, False, False, False, False],
|
||||
[False, True, True, False, False],
|
||||
[False, True, True, False, False],
|
||||
[False, False, False, False, False],
|
||||
[False, False, False, False, False],
|
||||
],
|
||||
[
|
||||
[False, False, False, False, False],
|
||||
[False, False, False, False, False],
|
||||
[False, False, False, True, True],
|
||||
[False, False, False, True, True],
|
||||
[False, False, False, False, False],
|
||||
],
|
||||
]
|
||||
),
|
||||
0.5,
|
||||
np.array([True, True]),
|
||||
DoesNotRaise(),
|
||||
), # two masks non-overlapping with no category
|
||||
(
|
||||
np.array([[0, 0, 0, 0, 0.8], [0, 0, 0, 0, 0.9]]),
|
||||
np.array(
|
||||
[
|
||||
[
|
||||
[False, False, False, False, False],
|
||||
[False, True, True, True, False],
|
||||
[False, True, True, True, False],
|
||||
[False, True, True, True, False],
|
||||
[False, False, False, False, False],
|
||||
],
|
||||
[
|
||||
[False, False, False, False, False],
|
||||
[False, False, True, True, True],
|
||||
[False, False, True, True, True],
|
||||
[False, False, True, True, True],
|
||||
[False, False, False, False, False],
|
||||
],
|
||||
]
|
||||
),
|
||||
0.4,
|
||||
np.array([False, True]),
|
||||
DoesNotRaise(),
|
||||
), # two masks partially overlapping with no category
|
||||
(
|
||||
np.array([[0, 0, 0, 0, 0.8, 0], [0, 0, 0, 0, 0.9, 1]]),
|
||||
np.array(
|
||||
[
|
||||
[
|
||||
[False, False, False, False, False],
|
||||
[False, True, True, True, False],
|
||||
[False, True, True, True, False],
|
||||
[False, True, True, True, False],
|
||||
[False, False, False, False, False],
|
||||
],
|
||||
[
|
||||
[False, False, False, False, False],
|
||||
[False, False, True, True, True],
|
||||
[False, False, True, True, True],
|
||||
[False, False, True, True, True],
|
||||
[False, False, False, False, False],
|
||||
],
|
||||
]
|
||||
),
|
||||
0.5,
|
||||
np.array([True, True]),
|
||||
DoesNotRaise(),
|
||||
), # two masks partially overlapping with different category
|
||||
(
|
||||
np.array(
|
||||
[
|
||||
[0, 0, 0, 0, 0.8],
|
||||
[0, 0, 0, 0, 0.85],
|
||||
[0, 0, 0, 0, 0.9],
|
||||
]
|
||||
),
|
||||
np.array(
|
||||
[
|
||||
[
|
||||
[False, False, False, False, False],
|
||||
[False, True, True, False, False],
|
||||
[False, True, True, False, False],
|
||||
[False, False, False, False, False],
|
||||
[False, False, False, False, False],
|
||||
],
|
||||
[
|
||||
[False, False, False, False, False],
|
||||
[False, True, True, False, False],
|
||||
[False, True, True, False, False],
|
||||
[False, False, False, False, False],
|
||||
[False, False, False, False, False],
|
||||
],
|
||||
[
|
||||
[False, False, False, False, False],
|
||||
[False, False, False, True, True],
|
||||
[False, False, False, True, True],
|
||||
[False, False, False, False, False],
|
||||
[False, False, False, False, False],
|
||||
],
|
||||
]
|
||||
),
|
||||
0.5,
|
||||
np.array([False, True, True]),
|
||||
DoesNotRaise(),
|
||||
), # three masks with no category
|
||||
(
|
||||
np.array(
|
||||
[
|
||||
[0, 0, 0, 0, 0.8, 0],
|
||||
[0, 0, 0, 0, 0.85, 1],
|
||||
[0, 0, 0, 0, 0.9, 2],
|
||||
]
|
||||
),
|
||||
np.array(
|
||||
[
|
||||
[
|
||||
[False, False, False, False, False],
|
||||
[False, True, True, False, False],
|
||||
[False, True, True, False, False],
|
||||
[False, False, False, False, False],
|
||||
[False, False, False, False, False],
|
||||
],
|
||||
[
|
||||
[False, False, False, False, False],
|
||||
[False, True, True, False, False],
|
||||
[False, True, True, False, False],
|
||||
[False, True, True, False, False],
|
||||
[False, False, False, False, False],
|
||||
],
|
||||
[
|
||||
[False, False, False, False, False],
|
||||
[False, True, True, False, False],
|
||||
[False, True, True, False, False],
|
||||
[False, False, False, False, False],
|
||||
[False, False, False, False, False],
|
||||
],
|
||||
]
|
||||
),
|
||||
0.5,
|
||||
np.array([True, True, True]),
|
||||
DoesNotRaise(),
|
||||
), # three masks with different category
|
||||
],
|
||||
)
|
||||
def test_mask_non_max_suppression(
|
||||
predictions: np.ndarray,
|
||||
masks: np.ndarray,
|
||||
iou_threshold: float,
|
||||
expected_result: Optional[np.ndarray],
|
||||
exception: Exception,
|
||||
) -> None:
|
||||
with exception:
|
||||
result = mask_non_max_suppression(
|
||||
predictions=predictions, masks=masks, iou_threshold=iou_threshold
|
||||
)
|
||||
assert np.array_equal(result, expected_result)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"xyxy, resolution_wh, expected_result",
|
||||
[
|
||||
@@ -911,6 +603,14 @@ def test_calculate_masks_centroids(
|
||||
{"test_1": []},
|
||||
DoesNotRaise(),
|
||||
), # single data dict with a single field name and empty list values
|
||||
(
|
||||
[
|
||||
{"test_1": []},
|
||||
{"test_1": []},
|
||||
],
|
||||
{"test_1": []},
|
||||
DoesNotRaise(),
|
||||
), # two data dicts with the same field name and empty list values
|
||||
(
|
||||
[
|
||||
{"test_1": np.array([])},
|
||||
@@ -918,6 +618,14 @@ def test_calculate_masks_centroids(
|
||||
{"test_1": np.array([])},
|
||||
DoesNotRaise(),
|
||||
), # single data dict with a single field name and empty np.array values
|
||||
(
|
||||
[
|
||||
{"test_1": np.array([])},
|
||||
{"test_1": np.array([])},
|
||||
],
|
||||
{"test_1": np.array([])},
|
||||
DoesNotRaise(),
|
||||
), # two data dicts with the same field name and empty np.array values
|
||||
(
|
||||
[
|
||||
{"test_1": [1, 2, 3]},
|
||||
@@ -932,7 +640,7 @@ def test_calculate_masks_centroids(
|
||||
],
|
||||
{"test_1": [3, 2, 1]},
|
||||
DoesNotRaise(),
|
||||
), # two data dicts with the same field name and empty and list values
|
||||
), # two data dicts with the same field name; one of with empty list as value
|
||||
(
|
||||
[
|
||||
{"test_1": [1, 2, 3]},
|
||||
@@ -1012,6 +720,49 @@ def test_calculate_masks_centroids(
|
||||
None,
|
||||
pytest.raises(ValueError),
|
||||
), # two data dicts with the same field name and different length arrays values
|
||||
(
|
||||
[{}, {"test_1": [1, 2, 3]}],
|
||||
None,
|
||||
pytest.raises(ValueError),
|
||||
), # two data dicts; one empty and one non-empty dict
|
||||
(
|
||||
[{"test_1": [], "test_2": []}, {"test_1": [1, 2, 3], "test_2": [1, 2, 3]}],
|
||||
{"test_1": [1, 2, 3], "test_2": [1, 2, 3]},
|
||||
DoesNotRaise(),
|
||||
), # two data dicts; one empty and one non-empty dict; same keys
|
||||
(
|
||||
[{"test_1": []}, {"test_1": [1, 2, 3], "test_2": [4, 5, 6]}],
|
||||
None,
|
||||
pytest.raises(ValueError),
|
||||
), # two data dicts; one empty and one non-empty dict; different keys
|
||||
(
|
||||
[
|
||||
{
|
||||
"test_1": [1, 2, 3],
|
||||
"test_2": [4, 5, 6],
|
||||
"test_3": [7, 8, 9],
|
||||
},
|
||||
{"test_1": [1, 2, 3], "test_2": [4, 5, 6]},
|
||||
],
|
||||
None,
|
||||
pytest.raises(ValueError),
|
||||
), # two data dicts; one with three keys, one with two keys
|
||||
(
|
||||
[
|
||||
{"test_1": [1, 2, 3]},
|
||||
{"test_1": [1, 2, 3], "test_2": [1, 2, 3]},
|
||||
],
|
||||
None,
|
||||
pytest.raises(ValueError),
|
||||
), # some keys missing in one dict
|
||||
(
|
||||
[
|
||||
{"test_1": [1, 2, 3], "test_2": ["a", "b"]},
|
||||
{"test_1": [4, 5], "test_2": ["c", "d", "e"]},
|
||||
],
|
||||
None,
|
||||
pytest.raises(ValueError),
|
||||
), # different value lengths for the same key
|
||||
],
|
||||
)
|
||||
def test_merge_data(
|
||||
@@ -1021,6 +772,9 @@ def test_merge_data(
|
||||
):
|
||||
with exception:
|
||||
result = merge_data(data_list=data_list)
|
||||
if expected_result is None:
|
||||
assert False, f"Expected an error, but got result {result}"
|
||||
|
||||
for key in result:
|
||||
if isinstance(result[key], np.ndarray):
|
||||
assert np.array_equal(
|
||||
@@ -1203,3 +957,184 @@ def test_get_data_item(
|
||||
assert (
|
||||
result[key] == expected_result[key]
|
||||
), f"Mismatch in non-array data for key {key}"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"mask, expected_result, exception",
|
||||
[
|
||||
(
|
||||
np.array([[0, 0, 0, 0], [0, 1, 1, 0], [0, 1, 0, 0], [0, 1, 1, 0]]).astype(
|
||||
bool
|
||||
),
|
||||
False,
|
||||
DoesNotRaise(),
|
||||
), # foreground object in one continuous piece
|
||||
(
|
||||
np.array([[1, 0, 0, 0], [1, 0, 0, 0], [0, 0, 0, 0], [0, 1, 1, 0]]).astype(
|
||||
bool
|
||||
),
|
||||
False,
|
||||
DoesNotRaise(),
|
||||
), # foreground object in 2 seperate elements
|
||||
(
|
||||
np.array([[0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0]]).astype(
|
||||
bool
|
||||
),
|
||||
False,
|
||||
DoesNotRaise(),
|
||||
), # no foreground pixels in mask
|
||||
(
|
||||
np.array([[1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 1, 1]]).astype(
|
||||
bool
|
||||
),
|
||||
False,
|
||||
DoesNotRaise(),
|
||||
), # only foreground pixels in mask
|
||||
(
|
||||
np.array([[1, 1, 1, 0], [1, 0, 1, 0], [1, 1, 1, 0], [0, 0, 0, 0]]).astype(
|
||||
bool
|
||||
),
|
||||
True,
|
||||
DoesNotRaise(),
|
||||
), # foreground object has 1 hole
|
||||
(
|
||||
np.array([[1, 1, 1, 0], [1, 0, 1, 1], [1, 1, 0, 1], [0, 1, 1, 1]]).astype(
|
||||
bool
|
||||
),
|
||||
True,
|
||||
DoesNotRaise(),
|
||||
), # foreground object has 2 holes
|
||||
],
|
||||
)
|
||||
def test_contains_holes(
|
||||
mask: npt.NDArray[np.bool_], expected_result: bool, exception: Exception
|
||||
) -> None:
|
||||
with exception:
|
||||
result = contains_holes(mask)
|
||||
assert result == expected_result
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"mask, connectivity, expected_result, exception",
|
||||
[
|
||||
(
|
||||
np.array([[0, 0, 0, 0], [0, 1, 1, 0], [0, 1, 0, 0], [0, 1, 1, 0]]).astype(
|
||||
bool
|
||||
),
|
||||
4,
|
||||
False,
|
||||
DoesNotRaise(),
|
||||
), # foreground object in one continuous piece
|
||||
(
|
||||
np.array([[1, 0, 0, 0], [1, 0, 0, 0], [0, 0, 0, 0], [0, 1, 1, 0]]).astype(
|
||||
bool
|
||||
),
|
||||
4,
|
||||
True,
|
||||
DoesNotRaise(),
|
||||
), # foreground object in 2 seperate elements
|
||||
(
|
||||
np.array([[0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0]]).astype(
|
||||
bool
|
||||
),
|
||||
4,
|
||||
False,
|
||||
DoesNotRaise(),
|
||||
), # no foreground pixels in mask
|
||||
(
|
||||
np.array([[1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 1, 1]]).astype(
|
||||
bool
|
||||
),
|
||||
4,
|
||||
False,
|
||||
DoesNotRaise(),
|
||||
), # only foreground pixels in mask
|
||||
(
|
||||
np.array([[1, 1, 1, 0], [1, 0, 1, 1], [1, 1, 0, 1], [0, 1, 1, 1]]).astype(
|
||||
bool
|
||||
),
|
||||
4,
|
||||
False,
|
||||
DoesNotRaise(),
|
||||
), # foreground object has 2 holes, but is in single piece
|
||||
(
|
||||
np.array([[1, 1, 0, 0], [1, 1, 0, 1], [1, 0, 1, 1], [0, 0, 1, 1]]).astype(
|
||||
bool
|
||||
),
|
||||
4,
|
||||
True,
|
||||
DoesNotRaise(),
|
||||
), # foreground object in 2 elements with respect to 4-way connectivity
|
||||
(
|
||||
np.array([[1, 1, 0, 0], [1, 1, 0, 1], [1, 0, 1, 1], [0, 0, 1, 1]]).astype(
|
||||
bool
|
||||
),
|
||||
8,
|
||||
False,
|
||||
DoesNotRaise(),
|
||||
), # foreground object in single piece with respect to 8-way connectivity
|
||||
(
|
||||
np.array([[1, 1, 0, 0], [1, 1, 0, 1], [1, 0, 1, 1], [0, 0, 1, 1]]).astype(
|
||||
bool
|
||||
),
|
||||
5,
|
||||
None,
|
||||
pytest.raises(ValueError),
|
||||
), # Incorrect connectivity parameter value, raises ValueError
|
||||
],
|
||||
)
|
||||
def test_contains_multiple_segments(
|
||||
mask: npt.NDArray[np.bool_],
|
||||
connectivity: int,
|
||||
expected_result: bool,
|
||||
exception: Exception,
|
||||
) -> None:
|
||||
with exception:
|
||||
result = contains_multiple_segments(mask=mask, connectivity=connectivity)
|
||||
assert result == expected_result
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"xywh, expected_result",
|
||||
[
|
||||
(np.array([[10, 20, 30, 40]]), np.array([[10, 20, 40, 60]])), # standard case
|
||||
(np.array([[0, 0, 0, 0]]), np.array([[0, 0, 0, 0]])), # zero size bounding box
|
||||
(
|
||||
np.array([[50, 50, 100, 100]]),
|
||||
np.array([[50, 50, 150, 150]]),
|
||||
), # large bounding box
|
||||
(
|
||||
np.array([[-10, -20, 30, 40]]),
|
||||
np.array([[-10, -20, 20, 20]]),
|
||||
), # negative coordinates
|
||||
(np.array([[50, 50, 0, 30]]), np.array([[50, 50, 50, 80]])), # zero width
|
||||
(np.array([[50, 50, 20, 0]]), np.array([[50, 50, 70, 50]])), # zero height
|
||||
(np.array([]).reshape(0, 4), np.array([]).reshape(0, 4)), # empty array
|
||||
],
|
||||
)
|
||||
def test_xywh_to_xyxy(xywh: np.ndarray, expected_result: np.ndarray) -> None:
|
||||
result = xywh_to_xyxy(xywh)
|
||||
np.testing.assert_array_equal(result, expected_result)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"xcycwh, expected_result",
|
||||
[
|
||||
(np.array([[50, 50, 20, 30]]), np.array([[40, 35, 60, 65]])), # standard case
|
||||
(np.array([[0, 0, 0, 0]]), np.array([[0, 0, 0, 0]])), # zero size bounding box
|
||||
(
|
||||
np.array([[50, 50, 100, 100]]),
|
||||
np.array([[0, 0, 100, 100]]),
|
||||
), # large bounding box centered at (50, 50)
|
||||
(
|
||||
np.array([[-10, -10, 20, 30]]),
|
||||
np.array([[-20, -25, 0, 5]]),
|
||||
), # negative coordinates
|
||||
(np.array([[50, 50, 0, 30]]), np.array([[50, 35, 50, 65]])), # zero width
|
||||
(np.array([[50, 50, 20, 0]]), np.array([[40, 50, 60, 50]])), # zero height
|
||||
(np.array([]).reshape(0, 4), np.array([]).reshape(0, 4)), # empty array
|
||||
],
|
||||
)
|
||||
def test_xcycwh_to_xyxy(xcycwh: np.ndarray, expected_result: np.ndarray) -> None:
|
||||
result = xcycwh_to_xyxy(xcycwh)
|
||||
np.testing.assert_array_equal(result, expected_result)
|
||||
|
||||
@@ -11,7 +11,7 @@ ALL_IMAGES_LIST = [os.path.join(ASSETS_DIR, f"{i}.jpg") for i in range(1, 6)]
|
||||
|
||||
|
||||
@fixture(scope="function")
|
||||
def empty_opencv_image() -> np.ndarray:
|
||||
def empty_cv2_image() -> np.ndarray:
|
||||
return np.zeros((128, 128, 3), dtype=np.uint8)
|
||||
|
||||
|
||||
|
||||
@@ -2,28 +2,28 @@ import numpy as np
|
||||
from PIL import Image, ImageChops
|
||||
|
||||
from supervision.utils.conversion import (
|
||||
convert_for_image_processing,
|
||||
cv2_to_pillow,
|
||||
ensure_cv2_image_for_processing,
|
||||
images_to_cv2,
|
||||
pillow_to_cv2,
|
||||
)
|
||||
|
||||
|
||||
def test_convert_for_image_processing_when_pillow_image_submitted(
|
||||
empty_opencv_image: np.ndarray, empty_pillow_image: Image.Image
|
||||
def test_ensure_cv2_image_for_processing_when_pillow_image_submitted(
|
||||
empty_cv2_image: np.ndarray, empty_pillow_image: Image.Image
|
||||
) -> None:
|
||||
# given
|
||||
param_a_value = 3
|
||||
param_b_value = "some"
|
||||
|
||||
@convert_for_image_processing
|
||||
@ensure_cv2_image_for_processing
|
||||
def my_custom_processing_function(
|
||||
image: np.ndarray,
|
||||
param_a: int,
|
||||
param_b: str,
|
||||
) -> np.ndarray:
|
||||
assert np.allclose(
|
||||
image, empty_opencv_image
|
||||
image, empty_cv2_image
|
||||
), "Expected conversion to OpenCV image to happen"
|
||||
assert (
|
||||
param_a == param_a_value
|
||||
@@ -48,21 +48,21 @@ def test_convert_for_image_processing_when_pillow_image_submitted(
|
||||
)
|
||||
|
||||
|
||||
def test_convert_for_image_processing_when_opencv_image_submitted(
|
||||
empty_opencv_image: np.ndarray,
|
||||
def test_ensure_cv2_image_for_processing_when_cv2_image_submitted(
|
||||
empty_cv2_image: np.ndarray,
|
||||
) -> None:
|
||||
# given
|
||||
param_a_value = 3
|
||||
param_b_value = "some"
|
||||
|
||||
@convert_for_image_processing
|
||||
@ensure_cv2_image_for_processing
|
||||
def my_custom_processing_function(
|
||||
image: np.ndarray,
|
||||
param_a: int,
|
||||
param_b: str,
|
||||
) -> np.ndarray:
|
||||
assert np.allclose(
|
||||
image, empty_opencv_image
|
||||
image, empty_cv2_image
|
||||
), "Expected conversion to OpenCV image to happen"
|
||||
assert (
|
||||
param_a == param_a_value
|
||||
@@ -74,22 +74,20 @@ def test_convert_for_image_processing_when_opencv_image_submitted(
|
||||
|
||||
# when
|
||||
result = my_custom_processing_function(
|
||||
empty_opencv_image,
|
||||
empty_cv2_image,
|
||||
param_a_value,
|
||||
param_b=param_b_value,
|
||||
)
|
||||
|
||||
# then
|
||||
assert (
|
||||
result is empty_opencv_image
|
||||
), "Expected to return OpenCV image without changes"
|
||||
assert result is empty_cv2_image, "Expected to return OpenCV image without changes"
|
||||
|
||||
|
||||
def test_cv2_to_pillow(
|
||||
empty_opencv_image: np.ndarray, empty_pillow_image: Image.Image
|
||||
empty_cv2_image: np.ndarray, empty_pillow_image: Image.Image
|
||||
) -> None:
|
||||
# when
|
||||
result = cv2_to_pillow(image=empty_opencv_image)
|
||||
result = cv2_to_pillow(image=empty_cv2_image)
|
||||
|
||||
# then
|
||||
difference = ImageChops.difference(result, empty_pillow_image)
|
||||
@@ -99,14 +97,14 @@ def test_cv2_to_pillow(
|
||||
|
||||
|
||||
def test_pillow_to_cv2(
|
||||
empty_opencv_image: np.ndarray, empty_pillow_image: Image.Image
|
||||
empty_cv2_image: np.ndarray, empty_pillow_image: Image.Image
|
||||
) -> None:
|
||||
# when
|
||||
result = pillow_to_cv2(image=empty_pillow_image)
|
||||
|
||||
# then
|
||||
assert np.allclose(
|
||||
result, empty_opencv_image
|
||||
result, empty_cv2_image
|
||||
), "Conversion to OpenCV image expected not to change the content of image"
|
||||
|
||||
|
||||
@@ -119,10 +117,10 @@ def test_images_to_cv2_when_empty_input_provided() -> None:
|
||||
|
||||
|
||||
def test_images_to_cv2_when_only_cv2_images_provided(
|
||||
empty_opencv_image: np.ndarray,
|
||||
empty_cv2_image: np.ndarray,
|
||||
) -> None:
|
||||
# given
|
||||
images = [empty_opencv_image] * 5
|
||||
images = [empty_cv2_image] * 5
|
||||
|
||||
# when
|
||||
result = images_to_cv2(images=images)
|
||||
@@ -131,13 +129,13 @@ def test_images_to_cv2_when_only_cv2_images_provided(
|
||||
assert len(result) == 5, "Expected the same number of output element as input ones"
|
||||
for result_element in result:
|
||||
assert (
|
||||
result_element is empty_opencv_image
|
||||
result_element is empty_cv2_image
|
||||
), "Expected CV images not to be touched by conversion"
|
||||
|
||||
|
||||
def test_images_to_cv2_when_only_pillow_images_provided(
|
||||
empty_pillow_image: Image.Image,
|
||||
empty_opencv_image: np.ndarray,
|
||||
empty_cv2_image: np.ndarray,
|
||||
) -> None:
|
||||
# given
|
||||
images = [empty_pillow_image] * 5
|
||||
@@ -149,16 +147,16 @@ def test_images_to_cv2_when_only_pillow_images_provided(
|
||||
assert len(result) == 5, "Expected the same number of output element as input ones"
|
||||
for result_element in result:
|
||||
assert np.allclose(
|
||||
result_element, empty_opencv_image
|
||||
result_element, empty_cv2_image
|
||||
), "Output images expected to be equal to empty OpenCV image"
|
||||
|
||||
|
||||
def test_images_to_cv2_when_mixed_input_provided(
|
||||
empty_pillow_image: Image.Image,
|
||||
empty_opencv_image: np.ndarray,
|
||||
empty_cv2_image: np.ndarray,
|
||||
) -> None:
|
||||
# given
|
||||
images = [empty_pillow_image, empty_opencv_image]
|
||||
images = [empty_pillow_image, empty_cv2_image]
|
||||
|
||||
# when
|
||||
result = images_to_cv2(images=images)
|
||||
@@ -166,8 +164,8 @@ def test_images_to_cv2_when_mixed_input_provided(
|
||||
# then
|
||||
assert len(result) == 2, "Expected the same number of output element as input ones"
|
||||
assert np.allclose(
|
||||
result[0], empty_opencv_image
|
||||
result[0], empty_cv2_image
|
||||
), "PIL image should be converted to OpenCV one, equal to example empty image"
|
||||
assert (
|
||||
result[1] is empty_opencv_image
|
||||
result[1] is empty_cv2_image
|
||||
), "Expected CV images not to be touched by conversion"
|
||||
|
||||
@@ -0,0 +1,193 @@
|
||||
from contextlib import ExitStack as DoesNotRaise
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Set
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from supervision.detection.core import Detections
|
||||
from supervision.utils.internal import get_instance_variables
|
||||
|
||||
|
||||
class MockClass:
|
||||
def __init__(self):
|
||||
self.public = 0
|
||||
self._protected = 1
|
||||
self.__private = 2
|
||||
|
||||
def public_method(self):
|
||||
pass
|
||||
|
||||
def _protected_method(self):
|
||||
pass
|
||||
|
||||
def __private_method(self):
|
||||
pass
|
||||
|
||||
@property
|
||||
def public_property(self):
|
||||
return 0
|
||||
|
||||
@property
|
||||
def _protected_property(self):
|
||||
return 1
|
||||
|
||||
@property
|
||||
def __private_property(self):
|
||||
return 2
|
||||
|
||||
|
||||
@dataclass
|
||||
class MockDataclass:
|
||||
public: int = 0
|
||||
_protected: int = 1
|
||||
__private: int = 2
|
||||
|
||||
public_field: int = field(default=0)
|
||||
_protected_field: int = field(default=1)
|
||||
__private_field: int = field(default=2)
|
||||
|
||||
public_field_with_factory: dict = field(default_factory=dict)
|
||||
_protected_field_with_factory: dict = field(default_factory=dict)
|
||||
__private_field_with_factory: dict = field(default_factory=dict)
|
||||
|
||||
def public_method(self):
|
||||
pass
|
||||
|
||||
def _protected_method(self):
|
||||
pass
|
||||
|
||||
def __private_method(self):
|
||||
pass
|
||||
|
||||
@property
|
||||
def public_property(self):
|
||||
return 0
|
||||
|
||||
@property
|
||||
def _protected_property(self):
|
||||
return 1
|
||||
|
||||
@property
|
||||
def __private_property(self):
|
||||
return 2
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"input_instance, include_properties, expected, exception",
|
||||
[
|
||||
(
|
||||
MockClass,
|
||||
False,
|
||||
None,
|
||||
pytest.raises(ValueError),
|
||||
),
|
||||
(
|
||||
MockClass(),
|
||||
False,
|
||||
{"public"},
|
||||
DoesNotRaise(),
|
||||
),
|
||||
(
|
||||
MockClass(),
|
||||
True,
|
||||
{"public", "public_property"},
|
||||
DoesNotRaise(),
|
||||
),
|
||||
(
|
||||
MockDataclass(),
|
||||
False,
|
||||
{"public", "public_field", "public_field_with_factory"},
|
||||
DoesNotRaise(),
|
||||
),
|
||||
(
|
||||
MockDataclass(),
|
||||
True,
|
||||
{"public", "public_field", "public_field_with_factory", "public_property"},
|
||||
DoesNotRaise(),
|
||||
),
|
||||
(
|
||||
Detections,
|
||||
False,
|
||||
None,
|
||||
pytest.raises(ValueError),
|
||||
),
|
||||
(
|
||||
Detections,
|
||||
True,
|
||||
None,
|
||||
pytest.raises(ValueError),
|
||||
),
|
||||
(
|
||||
Detections.empty(),
|
||||
False,
|
||||
{"xyxy", "class_id", "confidence", "mask", "tracker_id", "data"},
|
||||
DoesNotRaise(),
|
||||
),
|
||||
(
|
||||
Detections.empty(),
|
||||
True,
|
||||
{
|
||||
"xyxy",
|
||||
"class_id",
|
||||
"confidence",
|
||||
"mask",
|
||||
"tracker_id",
|
||||
"data",
|
||||
"area",
|
||||
"box_area",
|
||||
},
|
||||
DoesNotRaise(),
|
||||
),
|
||||
(
|
||||
Detections(xyxy=np.array([[1, 2, 3, 4]])),
|
||||
False,
|
||||
{
|
||||
"xyxy",
|
||||
"class_id",
|
||||
"confidence",
|
||||
"mask",
|
||||
"tracker_id",
|
||||
"data",
|
||||
},
|
||||
DoesNotRaise(),
|
||||
),
|
||||
(
|
||||
Detections(
|
||||
xyxy=np.array([[1, 2, 3, 4], [5, 6, 7, 8]]),
|
||||
class_id=np.array([1, 2]),
|
||||
confidence=np.array([0.1, 0.2]),
|
||||
mask=np.array([[[1]], [[2]]]),
|
||||
tracker_id=np.array([1, 2]),
|
||||
data={"key_1": [1, 2], "key_2": [3, 4]},
|
||||
),
|
||||
False,
|
||||
{
|
||||
"xyxy",
|
||||
"class_id",
|
||||
"confidence",
|
||||
"mask",
|
||||
"tracker_id",
|
||||
"data",
|
||||
},
|
||||
DoesNotRaise(),
|
||||
),
|
||||
(
|
||||
Detections.empty(),
|
||||
False,
|
||||
{"xyxy", "class_id", "confidence", "mask", "tracker_id", "data"},
|
||||
DoesNotRaise(),
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_get_instance_variables(
|
||||
input_instance: Any,
|
||||
include_properties: bool,
|
||||
expected: Set[str],
|
||||
exception: Exception,
|
||||
) -> None:
|
||||
with exception:
|
||||
result = get_instance_variables(
|
||||
input_instance, include_properties=include_properties
|
||||
)
|
||||
assert result == expected
|
||||
Reference in New Issue
Block a user