o
    *j                     @   s   d dl mZmZ d dlZd dlZd dlm  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mZ d dlmZ d d	lmZ e Zejejejd
G dd deZdS )    )AnyDictN)	Pipelines)
OutputKeys)InputPipeline)	PIPELINES)	LoadImagePreprocessor)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 )VisionEfficientTuningPipelinemodelc              	      s   t  jd
d|i| tj rdnd| _| j| j| _| j  t	j
| jjfi || _| jdu rPttdtdt tjg dg dd	g| _dS dS )a  
        use `model` to create a vision efficient tuning pipeline for prediction
        Args:
            model: model id on modelscope hub.
        Example:
            >>> from modelscope.pipelines import pipeline
            >>> petl_pipeline = pipeline('vision-efficient-tuning',
                'damo/cv_vitb16_classification_vision-efficient-tuning-adapter')
            >>> result = petl_pipeline(
                'data/test/images/vision_efficient_tuning_test_1.png')
            >>> print(f'Output: {result}.')
        r   cudacpuN   )   r   )g
ףp=
?gv/?gCl?)gZd;O?gy&1?g?)meanZstd )super__init__torchr   Zis_availabledevicer   toevalr
   Zfrom_pretrainedZ	model_dirpreprocessor
transformsZComposeZResizeZ
CenterCropZToTensorZ	Normalize)selfr   kwargs	__class__r   y/var/www/html/Deteccion_Ine/venv/lib/python3.10/site-packages/modelscope/pipelines/cv/vision_efficient_tuning_pipeline.pyr      s&   

z&VisionEfficientTuningPipeline.__init__inputsreturnc           
      K   s   d}dg}d}t | jtrGt |ts||i}n||vr3|D ]}||v r2t || tr2||| i} nq| |}||| d| ji}|S t |tr^|g| D ]}||v r]|| } nqQt	|}	| |	}||d| ji}|S )z9 Preprocess method build from transforms or Preprocessor zimg_path:FILEimageZimgsr   )

isinstancer   r
   dictstrZ	unsqueezer   r   r	   Zconvert_to_img)
r   r#   Zpreprocess_paramsZin_keyZother_in_keysZout_keyikdataresultZimgr   r   r"   
preprocess8   s2   






z(VisionEfficientTuningPipeline.preprocessc                 K   s:   t   | |}|W  d    S 1 sw   Y  d S )N)r   Zno_gradr   )r   r#   Zforward_paramsresultsr   r   r"   forwardR   s   

$z%VisionEfficientTuningPipeline.forwardc                    s   |t j   }tj|ddd ddd dd }tj|ddd ddd dd }dd |D  fd	d|D d
}t j|d t j|d i}|S )z  Postprocess for classification    )Zaxisr   N   c                 S   s   g | ]}|qS r   r   ).0Zscorer   r   r"   
<listcomp>`   s    z=VisionEfficientTuningPipeline.postprocess.<locals>.<listcomp>c                    s   g | ]} j j| qS r   )r   ZCLASSES)r2   labelr   r   r"   r3   a   s    )
pred_score
pred_classr6   r7   )r   ZSCORESr   numpynpsortZargsortZLABELS)r   r#   Zpost_paramsZscoresZpred_scoresZpred_labelsr+   Zoutputsr   r5   r"   postprocessX   s   $$

z)VisionEfficientTuningPipeline.postprocess)__name__
__module____qualname__r(   r   r   r   r   r,   r.   r;   __classcell__r   r   r    r"   r      s    


r   ) typingr   r   r8   r9   r   Ztorch.nn.functionalnnZ
functionalFZtorchvision.transformsr   Zmodelscope.metainfor   Zmodelscope.outputsr   Zmodelscope.pipelines.baser   r   Zmodelscope.pipelines.builderr   Zmodelscope.preprocessorsr	   r
   Zmodelscope.utils.constantr   Zmodelscope.utils.loggerr   loggerZregister_moduleZvision_efficient_tuningr   r   r   r   r"   <module>   s$   