refactor data generator using load function

Author: inzapp

generator.py

@@ -25,16 +25,8 @@
 from concurrent.futures.thread import ThreadPoolExecutor
 
 
-class DataGenerator:
+class DataGenerator(tf.keras.utils.Sequence):
     def __init__(self, root_path, image_paths, input_shape, batch_size, class_names, aug_brightness=0.0, aug_contrast=0.0, aug_rotate=0, aug_h_flip=False):
-        self.generator_flow = GeneratorFlow(root_path, image_paths, class_names, input_shape, batch_size, aug_brightness, aug_contrast, aug_rotate, aug_h_flip)
-
-    def flow(self):
-        return self.generator_flow
-
-
-class GeneratorFlow(tf.keras.utils.Sequence):
-    def __init__(self, root_path, image_paths, class_names, input_shape, batch_size, aug_brightness, aug_contrast, aug_rotate, aug_h_flip):
         assert 0.0 <= aug_brightness <= 1.0
         assert 0.0 <= aug_contrast <= 1.0
         assert type(aug_h_flip) == bool
@@ -59,17 +51,19 @@ def __init__(self, root_path, image_paths, class_names, input_shape, batch_size,
         self.transform = A.Compose(aug_methods)
         self.augmentation = len(aug_methods) > 0
 
-    def __getitem__(self, index):
+    def __len__(self):
+        return len(self.image_paths) // self.batch_size
+
+    def load(self):
         fs = []
-        for i in range(self.batch_size):
+        for _ in range(self.batch_size):
             fs.append(self.pool.submit(self.load_img, self.get_next_image_path()))
         batch_x = []
         batch_y = []
         for f in fs:
             img, path = f.result()
             x = self.preprocess(img, aug=self.augmentation)
             batch_x.append(x)
-
             dir_name = path.replace(self.root_path, '').split('/')[1]
             y = np.zeros((self.num_classes,), dtype=np.float32)
             if dir_name != 'unknown':
@@ -79,23 +73,6 @@ def __getitem__(self, index):
         batch_y = np.asarray(batch_y).reshape((self.batch_size, self.num_classes)).astype('float32')
         return batch_x, batch_y
 
-    def __len__(self):
-        return int(np.floor(len(self.image_paths) / self.batch_size))
-
-    def get_next_image_path(self):
-        path = self.image_paths[self.img_index]
-        self.img_index += 1
-        if self.img_index == len(self.image_paths):
-            self.img_index = 0
-            np.random.shuffle(self.image_paths)
-        return path
-
-    def random_blur(self, img):
-        if np.random.rand() > 0.5:
-            kernel_size = np.random.choice([3, 5])
-            img = cv2.GaussianBlur(img, (kernel_size, kernel_size), 0)
-        return img
-
     def preprocess(self, img, aug=False):
         img = cv2.resize(img, (self.input_shape[1], self.input_shape[0]))
         if aug:
@@ -105,6 +82,14 @@ def preprocess(self, img, aug=False):
         x = np.asarray(img).reshape(self.input_shape).astype('float32') / 255.0
         return x
 
+    def get_next_image_path(self):
+        path = self.image_paths[self.img_index]
+        self.img_index += 1
+        if self.img_index == len(self.image_paths):
+            self.img_index = 0
+            np.random.shuffle(self.image_paths)
+        return path
+
     def load_img(self, path):
         img = cv2.imdecode(np.fromfile(path, dtype=np.uint8), cv2.IMREAD_GRAYSCALE if self.input_shape[-1] == 1 else cv2.IMREAD_COLOR)
         return img, path

sigmoid_classifier.py

@@ -245,6 +245,7 @@ def train(self):
         if self.pretrained_iteration_count >= self.iterations:
             print(f'pretrained iteration count {self.pretrained_iteration_count} is greater or equal than target iterations {self.iterations}')
             exit(0)
+
         self.model.summary()
         print(f'\ntrain on {len(self.train_image_paths)} samples')
         print(f'validate on {len(self.validation_image_paths)} samples\n')
@@ -256,54 +257,51 @@ def train(self):
         eta_calculator = ETACalculator(iterations=self.iterations, start_iteration=iteration_count)
         eta_calculator.start()
         while True:
-            for idx, (batch_x, batch_y) in enumerate(self.train_data_generator.flow()):
-                lr_scheduler.update(optimizer, iteration_count)
-                loss = self.compute_gradient(self.model, optimizer, batch_x, batch_y, loss_function)
-                if self.show_class_activation_map and iteration_count % 100 == 0:
-                    try_count = 0
-                    while True:
-                        if try_count > len(batch_x):
-                            break
-                        rnum = random.randint(0, len(batch_x) - 1)
-                        if np.all(batch_y[rnum] < 0.3):  # skip cam view if unknown data
-                            continue
-                        else:
-                            new_input_tensor = batch_x[rnum]
-                            label_idx = np.argmax(batch_y[rnum]).item()
-                            break
-                    self.draw_cam(new_input_tensor, label_idx)
-                if self.live_loss_plot_flag:
-                    self.live_loss_plot.update(loss)
-                iteration_count += 1
-                progress_str = eta_calculator.update(iteration_count)
-                self.print_loss(progress_str, loss)
-                if iteration_count % 2000 == 0:
-                    self.save_last_model(self.model, iteration_count)
-                if iteration_count == self.iterations:
-                    self.save_last_model(self.model, iteration_count)
-                    self.save_model(iteration_count)
-                    self.remove_last_model()
-                    print('train end successfully')
-                    exit(0)
-                elif iteration_count >= int(self.iterations * self.warm_up) and self.checkpoint_interval > 0 and iteration_count % self.checkpoint_interval == 0:
-                    self.save_model(iteration_count)
+            batch_x, batch_y = self.train_data_generator.load()
+            lr_scheduler.update(optimizer, iteration_count)
+            loss = self.compute_gradient(self.model, optimizer, batch_x, batch_y, loss_function)
+            if self.show_class_activation_map and iteration_count % 100 == 0:
+                try_count = 0
+                while True:
+                    if try_count > len(batch_x):
+                        break
+                    rnum = random.randint(0, len(batch_x) - 1)
+                    if np.all(batch_y[rnum] < 0.3):  # skip cam view if unknown data
+                        continue
+                    else:
+                        new_input_tensor = batch_x[rnum]
+                        label_idx = np.argmax(batch_y[rnum]).item()
+                        break
+                self.draw_cam(new_input_tensor, label_idx)
+            if self.live_loss_plot_flag:
+                self.live_loss_plot.update(loss)
+            iteration_count += 1
+            progress_str = eta_calculator.update(iteration_count)
+            self.print_loss(progress_str, loss)
+            if iteration_count % 2000 == 0:
+                self.save_last_model(self.model, iteration_count)
+            if iteration_count == self.iterations:
+                self.save_last_model(self.model, iteration_count)
+                self.save_model(iteration_count)
+                self.remove_last_model()
+                print('train end successfully')
+                exit(0)
+            elif iteration_count >= int(self.iterations * self.warm_up) and self.checkpoint_interval > 0 and iteration_count % self.checkpoint_interval == 0:
+                self.save_model(iteration_count)
 
     def save_model(self, iteration_count):
         print()
-        if self.validation_data_generator.flow() is None:
-            self.save_last_model(self.model, iteration_count)
+        val_acc, val_class_score, val_unknown_score = self.evaluate(unknown_threshold=0.5, dataset='validation')
+        model_name = f'model_{iteration_count}_iter_acc_{val_acc:.4f}_class_score_{val_class_score:.4f}'
+        if self.include_unknown:
+            model_name += f'_unknown_score_{val_unknown_score:.4f}'
+        if val_acc > self.max_val_acc:
+            self.max_val_acc = val_acc
+            model_name = f'{self.checkpoint_path}/best_{model_name}.h5'
+            print(f'[best model saved]\n')
         else:
-            val_acc, val_class_score, val_unknown_score = self.evaluate(unknown_threshold=0.5, dataset='validation')
-            model_name = f'model_{iteration_count}_iter_acc_{val_acc:.4f}_class_score_{val_class_score:.4f}'
-            if self.include_unknown:
-                model_name += f'_unknown_score_{val_unknown_score:.4f}'
-            if val_acc > self.max_val_acc:
-                self.max_val_acc = val_acc
-                model_name = f'{self.checkpoint_path}/best_{model_name}.h5'
-                print(f'[best model saved]\n')
-            else:
-                model_name = f'{self.checkpoint_path}/{model_name}.h5'
-            self.model.save(model_name, include_optimizer=False)
+            model_name = f'{self.checkpoint_path}/{model_name}.h5'
+        self.model.save(model_name, include_optimizer=False)
 
     def evaluate(self, dataset, unknown_threshold=0.5):
         assert dataset in ['train', 'validation']
@@ -322,7 +320,8 @@ def graph_forward(model, x):
         hit_unknown_count = total_unknown_count = 0
         hit_scores = np.zeros(shape=(num_classes,), dtype=np.float32)
         unknown_score_sum = 0.0
-        for batch_x, batch_y in tqdm(data_generator.flow()):
+        for _ in tqdm(range(len(data_generator))):
+            batch_x, batch_y = data_generator.load()
             y = graph_forward(self.model, batch_x)[0]
             max_score_index = np.argmax(y)
             max_score = y[max_score_index]