
    ^jB7                        d Z ddlZddlZddlZddlmZ ddlZddlm	Z	m
Z
 ddlZddlZddlmZ ddlmZ 	 ddlZ	 ddlmZ ddlmZ  e       Zd Zd Zd	 Zd
 Zd ZddZ ddZ! G d de"      Z# G d dejH                        Z%	 	 	 	 dde&de&de'de(de(ddfdZ)e*dk(  rddl+m,Z,  e,e)       yy# e$ r dZY ~w xY w# e$ r dZY w xY w)zThis tool provides performance benchmarks by using ONNX Runtime and TensorRT to run inference on a given model with
the COCO validation set.

It offers reliable measurements of inference latency using ONNX Runtime or TensorRT on the device.
    N)OrderedDict
namedtuple)Image)tqdm)
get_loggerc                     t        | d      5 }t        j                  |      }d d d        |d   S # 1 sw Y   d   S xY w)Nrimages)openjsonload)ann_filefindatas      b/var/www/ramen.bs-engineer-server.com/venv/lib/python3.12/site-packages/rfdetr/export/benchmark.pyget_image_listr   *   s?    	h	 yy~>>s   0=c                 J    t        j                  |       j                  d      S )NRGB)r   r   convert)	file_paths    r   
load_imager   0   s    ::i ((//    c                      ddl m} m}m}m} ddlm}  |  |d       |        |t        j                  d       |       g      S )Nr   )ComposeResizeToDtypeToImage)	Normalize)  r   T)scale)	torchvision.transforms.v2r   r   r   r   rfdetr.datasets.transformsr   torchfloat32)r   r   r   r   r   s        r   infer_transformsr%   4   s=    KK4:IEMM.K		
 r   c                    | j                  d      \  }}}}|d|j                  d      z  z
  |d|j                  d      z  z
  |d|j                  d      z  z   |d|j                  d      z  z   g}t        j                  |d      S )Ng      ?g        )mindim)unbindclampr#   stack)xx_cy_cwhbs         r   box_cxcywh_to_xyxyr4   C   s    XXb\NCa	sQWWW%%	%	sQWWW%%	%	sQWWW%%	%	sQWWW%%	%		A ;;qb!!r   c           	         | d   | d   }}t        |      t        |      k(  sJ |j                  d   dk(  sJ |j                         }t        j                  |j                  |j                  d   d      dd      \  }}|}||j                  d   z  }||j                  d   z  }	t        |      }
t        j                  |
d|j                  d      j                  ddd	            }
|j                  d      \  }}t        j                  ||||gd      }|
|d d d d d f   z  }
t        ||	|
      D cg c]  \  }}}|||d
 }}}}|S c c}}}w )Nlabelsdets      r   r'   i,  r)      )scoresr6   boxes)lenshapesigmoidr#   topkviewr4   gather	unsqueezerepeatr+   r-   zip)outputstarget_sizes
out_logitsout_bboxprobtopk_valuestopk_indexesr;   
topk_boxesr6   r<   img_himg_w	scale_fctscorelabelboxresultss                     r   post_processrU   N   se   "8,gfoJz?c,////a A%%%D %

499Z5E5Ea5H"+MsXY ZKF!1!1!!44JJ,,Q//Fx(ELL:#7#7#;#B#B1a#KLE  &&q)LE5UE5%8a@IIaqj))EY\]cekmrYsttDUE5RU%53?tGtN us   7Ec                    g }t        |      D ]  }t        t        j                  j	                  ||d               }	|	j
                  \  }
}t        j                  ||
g      } t               |	d       \  }}|d    j                         }|j                          |5  t        |      D ]  }| j                  d d|i      } 	 d d d        |j                  |j                  |z         i }t        j                  d         j                  |      |d<   t        j                  |d         j                  |      |d<   t        j                   |gd      j                  |      }t#        ||      }|d   |d   i}|r|j%                  |        t&        j)                  d	j+                  d
t-        |      z  t/        |      z               i }|p|j1                          |j3                          |j5                          |j6                  d   j8                  j;                         |d<   t&        j)                  |       y y # 1 sw Y   xY w)N	file_nameinputr8   r6   r   r7   r)   idz%Model latency with ONNX Runtime: {}ms  bboxcoco_eval_bbox)r   r   ospathjoinsizer#   Tensorr%   numpyresetrangerunappendtotaltor-   rU   updateloggerinfoformatsumr=   synchronize_between_processes
accumulate	summarize	coco_evalstatstolist)sesscoco_evaluatortime_profileprefiximg_listdevicerepeats	time_listimg_dictimagewidthheightorig_target_sizesimage_tensor_samplesresrF   rT   rr   s                       r   
infer_onnxr   f   s   IN '277<<0EFG

v!LL&%9,*,UD9at$**, 	97^ 9hhtgw%789	9 	++g56!LLQ033F;,,s1v.11&9!KK):(;CFFvNw(9:~wqz*%!!#&+'. KK7>>tc)n?TWZ[cWd?def E!446!!#  ""0":":6"B"H"H"O"O"QE "%	9 	9s   %II	c                    g }t        |      D ]4  }t        t        j                  j	                  ||d               }	|	j
                  \  }
}t        j                  ||
g      } t               |	d       \  }}|d    j                  |      }|j                  \  }}}}|j                          |5  t        |      D ]  } | d|i      } 	 d d d        |j                  |j                  |z         t        j                  |gd      j                  |      }|t!        |      }|d   |d   i}|j#                  |       7 t$        j'                  dj)                  dt+        |      z  t-        |      z               i }|p|j/                          |j1                          |j3                          |j4                  d   j6                  j9                         |d	<   t$        j'                  |       y y # 1 sw Y   .xY w)
NrW   rX   r   r)   rY   z!Model latency with TensorRT: {}msrZ   r[   r\   )r   r   r]   r^   r_   r`   r#   ra   r%   rh   r>   rc   rd   rf   rg   r-   rU   ri   rj   rk   rl   rm   r=   rn   ro   rp   rq   rr   rs   )modelru   rv   rw   rx   ry   rz   r{   r|   r}   r~   r   r   r   r   r   r2   r1   rF   rT   r   rr   s                         r   infer_enginer      s   IN '277<<0EFG

v!LL&%9,*,UD9at$''/]]
1a 	 	47^ 4' 234	4 	++g56!KK):(;CFFvN%"7,=>GD>71:.C!!#&+'. KK3::4#i.;PSVW_S`;`ab E!446!!#  ""0":":6"B"H"H"O"O"QE "	4 	4s   /G77H	c                   l    e Zd ZdZ	 ddefdZdefdZd Zd Z	d Z
dd
Zd Zd Zd Zd Zd ZddZy	)TRTInferencezTensorRT inference engine.	sync_modec                 F   t         st        d      || _        || _        || _        || _        |r-t        j                  t         j                  j                        n,t        j                  t         j                  j                        | _	        | j                  |      | _        | j                  j                         | _        | j                  | j                  | j                  | j
                  | j                        | _        t!        d | j                  j#                         D              | _        | j'                         | _        | j+                         | _        d | _        | j                  s*t0        st        d      t1        j2                         | _        d | _        y )NzGTensorRT is not installed. Please install TensorRT to use TRTInference.c              3   >   K   | ]  \  }}||j                   f  y wN)ptr).0nvs      r   	<genexpr>z(TRTInference.__init__.<locals>.<genexpr>   s     (V1!QUU(Vs   zUpycuda is not installed. Please install `pycuda` to use TRTInference with async mode.)trtImportErrorengine_pathry   r   max_batch_sizeLoggerVERBOSEINFOrj   load_engineenginecreate_execution_contextcontextget_bindingsbindingsr   itemsbindings_addrget_input_namesinput_namesget_output_namesoutput_namesstreamcudaStreamrv   )selfr   ry   r   r   verboses         r   __init__zTRTInference.__init__   s4    ghh&",8?cjj!3!34SZZPSPZPZP_P_E`&&{3{{;;=))$++t||TEXEXZ^ZeZef((V@S@S@U(VV//1 113~~!k  ++-DK !r   
batch_sizec                    i }| j                   j                         D ]  \  }}| j                  j                  |      t        j
                  j                  k(  s=t        j                  d| d|j                          t        j                  |g|j                  dd   j                         j                  d      ||<    |S )Nzmake dummy input z with shape r8   cuda:0)r   r   r   get_tensor_moder   TensorIOModeINPUTrj   rk   r>   r#   randfloatrh   )r   r   blobnamebindings        r   get_dummy_inputzTRTInference.get_dummy_input   s    !]]002 	]MD'{{**40C4D4D4J4JJ/v\'--QR"ZZ
GW]]125FGMMORRS[\T
	] r   c                 @   t        j                  | j                  d       t        |d      5 }t        j                  | j                        5 }|j                  |j                               cddd       cddd       S # 1 sw Y   nxY wddd       y# 1 sw Y   yxY w)zLoad engine. rbN)r   init_libnvinfer_pluginsrj   r   Runtimedeserialize_cuda_engineread)r   r^   fruntimes       r   r   zTRTInference.load_engine   sv    ##DKK4$ 	=CKK$< 	=221668<	= 	= 	= 	= 	= 	= 	=s#    BA?,	B?B	BBc                     g }t        | j                        D ]M  \  }}| j                  j                  |      t        j                  j
                  k(  s=|j                  |       O |S r   )	enumerater   r   r   r   r   rf   r   namesr   r   s       r   r   zTRTInference.get_input_names   s[      - 	#GAt{{**40C4D4D4J4JJT"	# r   c                     g }t        | j                        D ]M  \  }}| j                  j                  |      t        j                  j
                  k(  s=|j                  |       O |S r   )r   r   r   r   r   OUTPUTrf   r   s       r   r   zTRTInference.get_output_names   s[      - 	#GAt{{**40C4D4D4K4KKT"	# r   Nc           	         t        dd      }t               }t        |      D ]  \  }}|j                  |      }	t	        j
                  |j                  |            }
|	d   dk(  rt        t        j                  t        j                  |	|
            j                  |      } |||
|	||j                               ||<    |S )zBuild binddings.Binding)r   dtyper>   r   r   r   r'   )r   )r   r   r   get_tensor_shaper   nptypeget_tensor_dtypeNotImplementedErrorr#   
from_numpynpemptyrh   data_ptr)r   r   r   r   ry   r   r   ir   r>   r   r   s               r   r   zTRTInference.get_bindings   s    Y(QR= ( 		TGAt++D1EJJv66t<=EQx2~)) ''e(DEHHP!(ueT4==?!S		T r   c                 ~   | j                   j                  | j                  D ci c]  }|||   j                          c}       | j                  j                  t        | j                   j                                      | j                  D ci c]  }|| j                  |   j                   }}|S c c}w c c}w r   )r   ri   r   r   r   
execute_v2listvaluesr   r   r   )r   r   r   rF   s       r   run_synczTRTInference.run_sync  s    !!$BRBR"SQ1d1g&6&6&8#8"STT%7%7%>%>%@ AB595F5FG1dmmA&+++GG #TGs   B5!B:c                    | j                   j                  | j                  D ci c]  }|||   j                          c}       | j                   j	                         D cg c]  \  }}t        |       }}}| j                  j                  || j                  j                         | j                  D ci c]  }|| j                  |   j                   }}| j                  j                          |S c c}w c c}}w c c}w )N)r   stream_handle)r   ri   r   r   r   intr   execute_async_v2r   handler   r   r   synchronize)r   r   r   r   r   r   rF   s          r   	run_asynczTRTInference.run_async  s    !!$BRBR"SQ1d1g&6&6&8#8"ST,0,>,>,D,D,FGDAqQGG%%}DKKL^L^%_595F5FG1dmmA&+++GG! #TGGs   C9#C>:!Dc                 ^    | j                   r| j                  |      S | j                  |      S r   )r   r   r   )r   r   s     r   __call__zTRTInference.__call__  s'    >>==&&>>$''r   c                 ^   | j                   r=t        j                  j                         rt        j                  j	                          y | j
                  | j
                  j	                          y t        j                  j                         rt        j                  j	                          y y r   )r   r#   r   is_availabler   r   r   s    r   r   zTRTInference.synchronize  sm     >>zz&&(

&&(;;"KK##%ZZ$$&JJ""$ 'r   c                     | j                   j                          | j                   5  t        |      D ]
  } | |      } 	 d d d        | j                   j                  |z  S # 1 sw Y   "xY wr   )rv   rc   rd   rg   )r   r   r   r   s       r   speedzTRTInference.speed)  sg    ! 	1X J	   &&**	 	s   A""A+c                    dt        t        j                  j                        z  }t        j                  | j
                        5 }|j                  |      5 }t        j                  || j
                        5 }|j                         5 }|j                  t        j                  j                  d       |j                  t        j                  j                         t        |d      5 }	|j!                  |	j#                               s}t
        j%                  d       t'        |j(                        D ]&  }
t
        j%                  |j+                  |
             ( 	 ddd       ddd       ddd       ddd       ddd       y	 ddd       |j-                  ||      }t        |d      5 }|j/                  |       ddd       |cddd       cddd       cddd       cddd       S # 1 sw Y   gxY w# 1 sw Y   ;xY w# 1 sw Y   nxY wddd       n# 1 sw Y   nxY wddd       n# 1 sw Y   nxY wddd       y# 1 sw Y   yxY w)zTakes an ONNX file and creates a TensorRT engine to run inference with
        http://gitlab.baidu.com/paddle-inference/benchmark/blob/main/backend_trt.py#L57
        r8   i   @r   z%ERROR: Failed to parse the ONNX file.Nwb)r   r   NetworkDefinitionCreationFlagEXPLICIT_BATCHBuilderrj   create_network
OnnxParsercreate_builder_configset_memory_pool_limitMemoryPoolType	WORKSPACEset_flagBuilderFlagFP16r   parser   errorrd   
num_errors	get_errorbuild_serialized_networkwrite)r   onnx_file_pathengine_file_pathr   explicit_batch_flagbuildernetworkparserconfigr   r   serialized_enginer   s                r   build_enginezTRTInference.build_engine0  s     C)J)J)Y)Y#ZZKK$	%(/""#67	%;BNN7DKK0	% 5;))+		% 06((););)E)EwOOOCOO001nd+  u||EJJL1LL!HI!&v'8'8!9 >V%5%5e%<=> 	% 	% 	% 	% 	% 2  !( @ @& Q&- +)*+ %)	% 	% 	% 	% 	%   + +#	% 	% 	% 	% 	% 	% 	% 	% 	% 	% 	%s   I!H09H
A H	*A4G.H	&H.H06I&H	.G: 	H			H	H0	I.G73H	:H?H	HH	H0H$ H0'	I0H9	5II)zdino.enginer   F    F)r  N)r  )__name__
__module____qualname____doc__boolr   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r      sX    $ ot!!EI!!F# =$(%+%r   r   c                   *    e Zd Zd Zd Zd Zd Zd Zy)TimeProfilerc                     d| _         y Nr   rg   r   s    r   r   zTimeProfiler.__init__M       
r   c                 0    | j                         | _        | S r   )timestartr   s    r   	__enter__zTimeProfiler.__enter__R  s     YY[
r   c                 d    | xj                   | j                         | j                  z
  z  c_         y r   )rg   r  r  )r   typevalue	tracebacks       r   __exit__zTimeProfiler.__exit__X  s    

diikDJJ..
r   c                     d| _         y r  r  r   s    r   rc   zTimeProfiler.reset[  r  r   c                     t         j                  j                         rt         j                  j                          t	        j
                         S r   )r#   r   r   r   r  perf_counterr   s    r   r  zTimeProfiler.time`  s3     ::""$JJ""$  ""r   N)r  r  r  r   r  r  rc   r  r
  r   r   r  r  L  s    
/
#r   r  r^   	coco_pathry   run_benchmarkdisable_evalreturnc           	         t         j                  | ||||d       t        j                  |d      }t	        |      }t        j                  |d      }|rd}t         j                  d       nd}|sddlm}	  |	|d	      }
nd
}
t               }| j                  d      r.dd
l	}|j                  | dg      }t        ||
|||d| |       y
| j                  d      r't        | dd|       }t        ||
|||d| |       y
t        d      )a'  Performance benchmark tool for ONNX/TRT models.

    Args:
        path: Engine file path (.onnx or .engine).
        coco_path: COCO dataset path.
        device: CUDA device index.
        run_benchmark: Repeat inference 10x to measure latency.
        disable_eval: Skip COCO evaluation.
    )r^   r  ry   r  r  z"annotations/instances_val2017.jsonval2017
   z6Inference for each image will be repeated 10 times ...r8   r   )CocoEvaluator)r[   Nz.onnxCUDAExecutionProvider)	providerszcuda:)ry   rz   z.engineT)r   ry   zFOnly model file names ending with ".onnx" and ".engine" are supported.)rj   rk   ospr_   r   rfdetr.evaluation.coco_evalr$  r  endswithonnxruntimeInferenceSessionr   r   r   r   )r^   r  ry   r  r  coco_gtrx   rw   rz   r$  ru   rv   nxrunrt   r   s                  r   mainr.  h  s     KK"*(	
 hhy"FGGg&HXXi+FLM=&w	:>L}}W#%%d7N6O%P4vxRWX^W_P`jqr	y	!TTE&:JKUNL&(UZ[aZbScmtu!"jkkr   __main__)CLI)r8   )z	data/cocor   FF)-r  
contextlibr   r]   os.pathr^   r'  r  collectionsr   r   rb   r   r#   PILr   	tqdm.autor   	_tensorrtr   r   pycuda.driverdriverr   rfdetr.utilities.loggerr   rj   r   r   r%   r4   rU   r   r   objectr   ContextDecoratorr  strr   r	  r.  r  jsonargparser0  r
  r   r   <module>r>     s$     	   /      /	0"0"J"JY%6 Y%x#:.. #< !3l
3l3l 3l 	3l
 3l 
3ll z I C  
C
  Ds"   B/ B< /B98B9<CC