U (‰d„ã@s`ddlZddlmZmZmZmZmZddlm Z ddl m Z Gdd„de ƒZ Gdd „d e ƒZ dS) éN)ÚAnyÚCallableÚOptionalÚTupleÚList)ÚImageé)Ú VisionDatasetcsŠeZdZdZdeeeeeeeeddœ‡fdd„ Zee j dœdd„Z ee e dœd d „Z eee e fd œd d „Zedœdd„Z‡ZS)Ú CocoDetectiona`MS Coco Detection `_ Dataset. It requires the `COCO API to be installed `_. Args: root (string): Root directory where images are downloaded to. annFile (string): Path to json annotation file. transform (callable, optional): A function/transform that takes in an PIL image and returns a transformed version. E.g, ``transforms.PILToTensor`` target_transform (callable, optional): A function/transform that takes in the target and transforms it. transforms (callable, optional): A function/transform that takes input sample and its target as entry and returns a transformed version. N)ÚrootÚannFileÚ transformÚtarget_transformÚ transformsÚreturncsBtƒ ||||¡ddlm}||ƒ|_tt|jj ¡ƒƒ|_ dS)Nr)ÚCOCO) ÚsuperÚ__init__Zpycocotools.cocorÚcocoÚlistÚsortedZimgsÚkeysÚids)Úselfr r r rrr©Ú __class__©ú=/tmp/pip-unpacked-wheel-vx7f76es/torchvision/datasets/coco.pyrs  zCocoDetection.__init__©ÚidrcCs0|j |¡dd}t tj |j|¡¡ d¡S)NrÚ file_nameZRGB) rZloadImgsrÚopenÚosÚpathÚjoinr Úconvert)rrr#rrrÚ _load_image'szCocoDetection._load_imagecCs|j |j |¡¡S©N)rZloadAnnsZ getAnnIds©rrrrrÚ _load_target+szCocoDetection._load_target)ÚindexrcCs@|j|}| |¡}| |¡}|jdk r8| ||¡\}}||fSr')rr&r)r)rr*rÚimageÚtargetrrrÚ __getitem__.s     zCocoDetection.__getitem__)rcCs t|jƒSr')Úlenr)rrrrÚ__len__8szCocoDetection.__len__)NNN)Ú__name__Ú __module__Ú __qualname__Ú__doc__ÚstrrrrÚintrr&rrr)rr-r/Ú __classcell__rrrrr s úù r cs,eZdZdZeeedœ‡fdd„ Z‡ZS)Ú CocoCaptionsaò`MS Coco Captions `_ Dataset. It requires the `COCO API to be installed `_. Args: root (string): Root directory where images are downloaded to. annFile (string): Path to json annotation file. transform (callable, optional): A function/transform that takes in an PIL image and returns a transformed version. E.g, ``transforms.PILToTensor`` target_transform (callable, optional): A function/transform that takes in the target and transforms it. transforms (callable, optional): A function/transform that takes input sample and its target as entry and returns a transformed version. Example: .. code:: python import torchvision.datasets as dset import torchvision.transforms as transforms cap = dset.CocoCaptions(root = 'dir where images are', annFile = 'json annotation file', transform=transforms.PILToTensor()) print('Number of samples: ', len(cap)) img, target = cap[3] # load 4th sample print("Image Size: ", img.size()) print(target) Output: :: Number of samples: 82783 Image Size: (3L, 427L, 640L) [u'A plane emitting smoke stream flying over a mountain.', u'A plane darts across a bright blue sky behind a mountain covered in snow', u'A plane leaves a contrail above the snowy mountain top.', u'A mountain that has a plane flying overheard in the distance.', u'A mountain view with a plume of smoke in the background'] rcsdd„tƒ |¡DƒS)NcSsg|] }|d‘qS)Úcaptionr)Ú.0ÚannrrrÚ hsz-CocoCaptions._load_target..)rr)r(rrrr)gszCocoCaptions._load_target) r0r1r2r3r5rr4r)r6rrrrr7<s*r7)Úos.pathr"ÚtypingrrrrrZPILrZvisionr r r7rrrrÚs   3