From 83bc94fa1bc331a4f4b53bc0478f406177505c50 Mon Sep 17 00:00:00 2001
From: facundoq <facundoq@gmail.com>
Date: Tue, 12 Jul 2022 16:44:26 -0300
Subject: [PATCH] changed name of package, added transform parameters

---
 setup.py        |  8 ++++----
 tinyimagenet.py | 17 +++++++++++------
 2 files changed, 15 insertions(+), 10 deletions(-)

diff --git a/setup.py b/setup.py
index 9747a8d..01f4596 100644
--- a/setup.py
+++ b/setup.py
@@ -12,8 +12,8 @@
     long_description = f.read()
 
 
-url="https://github.com/facundoq/torchvision-tinyimagenet"
-VERSION="0.2"
+url="https://github.com/facundoq/tinyimagenet"
+VERSION="0.1"
 
 class UploadCommand(Command):
     """Support setup.py upload."""
@@ -53,7 +53,7 @@ def run(self):
 
 
 setup(
-    name="torchvision-tinyimagenet",
+    name="tinyimagenet",
     version=VERSION,
     python_requires='>=3.6',
     packages=find_packages(),
@@ -77,7 +77,7 @@ def run(self):
     # metadata to display on PyPI
     author="Facundo Manuel Quiroga",
     author_email="facundoq@gmail.com",
-    description="Dataset class for PyTorch and the TinyImageNet dataset.",
+    description="Dataset class for PyTorch and the TinyImageNet dataset, with automated download and extraction.",
     keywords="TinyImageNet ImageNet Dataset PyTorch torch torchvision",
     url=url,   # project home page, if any
     project_urls={
diff --git a/tinyimagenet.py b/tinyimagenet.py
index f8fc6c9..b0a1131 100644
--- a/tinyimagenet.py
+++ b/tinyimagenet.py
@@ -103,7 +103,7 @@ class TinyImageNet(ImageFolder):
     mean = [0.485, 0.456, 0.406]
     std = [0.229, 0.224, 0.225]
 
-    def __init__(self, root: Path, split: str = "train") -> None:
+    def __init__(self, root: Path, split: str = "train",transform=None, target_transform=None) -> None:
         if isinstance(root,str):
             root = Path(root)
         assert split in ["train","val","test"]
@@ -112,7 +112,7 @@ def __init__(self, root: Path, split: str = "train") -> None:
         if not images_root.exists():
             download_resources(root,mirrors,resources)        
         preprocess_val(images_root)
-        super().__init__(images_root/split)
+        super().__init__(images_root/split,transform=transform,target_transform=target_transform)
         self.idx_to_words,self.idx_to_class = self.load_words_classes(images_root)
     
     def load_words_classes(self,root:Path):
@@ -131,16 +131,21 @@ def load_words_classes(self,root:Path):
 
 
 if __name__ == '__main__':
-    
+    from torchvision import transforms
+
     logging.basicConfig(level=logging.INFO)
-    
+    transform = transforms.Compose(
+        [transforms.ToTensor(),
+        transforms.Normalize(TinyImageNet.mean,TinyImageNet.std)]
+        )
+
     split ="val"
-    dataset = TinyImageNet(Path("~/.torchvision/tinyimagenet/"),split=split)
+    dataset = TinyImageNet(Path("~/.torchvision/tinyimagenet/"),split=split,transform=transform)
     n = len(dataset)
     print(f"TinyImageNet, split {split}, has  {n} samples.")
     print("Showing some samples")
     for i in range(0,n,n//5):
         image,klass = dataset[i]
-        print(f"Sample of class {klass:3d}, image {image}, words {dataset.idx_to_words[klass]}")
+        print(f"Sample of class {klass:3d}, image shape {image.shape}, words {dataset.idx_to_words[klass]}")