
    9iF                        S SK r S SKJr  S SKJrJrJr  S SKrS SK	J
r
  S SKJrJr  S SKJrJrJr  S SKJrJr  S SKJrJr  S SKJrJr  S S	KJr  S S
KJr  S SKJ r   S SK!J"r"  S SK#J$r$  S SK%J&r&  S SK'J(r(  S SK)J*r*  S SK+J,r,  S SK-J.r.  S SK/J0r0  S SK1J2r2  S SK3J4r4  S SK5J6r6J7r7J8r8  \4Rr                  " \*Rt                  S9 " S S\25      5       r;g)    N)CallableDictOptional)DetectionCheckpointer)build_detection_test_loaderbuild_detection_train_loader)SimpleTrainerhookslaunch)create_ddp_modeldefault_writers)inference_on_datasetprint_csv_format)LRMultiplierWarmupParamScheduler)get_default_optimizer_params)comm)PathManager)setup_logger)CosineParamScheduler)check_local_model_is_latest)snapshot_download)Trainers)DetEvaluator)build_action_detection_model)VideoDetMapper)BaseTrainer)TRAINERS)Invoke	ModelFileTasks)module_namec                      ^  \ rS rSr  SS\S\\   4U 4S jjjrS rS r	S\S\
\\4   4S	 jr  SS
 jrS rSS jrSrU =r$ )ActionDetectionTrainer!   cfg_filecfg_modify_fnc                   > U R                  U5      nUc.  [        R                  R                  U[        R
                  5      n[        T	U ]  U5        Ub  U" U R                  5      U l        U R                  R                  R                  U l        U R                  R                  R                  S   U l        U R                  R                  R                  R                  U l        [!        SU R                  R                  R"                  5      U R                  R                  R$                  S   -  U l        [)        U R                  R                  R*                  5      U l        UR/                  SS5      U l        X l        X0l        UR/                  S[6        R                  " U[        R8                  5      5      U l        g )Nwarmup_step   batch_size_per_gpuresumeFpretrained_model)get_or_download_model_dirospathjoinr    CONFIGURATIONsuper__init__cfgtrainmax_iter
total_steplr_schedulerr)   	optimizerlrmaxnum_gpus
dataloadertotal_batch_sizelenclasses_id_mapnum_classesgetr,   train_datasettest_datasetospTORCH_MODEL_FILEr-   )
selfmodel_idrD   rE   r&   r'   argskwargsmodel_cache_dir	__class__s
            o/var/www/html/land-doc-ocr/venv/lib/python3.13/site-packages/modelscope/trainers/cv/action_detection_trainer.pyr4   ActionDetectionTrainer.__init__$   sC    88Bww||OY5L5LMH"$$TXX.DH((..1188>>66}E((..**-- #txx~~&&!
HHNN%%&:;!< txx~~<<=jj51*( &

HH_i&@&@A!C    c                    [         R                  " 5       (       aM  U(       aF  [        R                  " U5        U R                  R                  [        R                  " US5      5        [         R                  " 5       n[        XSS9  [        XS9nUR                  SR                  U[         R                  " 5       5      5        g )Nz	config.pyfvcore)distributed_rankname)rS   z+Rank of current process: {}. World size: {})r   is_main_processr   mkdirsr5   dumprF   r1   get_rankr   infoformatget_world_size)rH   
output_dirrankloggers       rN   startActionDetectionTrainer.startA   s    !!jz*HHMM#((:{;<}}ZXFj@AHH$%%') 	*rP   c           	          U R                   R                  R                  S::  a  U R                  5         g [	        U R                  U R                   R                  R                  SSSSS9  g )Nr*   r   auto machine_rankdist_urlrJ   r5   r6   r=   do_trainr   )rH   rJ   rK   s      rN   r6   ActionDetectionTrainer.trainK   sM    88>>""a'MMO''rP   checkpoint_pathreturnc           	          U R                   R                  R                  S::  a  U R                  SUS9  g [	        U R                  U R                   R                  R                  SSSSU4S9  g )Nr*   T)	just_evalrj   r   rb   rd   rg   )rH   rj   rJ   rK   s       rN   evaluateActionDetectionTrainer.evaluateW   sY    88>>""a'MMD/MJ''O,.rP   c                   ^ ^ T R                  T R                  R                  R                  5        [	        T R
                  S9mT R                  R                  R                  S:  a  TR                  5         [        TSS9mU(       a/  [        T5      nUR                  U5        T R                  T5      nU$ [        R                  R                  [        TT R                   S9T R                   SS9n[#        U[%        ['        SS	5      ST R(                  T R*                  -  S
9T R*                  S9n[-        T R.                  [1        T R                  R                  R2                  SS9T R4                  T R                  R                  R6                  R8                  S9n[;        TXu5      n[        TT R                  R                  R                  US9n	UR=                  [>        R@                  " 5       [>        RB                  " US9[D        RF                  " 5       (       a3  [>        RH                  " U	T R                  R                  RJ                  S9OS [>        RL                  " T R                  RN                  RP                  UU 4S jS9[D        RF                  " 5       (       a3  [>        RR                  " [U        U	RV                  T R*                  5      SS9OS /5        U	RY                  T RZ                  SS9  T R\                  (       a)  U	RY                  T R\                  S9  UR^                  S-   n
OSn
UR                  U
T R*                  5        g )N)rB   r   F)broadcast_buffers)base_lrg?)paramsr;   weight_decayr*   gMbP?)warmup_factorwarmup_length)r7   Tis_train)mapperr?   num_workers)trainer)	scheduler)periodc                  &   > TR                  T 5      $ N)do_test)modelrH   s   rN   <lambda>1ActionDetectionTrainer.do_train.<locals>.<lambda>   s    dll5&9rP   )eval_periodeval_function   )r,   )0r_   r5   r6   work_dirr   rB   r=   cudar   r   loadr   torchoptimAdamWr   r;   r   r   r   r)   r8   r   rD   r   rA   r?   r>   workers_per_gpur	   register_hooksr
   IterationTimerLRSchedulerr   rU   PeriodicCheckpointercheckpoint_intervalEvalHook
evaluationintervalPeriodicWriterr   save_dirresume_or_loadr-   r,   iter)rH   rm   rj   
checkpointresultr   r9   train_loaderr{   checkpointer
start_iterr   s   `          @rN   rh   ActionDetectionTrainer.do_traind   s   
 	

488>>**+,9I9IJ88>>""Q&JJL$UeDE.u5JOOO,\\%(FM!!/twwGww "  $ $Q-"..@B __
 4!-->!2211AAC  |;,488>>**G= 	  "5 ##%% &&TXX^^%G%GI+/NN HH//889;
 #2244    5 5tG:> 
 	 	##D$9$9%#H;;''t{{'; )JJj$//2rP   c                    [        [        U R                  R                  R                  R                  5       5      U R                  R                  R                  U R                  R                  R                  S:  S9n[        U R                  [        U R                  R                  R                  SS9U R                  R                  R                  R                  S9n[        XU5      n[        U5        U$ )Nr*   )distributedFrw   )ry   rz   )r   listr5   r6   rA   keysr   r=   r   rE   r   r   r>   r   r   r   )rH   r   	evaluatortest_loaderr   s        rN   r   ActionDetectionTrainer.do_test   s     ..3356HHNN##//!35	 2!--?++66FF	H &e)D rP   c                 ~   [         R                  R                  U5      (       ao  [         R                  R                  U5      (       a  UO[         R                  R	                  U5      n[        U[        R                  [        R                  0S9  U$ [        UU[        R                  [        R                  0S9nU$ )N)
user_agent)revisionr   )r/   r0   existsisdirdirnamer   r   KEYLOCAL_TRAINERr   TRAINER)rH   r   model_revisionrL   s       rN   r.   0ActionDetectionTrainer.get_or_download_model_dir   s    77>>%  ')ww}}( (eGGOOE2 'VZZ9M9M,NP 	 0'"JJ79O rP   )
r5   r;   rB   r-   r,   rE   r?   r8   rD   r)   )NN)FNr   )__name__
__module____qualname____firstlineno__strr   r   r4   r_   r6   r   floatrn   rh   r   r.   __static_attributes____classcell__)rM   s   @rN   r$   r$   !   sv     "&59C 	C
 !) 2C C:*
. ."3:.. 83t rP   r$   )<r/   os.pathr0   rF   typingr   r   r   r   detectron2.checkpointr   detectron2.datar   r   detectron2.enginer	   r
   r   detectron2.engine.defaultsr   r   detectron2.evaluationr   r   detectron2.solverr   r   detectron2.solver.buildr   detectron2.utilsr   detectron2.utils.file_ior   detectron2.utils.loggerr   fvcore.common.param_schedulerr   modelscope.hub.check_modelr    modelscope.hub.snapshot_downloadr   modelscope.metainfor   -modelscope.metrics.action_detection_evaluatorr   Fmodelscope.models.cv.action_detection.modules.action_detection_pytorchr   3modelscope.preprocessors.cv.action_detection_mapperr   modelscope.trainers.baser   modelscope.trainers.builderr   modelscope.utils.constantr   r    r!   register_moduleaction_detectionr$   rc   rP   rN   <module>r      s    
  + +  7; : : H H @ @ ! 0 0 > B > ( F! N 0 0 > > 
h&?&?@V[ V AVrP   