dekker.one
Group work for a Monash Research Methods course
1'''
2Created by Tony Silvestre to prepare images for use from a Kaggle Where's Waldo dataset
3'''
4import os
5import numpy as np
6from matplotlib import pyplot as plt
7import math
8import cv2
9from skimage import color, exposure, transform
10import skimage as sk
11import random
12
13'''
14Defines some basic preprocessing for the images
15'''
16def preprocess(img):
17 # Histogram normalization in v channel
18 hsv = color.rgb2hsv(img)
19 hsv[:, :, 2] = exposure.equalize_hist(hsv[:, :, 2])
20 img = color.hsv2rgb(hsv)
21
22 return img
23
24'''
25Augments an image (horizontal reflection, hue, contrast, saturation adjustment) to expand on the dataset
26Code was adapted from Medium.com/@thimblot (Thomas Himblot)
27'''
28def augment(img):
29 # Randomly rotate the image
30 random_degree = random.uniform(-25, 25) # Random degree of rotation between 25% on the left and 25% on the right
31 img = sk.transform.rotate(img, random_degree)
32
33 # Add random noise to the image
34 img = sk.util.random_noise(img)
35
36 # Randomly (25% chance) flips the image horizontally (by inverting the image dimensions)
37 if random.randint(0, 4) == 1:
38 img = img[:, ::-1]
39
40 # Randomly (25% chance) inverts the intensity range of the image
41 if random.randint(0, 4) == 1:
42 img = sk.util.invert(img)
43
44 return img
45
46def gen_data(w_path, n_w_path):
47 waldo_file_list = os.listdir(os.path.join(w_path))
48 total_w = len(waldo_file_list)
49 not_waldo_file_list = os.listdir(os.path.join(n_w_path))
50 total_nw = len(not_waldo_file_list)
51 imgs_raw = [] # Images
52 imgs_lbl = [] # Image labels
53
54 w = 0
55 for image_name in waldo_file_list:
56 pic = cv2.imread(os.path.join(w_path, image_name)) # NOTE: cv2.imread() returns a numpy array in BGR not RGB
57 pic = preprocess(pic)
58
59 pic_roll = np.rollaxis(pic, -1) # rolls colour axis to 0
60 imgs_raw.append(pic_roll)
61 imgs_lbl.append(1) # Value of 1 as Waldo is present in the image
62
63 for j in range(0, 10):
64 pic = augment(pic)
65 pic_roll = np.rollaxis(pic, -1) # rolls colour axis to 0
66 imgs_raw.append(pic_roll)
67 imgs_lbl.append(1) # Value of 1 as Waldo is still present in the transformed image
68
69 print('Completed: {0}/{1} Waldo images'.format(w+1, total_w))
70 w += 1
71
72 nw = 0
73 for image_name in not_waldo_file_list:
74 pic = cv2.imread(os.path.join(n_w_path, image_name))
75 pic = preprocess(pic)
76
77 pic = np.rollaxis(pic, -1) # rolls colour axis to 0
78 imgs_raw.append(pic)
79 imgs_lbl.append(0)
80
81 print('Completed: {0}/{1} non-Waldo images'.format(nw+1, total_nw))
82 if nw > 10*w:
83 print("Non_Waldo files restricted")
84 break
85 else:
86 nw += 1
87
88 ## Randomise and split data into training and test sets
89 # Code was modified from code written by: Kyle O'Brien (medium.com/@kylepob61392)
90 n_images = len(imgs_raw)
91 TRAIN_TEST_SPLIT = 0.75 # Amount of trainingdata as a percentage of the total
92
93 # Split at the given index
94 split_index = int(TRAIN_TEST_SPLIT * n_images)
95 shuffled_indices = np.random.permutation(n_images)
96 train_indices = shuffled_indices[0:split_index]
97 test_indices = shuffled_indices[split_index:]
98
99 train_data = []
100 train_lbl = []
101 test_data = []
102 test_lbl = []
103
104 # Split the images and the labels
105 for index in train_indices:
106 train_data.append(imgs_raw[index])
107 train_lbl.append(imgs_lbl[index])
108
109 for index in test_indices:
110 test_data.append(imgs_raw[index])
111 test_lbl.append(imgs_lbl[index])
112
113 try:
114 # Save the data as numpy files
115 np.save('Waldo_train_data.npy', train_data)
116 np.save('Waldo_train_lbl.npy', train_lbl)
117 np.save('Waldo_test_data.npy', test_data)
118 np.save('Waldo_test_lbl.npy', test_lbl)
119 print("All data saved")
120 except:
121 print("ERROR: Data may not be completely saved")
122
123
124if __name__ == "__main__":
125 # Paths to the Waldo images
126 waldo_path = 'waldo_data/64/waldo'
127 n_waldo_path = 'waldo_data/64/notwaldo'
128
129 gen_data(waldo_path, n_waldo_path)