@@ -53,15 +53,15 @@ def get_char(self, layer):
5353# Modifications - Ishaan, Muktadir
5454
5555class MapsDataset (Dataset ):
56- def __init__ (self , window_size , step_size , sample_group_size , converter , outputDir = "../data/output" ):
56+ def __init__ (self , patch_size , stride , sample_group_size , converter , outputDir = "../data/output" ):
5757 self .outputDir = outputDir
5858 self .char_size = converter .char_size
5959 self .converter = converter
60- self .window_size = window_size
61- self .step_size = step_size
60+ self .patch_size = patch_size
61+ self .stride = stride
6262 self .sample_group_size = sample_group_size
6363 self .samples = []
64- self .block_size = self .window_size [0 ] * self .window_size [1 ] - 1
64+ self .block_size = self .patch_size [0 ] * self .patch_size [1 ] - 1
6565
6666 os .makedirs (self .outputDir , exist_ok = True )
6767
@@ -79,12 +79,12 @@ def __getitem__(self, idx):
7979 #return flat
8080
8181 def add (self , mapReader ):
82- for i in range (0 , mapReader .size [0 ] - self .window_size [0 ] + 1 , self .step_size ):
83- for j in range (0 , mapReader .size [1 ] - self .window_size [1 ] + 1 , self .step_size ):
82+ for i in range (0 , mapReader .size [0 ] - self .patch_size [0 ] + 1 , self .stride ):
83+ for j in range (0 , mapReader .size [1 ] - self .patch_size [1 ] + 1 , self .stride ):
8484 self .samples .append ([[
8585 (self .converter .get_char (mapReader .data [i + x ][j + y ]) / (len (self .converter .char_groups ) - 1 )) * - 2 + 1
86- for y in range (self .window_size [1 ])]
87- for x in range (self .window_size [0 ])])
86+ for y in range (self .patch_size [1 ])]
87+ for x in range (self .patch_size [0 ])])
8888
8989 #Generate image patches and write to data/output directory
9090 def generate_patches (self , mapReader , image_groups = 3 , outDirectory = None ):
@@ -98,12 +98,12 @@ def generate_patches(self, mapReader, image_groups=3, outDirectory=None):
9898 mapReader .standardize (converter = self .converter )
9999
100100 if outDirectory is None :
101- outDirectory = os .path .join (self .outputDir , mapReader .mapName , f"{ self .window_size [0 ]} x{ self .window_size [1 ]} " , f"stride-{ self .step_size } " )
101+ outDirectory = os .path .join (self .outputDir , mapReader .mapName , f"{ self .patch_size [0 ]} x{ self .patch_size [1 ]} " , f"stride-{ self .stride } " )
102102 os .makedirs (outDirectory , exist_ok = True )
103103
104104 img_group_number = 0
105- for i in range (0 , mapReader .size [0 ] - self .window_size [0 ] + 1 , self .step_size ):
106- for j in range (0 , mapReader .size [1 ] - self .window_size [1 ] + 1 , self .step_size ):
105+ for i in range (0 , mapReader .size [0 ] - self .patch_size [0 ] + 1 , self .stride ):
106+ for j in range (0 , mapReader .size [1 ] - self .patch_size [1 ] + 1 , self .stride ):
107107
108108 sample = self .extractSample (mapReader , topLeft = (i , j ))
109109
@@ -125,16 +125,16 @@ def extractSample(self, mapReader, topLeft):
125125 # sample = [
126126 # [
127127 # (self.converter.get_char(mapReader.data[i + x][j + y]) / (len(self.converter.char_groups) - 1)) * -2 + 1 # TODO this conversion should be done once in the original data instead of patches.
128- # for y in range(self.window_size [1])
128+ # for y in range(self.patch_size [1])
129129 # ]
130- # for x in range(self.window_size [0])
130+ # for x in range(self.patch_size [0])
131131 # ]
132132 sample = [
133133 [
134134 mapReader .data [i + x ][j + y ]
135- for y in range (self .window_size [1 ])
135+ for y in range (self .patch_size [1 ])
136136 ]
137- for x in range (self .window_size [0 ])
137+ for x in range (self .patch_size [0 ])
138138 ]
139139 return np .asarray (sample )
140140
@@ -143,8 +143,8 @@ def shuffle(self):
143143 random .shuffle (self .samples )
144144
145145 def get_train_test (self , train_ratio ):
146- train_dataset = MapsDataset (self .window_size , self .step_size , self .converter )
147- test_dataset = MapsDataset (self .window_size , self .step_size , self .converter )
146+ train_dataset = MapsDataset (self .patch_size , self .stride , self .converter )
147+ test_dataset = MapsDataset (self .patch_size , self .stride , self .converter )
148148 train_size = int (len (self .samples ) * train_ratio )
149149 train_dataset .samples = self .samples [:train_size ]
150150 test_dataset .samples = self .samples [train_size :]
0 commit comments