o
    * i                     @  sT   d dl mZ d dlZddlmZmZmZmZ deeeedZG dd dej	j
ZdS )    )annotationsN   )MFCCLogMelSpectrogramMelSpectrogramSpectrogram)rawZmelspectrogramZmfccZlogmelspectrogramspectrogramc                      sJ   e Zd ZdZ		dd fddZdddZdd Zdd Zdd Z  Z	S )AudioClassificationDatasetz5
    Base class of audio classification dataset.
    r   Nfiles	list[str]labels	list[int]	feat_typestrsample_rate
int | Nonec                   sT   t    |t vrtd| dtt  || _|| _|| _|| _	|| _
dS )a?  
        Args:
            files (:obj:`List[str]`): A list of absolute path of audio files.
            labels (:obj:`List[int]`): Labels of audio files.
            feat_type (:obj:`str`, `optional`, defaults to `raw`):
                It identifies the feature type that user wants to extract an audio file.
        zUnknown feat_type: z, it must be one in N)super__init__
feat_funcskeysRuntimeErrorlistr   r   r   r   feat_config)selfr   r   r   r   kwargs	__class__ i/home/app/PaddleOCR-VL-test/.venv_paddleocr/lib/python3.10/site-packages/paddle/audio/datasets/dataset.pyr   "   s   
z#AudioClassificationDataset.__init__
input_filec                 C  s   t N)NotImplementedError)r   r    r   r   r   	_get_dataA   s   z$AudioClassificationDataset._get_datac           	      C  s   | j | | j| }}tj|\}}|| _t| j }i }t|j	dkr)|
d}tj|tjd}|d ur]|d}| jdkrK|dd| ji| j}n|di | j}||
d|d< n||d< ||d< |S )	Nr   r   )Zdtyper	   srfeatlabelr   )r   r   paddleZaudioloadr   r   r   lenshapeZsqueezeZ	to_tensorZfloat32Z	unsqueezer   )	r   idxfiler&   Zwaveformr   Z	feat_funcrecordZfeature_extractorr   r   r   _convert_to_recordD   s*   



z-AudioClassificationDataset._convert_to_recordc                 C  s   |  |}|d |d fS )Nr%   r&   )r.   )r   r+   r-   r   r   r   __getitem__]   s   
z&AudioClassificationDataset.__getitem__c                 C  s
   t | jS r!   )r)   r   )r   r   r   r   __len__a   s   
z"AudioClassificationDataset.__len__)r   N)r   r   r   r   r   r   r   r   )r    r   )
__name__
__module____qualname____doc__r   r#   r.   r/   r0   __classcell__r   r   r   r   r
      s    
r
   )
__future__r   r'   featuresr   r   r   r   r   ioZDatasetr
   r   r   r   r   <module>   s   	