@@ -62,23 +62,32 @@ formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(messag
 handler.setFormatter(formatter)
 root.addHandler(handler)
 
-from mapdataset import ImageGroupReader, single_layer_converter, MapsDataset, MapReader
 
+from mapdataset import ImageGroupReader, single_layer_converter, MapsDataset, MapReader, ImageUtils
+
+dillFolder = "./data/output/SF_Layered/32x32/group-1280-stride-10"
+
+mapsDataset = MapsDataset(
+    patch_size=(32, 32), 
+    stride=10, 
+    sample_group_size=1280, 
+    converter=single_layer_converter,
+    outputDir="./data/output"
+    ) 
+
+mapsDataset.loadPatches("./data/output/SF_Layered/32x32/group-1280-stride-10")
+patchNo = randint(0, len(mapsDataset))
+logging.info(f"reading patch {patchNo}")
+patch = mapsDataset[patchNo]
 
 
-dillFolder = "./data/output/SF_Layered/32x32/stride-10"
-nGroups = 0
-# Iterate directory
-for path in os.listdir(dillFolder):
-    # check if current path is a file
-    if os.path.isfile(os.path.join(dillFolder, path)) and path.endswith(".dill"):
-        nGroups += 1
+im = ImageUtils.TorchNpPatchToPILImgGray(patch)
+path = os.path.join(dillFolder, f"{patchNo}.png")
+im.save(path)
+
+```
+
+# Using for training
 
-for i in range(nGroups):
-    reader = ImageGroupReader(dillFolder)
-    data = reader.load_group(groupNo=i)
-    patchImgArray = reader.asImg(data[0])
-    im = Image.fromarray(patchImgArray)
-    path = os.path.join(dillFolder, f"{i}-0.png")
-    im.save(path)
-```
+1. Create patches if you already do not have them
+2. Create a MapsDataset object and load patches. Now you can use the dataset object as a regular Pytorch dataset or use it with a Dataloader.
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "MapDatasetGenerator"
-version = "0.0.1"
+version = "0.0.2"
 description = "Map dataset generator for learning map representations and generation"
 authors = ["Ishaan <iparanja@ucsc.edu>", "Golam Md Muktadir <muktadir@ucsc.edu", "Bahar Bateni <bbateni@ucsc.edu>"]
 readme = "README.md"