Make it possible to specify star model paths outside of source

README.md

@@ -97,6 +97,20 @@ path_female_star = '/mypath/female/model.npz'
 path_neutral_star = '/mypath/neutral/model.npz'
 ```
 
+Alternatively, place a file called `star.ini` in your user home directory, i.e.
+- Linux: `/home/username/star.ini`
+- Windows: `C:\Users\username\star.ini`
+- MacOs: `/Users/username/star.ini`
+
+The `star.ini` file should specify the paths, e.g.
+```
+[DEFAULT]
+path_male_star = /Users/username/star_1_1/male/model.npz
+path_female_star = /Users/username/star_1_1/female/model.npz
+path_neutral_star = /Users/username/star_1_1/neutral/model.npz
+data_type = float32
+```
+
 7. Install with pip
 ```
 pip install .

demo/load_torch.py

@@ -34,11 +34,13 @@
 batch_size=1
 m = STAR(gender='male',num_betas=num_betas)
 
+device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
+
 # Zero pose
-poses = torch.cuda.FloatTensor(np.zeros((batch_size,72)))
-betas = torch.cuda.FloatTensor(betas)
+poses = torch.FloatTensor(np.zeros((batch_size,72)), device=device)
+betas = torch.FloatTensor(betas, device=device)
 
-trans = torch.cuda.FloatTensor(np.zeros((batch_size,3)))
+trans = torch.FloatTensor(np.zeros((batch_size,3)), device=device)
 model = star.forward(poses, betas,trans)
 shaped = model.v_shaped[-1, :, :]
 

star/config.py

@@ -19,13 +19,25 @@
 # Code Developed by:
 # Ahmed A. A. Osman
 
-import os
+from configparser import ConfigParser
+from pathlib import Path
+
 path_male_star = ''
 path_female_star = ''
 path_neutral_star = ''
-
 data_type = 'float32'
 
+# allow users to specify paths in a config file outside the source
+config_path = Path.home() / 'star.ini'
+if config_path.exists():
+    config = ConfigParser()
+    config.read(config_path)
+
+    path_male_star = config['DEFAULT'].get('path_male_star', '')
+    path_female_star = config['DEFAULT'].get('path_female_star', '')
+    path_neutral_star = config['DEFAULT'].get('path_neutral_star', '')
+    data_type = config['DEFAULT'].get('data_type', 'float32')
+
 if data_type not in ['float16','float32','float64']:
     raise RuntimeError('Invalid data type %s'%(data_type))
 

star/pytorch/star.py

@@ -33,7 +33,7 @@
 
 
 class STAR(nn.Module):
-    def __init__(self,gender='female',num_betas=10):
+    def __init__(self,gender='female',num_betas=10,device=None):
         super(STAR, self).__init__()
 
         if gender not in ['male','female','neutral']:
@@ -55,18 +55,21 @@ def __init__(self,gender='female',num_betas=10):
         rows,cols = np.where(J_regressor!=0)
         vals = J_regressor[rows,cols]
         self.num_betas = num_betas
+
+        if device:
+            device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
 
         # Model sparse joints regressor, regresses joints location from a mesh
-        self.register_buffer('J_regressor', torch.cuda.FloatTensor(J_regressor))
+        self.register_buffer('J_regressor', torch.FloatTensor(J_regressor).to(device))
 
         # Model skinning weights
-        self.register_buffer('weights', torch.cuda.FloatTensor(star_model['weights']))
+        self.register_buffer('weights', torch.FloatTensor(star_model['weights']).to(device))
         # Model pose corrective blend shapes
-        self.register_buffer('posedirs', torch.cuda.FloatTensor(star_model['posedirs'].reshape((-1,93))))
+        self.register_buffer('posedirs', torch.FloatTensor(star_model['posedirs'].reshape((-1,93))).to(device))
         # Mean Shape
-        self.register_buffer('v_template', torch.cuda.FloatTensor(star_model['v_template']))
+        self.register_buffer('v_template', torch.FloatTensor(star_model['v_template']).to(device))
         # Shape corrective blend shapes
-        self.register_buffer('shapedirs', torch.cuda.FloatTensor(np.array(star_model['shapedirs'][:,:,:num_betas])))
+        self.register_buffer('shapedirs', torch.FloatTensor(np.array(star_model['shapedirs'][:,:,:num_betas])).to(device))
         # Mesh traingles
         self.register_buffer('faces', torch.from_numpy(star_model['f'].astype(np.int64)))
         self.f = star_model['f']

star/pytorch/utils.py

@@ -80,8 +80,9 @@ def with_zeros(input):
     :param input: A tensor of dimensions batch size x 3 x 4
     :return: A tensor batch size x 4 x 4 (appended with 0,0,0,1)
     '''
+    device = input.device
     batch_size  = input.shape[0]
-    row_append     = torch.cuda.FloatTensor(([0.0, 0.0, 0.0, 1.0]))
+    row_append     = torch.FloatTensor(([0.0, 0.0, 0.0, 1.0])).to(device)
     row_append.requires_grad = False
     padded_tensor     = torch.cat([input, row_append.view(1, 1, 4).repeat(batch_size, 1, 1)], 1)
     return padded_tensor