# DISTRIBUTION STATEMENT A. Approved for public release: distribution unlimited.
#
# This material is based upon work supported by the Assistant Secretary of Defense for Research and
# Engineering under Air Force Contract No. FA8721-05-C-0002 and/or FA8702-15-D-0001. Any opinions,
# findings, conclusions or recommendations expressed in this material are those of the author(s) and
# do not necessarily reflect the views of the Assistant Secretary of Defense for Research and
# Engineering.
#
# © 2018 Massachusetts Institute of Technology.
#
# MIT Proprietary, Subject to FAR52.227-11 Patent Rights - Ownership by the contractor (May 2014)
#
# The software/firmware is provided to you on an As-Is basis
#
# Delivered to the U.S. Government with Unlimited Rights, as defined in DFARS Part 252.227-7013 or
# 7014 (Feb 2014). Notwithstanding any copyright notice, U.S. Government rights in this work are
# defined by DFARS 252.227-7013 or DFARS 252.227-7014 as detailed above. Use of this work other than
# as specifically authorized by the U.S. Government may violate any copyrights that exist in this
# work.

from torchvision.datasets import CIFAR10
from os import path
from torchvision import transforms

class CIFAR10(CIFAR10):
    """`CIFAR10 <https://www.cs.toronto.edu/~kriz/cifar.html>`_ Dataset.
        Args:
            root (string): Root directory of dataset where directory
                ``cifar-10-batches-py`` exists or will be saved to if download is set to True.
            train (bool, optional): If True, creates dataset from training set, otherwise
                creates from test set.
            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.
            download (bool, optional): If true, downloads the dataset from the internet and
                puts it in root directory. If dataset is already downloaded, it is not
                downloaded again.
    """

    def __init__(self, root, train=True):
        root = path.join(root, 'cifar10')
        if train:
            transform = transforms.Compose([
                transforms.RandomCrop(32, padding=4),
                transforms.RandomHorizontalFlip(),
                transforms.ToTensor(),
            ])
        else:
            transform = transforms.Compose([
                transforms.ToTensor(),
            ])
        super().__init__(root, train, transform, target_transform=False, download=False)
        if train:
            self.data = self.train_data
            self.targets = self.train_labels
        else:
            self.data = self.test_data
            self.targets = self.test_labels

# DISTRIBUTION STATEMENT A. Approved for public release: distribution unlimited.
#
# This material is based upon work supported by the Assistant Secretary of Defense for Research and
# Engineering under Air Force Contract No. FA8721-05-C-0002 and/or FA8702-15-D-0001. Any opinions,
# findings, conclusions or recommendations expressed in this material are those of the author(s) and
# do not necessarily reflect the views of the Assistant Secretary of Defense for Research and
# Engineering.
#
# © 2018 Massachusetts Institute of Technology.
#
# MIT Proprietary, Subject to FAR52.227-11 Patent Rights - Ownership by the contractor (May 2014)
#
# The software/firmware is provided to you on an As-Is basis
#
# Delivered to the U.S. Government with Unlimited Rights, as defined in DFARS Part 252.227-7013 or
# 7014 (Feb 2014). Notwithstanding any copyright notice, U.S. Government rights in this work are
# defined by DFARS 252.227-7013 or DFARS 252.227-7014 as detailed above. Use of this work other than
# as specifically authorized by the U.S. Government may violate any copyrights that exist in this
# work.

from torchvision.datasets import CIFAR10
from os import path
from torchvision import transforms

class CIFAR10(CIFAR10):
    """`CIFAR10 <https://www.cs.toronto.edu/~kriz/cifar.html>`_ Dataset.
        Args:
            root (string): Root directory of dataset where directory
                ``cifar-10-batches-py`` exists or will be saved to if download is set to True.
            train (bool, optional): If True, creates dataset from training set, otherwise
                creates from test set.
            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.
            download (bool, optional): If true, downloads the dataset from the internet and
                puts it in root directory. If dataset is already downloaded, it is not
                downloaded again.
    """

    def __init__(self, root, train=True):
        root = path.join(root, 'cifar10')
        if train:
            transform = transforms.Compose([
                transforms.RandomCrop(32, padding=4),
                transforms.RandomHorizontalFlip(),
                transforms.ToTensor(),
            ])
        else:
            transform = transforms.Compose([
                transforms.ToTensor(),
            ])
        super().__init__(root, train, transform, target_transform=False, download=False)
        if train:
            self.data = self.train_data
            self.targets = self.train_labels
        else:
            self.data = self.test_data
            self.targets = self.test_labels

# DISTRIBUTION STATEMENT A. Approved for public release: distribution unlimited. # # This material is based upon work supported by the Assistant Secretary of Defense for Research and # Engineering under Air Force Contract No. FA8721-05-C-0002 and/or FA8702-15-D-0001. Any opinions, # findings, conclusions or recommendations expressed in this material are those of the author(s) and # do not necessarily reflect the views of the Assistant Secretary of Defense for Research and # Engineering. # # © 2018 Massachusetts Institute of Technology. # # MIT Proprietary, Subject to FAR52.227-11 Patent Rights - Ownership by the contractor (May 2014) # # The software/firmware is provided to you on an As-Is basis # # Delivered to the U.S. Government with Unlimited Rights, as defined in DFARS Part 252.227-7013 or # 7014 (Feb 2014). Notwithstanding any copyright notice, U.S. Government rights in this work are # defined by DFARS 252.227-7013 or DFARS 252.227-7014 as detailed above. Use of this work other than # as specifically authorized by the U.S. Government may violate any copyrights that exist in this # work. `CIFAR10 <https://www.cs.toronto.edu/~kriz/cifar.html>`_ Dataset.         Args:             root (string): Root directory of dataset where directory                 ``cifar-10-batches-py`` exists or will be saved to if download is set to True.             train (bool, optional): If True, creates dataset from training set, otherwise                 creates from test set.             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.             download (bool, optional): If true, downloads the dataset from the internet and                 puts it in root directory. If dataset is already downloaded, it is not                 downloaded again.