o
    *j                     @   s   d dl mZmZmZmZ d dlZd dlZd dlZd dl	Z	d dl
mZ d dlmZ d dlmZmZmZ d dlmZ d dl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DictListUnionN)	Pipelines)
OutputKeys)InputModelPipeline)	PIPELINES	LoadImage)Tasks)
get_logger)module_namec                       s|   e Zd ZdZdef fddZdd Zdefdd	Zde	ee
f d
efddZded
e	ee
f fddZ fddZ  ZS )ImageMatchingFastPipelinea   Image Matching Pipeline.

    Examples:

    >>> from modelscope.outputs import OutputKeys
    >>> from modelscope.pipelines import pipeline
    >>> from modelscope.utils.constant import Tasks

    >>> task = 'image-matching'
    >>> model_id = 'Damo_XR_Lab/cv_transformer_image-matching_fast'

    >>> input_location =  [[
    >>>        'data/test/images/image_matching1.jpg',
    >>>         'data/test/images/image_matching1.jpg',
    >>>    ]]
    >>> estimator = pipeline(task, model=model_id)
    >>> result = estimator(input_location)
    >>> kpts0, kpts1, confidence = result[0][OutputKeys.MATCHES]
    >>> print(f'Found {len(kpts0)} matches')
    modelc                    s6   t  jdd|i| tj stdtd dS )z
        use `model` to create a image matching pipeline fast for prediction
        Args:
            model: model id on modelscope hub.
        r   z?Cuda is not available. Image matching model only supports cuda.z#image matching model, pipeline initN )super__init__torchcudaZis_availableRuntimeErrorloggerinfo)selfr   kwargs	__class__r   u/var/www/html/Deteccion_Ine/venv/lib/python3.10/site-packages/modelscope/pipelines/cv/image_matching_fast_pipeline.pyr   ,   s   
z"ImageMatchingFastPipeline.__init__c                 C   s   t dd}||d }|S )Ncv2)backendimgr   )r   Zimg_nameZimage_loaderr"   r   r   r   
load_image<   s   
z$ImageMatchingFastPipeline.load_imageinputc                 C   s:   t |dks
J d| |d }| |d }||dS )N   z$input should be a list of two imagesr      )Zimage0Zimage1)lenr#   )r   r$   Zimg1Zimg2r   r   r   
preprocessA   s   
z$ImageMatchingFastPipeline.preprocessreturnc                 C   s   | j |}|S )N)r   Z	inference)r   r$   resultsr   r   r   forwardH   s   z!ImageMatchingFastPipeline.forwardinputsc                 C   sf   | j |}|tj }|d    }|d    }|d    }tj|||gi}|S )Nkpts0kpts1
confidence)r   postprocessr   ZMATCHESdetachcpunumpy)r   r,   r*   matchesr-   r.   r/   Zoutputsr   r   r   r0   L   s   
z%ImageMatchingFastPipeline.postprocessc                    s   t  j|fi |S )a  
        Match two images and return the matched keypoints and confidence.

        Args:
            input (`List[List[str]]`): A list of two image paths.

        Return:
            A list of result.
            The list contain the following values:

            - kpts0 -- Matched keypoints in the first image
            - kpts1 -- Matched keypoints in the second image
            - confidence -- Confidence of the match
        )r   __call__)r   r$   r   r   r   r   r5   Z   s   z"ImageMatchingFastPipeline.__call__)__name__
__module____qualname____doc__strr   r#   r   r(   r   r   listr+   r0   r5   __classcell__r   r   r   r   r      s    r   )typingr   r   r   r   r    r3   npZPILr   Zmodelscope.metainfor   Zmodelscope.outputsr   Zmodelscope.pipelines.baser   r	   r
   Zmodelscope.pipelines.builderr   Zmodelscope.preprocessorsr   Zmodelscope.utils.constantr   Zmodelscope.utils.loggerr   r   Zregister_moduleZimage_matchingZimage_matching_fastr   r   r   r   r   <module>   s"   