Added support for raw dataset inference

Author: jackal02

experiments/mpii/hrnet/w32_256x256_adam_lr1e-3.yaml

@@ -4,10 +4,10 @@ CUDNN:
   DETERMINISTIC: false
   ENABLED: true
 DATA_DIR: ''
-GPUS: (0,1,2,3)
+GPUS: (0,)
 OUTPUT_DIR: 'output'
 LOG_DIR: 'log'
-WORKERS: 24
+WORKERS: 6
 PRINT_FREQ: 100
 
 DATASET:
@@ -107,7 +107,7 @@ TRAIN:
   MOMENTUM: 0.9
   NESTEROV: false
 TEST:
-  BATCH_SIZE_PER_GPU: 32
+  BATCH_SIZE_PER_GPU: 128
   MODEL_FILE: ''
   FLIP_TEST: true
   POST_PROCESS: true

lib/config/default.py

@@ -20,7 +20,7 @@
 _C.LOG_DIR = ''
 _C.DATA_DIR = ''
 _C.GPUS = (0,)
-_C.WORKERS = 4
+_C.WORKERS = 1
 _C.PRINT_FREQ = 20
 _C.AUTO_RESUME = False
 _C.PIN_MEMORY = True
@@ -89,14 +89,14 @@
 _C.TRAIN.RESUME = False
 _C.TRAIN.CHECKPOINT = ''
 
-_C.TRAIN.BATCH_SIZE_PER_GPU = 32
+_C.TRAIN.BATCH_SIZE_PER_GPU = 4
 _C.TRAIN.SHUFFLE = True
 
 # testing
 _C.TEST = CN()
 
 # size of images for each device
-_C.TEST.BATCH_SIZE_PER_GPU = 32
+_C.TEST.BATCH_SIZE_PER_GPU = 8
 # Test Model Epoch
 _C.TEST.FLIP_TEST = False
 _C.TEST.POST_PROCESS = False

lib/core/function.py

@@ -14,6 +14,7 @@
 
 import numpy as np
 import torch
+import cv2
 
 from core.evaluate import accuracy
 from core.inference import get_final_preds
@@ -94,8 +95,23 @@ def train(config, train_loader, model, criterion, optimizer, epoch,
                               prefix)
 
 
+def compute_joints(batch_image, batch_joints, batch_joints_vis):
+    for k in range(batch_image.size(0)):
+        image_tensor = batch_image[k]
+        image = image_tensor.mul(255).clamp(0, 255).byte().permute(1, 2, 0).cpu().numpy()
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+
+        joints = batch_joints[k]
+        joints_vis = batch_joints_vis[k]
+        for joint in joints:
+            cv2.circle(image, (int(joint[0]), int(joint[1])), 2, [0, 0, 255], 2)
+
+        cv2.imshow("im", image)
+        cv2.waitKey()
+
+
 def validate(config, val_loader, val_dataset, model, criterion, output_dir,
-             tb_log_dir, writer_dict=None):
+             tb_log_dir, writer_dict=None, predict_only=False):
     batch_time = AverageMeter()
     losses = AverageMeter()
     acc = AverageMeter()
@@ -116,6 +132,8 @@ def validate(config, val_loader, val_dataset, model, criterion, output_dir,
     with torch.no_grad():
         end = time.time()
         for i, (input, target, target_weight, meta) in enumerate(val_loader):
+            img = input.data[0].mul(255).clamp(0, 255).byte().permute(1, 2, 0).cpu().numpy()
+
             # compute output
             outputs = model(input)
             if isinstance(outputs, list):
@@ -147,12 +165,12 @@ def validate(config, val_loader, val_dataset, model, criterion, output_dir,
 
                 output = (output + output_flipped) * 0.5
 
+            num_images = input.size(0)
+
             target = target.cuda(non_blocking=True)
             target_weight = target_weight.cuda(non_blocking=True)
 
             loss = criterion(output, target, target_weight)
-
-            num_images = input.size(0)
             # measure accuracy and record loss
             losses.update(loss.item(), num_images)
             _, avg_acc, cnt, pred = accuracy(output.cpu().numpy(),
@@ -181,6 +199,7 @@ def validate(config, val_loader, val_dataset, model, criterion, output_dir,
             image_path.extend(meta['image'])
 
             idx += num_images
+            compute_joints(input, pred*4, meta['joints_vis'])
 
             if i % config.PRINT_FREQ == 0:
                 msg = 'Test: [{0}/{1}]\t' \
@@ -196,18 +215,20 @@ def validate(config, val_loader, val_dataset, model, criterion, output_dir,
                 )
                 save_debug_images(config, input, meta, target, pred*4, output,
                                   prefix)
+        name_values = None
+        perf_indicator = None
+        if not predict_only:
+            name_values, perf_indicator = val_dataset.evaluate(
+                config, all_preds, output_dir, all_boxes, image_path,
+                filenames, imgnums
+            )
 
-        name_values, perf_indicator = val_dataset.evaluate(
-            config, all_preds, output_dir, all_boxes, image_path,
-            filenames, imgnums
-        )
-
-        model_name = config.MODEL.NAME
-        if isinstance(name_values, list):
-            for name_value in name_values:
-                _print_name_value(name_value, model_name)
-        else:
-            _print_name_value(name_values, model_name)
+            model_name = config.MODEL.NAME
+            if isinstance(name_values, list):
+                for name_value in name_values:
+                    _print_name_value(name_value, model_name)
+            else:
+                _print_name_value(name_values, model_name)
 
         if writer_dict:
             writer = writer_dict['writer']
@@ -222,19 +243,22 @@ def validate(config, val_loader, val_dataset, model, criterion, output_dir,
                 acc.avg,
                 global_steps
             )
-            if isinstance(name_values, list):
-                for name_value in name_values:
+
+            if not predict_only:
+                if isinstance(name_values, list):
+                    for name_value in name_values:
+                        writer.add_scalars(
+                            'valid',
+                            dict(name_value),
+                            global_steps
+                        )
+                else:
                     writer.add_scalars(
                         'valid',
-                        dict(name_value),
+                        dict(name_values),
                         global_steps
                     )
-            else:
-                writer.add_scalars(
-                    'valid',
-                    dict(name_values),
-                    global_steps
-                )
+
             writer_dict['valid_global_steps'] = global_steps + 1
 
     return perf_indicator

lib/dataset/JointsDataset.py

@@ -174,6 +174,10 @@ def __getitem__(self, idx):
         if self.transform:
             input = self.transform(input)
 
+        input = cv2.resize(data_numpy, (256, 256), interpolation=cv2.INTER_LINEAR)
+        input = cv2.normalize(input, None, alpha=0, beta=1, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_32F)
+        input = torch.from_numpy(input).permute(2, 0, 1).float()
+
         for i in range(self.num_joints):
             if joints_vis[i, 0] > 0.0:
                 joints[i, 0:2] = affine_transform(joints[i, 0:2], trans)

lib/dataset/mpii.py

@@ -23,7 +23,7 @@
 
 
 class MPIIDataset(JointsDataset):
-    def __init__(self, cfg, root, image_set, is_train, transform=None):
+    def __init__(self, cfg, root, image_set, is_train, transform=None, unannot_imgs_path=None):
         super().__init__(cfg, root, image_set, is_train, transform)
 
         self.num_joints = 16
@@ -33,10 +33,13 @@ def __init__(self, cfg, root, image_set, is_train, transform=None):
         self.upper_body_ids = (7, 8, 9, 10, 11, 12, 13, 14, 15)
         self.lower_body_ids = (0, 1, 2, 3, 4, 5, 6)
 
-        self.db = self._get_db()
+        if unannot_imgs_path is None:
+            self.db = self._get_db()
 
-        if is_train and cfg.DATASET.SELECT_DATA:
-            self.db = self.select_data(self.db)
+            if is_train and cfg.DATASET.SELECT_DATA:
+                self.db = self.select_data(self.db)
+        else:
+            self.db = self._get_db_raw(unannot_imgs_path)
 
         logger.info('=> load {} samples'.format(len(self.db)))
 
@@ -93,6 +96,24 @@ def _get_db(self):
 
         return gt_db
 
+    def _get_db_raw(self, unannot_imgs_path):
+        gt_db = []
+
+        for image_name in os.listdir(unannot_imgs_path):
+            gt_db.append(
+                {
+                    'image': os.path.join(unannot_imgs_path, image_name),
+                    'center': np.array([0, 0]),
+                    'scale': np.array([1., 1.]),
+                    'joints_3d': np.ones((16, 3)),
+                    'joints_3d_vis': np.ones((16, 3)),
+                    'filename': '',
+                    'imgnum': 0,
+                }
+            )
+
+        return gt_db
+
     def evaluate(self, cfg, preds, output_dir, *args, **kwargs):
         # convert 0-based index to 1-based index
         preds = preds[:, :, 0:2] + 1.0

lib/utils/vis.py

@@ -43,8 +43,8 @@ def save_batch_image_with_joints(batch_image, batch_joints, batch_joints_vis,
             joints_vis = batch_joints_vis[k]
 
             for joint, joint_vis in zip(joints, joints_vis):
-                joint[0] = x * width + padding + joint[0]
-                joint[1] = y * height + padding + joint[1]
+                # joint[0] = x * width + padding + joint[0]
+                # joint[1] = y * height + padding + joint[1]
                 if joint_vis[0]:
                     cv2.circle(ndarr, (int(joint[0]), int(joint[1])), 2, [255, 0, 0], 2)
             k = k + 1

requirements.txt

@@ -1,5 +1,4 @@
 EasyDict==1.7
-opencv-python==3.4.1.15
 shapely==1.6.4
 Cython
 scipy

tools/test.py

@@ -21,7 +21,7 @@
 import torch.utils.data.distributed
 import torchvision.transforms as transforms
 
-import _init_paths
+#import _init_paths
 from config import cfg
 from config import update_config
 from core.loss import JointsMSELoss
@@ -62,6 +62,12 @@ def parse_args():
                         type=str,
                         default='')
 
+    parser.add_argument('--unnanotImgsPath',
+                        help='annotations availability',
+                        type=str,
+                        default='')
+                        # default='/home/andrettin/Repo/POSE_ESTIMATION/deep-high-resolution-net.pytorch/data/mpii/images')
+
     args = parser.parse_args()
     return args
 
@@ -95,7 +101,7 @@ def main():
         logger.info('=> loading model from {}'.format(model_state_file))
         model.load_state_dict(torch.load(model_state_file))
 
-    model = torch.nn.DataParallel(model, device_ids=cfg.GPUS).cuda()
+    model = torch.nn.DataParallel(model, device_ids=[0]).cuda()
 
     # define loss function (criterion) and optimizer
     criterion = JointsMSELoss(
@@ -111,7 +117,8 @@ def main():
         transforms.Compose([
             transforms.ToTensor(),
             normalize,
-        ])
+        ]),
+        args.unnanotImgsPath if args.unnanotImgsPath != '' else None
     )
     valid_loader = torch.utils.data.DataLoader(
         valid_dataset,
@@ -123,7 +130,8 @@ def main():
 
     # evaluate on validation set
     validate(cfg, valid_loader, valid_dataset, model, criterion,
-             final_output_dir, tb_log_dir)
+             final_output_dir, tb_log_dir,
+             predict_only=True if args.unnanotImgsPath != '' else False)
 
 
 if __name__ == '__main__':