rsbox.ml

  1"""
  2File: ml.py 
  3-------------- 
  4Contains useful functions that I may import while programming 
  5machine learning projects.  
  6""" 
  7
  8from matplotlib import pyplot as plt
  9from PIL import Image
 10import requests
 11import torch 
 12import numpy as np 
 13from glob import glob
 14from tqdm.auto import tqdm
 15import random
 16import torchvision.transforms as T
 17
 18
 19def plot_png(file_path, color):
 20	"""
 21	Plots a png file (or jpeg really) using matplotlib. Color is a bool value. 
 22	If true, image will be displayed in color. Else, it will be displayed in 
 23	greyscale.
 24	"""
 25	image = Image.open(file_path)
 26	if color == True:
 27		plt.imshow(image, interpolation='nearest')
 28		plt.show()
 29	else:
 30		plt.imshow(image, interpolation='nearest', cmap='Greys_r')
 31		plt.show()
 32
 33
 34def plot_np_img(image_array, color):
 35	"""
 36    (Here for longevity purposes... use plot instead)
 37	Plots a numpy array using matplotlib. Color is a bool value. If true, 
 38	image will be displayed in color. Else, it will be displayed in 
 39	greyscale.
 40
 41    input: image_array must be of shape (1, H, W) or (H, W) for greyscale
 42    and (3, H, W) for color. 
 43	"""
 44	if color == True:
 45		plt.imshow(image_array, interpolation='nearest')
 46		plt.show()
 47	else:
 48		plt.imshow(image_array, interpolation='nearest', cmap='Greys_r')
 49		plt.show()
 50
 51
 52def plot(image_array, color=True):
 53    """
 54    Plots a numpy array using matplotlib. Color is a bool value which
 55    is true by default. If passed in as False, image will be displayed in 
 56    greyscale. 
 57
 58    input: image_array must be of shape (1, H, W) or (H, W, 1) or (H, W) for greyscale
 59    and (3, H, W) or (H, W, 3) for color. 
 60    """
 61
 62    # move to numpy is not already in numpy array format 
 63    image_array = np.array(image_array)
 64    image_array = np.squeeze(image_array)
 65
 66    if color == True:
 67        if image_array.shape[0] == 3:
 68            image_array = np.moveaxis(image_array, 0, -1)
 69        plt.imshow(image_array, interpolation='nearest')
 70        plt.show()
 71    else:
 72        plt.imshow(image_array, interpolation='nearest', cmap='Greys_r')
 73        plt.show()
 74
 75
 76def image_dir_to_data(dirpath, extension):
 77    """
 78    Takes in a directory containing images
 79    and returns a list of numpy arrays representing
 80    those images. 
 81    Args:
 82        - dirpath: path to directory 
 83        - extension: image extension type (e.g. png) (string)
 84    """
 85    data_subset = []
 86    sub_set = glob(dirpath + '/*.' + extension)
 87    for elem in sub_set:
 88        image = Image.open(elem)
 89        image_array = np.array(image)
 90        data_subset.append(image_array)
 91    return data_subset  
 92
 93
 94def image_dir_to_data_norm(dirpath, extension):
 95    """
 96    (Same as image_dir_to_data but with max value 
 97    (255) normalization). 
 98    Takes in a directory containing images
 99    and returns a list of numpy arrays representing
100    those images. 
101    Args:
102        - dirpath: path to directory 
103        - extension: image extension type (e.g. png) (string)
104    """
105    data_subset = []
106    sub_set = glob(dirpath + '/*.' + extension)
107    for elem in sub_set:
108        image = Image.open(elem)
109        image_array = np.array(image)
110        image_array = image_array.astype(float)
111        image_array = image_array/255.0
112        data_subset.append(image_array)
113    return data_subset    
114
115
116def gen_label_pair(data_samples, label):
117	"""
118	Takes in a list of data samples (i.e. images, x)
119	and a corresponding label and returns a list of tuples
120	where each tuple is an (x, y) pair. 
121	"""
122	d_set = []
123	for elem in data_samples:
124	    sample_pair = (elem, label)
125	    d_set.append(sample_pair)
126	return d_set 
127
128
129def gen_distro(master_list):
130    """
131    Takes in a list of lists who each contain
132    (x, y) tuple samples for that class and returns 
133    a randomly concatenated version. 
134    """
135    generated_distro = []
136    for class_list in master_list:
137        generated_distro = generated_distro + class_list
138    random.shuffle(generated_distro)
139    return generated_distro
140
141
142def get_img(url='https://stanford.edu/~rsikand/assets/images/seal.png', resize=True, size=(128, 128)):
143    """
144    Returns a sample numpy image to play with. 
145    Arguments:
146    -----------
147    - url: Can pass in custom url if you want a different image. 
148    Defaults to a baby seal!  
149    - resize: by default, all images are resized to 128x128. If you'd
150    like to keep the default size, specify False. 
151    - size: size to resize to. Default is 128x128. 
152    """
153
154    image = np.array(Image.open(requests.get(url, stream=True).raw))
155
156
157    if resize:
158        if len(image.shape) == 3:
159            image = np.array(T.Resize(size=size)(torch.movedim(torch.tensor(image), -1, 0)))
160        else:
161            image = np.array(T.Resize(size=size)(torch.tensor(image)))
162    else:
163        if len(image.shape) == 3:
164            if image.shape[0] != 3 and image.shape[2] == 3:
165                image = np.moveaxis(image, -1, 0)
166            elif image.shape[0] != 1 and image.shape[2] == 1:
167                image = np.moveaxis(image, -1, 0)
168            else:
169                image = image
170        else:
171            image = image
172
173    return image
174
175
176def img_dataset_from_dir(dir_path):
177    """
178    Given a directory containing folders
179    representing classes of images, this
180    functions builds a valid numpy
181    dataset distribution. 
182    Input (dir_path) structure: 
183    dir_path/class_1, class_n/1.png 
184    Note: 'dir_path' must be the raw
185    dir name (no trailing dash) 
186    Output: [(x,y), ..., (x,y)]
187    """
188
189    dir_path = dir_path + "/*/"
190    class_list = glob(dir_path, recursive = True)
191    
192    master_list = []
193    idx = 0
194    for class_ in class_list:
195        curr_class = image_dir_to_data_norm(class_, "png")
196        new_arrays = []
197        for elem in curr_class:
198            if len(elem.shape) == 2:
199                elem = np.expand_dims(elem, axis=0)
200            assert len(elem.shape) == 3
201            if elem.shape[0] != 3 and elem.shape[2] == 3:
202                elem = np.moveaxis(elem, -1, 0)
203            elif elem.shape[0] != 1 and elem.shape[2] == 1:
204                elem = np.moveaxis(elem, -1, 0)
205            new_arrays.append(elem)
206
207        labeled_list = gen_label_pair(new_arrays, idx)
208        master_list.append(labeled_list)
209        idx += 1
210
211    return gen_distro(master_list)
212
213
214def print_model_size(net):
215    """
216    Function that calculates PyTorch model size
217    and prints it out in mb, bytes, and num params.  
218    """
219    # taken from 
220    # https://stackoverflow.com/questions/49201236/check-the-total-number-of-parameters-in-a-pytorch-model
221    total_params = sum(p.numel() for p in net.parameters())
222    
223    # taken from 
224    # https://discuss.pytorch.org/t/finding-model-size/130275/2 
225    param_size = 0
226    for param in net.parameters():
227        param_size += param.nelement() * param.element_size()
228    buffer_size = 0
229    for buffer in net.buffers():
230        buffer_size += buffer.nelement() * buffer.element_size()
231
232    size_all_mb = (param_size + buffer_size) / 1024**2
233    
234    print("Total number of parameters: ", total_params)
235    print("model size (bytes): ", (param_size + buffer_size))
236    print('model size (mb): {:.3f}MB'.format(size_all_mb))
237
238
239class MeanMetric:
240  """
241  Scalar metric designed to use on a per-epoch basis 
242  and updated on per-batch basis. For getting average across
243  the epoch. 
244  """
245
246  def __init__(self):
247    self.vals = []
248
249  def update(self, new_val):
250    self.vals.append(new_val)
251
252  def reset(self):
253    self.vals = []
254
255  def get(self):
256    mean_value = sum(self.vals)/len(self.vals)
257    return mean_value
258    
259
260
261def numpy_collate(batch):
262    """batch collate function but keep as numpy array instead 
263    of converting to torch tensor like pytorch default does. inputs
264    must be np.arrays 
265    https://github.com/pytorch/pytorch/blob/35c8f93fd238d42aaea8fd6e730c3da9e18257cc/torch/utils/data/dataloader.py#L218
266    """
267
268    r_tuple = []
269
270    for i in range(len(batch[0])):
271        arrays_to_stack = []
272        for j in range(len(batch)):
273            arrays_to_stack.append(batch[j][i])
274        stacked = np.stack(arrays_to_stack, axis=0)
275        r_tuple.append(stacked)
276    
277    return tuple(r_tuple)
278
279
280def load_image(img_path, resize=None, normalize=True):
281    """
282    Loads an image file into a numpy array in (C, H, W) form. 
283    Arguments:
284    -----------
285    - img_path (str): image filename (relative) 
286    - resize (tuple): (x, y) size to resize to. If None, won't resize. Default = None. 
287    - normalize: if True, divide by 255. 
288    Returns:
289    -----------
290    - image_array: loaded numpy array of the image file     
291    """
292    
293    x = Image.open(img_path)
294    image_array = np.array(x)
295
296    # remove transparency dimension if it exists 
297    if image_array.shape[-1] == 4 and len(image_array.shape) == 3:
298        image_array = image_array[:, :, :3]
299    elif image_array.shape[0] == 4 and len(image_array.shape) == 3:
300        image_array = image_array[:3, :, :]
301    
302    # if image has only 2 channels, add a third channel
303    if len(image_array.shape) == 3 and image_array.shape[-1] == 2:
304        image_array = np.concatenate([np.expand_dims(image_array[:, :, -1], axis=-1), image_array], axis=-1)
305    elif len(image_array.shape) == 3 and image_array.shape[0] == 2:
306        image_array = np.concatenate([np.expand_dims(image_array[-1, :, :], axis=0), image_array], axis=0)
307
308    if normalize:
309        image_array = image_array / 255.0
310
311    # if image is grayscale, add a channel dimension
312    if len(image_array.shape) == 2:
313        image_array = np.expand_dims(image_array, axis=0)
314    
315    assert len(image_array.shape) == 3
316
317    # switch dim 
318    if (image_array.shape[0] != 3 and image_array.shape[2] == 3) or (image_array.shape[0] != 1 and image_array.shape[2] == 1):
319        image_array = np.moveaxis(image_array, -1, 0)
320
321    assert image_array.shape[0] == 3 or image_array.shape[0] == 1
322
323    # resize 
324    if resize is not None:
325        image_array = np.array(T.Resize(size=resize)(torch.tensor(image_array)))
326
327    return image_array
328
329
330
331# v2 functions 
332
333def img_dir_to_data(dirpath, resize=None, normalize=True, extension=None):
334    """
335    (New version of image_dir_to_data). 
336    Takes in a directory containing images
337    and returns a list of numpy arrays representing
338    those images. 
339    Args:
340        - dirpath: path to directory 
341        - extension: include only files with this extension. Default = None. 
342        - resize: tuple of (h, w) for what to resize the images to. If None, won't resize. Default = None. 
343        - normalize: if True, divide by 255. 
344    """
345    data_subset = []
346    if extension is not None:
347        dirpath = dirpath + "/*." + extension
348    else:
349        dirpath = dirpath + "/*"
350    sub_set = glob(dirpath)
351
352    for elem in sub_set:
353        image = load_image(elem, resize=resize, normalize=normalize)
354        data_subset.append(image)
355
356    return data_subset  
357
358
359
360def classification_dataset(dir_path, resize=None, normalize=True, extension=None):
361    """
362    (New version of img_dataset_from_dir).  
363    Given a directory containing folders
364    representing classes of images, this
365    functions builds a valid numpy
366    dataset distribution. 
367    Input (dir_path) structure: 
368    dir_path/class_1, class_n/1.png 
369    Note: 'dir_path' must be the raw
370    dir name (no trailing dash) 
371    Output: [(x,y), ..., (x,y)]
372    Arguments:
373    - dir_path (str): path to the directory containing the class folders. 
374    - resize: tuple of (h, w) for what to resize the images to. If None, won't resize. Default = None. 
375    - normalize: if True, divide by 255. Default = True. 
376    - extension: include only files with this extension. Default = None. 
377    """
378
379    dir_path = dir_path + "/*/"
380    class_list = glob(dir_path, recursive = True)
381    
382    master_list = []
383    idx = 0
384    tqdm_class_idx = 1
385    tqdm_class_list = tqdm(class_list)
386    for class_ in tqdm_class_list:
387        tqdm_class_list.set_description(f"Loading class {tqdm_class_idx}")
388        tqdm_class_idx += 1
389        curr_class = img_dir_to_data(class_, resize, normalize, extension)
390        new_arrays = []
391        for elem in curr_class:
392            if len(elem.shape) == 2:
393                elem = np.expand_dims(elem, axis=0)
394            assert len(elem.shape) == 3
395            if elem.shape[0] != 3 and elem.shape[2] == 3:
396                elem = np.moveaxis(elem, -1, 0)
397            elif elem.shape[0] != 1 and elem.shape[2] == 1:
398                elem = np.moveaxis(elem, -1, 0)
399            new_arrays.append(elem)
400
401        labeled_list = gen_label_pair(new_arrays, idx)
402        master_list.append(labeled_list)
403        idx += 1
404
405    return gen_distro(master_list)