Add initial file for running the network and visualizing the results in real time for single images or videos

Author: Turkka Helinoja

tools/run_and_visualize.py

@@ -0,0 +1,147 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import argparse
+import os
+import pprint
+
+import torch
+import torch.nn.parallel
+import torch.backends.cudnn as cudnn
+import torch.optim
+import torch.utils.data
+import torch.utils.data.distributed
+import torchvision.transforms as transforms
+
+import _init_paths
+from config import cfg
+from config import update_config
+from core.loss import JointsMSELoss
+from core.function import validate
+from utils.utils import create_logger
+
+import dataset
+import models
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Run and visualize keypoints network for image or video.")
+    # general
+    parser.add_argument("--cfg",
+                        help="Configuration file name",
+                        required=True,
+                        type=str)
+
+    parser.add_argument("opts",
+                        help="Modify config options using the command-line",
+                        default=None,
+                        nargs=argparse.REMAINDER)
+
+    parser.add_argument("--modelDir",
+                        help="model directory",
+                        type=str,
+                        default="")
+    parser.add_argument("--logDir",
+                        help="log directory",
+                        type=str,
+                        default="")
+    parser.add_argument("--dataDir",
+                        help="data directory",
+                        type=str,
+                        default="")
+    parser.add_argument("--prevModelDir",
+                        help="prev Model directory",
+                        type=str,
+                        default="")
+    parser.add_argument("--visualize",
+                        help="Visualize the results",
+                        type=bool,
+                        default=False)
+    parser.add_argument("--input",
+                        help="Input image file",
+                        type=str,
+                        default="")
+    parser.add_argument("--video",
+                        help="Input video file",
+                        type=str,
+                        default="")
+
+    args = parser.parse_args()
+    return args
+
+def main():
+    args = parse_args()
+    update_config(cfg, args)
+
+    # Create a logger
+    logger, final_output_dir, tb_log_dir = create_logger(
+        cfg, args.cfg, 'valid')
+    logger.info(pprint.pformat(args))
+    logger.info(cfg)
+
+    # cudnn related setting
+    cudnn.benchmark = cfg.CUDNN.BENCHMARK
+    torch.backends.cudnn.deterministic = cfg.CUDNN.DETERMINISTIC
+    torch.backends.cudnn.enabled = cfg.CUDNN.ENABLED
+
+    # Configure model
+    model = eval('models.'+cfg.MODEL.NAME+'.get_pose_net')(
+        cfg, is_train=False
+    )
+
+    if cfg.TEST.MODEL_FILE:
+        logger.info('=> loading model from {}'.format(cfg.TEST.MODEL_FILE))
+        model.load_state_dict(torch.load(cfg.TEST.MODEL_FILE), strict=False)
+    else:
+        model_state_file = os.path.join(
+            final_output_dir, 'final_state.pth'
+        )
+        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()
+
+    # define loss function (criterion) and optimizer
+    criterion = JointsMSELoss(
+        use_target_weight=cfg.LOSS.USE_TARGET_WEIGHT
+    ).cuda()
+
+    # Data loading code
+    normalize = transforms.Normalize(
+        mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
+    )
+
+    # Load data
+    if args.input != "":
+        with open(args.input, "r") as image:
+            # TODO: Write a way to handle single images
+            # TODO: Handle visualization
+            pass
+    elif args.video:
+        # TODO: Write a way to handle videos image by image
+        # TODO: Handle visualization
+        pass
+    else:
+        # Original dataset way
+        valid_dataset = eval('dataset.'+cfg.DATASET.DATASET)(
+            cfg, cfg.DATASET.ROOT, cfg.DATASET.TEST_SET, False,
+            transforms.Compose([
+                transforms.ToTensor(),
+                normalize,
+            ])
+        )
+
+        valid_loader = torch.utils.data.DataLoader(
+            valid_dataset,
+            batch_size=cfg.TEST.BATCH_SIZE_PER_GPU*len(cfg.GPUS),
+            shuffle=False,
+            num_workers=cfg.WORKERS,
+            pin_memory=True
+        )
+
+        # evaluate on validation set
+        validate(cfg, valid_loader, valid_dataset, model, criterion,
+             final_output_dir, tb_log_dir)
+
+if __name__ == '__main__':
+    main()