o Niz,@sdZddlmZddlmZddlmZddlZddlZddlZddlZ ddl m m Z ddlmZdZdZdZed ZGd d d ejjZGd d d eZGdddedgdZdS)zCaltech birds dataset.)absolute_import)division)print_functionNa.Caltech-UCSD Birds 200 (CUB-200) is an image dataset with photos of 200 bird species (mostly North American). The total number of categories of birds is 200 and there are 6033 images in the 2010 dataset and 11,788 images in the 2011 dataset. Annotations include bounding boxes, segmentation labels. z4http://www.vision.caltech.edu/visipedia/CUB-200.htmla @techreport{WelinderEtal2010, Author = {P. Welinder and S. Branson and T. Mita and C. Wah and F. Schroff and S. Belongie and P. Perona}, Institution = {California Institute of Technology}, Number = {CNS-TR-2010-001}, Title = {{Caltech-UCSD Birds 200}}, Year = {2010} } z((\w*)/)*(\d*).(\w*)/(\w*.jpg)$c@sHeZdZdZejdZeddZ ddZ ddZ d d Z d d Z d S)CaltechBirds2010zCaltech Birds 2010 dataset.0.1.0cCst|jddddS)Nz?http://www.vision.caltech.edu/visipedia-data/CUB-200/images.tgzz>http://www.vision.caltech.edu/visipedia-data/CUB-200/lists.tgzzDhttp://www.vision.caltech.edu/visipedia-data/CUB-200/annotations.tgzname images_url split_urlannotations_urlCaltechBirdsInforselfrj/home/ubuntu/.local/lib/python3.10/site-packages/tensorflow_datasets/image_classification/caltech_birds.py_caltech_birds_info9 z$CaltechBirds2010._caltech_birds_infoc CZtjj|ttjtjtjtjjddtjtj tjjddddt t dSN) num_classes)NN)shapeimagezimage/filenamelabel label_namebboxsegmentation_mask)rr)builder descriptionfeaturessupervised_keyshomepagecitation tfdscore DatasetInfo _DESCRIPTIONr" FeaturesDictImageText ClassLabel BBoxFeature_URL _CITATIONrrrr_infoB   zCaltechBirds2010._infoc Cs||jj|jj|jjg}||jj|jjg}tj|dd}tj|dd}t j j |}| }Wdn1sCwYt j j |}| }Wdn1s`wYtt} t j j |dD]X\} } } | D]P} | drtj| | }t j j |d}tjjjj j|ddd gd }Wdn1swY| | d d|d| | d d|d qzqstjjtjj||d || d dtjjtjj||d || d dgS)Nrzlists/train.txtzlists/test.txtrz.matrbTrseg) squeeze_mevariable_names.archive file_names annotationsr gen_kwargs) downloadrr r r download_and_extractospathjointfiogfileGFileread splitlines collections defaultdictlistwalkendswithr'r( lazy_importsscipyloadmatsplitappendSplitGeneratorSplitTRAIN iter_archiveTEST)r dl_manager download_pathextracted_path train_path test_pathf train_list test_list attributesroot_filesfnamerCmatrrr_split_generatorsTs\       z"CaltechBirds2010._split_generatorscCs<|d|}|d|}|d|}|d|}||||fS)aFunction to get normalized bounding box values. Args: bbox_annotations: list of bbox values in kitti format img_width: image width img_height: image height Returns: Normalized bounding box xmin, ymin, xmax, ymax values topleftbottomrightr)rbbox_annotations img_width img_heightyminxminymaxxmaxrrr_get_bounding_box_valuess    z)CaltechBirds2010._get_bounding_box_valuesc cs|D]\}}|dd}t|}|r"d|ddd|vr#q|}|d}t|dd} |ddd } || d} | j\} } | || d | | }|||| |t j j t |d d t |dd t |d d t |d d d | ddddtjfdfVqdS)aGenerate birds images, labels and bounding box given the directory path. Args: archive: object that iterates over the zip file_names : list of train/test image file names obtained from mat file annotations : dict of image file names and bbox attributes, segmentation labels Yields: Image path, Image file name, its corresponding label and bounding box values \/Nrr8rg?r9)rprqrrrsr)replace_NAME_REmatchrDrSgroupslowerintrrtr'r"BBoxminnpnewaxis)rr;r<r=rffobjresmatchesr label_key file_namerheightwidthrrrr_generate_exampless:            z#CaltechBirds2010._generate_examplesN)__name__ __module__ __qualname____doc__r'r(VersionVERSIONpropertyrr2rhrtrrrrrr4s  3 rc@s@eZdZdZejdZeddZ ddZ ddZ d d Z d S) CaltechBirds2011zCaltech Birds 2011 dataset.rcCst|jddddS)NzJhttp://www.vision.caltech.edu/visipedia-data/CUB-200-2011/CUB_200_2011.tgzzKhttp://www.vision.caltech.edu/visipedia-data/CUB-200-2011/segmentations.tgzrr rrrrrrz$CaltechBirds2011._caltech_birds_infoc Crrr&rrrrr2r3zCaltechBirds2011._infoc Cs||jjg}||jj|jjg}tj|dd}tj|dd}tj|dd}gg}}t t } t j j |} t j j |i} t j j |R} t| | | D]C\} }}| \}}|\}}t|}|}| |ddd|dd||krt|dkr||qV||qVWdn1swYWdn1swYWdn1swYt j j |dD]L\}}}|D]D}|drt j j tj||d }tjjjjtj|tj d dd }Wdn 1s wY| |dd|qqtjj!tj"j#|$|d|| d d tjj!tj"j%|$|d|| d d gS)NrzCUB_200_2011/images.txtz!CUB_200_2011/train_test_split.txtzCUB_200_2011/bounding_boxes.txtryr8rz.pngr4)dtype)flagsr:r>)&r@rr rAr rBrCrDrKrLrMrErFrGrHziprSr|r}r~rTrrNrOr'r(rPcv2imdecoder fromstringrIuint8rUrVrWrXrY)rrZr[r\image_names_path split_path bbox_pathr`rarbr_f1f2lineline1line2img_idxvalidximg_namerrrcrdrerfpng_fmaskrrrrhs        (       z"CaltechBirds2011._split_generatorscCsdt|d|}t|d|}t|dt|d|}t|dt|d|}||||fS)aGets normalized bounding box values (Conversion to KITTI format). Args: bbox_annotations: list of bbox values in kitti format img_width: image width img_height: image height Returns: Normalized bounding box xmin, ymin, xmax, ymax values rrr9rz)float)rrmrnrorqrprsrrrrrrt+s    z)CaltechBirds2011._get_bounding_box_valuesN) rrrrr'r(rrrrr2rhrtrrrrrs   s.      p