before debugging

Author: yoojy31

src/lib/tester.py

@@ -1,8 +1,13 @@
 import os
 import abc
+import sys
+import tqdm
+import pickle
 import numpy as np
 import scipy.misc
 from lib import util
+from lib.external.dataset import factory
+from . import tester_util
 
 
 class TesterABC(abc.ABC):
@@ -12,7 +17,7 @@ def __init__(self, global_args, tester_args):
         self.result_root = tester_args['result_root']
 
     @ abc.abstractmethod
-    def test(self, framework, data_loader, result_dir):
+    def run(self, framework, data_loader, result_dir):
         pass
 
 
@@ -22,14 +27,12 @@ def __init__(self, global_args, tester_args):
         self.n_samples = tester_args['n_samples']
         self.max_boxes = tester_args['max_boxes']
         self.conf_thresh = tester_args['conf_thresh']
-        self.coord_size = (global_args['coord_h'], global_args['coord_w'])
 
-    def test(self, framework, data_loader, result_dir):
+    def run(self, framework, data_loader, result_dir):
         assert data_loader.batch_size == 1
         pre_proc = data_loader.dataset.pre_proc
         class_map = data_loader.dataset.get_number2name_map()
-        save_dir = os.path.join(self.result_root, result_dir)
-        util.make_dir(save_dir)
+        util.make_dir(result_dir)
 
         for i, data_dict in enumerate(data_loader):
             if i >= self.n_samples:
@@ -44,28 +47,109 @@ def test(self, framework, data_loader, result_dir):
             img_s = data_dict['img'][0]
             gt_boxes_s = data_dict['boxes'][0]
             gt_labels_s = data_dict['labels'][0]
-            img_size = data_dict['img_size'][0]
-            input_size = img_s.shape[:2]
 
             sort_idx = 0
-            gt_img_path = os.path.join(save_dir, '%03d_%d_%s.png' % (i, sort_idx, 'gt'))
-            gt_img_s = testutil.draw_boxes(img_s, gt_boxes_s, None, gt_labels_s,
-                                           class_map, self.conf_thresh, self.max_boxes)
+            gt_img_path = os.path.join(result_dir, '%03d_%d_%s.png' % (i, sort_idx, 'gt'))
+            gt_img_s = tester_util.draw_boxes(
+                img_s, gt_boxes_s, None, gt_labels_s,
+                class_map, self.conf_thresh, self.max_boxes)
             scipy.misc.imsave(gt_img_path, gt_img_s)
-            del gt_img_s
+            sort_idx += 1
 
+            # draw_boxes
+            pred_img_path = os.path.join(result_dir, '%03d_%d_%s.png' % (i, sort_idx, 'pred'))
+            pred_img_s = tester_util.draw_boxes(
+                img_s, pred_boxes_s, pred_confs_s, pred_labels_s,
+                class_map, self.conf_thresh, self.max_boxes)
+            scipy.misc.imsave(pred_img_path, pred_img_s)
+            sort_idx += 1
+
+
+class QuantTester(TesterABC):
+    def __init__(self, global_args, tester_args):
+        super(QuantTester, self).__init__(global_args, tester_args)
+        self.n_classes = global_args['n_classes']
+        self.imdb_name = tester_args['dataset']
+        self.iou_thresh = tester_args['iou_thresh']
+        assert self.imdb_name in ('voc_2007_test', 'coco_2017_val', 'coco_2017_test-dev')
+
+    def test(self, framework, data_loader, result_dir_name):
+        assert data_loader.batch_size == 1
+        num_samples = data_loader.dataset.__len__()
+        all_boxes = [[[] for _ in range(num_samples)] for _ in range(self.n_classes)]
+
+        # import cv2
+        s = 0
+        times = list()
+        data_loader_pbar = tqdm(data_loader)
+        for idx, data_dict in enumerate(data_loader_pbar):
+            output_dict, result_dict = framework.infer_forward(data_dict, flip=self.flip_test)
+            times.append(result_dict['time'])
+            mean_time = np.mean(times[10:]) if len(times) > 10 else np.mean(times)
+            data_loader_pbar.set_description('infer time: %.4f sec' % mean_time)
+
+            # total predict boxes shape : (batch, # pred box, 4)
+            # total predict boxes confidence shape : (batch, # pred box, 1)
+            # total predict boxes label shape : (batch, # pred box, 1)
+            img_size_s = data_dict['img_size'].float()[0]
+            input_size = data_dict['img'].shape[2:]
+
+            boxes_s = result_dict['boxes_l'][0]
+            confs_s = result_dict['confs_l'][0]
+            labels_s = result_dict['labels_l'][0]
+
+            boxes_s[:, [0, 2]] *= (img_size_s[1] / input_size[1])
+            boxes_s[:, [1, 3]] *= (img_size_s[0] / input_size[0])
+            boxes_s, confs_s, labels_s = \
+                util.sort_boxes_s(boxes_s, confs_s, labels_s)
+
+            boxes_s = util.cvt_torch2numpy(boxes_s)
+            confs_s = util.cvt_torch2numpy(confs_s)
+            labels_s = util.cvt_torch2numpy(labels_s)
+
+            if len(confs_s.shape) == 1:
+                confs_s = np.expand_dims(confs_s, axis=1)
+            for i, (cls_box, cls_conf, cls_label) in \
+                    enumerate(zip(boxes_s, confs_s, labels_s)):
+
+                cls_box_with_conf = np.concatenate((cls_box, cls_conf), axis=0)
+                cls_box_with_conf = np.expand_dims(cls_box_with_conf, axis=0)
+                all_boxes[int(cls_label)][s].append(cls_box_with_conf)
+
+            for c in range(self.n_classes):
+                all_boxes[c][s] = np.concatenate(all_boxes[c][s], axis=0) \
+                    if 0 < len(all_boxes[c][s]) else np.concatenate([[]], axis=0)
             data_dict.clear()
-            output_dict.clear()
-            result_dict.clear()
-        return save_dir
-
-
-def save_img(save_dir, img_dict, key, size, range_ratio, idx, sort_idx, tag):
-    if key in img_dict.keys():
-        img = util.cvt_torch2numpy(img_dict[key])[0] * range_ratio
-        if len(img.shape) == 3:
-            img = np.squeeze(np.transpose(img, (1, 2, 0)))
-        img = scipy.misc.imresize(img, size, interp='bilinear')
-        img_path = os.path.join(save_dir, '%03d_%d_%s.png' % (idx, sort_idx, tag))
-        scipy.misc.imsave(img_path, img)
-    return sort_idx + 1
+
+        # create result directories
+        if not os.path.exists(self.result_dir):
+            os.mkdir(self.result_dir)
+        result_dir_path = os.path.join(self.result_dir, result_dir_name)
+        if not os.path.exists(result_dir_path):
+            os.mkdir(result_dir_path)
+
+        dataset_root = data_loader.dataset.get_dataset_root()
+        if isinstance(dataset_root, list):
+            dataset_root = dataset_root[0]
+        imdb = factory.get_imdb(self.imdb_name, dataset_root)
+
+        if 'coco' in self.imdb_name:
+            sys_stdout = sys.stdout
+            result_file_path = open(os.path.join(result_dir_path, 'ap_ar.txt'), 'w')
+            sys.stdout = result_file_path
+            imdb.evaluate_detections(all_boxes, result_dir_path)
+            sys.stdout = sys_stdout
+            result_file_path.close()
+
+        else:
+            det_file_path = os.path.join(result_dir_path, 'detection_results.pkl')
+            with open(det_file_path, 'wb') as det_file:
+                pickle.dump(all_boxes, det_file, pickle.HIGHEST_PROTOCOL)
+            result_msg = imdb.evaluate_detections(all_boxes, result_dir_path)
+            result_file_path = os.path.join(result_dir_path, 'mean_ap.txt')
+            with open(result_file_path, 'w') as file:
+                file.write(result_msg)
+            os.remove(det_file_path)
+
+        all_boxes.clear()
+

src/lib/tester_util.py

@@ -0,0 +1,67 @@
+import numpy as np
+import cv2
+
+colors = ((64, 64, 64), (31, 119, 180), (174, 199, 232), (255, 127, 14),
+          (255, 187, 120), (44, 160, 44), (152, 223, 138), (214, 39, 40),
+          (255, 152, 150), (148, 103, 189), (197, 176, 213), (140, 86, 75),
+          (196, 156, 148), (227, 119, 194), (247, 182, 210), (127, 127, 127),
+          (199, 199, 199), (188, 189, 34), (219, 219, 141), (23, 190, 207),
+          (158, 218, 229), (180, 119, 31))
+
+
+def draw_boxes(img_s, boxes_s, confs_s=None, labels_s=None,
+               class_map=None, conf_thresh=0.0, max_boxes=100):
+
+    box_img_s = img_s.copy()
+    n_draw_boxes = 0
+    n_wrong_boxes = 0
+    n_thresh_boxes = 0
+    for i, box in enumerate(boxes_s):
+        try:
+            l, t = int(round(box[0])), int(round(box[1]))
+            r, b = int(round(box[2])), int(round(box[3]))
+        except IndexError:
+            print(boxes_s)
+            print(i, box)
+            print('IndexError')
+            exit()
+
+        if confs_s is not None:
+            if conf_thresh > confs_s[i]:
+                n_thresh_boxes += 1
+                continue
+        if (r - l <= 0) or (b - t <= 0):
+            n_wrong_boxes += 1
+            continue
+        if n_draw_boxes >= max_boxes:
+            continue
+
+        conf_str = '-' if confs_s is None else '%0.3f' % confs_s[i]
+        if labels_s is None:
+            lab_str, color = '-', colors[i % len(colors)]
+        else:
+            lab_i = int(labels_s[i])
+            lab_str = str(lab_i) if class_map is None else class_map[lab_i]
+            color = colors[lab_i % len(colors)]
+
+        box_img_s = cv2.rectangle(box_img_s, (l, t), (r, b), color, 2)
+        l = int(l - 1 if l > 1 else r - 60)
+        t = int(t - 8 if t > 8 else b)
+        r, b = int(l + 60), int(t + 8)
+        box_img_s = cv2.rectangle(box_img_s, (l, t), (r, b), color, cv2.FILLED)
+        box_img_s = cv2.putText(box_img_s, '%s %s' % (conf_str, lab_str), (l + 1, t + 7),
+                                cv2.FONT_HERSHEY_SIMPLEX, 0.25, (255, 255, 255),
+                                1, cv2.LINE_AA)
+        n_draw_boxes += 1
+
+    info_text = 'n_draw_b: %d, n_thr_b: %d, n_wrong_b: %d' % \
+                (n_draw_boxes, n_thresh_boxes, n_wrong_boxes)
+    if confs_s is not None:
+        info_text += ', sum_of_conf: %.3f' % (np.sum(confs_s))
+    else:
+        info_text += ', sum_of_conf: -'
+
+    box_img_s = cv2.rectangle(box_img_s, (0, 0), (350, 11), (0, 0, 0), cv2.FILLED)
+    box_img_s = cv2.putText(box_img_s, info_text, (5, 10), cv2.FONT_HERSHEY_SIMPLEX,
+                        0.25, (255, 255, 255), 1, cv2.LINE_AA)
+    return box_img_s

src/run.py

@@ -1,104 +1,79 @@
 import os
 import time
 import torch
-import traceback
+# import traceback
 from tensorboardX import SummaryWriter
-import util as run_util
 import option
+import util
 
 
 def main():
     torch.multiprocessing.set_sharing_strategy('file_system')
+    print('[RUN] parse arguments')
     args, framework, optimizer, data_loader_dict, tester_dict = option.parse_options()
 
-    print('CREATE RESULT DIRECTORIES')
-    result_dir_dict = run_util.create_result_dir(args.result_dir, ['src', 'log', 'snapshot'])
-    run_util.copy_file(args.bash_file, args.result_dir)
-    run_util.copy_dir('./src', result_dir_dict['src'])
+    print('[RUN] create result directories')
+    result_dir_dict = util.create_result_dir(args.result_dir, ['src', 'log', 'snapshot', 'test'])
+    util.copy_file(args.bash_file, args.result_dir)
+    util.copy_dir('./src', result_dir_dict['src'])
 
-    print('CREATE LOGGERS')
+    print('[RUN] create loggers')
     train_log_dir = os.path.join(result_dir_dict['log'], 'train')
     train_logger = SummaryWriter(train_log_dir)
 
-    print('')
-    print('START TRAINING')
     print('[OPTIMIZER] learning rate:', optimizer.param_groups[0]['lr'])
     n_batches = data_loader_dict['train'].__len__()
     global_step = args.train_info_dict['init_iter']
     max_grad = args.train_info_dict['max_grad']
-    skip_flag = False
 
+    print('')
     while True:
         start_time = time.time()
         for train_data_dict in data_loader_dict['train']:
             batch_time = time.time() - start_time
 
-            try:
-                if not skip_flag:
-                    if global_step in args.snapshot_iters:
-                        snapshot_dir = os.path.join(result_dir_dict['snapshot'], '%07d' % global_step)
-                        run_util.save_snapshot(framework.network, optimizer, snapshot_dir)
-
-                    # if global_step in args.test_iters:
-                    #     snapshot_dir = os.path.join(result_dir_dict['snapshot'], '%07d' % global_step)
-                    #     test_networks(args, tester_dict, framework, test_set_loader, global_step)
-
-                    if args.train_info_dict['max_iter'] <= global_step:
-                        break
-
-                    # update_learning_rate(optimizer, args.lr_decay_schd_dict, global_step)
-                    # update_loss_weights({'loss_func': loss_func}, args.lw_schd_dict, global_step)
-
-                try:
-                    start_time = time.time()
-                    _, train_loss_dict = framework.train_forward(train_data_dict)
-                    run.update_networks(network, optimizer, train_loss_dict, max_grad)
-                    train_time = time.time() - start_time
-
-                    # if valid_loss_dict is not None:
-                    if global_step % args.train_info_dict['print_intv'] == 0:
-                        iter_str = '[%d/%d] ' % (global_step, args.train_info_dict['max_iter'])
-                        info_str = 'n_batches:%d batch_time:%0.3f train_time:%0.3f' % \
-                                   (n_batches, batch_time, train_time)
-                        train_str = util.cvt_dict2str(train_loss_dict)
-
-                        print(iter_str + info_str)
-                        print('[train] ' + train_str + '\n')
+            if global_step in args.snapshot_iters:
+                snapshot_dir = os.path.join(result_dir_dict['snapshot'], '%07d' % global_step)
+                util.save_snapshot(framework.network, optimizer, snapshot_dir)
 
-                        for key, value in train_loss_dict.items():
-                            train_logger.add_scalar(key, value, global_step)
+            if global_step in args.test_iters:
+                test_dir = os.path.join(result_dir_dict['test'], '%07d' % global_step)
+                util.run_testers(tester_dict, framework, data_loader_dict['test'], test_dir)
 
-                    train_loss_dict.clear()
-                    del train_loss_dict
+            if args.train_info_dict['max_iter'] <= global_step:
+                break
 
-                    skip_flag = False
-                    global_step += 1
+            if global_step in args.lr_decay_schd_dict.keys():
+                util.update_learning_rate(optimizer, args.lr_decay_schd_dict[global_step])
 
-                except Exception as e:
-                    print('[WARNING] %s' % (str(e)))
-                    print('[%d/%d] skip this mini-batch\n' % (global_step, args.train_info_dict['max_iter']))
+            start_time = time.time()
+            _, train_loss_dict = framework.train_forward(train_data_dict)
+            util.update_network(framework.network, optimizer, train_loss_dict, max_grad)
+            train_time = time.time() - start_time
 
-                    if 'memory' in str(e):
-                        optimizer.zero_grad()
-                        optimizer.step()
+            if global_step % args.train_info_dict['print_intv'] == 0:
+                iter_str = '[TRAINING] %d/%d:' % (global_step, args.train_info_dict['max_iter'])
+                info_str = 'n_batches: %d, batch_time: %0.3f, train_time: %0.3f' % \
+                           (n_batches, batch_time, train_time)
+                train_str = util.cvt_dict2str(train_loss_dict)
+                print(iter_str + '\n' + info_str + '\n' + train_str + '\n')
 
-                        train_loss_dict.clear()
-                        del train_loss_dict
-                        torch.cuda.empty_cache()
+                for key, value in train_loss_dict.items():
+                    train_logger.add_scalar(key, value, global_step)
 
-                    skip_flag = True
+            train_loss_dict.clear()
+            train_data_dict.clear()
+            del train_loss_dict, train_data_dict
 
-                train_data_dict.clear()
-                start_time = time.time()
+            global_step += 1
+            start_time = time.time()
 
-            except Exception as e:
-                traceback.print_tb(e.__traceback__)
-                print('[ERROR] %s' % (str(e)))
-                snapshot_dir = os.path.join(result_dir_dict['snapshot'], '%07d' % global_step)
-                run_util.save_snapshot(framework.network, optimizer, snapshot_dir)
-                exit()
+        if args.train_info_dict['max_iter'] <= global_step:
+            break
     train_logger.close()
 
 
 if __name__ == '__main__':
     main()
+
+