From 8c277d79c6f9462e497fb77fed3e85c6b7630cb3 Mon Sep 17 00:00:00 2001
From: D-X-Y <280835372@qq.com>
Date: Tue, 14 Jan 2020 19:53:27 +1100
Subject: [PATCH] update CPU mode for SBR

---
 .gitignore       |  2 ++
 SBR/README.md    |  7 ++++++-
 SBR/exps/eval.py | 16 ++++++++++------
 3 files changed, 18 insertions(+), 7 deletions(-)

diff --git a/.gitignore b/.gitignore
index 787a0dd0..ab68cf18 100644
--- a/.gitignore
+++ b/.gitignore
@@ -16,3 +16,5 @@ cache_data/full.pdf
 cache_data/cache/
 *.zip
 */*.zip
+.DS_Store
+*/.DS_Store
diff --git a/SBR/README.md b/SBR/README.md
index f31ac11f..f39504d8 100644
--- a/SBR/README.md
+++ b/SBR/README.md
@@ -44,7 +44,7 @@ image_path annotation_path x1 y1 x2 y2 (face_size)
 
 ## Training
 
-See the `configs` directory for some example configurations, and a pre-trained model on 300-W is at [Google Driver](https://drive.google.com/drive/folders/1ylMoVuUaNPqP7GSeWS3yE-wfU9JEJSSu).
+See the `configs` directory for some example configurations, and a pre-trained model on 300-W is at [Google Driver](https://drive.google.com/drive/folders/1ylMoVuUaNPqP7GSeWS3yE-wfU9JEJSSu) and [Google Driver](https://drive.google.com/open?id=1znVYaCheFJ6VbZmzS1fcfkHoOV4K9H6L).
 
 ### Basic Training
 ```
@@ -88,6 +88,11 @@ CUDA_VISIBLE_DEVICES=0 python ./exps/eval.py --image ./cache_data/cache/self.jpe
 - face  : the face bounding box
 - save  : save the visualized results
 
+Or if you want to run on CPU:
+```
+python ./exps/eval.py --cpu --image ./cache_data/cache/self.jpeg --model ./snapshots/300W-CPM-DET/checkpoint/cpm_vgg16-epoch-049-050.pth --face 250 150 900 1100 --save ./cache_data/cache/test.jpeg
+```
+
 
 ## Visualization
 
diff --git a/SBR/exps/eval.py b/SBR/exps/eval.py
index 744f69cb..f230725c 100644
--- a/SBR/exps/eval.py
+++ b/SBR/exps/eval.py
@@ -23,9 +23,10 @@
 
 
 def evaluate(args):
-  assert torch.cuda.is_available(), 'CUDA is not available.'
-  torch.backends.cudnn.enabled   = True
-  torch.backends.cudnn.benchmark = True
+  if not args.cpu:
+    assert torch.cuda.is_available(), 'CUDA is not available.'
+    torch.backends.cudnn.enabled   = True
+    torch.backends.cudnn.benchmark = True
 
   print ('The image is {:}'.format(args.image))
   print ('The model is {:}'.format(args.model))
@@ -33,7 +34,8 @@ def evaluate(args):
   assert snapshot.exists(), 'The model path {:} does not exist'
   print ('The face bounding box is {:}'.format(args.face))
   assert len(args.face) == 4, 'Invalid face input : {:}'.format(args.face)
-  snapshot = torch.load(snapshot)
+  if args.cpu: snapshot = torch.load(snapshot, map_location='cpu')
+  else       : snapshot = torch.load(snapshot)
 
   # General Data Argumentation
   mean_fill   = tuple( [int(x*255) for x in [0.485, 0.456, 0.406] ] )
@@ -47,7 +49,7 @@ def evaluate(args):
   dataset.reset(param.num_pts)
   
   net = obtain_model(model_config, param.num_pts + 1)
-  net = net.cuda()
+  if not args.cpu: net = net.cuda()
   #import pdb; pdb.set_trace()
   try:
     weights = remove_module_dict(snapshot['detector'])
@@ -58,7 +60,8 @@ def evaluate(args):
   [image, _, _, _, _, _, cropped_size], meta = dataset.prepare_input(args.image, args.face)
   # network forward
   with torch.no_grad():
-    inputs = image.unsqueeze(0).cuda()
+    if args.cpu: inputs = image.unsqueeze(0)
+    else       : inputs = image.unsqueeze(0).cuda()
     batch_heatmaps, batch_locs, batch_scos = net(inputs)
     flops, params = get_model_infos(net, inputs.shape)
     print ('IN-shape : {:}, FLOPs : {:} MB, Params : {:} MB'.format(list(inputs.shape), flops, params))
@@ -92,5 +95,6 @@ def evaluate(args):
   parser.add_argument('--model',            type=str,   help='The snapshot to the saved detector.')
   parser.add_argument('--face',  nargs='+', type=float, help='The coordinate [x1,y1,x2,y2] of a face')
   parser.add_argument('--save',             type=str,   help='The path to save the visualized results.')
+  parser.add_argument('--cpu',     action='store_true', help='Use CPU or not.')
   args = parser.parse_args()
   evaluate(args)