o Ni@s~dZddlmZddlmZddlmZddlZddlmZddl m m Z ddl mZdZdZd ZGd d d ejjZdS) zdSprites dataset.)absolute_import)division)print_functionN)movesz@misc{dsprites17, author = {Loic Matthey and Irina Higgins and Demis Hassabis and Alexander Lerchner}, title = {dSprites: Disentanglement testing Sprites dataset}, howpublished= {https://github.com/deepmind/dsprites-dataset/}, year = "2017", } zqhttps://github.com/deepmind/dsprites-dataset/blob/master/dsprites_ndarray_co1sh3sc6or40x32y32_64x64.hdf5?raw=trueadSprites is a dataset of 2D shapes procedurally generated from 6 ground truth independent latent factors. These factors are *color*, *shape*, *scale*, *rotation*, *x* and *y* positions of a sprite. All possible combinations of these latents are present exactly once, generating N = 737280 total images. ### Latent factor values * Color: white * Shape: square, ellipse, heart * Scale: 6 values linearly spaced in [0.5, 1] * Orientation: 40 values in [0, 2 pi] * Position X: 32 values in [0, 1] * Position Y: 32 values in [0, 1] We varied one latent at a time (starting from Position Y, then Position X, etc), and sequentially stored the images in fixed order. Hence the order along the first dimension is fixed and allows you to map back to the value of the latents corresponding to that image. We chose the latents values deliberately to have the smallest step changes while ensuring that all pixel outputs were different. No noise was added. c@sDeZdZdZejddZejdgZddZ ddZ d d Z d S) DspriteszdSprites data set.2.0.0z6New split API (https://tensorflow.org/datasets/splits)z2.1.0cCstjjddtjjddtjjddtjjddtjjddtjjddtjjgtjdtjjgtjdtjjgtjdtjjgtjdtjjgtjdd }|jd kr[tj|d <tj j |t tj |d t d S)N)@r)shape) num_classes( )r dtype image label_shape label_scalelabel_orientationlabel_x_positionlabel_y_position value_shape value_scalevalue_orientationvalue_x_positionvalue_y_positionridz,https://github.com/deepmind/dsprites-dataset)builder descriptionfeatureshomepagecitation)tfdsr Image ClassLabelTensortffloat32versionTextcore DatasetInfo _DESCRIPTION FeaturesDict _CITATION)self features_dictr2V/home/ubuntu/.local/lib/python3.10/site-packages/tensorflow_datasets/image/dsprites.py_infoMs*        zDsprites._infocCs&|t}tjjtjjt|ddgS)N)filepath)name gen_kwargs)download_URLr#r+SplitGeneratorSplitTRAINdict)r0 dl_managerr5r2r2r3_split_generatorses zDsprites._split_generatorsc cstjjj|d!}t|d}t|dd}t|dd}Wdn1s.wYtt |||D]?\}\}}} t t |d|d|d |d |d |d | d| d | d | d | d d } |j dkrvd || d<|| fVqs