cherry pick from human pose estimation

Author: longcw

.gitmodules

@@ -0,0 +1,3 @@
+[submodule "lib/third_part/poseval"]
+	path = lib/third_part/poseval
+	url = [email protected]:longcw/poseval.git

experiments/coco/hrnet/w32_256x192_adam_lr1e-3.yaml

@@ -4,7 +4,7 @@ CUDNN:
   DETERMINISTIC: false
   ENABLED: true
 DATA_DIR: ''
-GPUS: (0,1,2,3)
+GPUS: (0,)
 OUTPUT_DIR: 'output'
 LOG_DIR: 'log'
 WORKERS: 24

experiments/mpii/hrnet/w32_256x256_adam_lr1e-3.yaml

@@ -4,7 +4,7 @@ CUDNN:
   DETERMINISTIC: false
   ENABLED: true
 DATA_DIR: ''
-GPUS: (0,1,2,3)
+GPUS: (0,)
 OUTPUT_DIR: 'output'
 LOG_DIR: 'log'
 WORKERS: 24

lib/core/evaluate.py

@@ -69,5 +69,3 @@ def accuracy(output, target, hm_type='gaussian', thr=0.5):
     if cnt != 0:
         acc[0] = avg_acc
     return acc, avg_acc, cnt, pred
-
-

lib/core/function.py

@@ -7,16 +7,18 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
- 
-import time
+
 import logging
+import time
 import os
+from collections import defaultdict
+import ujson as json
 
 import numpy as np
 import torch
 
 from core.evaluate import accuracy
-from core.inference import get_final_preds
+from core.inference import get_final_preds, get_max_preds
 from utils.transforms import flip_back
 from utils.vis import save_debug_images
 
@@ -110,8 +112,7 @@ def validate(config, val_loader, val_dataset, model, criterion, output_dir,
     )
     all_boxes = np.zeros((num_samples, 6))
     image_path = []
-    filenames = []
-    imgnums = []
+    image_ids = []
     idx = 0
     with torch.no_grad():
         end = time.time()
@@ -179,6 +180,8 @@ def validate(config, val_loader, val_dataset, model, criterion, output_dir,
             all_boxes[idx:idx + num_images, 4] = np.prod(s*200, 1)
             all_boxes[idx:idx + num_images, 5] = score
             image_path.extend(meta['image'])
+            if config.DATASET.DATASET == 'posetrack':
+                image_ids.extend(meta['image_id'].numpy())
 
             idx += num_images
 
@@ -198,8 +201,7 @@ def validate(config, val_loader, val_dataset, model, criterion, output_dir,
                                   prefix)
 
         name_values, perf_indicator = val_dataset.evaluate(
-            config, all_preds, output_dir, all_boxes, image_path,
-            filenames, imgnums
+            config, all_preds, output_dir, all_boxes, image_path, image_ids=image_ids
         )
 
         model_name = config.MODEL.NAME
@@ -240,6 +242,107 @@ def validate(config, val_loader, val_dataset, model, criterion, output_dir,
     return perf_indicator
 
 
+def inference(config, image_loader, image_dataset, model, output_dir):
+    batch_time = AverageMeter()
+    losses = AverageMeter()
+    acc = AverageMeter()
+
+    # switch to evaluate mode
+    model.eval()
+
+    num_samples = len(image_dataset)
+    all_preds = np.zeros((num_samples, config.MODEL.NUM_JOINTS, 3),
+                         dtype=np.float32)
+    all_boxes = np.zeros((num_samples, 5))
+    all_image_pathes = []
+    all_image_ids = []
+    idx = 0
+    with torch.no_grad():
+        end = time.time()
+        for i, (input, target, target_weight, meta) in enumerate(image_loader):
+            num_images = input.size(0)
+            # compute output
+            outputs = model(input)
+            if isinstance(outputs, list):
+                output = outputs[-1]
+            else:
+                output = outputs
+
+            if config.TEST.FLIP_TEST:
+                # this part is ugly, because pytorch has not supported negative index
+                # input_flipped = model(input[:, :, :, ::-1])
+                input_flipped = np.flip(input.cpu().numpy(), 3).copy()
+                input_flipped = torch.from_numpy(input_flipped).cuda()
+                outputs_flipped = model(input_flipped)
+                if isinstance(outputs_flipped, list):
+                    output_flipped = outputs_flipped[-1]
+                else:
+                    output_flipped = outputs_flipped
+
+                output_flipped = flip_back(output_flipped.cpu().numpy(),
+                                           image_dataset.flip_pairs)
+                output_flipped = torch.from_numpy(output_flipped.copy()).cuda()
+
+                # feature is not aligned, shift flipped heatmap for higher accuracy
+                if config.TEST.SHIFT_HEATMAP:
+                    output_flipped[:, :, :, 1:] = \
+                        output_flipped.clone()[:, :, :, 0:-1]
+                    # output_flipped[:, :, :, 0] = 0
+
+                output = (output + output_flipped) * 0.5
+
+            # measure elapsed time
+            batch_time.update(time.time() - end)
+            end = time.time()
+
+            c = meta['center'].numpy()
+            s = meta['scale'].numpy()
+            score = meta['score'].numpy()
+            tlwhs = meta['bbox_tlwh'].numpy()
+            output = output.data.cpu()
+
+            preds, maxvals = get_final_preds(config, output.numpy(), c, s)
+
+            all_preds[idx:idx + num_images, :, 0:2] = preds[:, :, 0:2]
+            all_preds[idx:idx + num_images, :, 2:3] = maxvals
+            # double check this all_boxes parts
+            all_boxes[idx:idx + num_images, 0:4] = tlwhs
+            all_boxes[idx:idx + num_images, 4] = score
+            all_image_pathes.extend(meta['image'])
+            if config.DATASET.DATASET == 'mot':
+                seq_names, frame_ids = meta['image_id']
+                frame_ids = frame_ids.numpy().astype(int)
+                all_image_ids.extend(list(zip(seq_names, frame_ids)))
+
+            idx += num_images
+
+            if i % config.PRINT_FREQ == 0:
+                msg = 'Test: [{0}/{1}]\t' \
+                      'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'.format(
+                    i, len(image_loader), batch_time=batch_time)
+                logger.info(msg)
+
+                prefix = '{}_{}'.format(os.path.join(output_dir, 'inference'), i)
+                pred, _ = get_max_preds(output.numpy())
+                save_debug_images(config, input, meta, target, pred * 4, output, prefix)
+
+    # write output
+    frame_results = defaultdict(list)
+    for image_id, pred, box in zip(all_image_ids, all_preds, all_boxes):
+        frame_results[image_id].append((pred.astype(float).tolist(), box.astype(float).tolist()))
+
+    final_results = {}
+    for image_id, results in frame_results.items():
+        keypoints, boxes = zip(*results)
+        final_results[image_id] = {'keypoints': keypoints, 'boxes': boxes}
+
+    if not os.path.isdir(output_dir):
+        os.makedirs(output_dir)
+    with open(os.path.join(output_dir, 'box_keypoints.json'), 'w') as f:
+        json.dump(final_results, f)
+    logger.info('Save results to {}'.format(os.path.join(output_dir, 'box_keypoints.json')))
+
+
 # markdown format output
 def _print_name_value(name_value, full_arch_name):
     names = name_value.keys()

lib/dataset/JointsDataset.py

@@ -46,8 +46,8 @@ def __init__(self, cfg, root, image_set, is_train, transform=None):
         self.prob_half_body = cfg.DATASET.PROB_HALF_BODY
         self.color_rgb = cfg.DATASET.COLOR_RGB
 
-        self.target_type = cfg.MODEL.TARGET_TYPE
         self.image_size = np.array(cfg.MODEL.IMAGE_SIZE)
+        self.target_type = cfg.MODEL.TARGET_TYPE
         self.heatmap_size = np.array(cfg.MODEL.HEATMAP_SIZE)
         self.sigma = cfg.MODEL.SIGMA
         self.use_different_joints_weight = cfg.LOSS.USE_DIFFERENT_JOINTS_WEIGHT
@@ -114,8 +114,8 @@ def __getitem__(self, idx):
         db_rec = copy.deepcopy(self.db[idx])
 
         image_file = db_rec['image']
-        filename = db_rec['filename'] if 'filename' in db_rec else ''
-        imgnum = db_rec['imgnum'] if 'imgnum' in db_rec else ''
+        image_id = db_rec.get('image_id', -1)
+        bbox_tlwh = db_rec.get('bbox_tlwh', (0, 0, 0, 0))
 
         if self.data_format == 'zip':
             from utils import zipreader
@@ -185,8 +185,8 @@ def __getitem__(self, idx):
 
         meta = {
             'image': image_file,
-            'filename': filename,
-            'imgnum': imgnum,
+            'image_id': image_id,
+            'bbox_tlwh': bbox_tlwh,
             'joints': joints,
             'joints_vis': joints_vis,
             'center': c,

lib/dataset/__init__.py

@@ -10,3 +10,5 @@
 
 from .mpii import MPIIDataset as mpii
 from .coco import COCODataset as coco
+from .posetrack import PoseTrackDataset as posetrack
+from .mot import MOTDataset as mot

lib/dataset/coco.py

@@ -8,15 +8,17 @@
 from __future__ import division
 from __future__ import print_function
 
-from collections import defaultdict
-from collections import OrderedDict
 import logging
 import os
+import pickle
+from collections import defaultdict
+from collections import OrderedDict
 
+# import json_tricks as json
+import ujson as json
+import numpy as np
 from pycocotools.coco import COCO
 from pycocotools.cocoeval import COCOeval
-import json_tricks as json
-import numpy as np
 
 from dataset.JointsDataset import JointsDataset
 from nms.nms import oks_nms
@@ -51,8 +53,15 @@ class COCODataset(JointsDataset):
         [16,14],[14,12],[17,15],[15,13],[12,13],[6,12],[7,13], [6,7],[6,8],
         [7,9],[8,10],[9,11],[2,3],[1,2],[1,3],[2,4],[3,5],[4,6],[5,7]]
     '''
+
+    num_joints = 17
+    flip_pairs = [[1, 2], [3, 4], [5, 6], [7, 8],
+                  [9, 10], [11, 12], [13, 14], [15, 16]]
+    upper_body_ids = (0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10)
+    lower_body_ids = (11, 12, 13, 14, 15, 16)
+
     def __init__(self, cfg, root, image_set, is_train, transform=None):
-        super().__init__(cfg, root, image_set, is_train, transform)
+        super(COCODataset, self).__init__(cfg, root, image_set, is_train, transform)
         self.nms_thre = cfg.TEST.NMS_THRE
         self.image_thre = cfg.TEST.IMAGE_THRE
         self.soft_nms = cfg.TEST.SOFT_NMS
@@ -87,12 +96,11 @@ def __init__(self, cfg, root, image_set, is_train, transform=None):
         self.num_images = len(self.image_set_index)
         logger.info('=> num_images: {}'.format(self.num_images))
 
-        self.num_joints = 17
-        self.flip_pairs = [[1, 2], [3, 4], [5, 6], [7, 8],
-                           [9, 10], [11, 12], [13, 14], [15, 16]]
+        self.num_joints = COCODataset.num_joints
+        self.flip_pairs = COCODataset.flip_pairs
+        self.upper_body_ids = COCODataset.upper_body_ids
+        self.lower_body_ids = COCODataset.lower_body_ids
         self.parent_ids = None
-        self.upper_body_ids = (0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10)
-        self.lower_body_ids = (11, 12, 13, 14, 15, 16)
 
         self.joints_weight = np.array(
             [
@@ -201,8 +209,6 @@ def _load_coco_keypoint_annotation_kernal(self, index):
                 'scale': scale,
                 'joints_3d': joints_3d,
                 'joints_3d_vis': joints_3d_vis,
-                'filename': '',
-                'imgnum': 0,
             })
 
         return rec
@@ -378,7 +384,7 @@ def _write_coco_keypoint_results(self, keypoints, res_file):
         ]
 
         results = self._coco_keypoint_results_one_category_kernel(data_pack[0])
-        logger.info('=> writing results json to %s' % res_file)
+        logger.info('=> Writing results json to %s' % res_file)
         with open(res_file, 'w') as f:
             json.dump(results, f, sort_keys=True, indent=4)
         try: