o
    0j                     @   s|   d dl mZ ddlmZ ddlmZ ddlmZ ddlm	Z	 d	d
l
mZmZmZmZmZmZ d	dlmZ G dd de	ZdS )    )Union   )FuncRegister   VideoBatchSampler	ReadVideo   )RunnerPredictor   )
CenterCropImage2ArrayNormalizeVideoScaleToBatchVideoClasTopkTopkVideoResultc                       s   e Zd Zi ZeeZd(deedf f fddZdd Z	dd Z
d	d
 Zd(deedf fddZed				d)ddZedd*ddZedd*ddZedd+ddZedg dg dfd d!Zed"d(d#d$Zed%d&d' Z  ZS ),VideoClasRunnerPredictorNtopkc                    s,   t  j|i | || _|  \| _| _d S N)super__init__r   _buildpre_tfspost_op)selfr   argskwargs	__class__ x/var/www/html/Deteccion_Ine/venv/lib/python3.10/site-packages/paddlex/inference/models/video_classification/predictor.pyr   %   s   z!VideoClasRunnerPredictor.__init__c                 C   s   t  S r   r   r   r"   r"   r#   _build_batch_sampler*   s   z-VideoClasRunnerPredictor._build_batch_samplerc                 C   s   t S r   r   r$   r"   r"   r#   _get_result_class-   s   z*VideoClasRunnerPredictor._get_result_classc           
      C   s   i }| j d d D ]2}t| d }|| jv sJ | j| }||i }|r/|| fi |n|| \}}|r;|||< q	t |d< i }| j d D ]%}	| j|	}| j d |	i }|rc|| fi |n|| \}}|||< qH||fS )NZ
PreProcessZtransform_opsr   r   ZPostProcess)configlistkeys	_FUNC_MAPgetr   )
r   r   cfgZtf_keyfuncr   nameopr   keyr"   r"   r#   r   0   s$   
 
 
zVideoClasRunnerPredictor._buildc           
      C   s   | j d |d}| j d |d}| j d |d}| j d |d}| j d |d}| j d |d}| j|d}| jd	 ||p?| jd
\}}}	||||	dS )Nr	   )Zvideosr   r   r   r   r   )xTopk)r   )Z
input_path	class_idsZscoresZlabel_names)r   runnerr   r   )
r   Z
batch_datar   Zbatch_raw_videosZbatch_videosr1   Zbatch_predsZbatch_class_idsZbatch_scoresZbatch_label_namesr"   r"   r#   processE   s   
z VideoClasRunnerPredictor.processr	         r   c                 C   s   dt d|||dfS )Nr	   Zdecord)backendnum_segseg_lensample_typer   )r   r9   target_sizer:   r;   r"   r"   r#   build_readvideoX   s   z(VideoClasRunnerPredictor.build_readvideor   c                 C   s   dt |dd ddfS )Nr   TF)
short_sizeZfixed_ratioZ
keep_ratioZdo_round)r   )r   r>   r"   r"   r#   build_scaleg   s   z$VideoClasRunnerPredictor.build_scaler   c                 C   s   dt |dfS )Nr   )r<   )r   )r   r<   r"   r"   r#   build_center_cropp   s   z*VideoClasRunnerPredictor.build_center_cropr   tchwc                 C   s   dt dddfS )Nr   TrA   )Z	transposedata_format)r   )r   rB   r"   r"   r#   build_image2arrayt   s   z*VideoClasRunnerPredictor.build_image2arrayr   )g
ףp=
?gv/?gCl?)gZd;O?gy&1?g?c                 C   s   dt ||dfS )Nr   )meanstd)r   )r   rD   rE   r"   r"   r#   build_normalizex   s   z(VideoClasRunnerPredictor.build_normalizer2   c                 C   s   | j st|| _ dt|dfS )Nr2   )r3   )r   intr   )r   r   Z
label_listr"   r"   r#   
build_topk   s   
z#VideoClasRunnerPredictor.build_topkZKeepKeysc                 O   s   dS )N)NNr"   )r   r   r   r"   r"   r#   foo   s   zVideoClasRunnerPredictor.foor   )r6   r7   r   N)r7   )rA   )__name__
__module____qualname__r*   r   registerr   rG   r   r%   r&   r   r5   r=   r?   r@   rC   rF   rH   rI   __classcell__r"   r"   r    r#   r       s8    r   N)typingr   Zutils.func_registerr   Zcommon.batch_samplerr   Zcommon.readerr	   Z
predictorsr   Z
processorsr   r   r   r   r   r   resultr   r   r"   r"   r"   r#   <module>   s    