o .wi@s<ddlZddlZddlmZddlmZGdddeZdS)N) BaseDataset) _urlretrievec@s(eZdZdZddZddZddZdS) StanfordOnlineProductszwhttps://drive.usercontent.google.com/download?id=1TclrpQOF_ullUP99wk_gjGN8pKvtErG8&export=download&authuser=0&confirm=tcCsb|jdvr||j\}}n|jdkr)|d\}}|d\}}||}||}||_||_dS)N)traintestz train+testrr)split_load_split_txtpathslabels)selfr r train_paths train_labels test_paths test_labelsra/home/ubuntu/sommelier/.venv/lib/python3.10/site-packages/pytorch_metric_learning/datasets/sop.pygenerate_split s   z%StanfordOnlineProducts.generate_splitc Csgg}}ttj|jdd|d;}t|D]+\}}|dkr"q|}t|d|d}} | tj|jd| | |qWd||fS1sRwY||fS)NStanford_Online_ProductsEbay_z.txtr) openospathjoinroot enumeratestriprintappend) r rr r fill_splitlabelrrrrr s"      z&StanfordOnlineProducts._load_split_txtcCstj|jddtj|jtjdd}ttj|dt |d}| |jWdn1s4wYt |dS)NT)exist_ok/)urlfilenamer) rmakedirsrrrr DOWNLOAD_URLrrzipfileZipFile extractallremove)r download_folder_pathzip_refrrrdownload_and_remove'sz*StanfordOnlineProducts.download_and_removeN)__name__ __module__ __qualname__r-rr r4rrrrrs  r)rr.datasets.base_datasetrutils.common_functionsrrrrrrs