o
    *j                     @   s   d dl Z d dlZd dlmZmZ d dl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 d dlmZ d dlmZ d d	lmZ d d
lmZmZ d dlmZ e ZejejejdG dd deZ dS )    N)AnyDict)Image)	Pipelines)
OutputKeys)InputPipeline)	PIPELINES)	LoadImage)Config)	ModelFileTasks)
get_logger)module_namec                       s   e Zd Zdef fddZdedeeef fddZdeeef deeef fdd	Z	d
eeef deeef fddZ
  ZS )ImageReidPersonPipelinemodelc                    s   t |ts	J dt jd|dd| td|  tj|t	j
}t|}|jj}| j| j| _| j  tt|jjt tj|jj|jjdg| _dS )z8
            model: model id on modelscope hub.
        zmodel must be a single strF)r   Zauto_collatezloading model config from dir )meanZstdN )
isinstancestrsuper__init__loggerinfoospathjoinr   ZCONFIGURATIONr   	from_filer   cfgtodeviceevalTZComposeZResizeZINPUTZ	SIZE_TESTZToTensorZ	NormalizeZ
PIXEL_MEANZ	PIXEL_STDval_transforms)selfr   kwargsZcfg_pathr   	__class__r   s/var/www/html/Deteccion_Ine/venv/lib/python3.10/site-packages/modelscope/pipelines/cv/image_reid_person_pipeline.pyr      s   

z ImageReidPersonPipeline.__init__inputreturnc                 C   s2   t |}| |}|d}|| j}d|iS )Nr   img)r
   Zconvert_to_imgr#   Z	unsqueezer   r    )r$   r)   r+   r   r   r(   
preprocess.   s
   


z"ImageReidPersonPipeline.preprocessc                 C   s,   |d }|  |}|   }tj|iS )Nr+   )r   detachcpunumpyr   ZIMG_EMBEDDING)r$   r)   r+   Zimg_embeddingr   r   r(   forward5   s   

zImageReidPersonPipeline.forwardinputsc                 C   s   |S )Nr   )r$   r1   r   r   r(   postprocess;   s   z#ImageReidPersonPipeline.postprocess)__name__
__module____qualname__r   r   r   r   r   r,   r0   r2   __classcell__r   r   r&   r(   r      s
    "*r   )!mathr   typingr   r   ZtorchZtorchvision.transformsZ
transformsr"   ZPILr   Zmodelscope.metainfor   Zmodelscope.outputsr   Zmodelscope.pipelines.baser   r   Zmodelscope.pipelines.builderr	   Zmodelscope.preprocessors.imager
   Zmodelscope.utils.configr   Zmodelscope.utils.constantr   r   Zmodelscope.utils.loggerr   r   Zregister_moduleZimage_reid_personr   r   r   r   r(   <module>   s&   