Add scripts and config files for running on real data.

sfwang · sfwang · commit 03bf7a8c6ca8 · 2019-07-25T16:59:15.000-04:00
diff --git a/experiments/real_data/128x128_d256x3_adam_lr1e-3.yaml b/experiments/real_data/128x128_d256x3_adam_lr1e-3.yaml
@@ -0,0 +1,77 @@
+GPUS: '0'
+DATA_DIR: ''
+OUTPUT_DIR: 'output'
+LOG_DIR: 'log'
+WORKERS: 0
+PRINT_FREQ: 1
+CUDNN:
+  BENCHMARK: False
+  DETERMINISTIC: True
+  ENABLED: True
+DATASET:
+  DATASET: 'real_dataset'
+  ROOT: None
+  TEST_SET: ['srep31332-s1.mat']
+  TRAIN_SET: None
+  X_MIN: 20
+  Y_MIN: 190
+  X_MAX: 148
+  Y_MAX: 318
+  # FLIP: true
+  # ROT_FACTOR: 30
+  # SCALE_FACTOR: 0.25
+MODEL:
+  NAME: neuron_resnet
+  PRETRAINED: 'models/pytorch/imagenet/resnet18-5c106cde.pth'
+  IMAGE_SIZE:
+  - 128 
+  - 128 
+  EXTRA:
+    TARGET_TYPE: gaussian
+    SIGMA: 2
+    HEATMAP_SIZE:
+    - 64
+    - 64
+    FINAL_CONV_KERNEL: 1
+    DECONV_WITH_BIAS: false
+    NUM_DECONV_LAYERS: 4
+    NUM_DECONV_FILTERS:
+    - 256
+    - 256
+    - 256
+    - 256
+    NUM_DECONV_KERNELS:
+    - 4
+    - 4
+    - 4
+    - 4
+    NUM_LAYERS: 18
+LOSS:
+  USE_TARGET_WEIGHT: False
+TRAIN:
+  BATCH_SIZE: 32
+  SHUFFLE: true
+  BEGIN_EPOCH: 0
+  END_EPOCH: 140
+  RESUME: false
+  OPTIMIZER: adam
+  LR: 0.001
+  LR_FACTOR: 0.1
+  LR_STEP:
+  - 90
+  - 120
+  WD: 0.0001
+  GAMMA1: 0.99
+  GAMMA2: 0.0
+  MOMENTUM: 0.9
+  NESTEROV: false
+TEST:
+  BATCH_SIZE: 1
+  FLIP_TEST: false
+  MODEL_FILE: 'output/simulation/neuron_resnet_18/128x128_d256x3_adam_lr1e-3/model_best.pth.tar'
+DEBUG:
+  DEBUG: false
+  SAVE_BATCH_IMAGES_GT: true
+  SAVE_BATCH_IMAGES_PRED: true
+  SAVE_HEATMAPS_GT: true
+  SAVE_HEATMAPS_PRED: true
diff --git a/lib/core/config.py b/lib/core/config.py
@@ -76,6 +76,11 @@
 config.DATASET.TEST_SET = 'valid'
 config.DATASET.DATA_FORMAT = 'jpg'
 
+config.DATASET.X_MIN = 250
+config.DATASET.Y_MIN = 50
+config.DATASET.X_MAX = 378
+config.DATASET.Y_MAX = 178
+
 # training data augmentation
 config.DATASET.FLIP = True
 config.DATASET.SCALE_FACTOR = 0.25
diff --git a/lib/core/evaluate.py b/lib/core/evaluate.py
@@ -16,7 +16,7 @@
 
 
 def calc_dists(preds, target):
-    preds = preds.astype(np.float32)
+    preds = preds.astype(np.float32)[:, :2]
     target = target.astype(np.float32)
 
     dists = np.sqrt(((preds.reshape(preds.shape[0], 1, preds.shape[1]) - \
@@ -32,8 +32,8 @@ def calc_tp_fp_fn(preds, target, hit_thr=2):
 
     row_ind, col_ind = linear_sum_assignment(dists)
 
-    tp = np.sum(dists[row_ind, col_ind] <= hit_thr)
-    fp = dists.shape[0] - tp
-    fn = dists.shape[1] - tp
+    tp = dists[row_ind, col_ind] <= hit_thr
+    fp = np.logical_not(tp)
+    fn = dists.shape[1] - np.sum(tp)
 
     return tp, fp, fn
diff --git a/lib/core/function.py b/lib/core/function.py
@@ -20,7 +20,7 @@
 from lib.core.evaluate import calc_tp_fp_fn
 from lib.core.inference import get_final_preds
 # from utils.transforms import flip_back
-# from utils.vis import save_debug_images
+from lib.utils.vis import vis_preds
 # from utils.vis_plain_keypoint import vis_mpii_keypoints
 # from utils.integral import softmax_integral_tensor
 
@@ -104,7 +104,7 @@ def train(config, train_loader, model, criterion, optimizer, epoch,
             #                   prefix)
 
 
-def validate(config, val_loader, val_dataset, model, criterion, output_dir,
+def validate(config, val_loader, model, criterion, output_dir,
              tb_log_dir, writer_dict=None):
     batch_time = AverageMeter()
     losses = AverageMeter()
@@ -167,8 +167,16 @@ def validate(config, val_loader, val_dataset, model, criterion, output_dir,
                     heatmap_pred[heatmap_pred < 0.0] = 0
                     heatmap_pred[heatmap_pred > 1.0] = 1.0
 
-                    writer.add_image('input_recording', input_image, global_steps,
-                        dataformats='CHW')
+                    input_image = (input_image * 255).astype(np.uint8)
+                    input_image = np.transpose(input_image, (1, 2, 0))
+                    pred = preds[idx]
+                    gt = sources[idx][:valid_source_nums[idx], :]
+
+                    tp, _, _ = calc_tp_fp_fn(pred, gt)
+                    final_preds = vis_preds(input_image, pred, tp)
+
+                    writer.add_image('final_preds', final_preds, global_steps,
+                        dataformats='HWC')
                     writer.add_image('heatmap_target', heatmap_target, global_steps,
                         dataformats='CHW')
                     writer.add_image('heatmap_pred', heatmap_pred, global_steps,
@@ -182,8 +190,8 @@ def validate(config, val_loader, val_dataset, model, criterion, output_dir,
         total_tp = total_fp = total_fn = 0
         for preds, target in zip(all_preds, all_gts):
             tp, fp, fn = calc_tp_fp_fn(preds, target)
-            total_tp += tp
-            total_fp += fp
+            total_tp += np.sum(tp)
+            total_fp += np.sum(fp)
             total_fn += fn
 
         recall = total_tp / (total_tp + total_fn)
@@ -204,6 +212,68 @@ def validate(config, val_loader, val_dataset, model, criterion, output_dir,
     return perf_indicator
 
 
+def test(config, test_loader, model, output_dir, tb_log_dir,
+         writer_dict=None):
+    batch_time = AverageMeter()
+
+    # switch to evaluate mode
+    model.eval()
+
+    all_preds = []
+
+    with torch.no_grad():
+        end = time.time()
+        for i, input in enumerate(test_loader):
+            # compute output
+            output = model(input)
+
+            num_images = input.size(0)
+
+            preds = get_final_preds(output.detach().cpu().numpy())
+            all_preds.extend(preds)
+
+            # measure elapsed time
+            batch_time.update(time.time() - end)
+            end = time.time()
+
+            if i % config.PRINT_FREQ == 0:
+                msg = 'Test: [{0}/{1}]\t' \
+                      'Time {batch_time.val:.3f} ({batch_time.avg:.3f})'.format(
+                          i, len(test_loader), batch_time=batch_time)
+                logger.info(msg)
+
+            if writer_dict:
+                writer = writer_dict['writer']
+                global_steps = writer_dict['vis_global_steps']
+
+                idx = np.random.randint(0, num_images)
+
+                input_image = input.detach().cpu().numpy()[idx]
+                min_val = input_image.min()
+                max_val = input_image.max()
+                input_image = (input_image - min_val) / (max_val - min_val)
+                heatmap_pred = output.detach().cpu().numpy()[idx]
+                heatmap_pred[heatmap_pred < 0.0] = 0
+                heatmap_pred[heatmap_pred > 1.0] = 1.0
+
+                input_image = (input_image * 255).astype(np.uint8)
+                input_image = np.transpose(input_image, (1, 2, 0))
+                pred = preds[idx]
+                tp = np.ones(pred.shape[0], dtype=bool)
+                final_preds = vis_preds(input_image, pred, tp)
+
+                writer.add_image('final_preds', final_preds, global_steps,
+                    dataformats='HWC')
+                writer.add_image('input_recording', input_image, global_steps,
+                    dataformats='HWC')
+                writer.add_image('heatmap_pred', heatmap_pred, global_steps,
+                    dataformats='CHW')
+
+                writer_dict['vis_global_steps'] = global_steps + 1
+
+                # prefix = '{}_{}'.format(os.path.join(output_dir, 'val'), i)
+
+
 # markdown format output
 def _print_name_value(name_value, full_arch_name):
     names = name_value.keys()
diff --git a/lib/core/inference.py b/lib/core/inference.py
@@ -59,7 +59,9 @@ def get_final_preds(batch_heatmaps, score_thresh=0.5):
 
         rows, cols = np.where(np.logical_and(local_max, heatmaps > score_thresh))
 
-        preds = np.vstack((cols, rows)).transpose()
+        scores = heatmaps[rows, cols]
+
+        preds = np.vstack((cols, rows, scores)).transpose()
 
         batch_preds.append(preds)
 
diff --git a/lib/dataset/RealDataset.py b/lib/dataset/RealDataset.py
@@ -0,0 +1,78 @@
+import logging
+import os
+import json
+import numpy as np
+# from scipy.io import loadmat, savemat
+import copy
+
+import h5py
+
+from torch.utils.data import Dataset
+
+# from lib.utils.utils import get_source, get_mesh, evalpotential, generate_heatmaps
+
+logger = logging.getLogger(__name__)
+
+class RealDataset(Dataset):
+    def __init__(self, cfg, root, image_set, is_train, transform):
+        self.cfg = cfg
+        self.is_train = is_train
+
+        self.root = root
+        self.image_set = image_set
+
+        if is_train:
+            raise ValueError('Real dataset has no labels and thus cannot be trained on')
+        # self.is_train = is_train
+
+        self.transform = transform
+
+        self.patch_width = cfg.MODEL.IMAGE_SIZE[0]
+        self.patch_height = cfg.MODEL.IMAGE_SIZE[1]
+
+        self.x_min = cfg.DATASET.X_MIN
+        self.y_min = cfg.DATASET.Y_MIN
+        self.x_max = cfg.DATASET.X_MAX
+        self.y_max = cfg.DATASET.Y_MAX
+
+        self.db = self._get_db()
+        self.db_length = len(self.db)
+
+        logger.info('=> load {} samples'.format(len(self.db)))
+
+    def __getitem__(self, idx):
+        the_db = copy.deepcopy(self.db[idx])
+
+        image = self.transform(the_db['image'])
+
+        return image
+
+    def _get_db(self):
+        gt_db = []
+
+        all_frames = []
+        for i, file_name in enumerate(self.image_set):
+            dataset = {}
+            file_path = os.path.join('data', self.cfg.DATASET.DATASET, file_name)
+            with h5py.File(file_path, 'r') as f:
+                for k, v in f.items():
+                    dataset[k] = np.array(v)
+
+            # assert ('D1' in dataset.keys())
+            assert ('D2' in dataset.keys())
+
+            all_frames = np.expand_dims(dataset['D2'].copy().astype(np.float32), axis=-1).repeat(3, axis=-1)
+            # all_frames_vis = np.transpose(dataset['D1'], (0, 2, 3, 1)).astype(np.uint8)
+
+            all_frames = all_frames[:, self.y_min:self.y_max, self.x_min:self.x_max, :]
+            # all_frames_vis = all_frames_vis[269:270, self.y_min:self.y_max, self.x_min:self.x_max, :]
+
+            for frame in all_frames:
+                gt_db.append({
+                    'image': frame
+                })
+
+        return gt_db
+
+    def __len__(self):
+        return self.db_length
diff --git a/lib/dataset/__init__.py b/lib/dataset/__init__.py
@@ -9,3 +9,4 @@
 from __future__ import print_function
 
 from lib.dataset.SimulatedDataset import SimulatedDataset as simulation
+from lib.dataset.RealDataset import RealDataset as real_dataset
diff --git a/lib/utils/vis.py b/lib/utils/vis.py
@@ -0,0 +1,59 @@
+# Copyright (c) 2017-present, Facebook, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+##############################################################################
+
+"""Detection output visualization module."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+from __future__ import unicode_literals
+
+import cv2
+import numpy as np
+import os
+
+# from lib.utils.colormap import colormap
+# import lib.utils.env as envu
+# import lib.utils.keypoints as keypoint_utils
+#
+# # Matplotlib requires certain adjustments in some environments
+# # Must happen before importing matplotlib
+# envu.set_up_matplotlib()
+# import matplotlib.pyplot as plt
+# from matplotlib.patches import Polygon
+#
+# plt.rcParams['pdf.fonttype'] = 42  # For editing in Adobe Illustrator
+#
+#
+# _GRAY = (218, 227, 218)
+# _GREEN = (18, 127, 15)
+# _WHITE = (255, 255, 255)
+
+def vis_preds(img, preds, tp):
+    # Draw the detections.
+    img = img.copy()
+    for idx in range(preds.shape[0]):
+        pt = preds[idx, 0].astype(np.int32) * 2, preds[idx, 1].astype(np.int32) * 2
+
+        if tp[idx]:
+            cv2.circle(
+                img, pt,
+                radius=1, color=(255, 0, 0), thickness=-1, lineType=cv2.LINE_AA)
+        else:
+            cv2.circle(
+                img, pt,
+                radius=1, color=(0, 0, 255), thickness=-1, lineType=cv2.LINE_AA)
+
+    return img
diff --git a/source_detection/test.py b/source_detection/test.py
diff --git a/source_detection/train.py b/source_detection/train.py
diff --git a/source_detection/validate.py b/source_detection/validate.py

Original file line number	Diff line number	Diff line change
`@@ -9,3 +9,4 @@`
`9`	`9`	`from __future__ import print_function`
`10`	`10`
`11`	`11`	`from lib.dataset.SimulatedDataset import SimulatedDataset as simulation`
	`12`	`+from lib.dataset.RealDataset import RealDataset as real_dataset`