U
    (d	                     @   sL   d dl mZmZmZmZ d dlZddlmZ ddlm	Z	 G dd de	Z
dS )	    )AnyCallableOptionalTupleN   )
transforms   )VisionDatasetc                	       sl   e Zd ZdZdeeeeef eee ee edd fdd	Zeee	e	f d
ddZ
edddZ  ZS )FakeDataa  A fake dataset that returns randomly generated images and returns them as PIL images

    Args:
        size (int, optional): Size of the dataset. Default: 1000 images
        image_size(tuple, optional): Size if the returned images. Default: (3, 224, 224)
        num_classes(int, optional): Number of classes in the dataset. Default: 10
        transform (callable, optional): A function/transform that  takes in an PIL image
            and returns a transformed version. E.g, ``transforms.RandomCrop``
        target_transform (callable, optional): A function/transform that takes in the
            target and transforms it.
        random_offset (int): Offsets the index-based random seed used to
            generate each image. Default: 0

            r   
   Nr   )size
image_sizenum_classes	transformtarget_transformrandom_offsetreturnc                    s.   t  jd ||d || _|| _|| _|| _d S )N)r   r   )super__init__r   r   r   r   )selfr   r   r   r   r   r   	__class__ A/tmp/pip-unpacked-wheel-vx7f76es/torchvision/datasets/fakedata.pyr      s
    	zFakeData.__init__)indexr   c                 C   s   |t | krt| jj dt }t|| j  tj| j	 }tj
d| jdtjdd }t| t |}| jdk	r| |}| jdk	r| |}|| fS )z
        Args:
            index (int): Index

        Returns:
            tuple: (image, target) where target is class_index of the target class.
        z index out of ranger   )r   )r   ZdtypeN)len
IndexErrorr   __name__torchZget_rng_stateZmanual_seedr   Zrandnr   randintr   longZset_rng_stater   Z
ToPILImager   r   item)r   r   Z	rng_stateimgtargetr   r   r   __getitem__(   s    	




zFakeData.__getitem__)r   c                 C   s   | j S )N)r   )r   r   r   r   __len__B   s    zFakeData.__len__)r   r   r   NNr   )r!   
__module____qualname____doc__intr   r   r   r   r   r(   r)   __classcell__r   r   r   r   r
   	   s$         r
   )typingr   r   r   r   r"    r   Zvisionr	   r
   r   r   r   r   <module>   s   