o
    oi                     @   s   d dl Z d dlZd dlZd dlZd dlZd dlZd dlZd dlZ	dej
dedejfddZdd Zed	krne  Zejd
ed ejdedd ejdedd e Zejejs`J ejejdshJ ee dS dS )    Nbuffercodecreturnc                 C   s$   t jt| d  |d\}}|S )N.)format)taloadioBytesIOtobytes)r   r   wav_ r   [/home/ubuntu/SoloSpeech/.venv/lib/python3.10/site-packages/df/scripts/fix_n_samples_hdf5.pyload_encoded   s    r   c           
      C   s  t | jd}|jd| j}|d u rtdtjd t	d ||jd< |jd| j
}|d u r;tdtjd |d }||jd< |jd	d
}| D ]}| D ]\}}|jdd }t||dd |d
krot|d }	nt||d}	t|	j |	 dksJ |d ur||	jd ks|	 |k rtd| d|	j d |	 dks|	jd dksJ t|	jd |jd< |	 dkrdnt|	jd |jd< z|jd= W qQ ty   Y qQw qKW d    d S 1 sw   Y  d S )Nzr+srzsr not found.)file   max_freqzmax_freq not found.   r   pcm	n_samples )end.)r   zWarning: Number of samples (z#) does not match with audio shape ()r      
n_channelsn_ch)h5pyFilehdf5attrsgetr   printsysstderrexitr   valuesitemstorch
from_numpyr   shapedimnumelintKeyError)
argsh5fr   r   r   groupkeydsr   audior   r   r   main   sN   


$"r6   __main__r    )typez--sr)r8   defaultz
--max-freqz.hdf5)argparser   osr$   r   numpynpr)   
torchaudior   ndarraystrTensorr   r6   __name__ArgumentParserparseradd_argumentr.   
parse_argsr0   pathisfiler    endswithr   r   r   r   <module>   s(   &