From 0afbc435dff14054f0ff661fad51448d8c55afd5 Mon Sep 17 00:00:00 2001 From: Laurent Mazare Date: Fri, 25 Aug 2023 13:50:31 +0100 Subject: [PATCH] Add some configurable legend for yolo detection. (#603) * Add some configurable legend for yolo detection. * Clippyness. --- Cargo.toml | 1 + candle-examples/Cargo.toml | 1 + candle-examples/examples/yolo-v8/main.rs | 44 +++++++++++++++++- .../examples/yolo-v8/roboto-mono-stripped.ttf | Bin 0 -> 12172 bytes 4 files changed, 45 insertions(+), 1 deletion(-) create mode 100644 candle-examples/examples/yolo-v8/roboto-mono-stripped.ttf diff --git a/Cargo.toml b/Cargo.toml index 89b3c63a..f1ad66eb 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -46,6 +46,7 @@ num-traits = "0.2.15" rand = "0.8.5" rand_distr = "0.4.3" rayon = "1.7.0" +rusttype = { version = "0.9", default-features = false } safetensors = "0.3.1" serde = { version = "1.0.171", features = ["derive"] } serde_json = "1.0.99" diff --git a/candle-examples/Cargo.toml b/candle-examples/Cargo.toml index 731052ea..8383ab37 100644 --- a/candle-examples/Cargo.toml +++ b/candle-examples/Cargo.toml @@ -33,6 +33,7 @@ hf-hub = { workspace = true, features=["tokio"]} imageproc = { workspace = true } memmap2 = { workspace = true } rand = { workspace = true } +rusttype = { workspace = true } tokenizers = { workspace = true, features = ["onig"] } tracing = { workspace = true } tracing-chrome = { workspace = true } diff --git a/candle-examples/examples/yolo-v8/main.rs b/candle-examples/examples/yolo-v8/main.rs index 1a378680..ab047304 100644 --- a/candle-examples/examples/yolo-v8/main.rs +++ b/candle-examples/examples/yolo-v8/main.rs @@ -59,6 +59,7 @@ pub fn report_detect( h: usize, confidence_threshold: f32, nms_threshold: f32, + legend_size: u32, ) -> Result { let (pred_size, npreds) = pred.dims2()?; let nclasses = pred_size - 4; @@ -96,6 +97,8 @@ pub fn report_detect( let w_ratio = initial_w as f32 / w as f32; let h_ratio = initial_h as f32 / h as f32; let mut img = img.to_rgb8(); + let font = Vec::from(include_bytes!("roboto-mono-stripped.ttf") as &[u8]); + let font = rusttype::Font::try_from_vec(font); for (class_index, bboxes_for_class) in bboxes.iter().enumerate() { for b in bboxes_for_class.iter() { println!( @@ -114,6 +117,29 @@ pub fn report_detect( image::Rgb([255, 0, 0]), ); } + if legend_size > 0 { + if let Some(font) = font.as_ref() { + imageproc::drawing::draw_filled_rect_mut( + &mut img, + imageproc::rect::Rect::at(xmin, ymin).of_size(dx as u32, legend_size), + image::Rgb([170, 0, 0]), + ); + let legend = format!( + "{} {:.0}%", + candle_examples::coco_classes::NAMES[class_index], + 100. * b.confidence + ); + imageproc::drawing::draw_text_mut( + &mut img, + image::Rgb([255, 255, 255]), + xmin, + ymin, + rusttype::Scale::uniform(legend_size as f32 - 1.), + font, + &legend, + ) + } + } } } Ok(DynamicImage::ImageRgb8(img)) @@ -255,6 +281,10 @@ pub struct Args { /// The task to be run. #[arg(long, default_value = "detect")] task: YoloTask, + + /// The size for the legend, 0 means no legend. + #[arg(long, default_value_t = 14)] + legend_size: u32, } impl Args { @@ -291,6 +321,7 @@ pub trait Task: Module + Sized { h: usize, confidence_threshold: f32, nms_threshold: f32, + legend_size: u32, ) -> Result; } @@ -306,8 +337,17 @@ impl Task for YoloV8 { h: usize, confidence_threshold: f32, nms_threshold: f32, + legend_size: u32, ) -> Result { - report_detect(pred, img, w, h, confidence_threshold, nms_threshold) + report_detect( + pred, + img, + w, + h, + confidence_threshold, + nms_threshold, + legend_size, + ) } } @@ -323,6 +363,7 @@ impl Task for YoloV8Pose { h: usize, confidence_threshold: f32, nms_threshold: f32, + _legend_size: u32, ) -> Result { report_pose(pred, img, w, h, confidence_threshold, nms_threshold) } @@ -385,6 +426,7 @@ pub fn run(args: Args) -> anyhow::Result<()> { height, args.confidence_threshold, args.nms_threshold, + args.legend_size, )?; image_name.set_extension("pp.jpg"); println!("writing {image_name:?}"); diff --git a/candle-examples/examples/yolo-v8/roboto-mono-stripped.ttf b/candle-examples/examples/yolo-v8/roboto-mono-stripped.ttf new file mode 100644 index 0000000000000000000000000000000000000000..6d807fc4ff348c8d1a86285c1b6a6e3b6ba00543 GIT binary patch literal 12172 zcmbVy33ycHx&QgjIg_1CCX<;-W+thqj4?|7m?g8+ylwKFfgi=NP7axxc5g8`%$MW0-y~ z^f~&WBjX~{!}`16TG>CmdcE$lj1BkSfa_O=MwfSfarILjhUtFFTn~4yAH%jF4_u#x zYuQNWa8JqqYyAuC|1yT*%CXV$)w9D`4~7|k&)gqd)ic(gvTHlszXI2(7>+e#7`c}< z1#O@z$h9E~*|<(Fm*^S4LR+7r_0wyt*Iq^&=m%L-H=EFgWBB^$L%GA!a5x{7M4!K-~@MqCGNGKA|^Sb91N~8!w*x>BvWIO9s%n#?K zfS{BVN)aPet+v}sm{Y*7&B~B0u)$%}YDy8K9LDlul*1jHl^I30Pp22FRD;#wZk5W+ zX3v?!W?NLso>29mN>!XbomK1eogH4h-|cA?iA4y(pH`3i(M1Di{JsYCLTIHr^r?X3?DzX$6tqS#@^?AcKgJDxu)kO34=83ARjRu1)>@7P5GJ6VmKL)(# zVMTX157gN0F=~Neja(4tF>?l9?2HWSsi#o$!1Gm+5hbNU8kO#Q3j+t+=nFqQP`=s} zS}fx65X#FFG&*eCnhqedp=e`m)!B~0?x~vj+jSB-I)n;^)XECS=E!h(lRY?~q|^*Q zUjS-KfSOF;r(MpGC%(ix(MkM=GdqdTZ{7jq{9Swae9`RvU`c>Op_+iBgZ>VKi;2sX5kss^wU9 z&3#6b9U+}QKL}^q4d?gAIo8RcICA*f;4kUF5IuALJl@8diavs`0uCIDE@OCb9>z=J z<6C*aVP2fa=TRGdh8UxtM(wPrW4HEAPLiX?nDJep;d~e`0{(Jh{7Psz6;t21?m9Wd zLR_O%^@jp|pmoF#(`$U*=X!h2`2E#5YicIbH>%fL^cr=)x94fM$5-z2p6($4`cTaF zpMmwWV0~mm02!&VBjB0n$|uoJQ1lbJZTF9Ugg+BqLfG)l(LYS#{|4CcIWzZlI7=nW zWq~}Iv)j>@Y2v#_SX|H!=Kf=FzYgvLCI%U{5CBO8{0C6j@Sm62e_Mg(urGf)i!o?~ zm*cfj7LgO3jJ^mw(=zk-zATrP25vDB ziLC%$b%K7Wuv$hNZ|l`RHF>Zz%^i!n0~~h{?>2Do`D7<55Q>^zuKkM^?Q^-Bg`xrk zNF;1=x%MS`gat34%xt!aQZ5e$yOc@>K@glT+pJV}1>!x~S@;9bHcZykeAjGtI?a~v z*3=wrc(!5Uc8B?X==k50C99mEej1}uqwaM%SJ_V5R=S)$8nw}&*7P`CE8(P!)olRI zzuU78PW(LZcNOqA52N73$mT?`BPFreey``l_aezPL|-rKRw{M%;mGA) zy6wte!Yg$;qgkgP@!>7e^I7_uZ`^TmZCl}b4f zS^9)KRAI50_jh$ZW^+~+u(cA&e7|RbL?Y3Fi2U8h-EP0#W_xVu(w#PIWdTPol`VnO z)d6200KTR&dIZ=9x=KE9>kahEWmfLQ#LX{Qxs2WK1RjK7R3VI_;!YyL1%e`yV9!nl zzE9w4rD~|Ma#*R-AaV*v6+Ykj75(S^{we@Ib#T4Wx0BcGwR;&hr)IB;n1aC z*rgI^*#f(y0lF}uQmO_~IQJO3{u=#bK97F+b%2*^)s>_l#)34+2^b6HjYCK!;M%wp zk*Wc5m_Ijvd?_>I;*Ty~7)Z&ycx);&WnhZT^m%Xn6)51Bl@fo1zaHXqSZ6-0vj^5; z{3~q904{^M_)7Y%57BdU!^bF$LcgFJ(Q_XWYw*&jF{;2TqF3=> zvZpx0rX%voF=_%EiPS z`c(&jfJiFQ@98+3<5Fvu<`_0O(BBq6WHcCD2IHzygL4HK559KT5AInN%9CNMS+J77Tl#5Vw$*Z z@ASD`N4vTocch)3w(sj)KIw9Mrxi+Sv|;f^rE+6?{RpL0;A-CLH*0FV=WS_O z5NdK+jO!ayTUWkPRoPruS@q5~{mPdsDjJBlO7!}1YkKXtR-#sK{`RgAg5aupq;mr)8?r{6cM;|r|?>IQAcVW-PWe}MGFnLTxpUYS(N z%YGs?la-ndS;LY3x;=3EUGzM%FlNb(_ecTMI35%v@H90^@9i8aHHm^G8Ns0B&~sOr z%Q?Bcv{x@O@w6)Sijb;CrK(bNsQO&aA${*jZ&~2j4$%NZ#e>!eb1 zZmwA(tv55Wl;4(Rfq<>eZ?70qs?=-_$E_Q+SMm7VGLPrkL6TaVuVoiXlB6uihP{#KZK_cbzV$D<@};Rw&fX$zhY8W6{*fN#aQ)Hk6ea3iR2KX3k;jVeR+n zT4DjLU5J6A1i?2ZyZ?Iyb!R_N7F2`=Gz!I%l5C4mRKl54U=)Zw((WTphx>hpEB0}~ zt}`9=12UNi;fO^XB2#f3cRL-E{bUdo8oXYE!S83Za|+J40yz2i*#P=8+Wkj56}|8Y zePqMctLPg#6<>$0(TBGZXK!L?Ei*U98<-2;<}PRM@F4{|b}KkJ$?AZS?Xb%#&`zjz z#PWh9aE*zpBwllf2Q_~uz=Vo|FcI5g>;b% z+zDTuohjuC%Zus;<)r7<>qL8w7{W+)4yQyh^K(LGxPB0DKD61LF3) zF%t0qV0_gJ0e>Xo4=^o%ATkZ<^{V>%aZN1E9D-!{H`egi0H(IJ z)zrMXli#?{=?n&)&V7yiCTIo%&}3*kM|Y5Zpsf%CLuj>el7S5pd=<7Y|0M&GW{+0IWB_;gyQ}~CJQaM&{Zj_0OIb5r(u)3&WOr@mg4m^#oV&@Bi z&^K4}He9c+ZYjvmF^JbjKj$^>1zPiYxg`=jUE$(#^8yZM9Q`^-@f}+4`XJCVCj0ys zQLfbGI@V`jqLJG3r%$Kx%p!e{r~he}%QH=>fYhRDsl>v~v&tk@MfIafC3PFm^0?XB zf+)@UXb`X~KRb==R zEE$MLX>n<3Vo-nA`2bSdl5{N2G{8^zB{t`Ar6V<_4vp4MPW#JB&yI8)_jrPChhw6< zZh4_3Um_`7Sn6IZlPRaYF4tst`%^BLugvLqqP?zHDitEB(8!fidHMMb?(+N1mQtf$ zzqTset5G|7Y*`-gQ0n3Hg)*TqOjT|&8cVE2MdP7xk4o*}vQ>P(9ZvQL*oKZ3h~;BY zYGB}zp)C=3YGU9XHv&83Rq^ht^zV__ZZvMLDw?m+gaks4T_O!gbk$<93>~JsSpoXy zp6KQ7-Ngo7T3RYEdmEX_Vx{I4^-8C$F*% z5@naq-Ygc28KZx=ZN+hb-I=ecQl79-Oet0`z_&*?Z#d_2*%vm!+Fc;`9$5RH>_fU6 zP5+2~H*YmCXW0?Nk2k23{>lffniawyb(4!yXr_jFW zAELivc!g8A0A@$%T5>zk!PKH+HB%>wr%fpJ4*$VXAy)O6USG<^RqxP$u$k!_@1XZ; zGN}9lzdGAeAjz$uYaLqcQa$-eG_++Ke)ZNjd-HsSg?a)%NP^=*7=R(^g?+%2^dKfp zL{KF|yYSsJd!YigkvMXYkZpTn=Cg+mVo<%_2T}^KUIWzBW2Hb@a>~NHDzx`MJ6~27h-T{aR8J_lTtVs5$ReYn%U{-^ zP#M#42W~)ZJ8#s~)YsS4+}KH<*?FV7wzjUO`oblT;~ zMTkV0cuz_ePg)REUnraxps9oCWZ4Fv zZ%c6gV-AP6%Ikf4Nn}K$*6S#0nWMPXXd0=Eez6^>SUUST@eNP`*?W+e@XAT7vr|c? zm_U~pf+hO}Ko>_M&IPE6&qg~QC?D5p4H~7g*Y9V_P1Dsr->HtSQ$AmXKwP-2va}8A zK>?5FcvsJf($e5n;9zH|cbQDCfUJM8g0MU;lgpP!%6k+Fxl%4)>h5^Fq{IP{%|x%Y z*J!K}6x0^!7wLPR1kw8)&PQ9_OXPA&DwTBxozPMVh}9{fI2kKT*nr(o zRz3V$bTxix`ufa8bWVc)K>XVbdze0tS`Ok*xS}TD&5NLYjW8Q*v{i5ix9=Pg(zE13 z?Ol#=0i?cstiU7^g^d+MDwT#(P`#lw2RKt1jXv+h!qsQpE>{_VX;0(&A|2is<>4Qx z`HHNpIo6`$wUOG5Ma8E3zteNbYPD7e1DEH_q@W#&w~fhra2BBsFtAhVdR z;P<40cL@Q@jgxI7;wezFWaEJtu^{L0|IXVyFA(&c9veO7cDv0s+r+SSu|^varUdnxg)eBe`=g-?rPrfjJ#rus_iy9)~%e@QFMk1{OZk}D* z@j}R7OY>bO(}qCRrs85N)Qb*u_>6TTk%ZCxF^~(?a2Wdt#z4SER?TRMyy$bUu&Vlx z-~2P!{>2#gU&TyN%~=NHzU{$cz?2*UBr7P9eAXoW{Ac~7G30?VM*9($J0KB>8qBUm zGMS>v?>p7iankRrreyNPE~p2X9S3zvRj)U|=yXiGy-=gnDIm67WOpo-$sxUvEpRv% z$z>|CWZ40m4KhI6fn|;^oj#nOUu7^3d%K^6FxY7?d2+dT$Y`w0&(R9ZTcXGP%M}Wx zQlVJx_x7rkS}kmY_qmoZOUh-7-R^eKZG#Xz%%Ute$u{pe4CYLEG zMR(8*6>9^hRHa_{%!%bq51_VF;dTGh6lraal>eY_0TdtG1bKOO3BKfPtv8H*S!y+w9YwSB9)2d^3k~~kJ}xtw3l8^FSA-6p1*L1%~mHBI%Tp3tF2Wc zR~ogtmErM=Wd?V~FMpZgv)Ue-yYQhBdtG6^TQ0B1aP;ZySNKms-wBKjXNCA8Im^Q@ zUCjl%D=mdT=mzoOjD~*L=<^xP0WffQ9x(eADFR7bm>OX)4z+-`@OzJ;e;b42jyriL zG&@bErO73xPL0M2j|PtREI#7)M9{RJQY*0P+O4E`OOp450nqLM;= zm)o_aYSAu>#bz;@cFZlQEyx$hq~hih>mseLhBAhgs&;zIZPF3yqNm{D;Kwmv4?+Ti z1Qf-$y{R>iH!Uri`hMF)g})M9B|=}${)xy#sZ8wBb35N=^4~A3t19bitE#Kkdzza{ z%Ub8o&f*FfTS1PKYAj-9o{dK^D={=Xi$0~N2ru+&phrgx!JRPn0LGxrQ?mmkU^s9e zD+jb8%`<}s4?=u_u#a#HaRYyoxx?TPaU)9LmP3btqAy4}z8l~Po__^_c6dIvdjdVo z+Bdp|o&>3&E>cb`i`@sB<0Fuh8Hc`k9BoA#(c|d5tVf48(-UYF^Oreg53MKGqznRb zgRi|)!jG4RAZEY&1cfPR#ifUO3~JSKr(@Kr(KV}85n(=49i3~G24x!Shy&sPwE=I; z*fTRY1O=>4@k)7nng3|}l8-iLZ~fiU{tH2UB!8p0tGxWu;Lr?1LjY)KPS9}Iu;L{M zh`AFfC$o9nM(Th0VD`Fq7j-@z40K6XiaRUHFAfdHXgE#3K{!(GPdq7rkRI&8L9a%; ze#QE)Ltj@NN_mK(^Z<-fBu4QdLA#? zG4^mhmQ*tql3*lrDhl2SNVrMjJ)W;QzaS5dtX@5u#pktE@p>L9HWhtu_wMf*4W>uv z{o{DRYT4vFweESN>G}1iz2C7~0{A%590;FkKE3a|`_R_Cr<$G)mCsG0!2#M8=hxRde~gIYx6qSgQ#N zEhb)mLCdXid|o*GbknJdUIisb`-)8R#bu}Go-Oy*5kR!X`W^4-hV52SaZZk2XxY(t z%D2g44S=Sv!}gQyS(k|Mncet>k-))8GET7>W&qXI$Q{TLwVgQ>h zhPM$WVtIoNtH=7V!T-OS3sgJ!U=)emIvH;Nd7XxZG^m(TD0+gyZivSzxdPIQZn=Vb zI6H@Jp_G;!HuH(Ka@Qhi^)!boKjeHeS62ucz`B;azLNTp$M zDXuCNFQ9Mj|JeuoXEAi2gFaArsepdeK|D%NuK-7IU<5^$UV0pzp%)*#z<4$AW#kjE zhbUIrC2U-qor--aLJrU0KtJJYgMKGhPT<#?_B74>uKd5=ZrRg9f6e^)4}_P~fBPTk z7nY;6a`YtKi=O-*>e`7W=?><*lb-lqqK`peBKAJfDbNWAR)Z}EeMl0mQ(%MV+IHp< z3Sv;>Ue}s{>I5JsMNm#BikUHy6kf(HIyE&GnhY`(ljYuVyUZ(Y{nhgo)z2XhXUx)CV51yo^SV!73rY&i^o$ImShyfuxwhEgN>AYk_*E8#PsG>cr1IH3kFus$E|W!D(&Qphqn%xt ztwx!`Y=KjfPi1DN5xj$cTd-i+s_p^$x06rMv+eDD=^5!_?izM}T6%hYP4&+n+V_`v zP4jCiD&E>sz13`SXRxzvI`p;D9`^aaKiaU>WGcza&CwU+YuG&gyh!9Zs%5*$XiG_@ z>of9I+5Eghu?Sx310=&w5dR650I1j+uK*$Hv$NOO?AOoGKiPtxK%Yh;w2*)>1FA)6 zKY5r?KrDy_HSq)|Ma?*u*l$cYDLg@vI)a$)NQykXL&z6+DTzxYFbg%c@Me`UJfhJS zLpR?`35s|UuUuQJQ8$~yBU)`S{`*{Yg-|3zY<3pKmzCt^a_t7g`VdDQF32xTOV84Y z6wsLqS>8tQgGb4i2m>n~OA!>8FkJY}*W^nmgZ^KJat!OhE@FJ_K|p8b>h6cQb1lX; zU>_kl@}MSk6}^Z48yDg(ybABf@5jF)iis1%@5pAdn>LYnVT|iWd0$m zIP0mbt8-X$`sQrNPRm}A{hRD4dl7q-y`8;}J;6T1zQTTk{kt4xPH)bJoSivub7UMB zr;gLb*}!>-b1>JD8_Dg@J;wEL|CC4MHRpA}ahOv<9QpA(Pk+9QUG^202?B-Bt6y@A z%>R#m;XgCOivB#MGOY_P)4{%CAM)+kzh-z_YY^TTd;#A7>4rDY&O;UPEWCyF7`!R9 z0{WK1dn;?9O0yHP;c9r-zzc5x9L6rd-JQVGL)b2uaUFXfp7LCWXE5injo3MC7gV4x zV2@yzVJ$P%2P05@+XuC&Hdw6|dl7aS1DPDaj-yYpmrxh_5?4U$DUiVzaMgnX=rlS2 zZ!Lr&dN>E3{t`BWeTEg|Vpzuu)CSnyMWBfrvHNCcp+^RM@jwQ>80*9~#d;v5 zpZy7ZUxd2|=D{nt@P*h8^Tm3C;5`<>D)V7ZHE^*GTL|ywhOsK(ZW~sPRbZ`H3pN)E zVM`$Lmz2v9!?ny>*h zY5-!Terz@H#R{CUz}q4+I7>g)1OHoKR2O_#!&e4p8wHJGT1_x-6m+x)yecy@2;O2f z%o&CLak$pTR$L7`jKc{`@Y4tTFeBE$dPbNviVeqlN1@e+4Z%vi&>Dd;J6l|1wIDGbn?JTrMr@>%a4H4o-6 zXB>!8B}>{bdHz4!V?NBj&F^YBYbQ`W0NiMR)m8yVm}^FB;WN7eFoDs>zyH8nWuUc8 z#TkPdrxh!~$Vf%qTuiup!>S=n*|(}^5K|6yt{%aZU<>ex>G<7Mn77i>Vs?`k`>Vw8 ob$$J-ndJsncMgHxfZfLIiOG4h!G1CPhj5dzFC4Uo`M}`+0}6(v{{R30 literal 0 HcmV?d00001