@@ -5,10 +5,10 @@ use rayon::prelude::*;
55
66use crate :: {
77 elapsed_module, inputs, Config , DynConf , Engine , Engines , FromConfig , Hbb , Image ,
8- ImageProcessor , Model , Module , Version , Xs , X , Y ,
8+ ImageProcessor , Mask , Model , Module , Ops , Version , Xs , X , Y ,
99} ;
1010
11- /// RT-DETR: Real-Time Detection Transformer
11+ /// PP-DocLayoutsss-v1/v2/v3
1212#[ derive( Debug , Builder ) ]
1313pub struct PPDocLayout {
1414 pub height : usize ,
@@ -158,6 +158,68 @@ impl PPDocLayout {
158158
159159 Ok ( ys)
160160 }
161+ Version ( 3 , _, _) => {
162+ // preds: [batch * max_det, 7]，[label_index, score, xmin, ymin, xmax, ymax, reading_order]
163+ let preds_reshaped = preds. to_shape ( [ self . batch , self . max_det , 7 ] ) ?;
164+ let preds_masks = xs
165+ . get :: < f32 > ( 2 )
166+ . ok_or_else ( || anyhow:: anyhow!( "Failed to get masks" ) ) ?;
167+ let preds_masks = preds_masks. to_shape ( [ self . batch , self . max_det , 200 , 200 ] ) ?;
168+
169+ let ys: Vec < Y > = preds_reshaped
170+ . axis_iter ( Axis ( 0 ) )
171+ . into_par_iter ( )
172+ . zip ( preds_masks. axis_iter ( Axis ( 0 ) ) . into_par_iter ( ) )
173+ . enumerate ( )
174+ . filter_map ( |( idx, ( preds, preds_masks) ) | {
175+ let info = & self . processor . images_transform_info [ idx] ;
176+ let ( image_height, image_width) = ( info. height_src , info. width_src ) ;
177+
178+ let mut items: Vec < ( Hbb , Mask , usize ) > = preds
179+ . outer_iter ( )
180+ . zip ( preds_masks. outer_iter ( ) )
181+ . filter_map ( |( pred_slice, pred_mask) | {
182+ let slice = pred_slice. as_slice ( ) ?;
183+ let hbb = f ( slice) ?;
184+ let order = slice[ 6 ] as usize ;
185+ let ( mh, mw) = ( pred_mask. shape ( ) [ 0 ] , pred_mask. shape ( ) [ 1 ] ) ;
186+ let ( ih, iw) = ( image_height, image_width) ;
187+ let mask_f32: Vec < f32 > =
188+ pred_mask. iter ( ) . map ( |& x| 1. / ( 1. + ( -x) . exp ( ) ) ) . collect ( ) ;
189+ let mask_f32: Vec < f32 > = Ops :: interpolate_1d (
190+ & mask_f32, mw as _ , mh as _ , iw as _ , ih as _ , false ,
191+ )
192+ . ok ( ) ?;
193+ let mask_u8: Vec < u8 > = mask_f32
194+ . into_iter ( )
195+ . map ( |x| if x <= 0.5 { 0 } else { 1 } )
196+ . collect ( ) ;
197+
198+ let mut mask = Mask :: new ( & mask_u8, iw as _ , ih as _ ) . ok ( ) ?;
199+
200+ if let Some ( id) = hbb. id ( ) {
201+ mask = mask. with_id ( id) ;
202+ }
203+ if let Some ( name) = hbb. name ( ) {
204+ mask = mask. with_name ( name) ;
205+ }
206+ if let Some ( confidence) = hbb. confidence ( ) {
207+ mask = mask. with_confidence ( confidence) ;
208+ }
209+
210+ Some ( ( hbb, mask, order) )
211+ } )
212+ . collect ( ) ;
213+
214+ items. sort_by ( |a, b| a. 2 . cmp ( & b. 2 ) ) ;
215+ let hbbs: Vec < Hbb > = items. iter ( ) . map ( |( h, _, _) | h. clone ( ) ) . collect ( ) ;
216+ let masks: Vec < Mask > = items. into_iter ( ) . map ( |( _, m, _) | m) . collect ( ) ;
217+ Some ( Y :: default ( ) . with_hbbs ( & hbbs) . with_masks ( & masks) )
218+ } )
219+ . collect ( ) ;
220+
221+ Ok ( ys)
222+ }
161223 _ => anyhow:: bail!( "Unsupported version" ) ,
162224 }
163225 }
0 commit comments